This commit is contained in:
Kijai 2024-05-29 15:04:10 +03:00
commit a15d79d5dc

View File

@ -955,7 +955,7 @@ class CrossFadeImages:
class GetImageRangeFromBatch: class GetImageRangeFromBatch:
RETURN_TYPES = ("IMAGE",) RETURN_TYPES = ("IMAGE", "MASK", )
FUNCTION = "imagesfrombatch" FUNCTION = "imagesfrombatch"
CATEGORY = "KJNodes/image" CATEGORY = "KJNodes/image"
DESCRIPTION = """ DESCRIPTION = """
@ -971,9 +971,12 @@ batch, starting from start_index.
"start_index": ("INT", {"default": 0,"min": -1, "max": 4096, "step": 1}), "start_index": ("INT", {"default": 0,"min": -1, "max": 4096, "step": 1}),
"num_frames": ("INT", {"default": 1,"min": 1, "max": 4096, "step": 1}), "num_frames": ("INT", {"default": 1,"min": 1, "max": 4096, "step": 1}),
}, },
"optional": {
"masks": ("MASK",),
}
} }
def imagesfrombatch(self, images, start_index, num_frames): def imagesfrombatch(self, images, start_index, num_frames, masks=None):
if start_index == -1: if start_index == -1:
start_index = len(images) - num_frames start_index = len(images) - num_frames
if start_index < 0 or start_index >= len(images): if start_index < 0 or start_index >= len(images):
@ -982,7 +985,8 @@ batch, starting from start_index.
if end_index > len(images): if end_index > len(images):
raise ValueError("GetImageRangeFromBatch: End index is out of range") raise ValueError("GetImageRangeFromBatch: End index is out of range")
chosen_images = images[start_index:end_index] chosen_images = images[start_index:end_index]
return (chosen_images, ) chosen_masks = masks[start_index:end_index] if masks is not None else None
return (chosen_images, chosen_masks,)
class GetImagesFromBatchIndexed: class GetImagesFromBatchIndexed: