From a7141dc5ca567df698efa66e9357529b38967744 Mon Sep 17 00:00:00 2001
From: kijai <jukka.seppanen@poke.fi>
Date: Fri, 17 May 2024 12:21:59 +0300
Subject: [PATCH] Add optional mask batch slicing to GetImageRangeFromBatch

---
 nodes/image_nodes.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/nodes/image_nodes.py b/nodes/image_nodes.py
index 0fa04b7..38c6fc2 100644
--- a/nodes/image_nodes.py
+++ b/nodes/image_nodes.py
@@ -955,7 +955,7 @@ class CrossFadeImages:
 
 class GetImageRangeFromBatch:
     
-    RETURN_TYPES = ("IMAGE",)
+    RETURN_TYPES = ("IMAGE", "MASK", )
     FUNCTION = "imagesfrombatch"
     CATEGORY = "KJNodes/image"
     DESCRIPTION = """
@@ -971,9 +971,12 @@ batch, starting from start_index.
                  "start_index": ("INT", {"default": 0,"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:
             start_index = len(images) - num_frames
         if start_index < 0 or start_index >= len(images):
@@ -982,7 +985,8 @@ batch, starting from start_index.
         if end_index > len(images):
             raise ValueError("GetImageRangeFromBatch: End index is out of range")
         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: