Update GetImageRangeFromBatch to work with only masks too

This commit is contained in:
kijai 2024-07-06 14:35:17 +03:00
parent 2ead4fae1d
commit 3d7577f316
2 changed files with 29 additions and 12 deletions

View File

@ -43,7 +43,7 @@ NODE_CONFIG = {
"ColorMatch": {"class": ColorMatch, "name": "Color Match"}, "ColorMatch": {"class": ColorMatch, "name": "Color Match"},
"CrossFadeImages": {"class": CrossFadeImages, "name": "Cross Fade Images"}, "CrossFadeImages": {"class": CrossFadeImages, "name": "Cross Fade Images"},
"GetImagesFromBatchIndexed": {"class": GetImagesFromBatchIndexed, "name": "Get Images From Batch Indexed"}, "GetImagesFromBatchIndexed": {"class": GetImagesFromBatchIndexed, "name": "Get Images From Batch Indexed"},
"GetImageRangeFromBatch": {"class": GetImageRangeFromBatch, "name": "Get Image Range From Batch"}, "GetImageRangeFromBatch": {"class": GetImageRangeFromBatch, "name": "Get Image or Mask Range From Batch"},
"GetImageSizeAndCount": {"class": GetImageSizeAndCount, "name": "Get Image Size & Count"}, "GetImageSizeAndCount": {"class": GetImageSizeAndCount, "name": "Get Image Size & Count"},
"ImageAndMaskPreview": {"class": ImageAndMaskPreview}, "ImageAndMaskPreview": {"class": ImageAndMaskPreview},
"ImageAddMulti": {"class": ImageAddMulti, "name": "Image Add Multi"}, "ImageAddMulti": {"class": ImageAddMulti, "name": "Image Add Multi"},

View File

@ -988,25 +988,42 @@ batch, starting from start_index.
def INPUT_TYPES(s): def INPUT_TYPES(s):
return { return {
"required": { "required": {
"images": ("IMAGE",),
"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": { "optional": {
"images": ("IMAGE",),
"masks": ("MASK",), "masks": ("MASK",),
} }
} }
def imagesfrombatch(self, images, start_index, num_frames, masks=None): def imagesfrombatch(self, start_index, num_frames, images=None, masks=None):
chosen_images = None
chosen_masks = None
# Process images if provided
if images is not 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):
raise ValueError("GetImageRangeFromBatch: Start index is out of range") raise ValueError("Start index is out of range")
end_index = start_index + num_frames end_index = start_index + num_frames
if end_index > len(images): if end_index > len(images):
raise ValueError("GetImageRangeFromBatch: End index is out of range") raise ValueError("End index is out of range")
chosen_images = images[start_index:end_index] chosen_images = images[start_index:end_index]
chosen_masks = masks[start_index:end_index] if masks is not None else None
# Process masks if provided
if masks is not None:
if start_index == -1:
start_index = len(masks) - num_frames
if start_index < 0 or start_index >= len(masks):
raise ValueError("Start index is out of range for masks")
end_index = start_index + num_frames
if end_index > len(masks):
raise ValueError("End index is out of range for masks")
chosen_masks = masks[start_index:end_index]
return (chosen_images, chosen_masks,) return (chosen_images, chosen_masks,)
class GetImagesFromBatchIndexed: class GetImagesFromBatchIndexed: