From aa54cbb6cf2b34a7d4d7517e65af5ea3e20be3d6 Mon Sep 17 00:00:00 2001
From: Kijai <40791699+kijai@users.noreply.github.com>
Date: Thu, 21 Dec 2023 14:52:48 +0200
Subject: [PATCH 1/2] Add GetImagesFromBatchIndexed:

---
 nodes.py | 32 +++++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/nodes.py b/nodes.py
index cddcb6c..16adc51 100644
--- a/nodes.py
+++ b/nodes.py
@@ -493,7 +493,35 @@ class GetImageRangeFromBatch:
             raise ValueError("GetImageRangeFromBatch: End index is out of range")
         chosen_images = images[start_index:end_index]
         return (chosen_images, )
-  
+    
+class GetImagesFromBatchIndexed:
+    
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "indexedimagesfrombatch"
+    CATEGORY = "KJNodes"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                 "images": ("IMAGE",),
+                 "indexes": ("STRING", {"default": "0, 1, 2", "multiline": True}),
+        },
+    } 
+    
+    def indexedimagesfrombatch(self, images, indexes):
+        # Parse the indexes string into a list of integers
+        index_list = [int(index.strip()) for index in indexes.split(',')]
+        
+        # Fetch images based on parsed indexes
+        chosen_images = [images[index] for index in index_list if 0 <= index < len(images)]
+        
+        # Check if any index was out of range and raise an error if so
+        if len(chosen_images) != len(index_list):
+            raise ValueError("One or more indexes are out of range")
+        
+        return (chosen_images,)
+    
 class ReplaceImagesInBatch:
     
     RETURN_TYPES = ("IMAGE",)
@@ -2945,6 +2973,7 @@ NODE_CLASS_MAPPINGS = {
     "SoundReactive": SoundReactive,
     "GenerateNoise": GenerateNoise,
     "StableZero123_BatchSchedule": StableZero123_BatchSchedule,
+    "GetImagesFromBatchIndexed": GetImagesFromBatchIndexed,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "INTConstant": "INT Constant",
@@ -2998,4 +3027,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "SoundReactive": "SoundReactive",
     "GenerateNoise": "GenerateNoise",
     "StableZero123_BatchSchedule": "StableZero123_BatchSchedule",
+    "GetImagesFromBatchIndexed": "GetImagesFromBatchIndexed",
 }
\ No newline at end of file

From ccce5cc2b17d6b6954ea65272874db3f63b74db1 Mon Sep 17 00:00:00 2001
From: Kijai <40791699+kijai@users.noreply.github.com>
Date: Thu, 21 Dec 2023 15:08:06 +0200
Subject: [PATCH 2/2] Update nodes.py

---
 nodes.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/nodes.py b/nodes.py
index 16adc51..75deebe 100644
--- a/nodes.py
+++ b/nodes.py
@@ -510,15 +510,15 @@ class GetImagesFromBatchIndexed:
     } 
     
     def indexedimagesfrombatch(self, images, indexes):
+        
         # Parse the indexes string into a list of integers
         index_list = [int(index.strip()) for index in indexes.split(',')]
         
-        # Fetch images based on parsed indexes
-        chosen_images = [images[index] for index in index_list if 0 <= index < len(images)]
+        # Convert list of indices to a PyTorch tensor
+        indices_tensor = torch.tensor(index_list, dtype=torch.long)
         
-        # Check if any index was out of range and raise an error if so
-        if len(chosen_images) != len(index_list):
-            raise ValueError("One or more indexes are out of range")
+        # Select the images at the specified indices
+        chosen_images = images[indices_tensor]
         
         return (chosen_images,)