From 3d7577f31617380cf4f32dbb1858918d51d7e1c6 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sat, 6 Jul 2024 14:35:17 +0300
Subject: [PATCH] Update GetImageRangeFromBatch to work with only masks too

---
 __init__.py          |  2 +-
 nodes/image_nodes.py | 39 ++++++++++++++++++++++++++++-----------
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/__init__.py b/__init__.py
index 9b5d6bf..8f422f8 100644
--- a/__init__.py
+++ b/__init__.py
@@ -43,7 +43,7 @@ NODE_CONFIG = {
     "ColorMatch": {"class": ColorMatch, "name": "Color Match"},
     "CrossFadeImages": {"class": CrossFadeImages, "name": "Cross Fade Images"},
     "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"},
     "ImageAndMaskPreview": {"class": ImageAndMaskPreview},
     "ImageAddMulti": {"class": ImageAddMulti, "name": "Image Add Multi"},
diff --git a/nodes/image_nodes.py b/nodes/image_nodes.py
index 4d2a320..1ccdbf5 100644
--- a/nodes/image_nodes.py
+++ b/nodes/image_nodes.py
@@ -988,25 +988,42 @@ batch, starting from start_index.
     def INPUT_TYPES(s):
         return {
             "required": {
-                 "images": ("IMAGE",),
                  "start_index": ("INT", {"default": 0,"min": -1, "max": 4096, "step": 1}),
                  "num_frames": ("INT", {"default": 1,"min": 1, "max": 4096, "step": 1}),
         },
         "optional": {
+            "images": ("IMAGE",),
             "masks": ("MASK",),
         }
     } 
     
-    def imagesfrombatch(self, images, start_index, num_frames, masks=None):
-        if start_index == -1:
-            start_index = len(images) - num_frames
-        if start_index < 0 or start_index >= len(images):
-            raise ValueError("GetImageRangeFromBatch: Start index is out of range")
-        end_index = start_index + num_frames
-        if end_index > len(images):
-            raise ValueError("GetImageRangeFromBatch: End index is out of range")
-        chosen_images = images[start_index:end_index]
-        chosen_masks = masks[start_index:end_index] if masks is not None else 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:
+                start_index = len(images) - num_frames
+            if start_index < 0 or start_index >= len(images):
+                raise ValueError("Start index is out of range")
+            end_index = start_index + num_frames
+            if end_index > len(images):
+                raise ValueError("End index is out of range")
+            chosen_images = images[start_index:end_index]
+
+        # 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,)
     
 class GetImagesFromBatchIndexed: