From 0a32019ed098c159697a4d76c35dcdef8cf2cbed Mon Sep 17 00:00:00 2001
From: Nourepide <nourepide@gmail.com>
Date: Tue, 28 Nov 2023 06:29:13 +0300
Subject: [PATCH] Add ImageBatchPermute node.

---
 modules/ImageBatch.py | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/modules/ImageBatch.py b/modules/ImageBatch.py
index e0065d2..4203915 100644
--- a/modules/ImageBatch.py
+++ b/modules/ImageBatch.py
@@ -1,3 +1,5 @@
+import re
+
 import torch
 
 
@@ -128,9 +130,36 @@ class ImageBatchJoin:
         return (torch.cat((images_a, images_b)),)
 
 
+class ImageBatchPermute:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "permute": ("STRING", {"multiline": False}),
+                "start_with_zero": ("BOOLEAN",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "node"
+    CATEGORY = "image/batch"
+
+    def node(self, images, permute, start_with_zero):
+        order = [int(num) - 1 if not start_with_zero else int(num) for num in re.findall(r'\d+', permute)]
+        order = torch.tensor(order)
+        order = order.clamp(0, images.shape[0] - 1)
+
+        return (images.index_select(0, order),)
+
+
 NODE_CLASS_MAPPINGS = {
     "ImageBatchGet": ImageBatchGet,
     "ImageBatchRemove": ImageBatchRemove,
     "ImageBatchFork": ImageBatchFork,
-    "ImageBatchJoin": ImageBatchJoin
+    "ImageBatchJoin": ImageBatchJoin,
+    "ImageBatchPermute": ImageBatchPermute
 }