Add ImageBatch module.

Loader.py

@@ -72,6 +72,10 @@ class Loader:
             from .modules import Clamp
             modules.update(Clamp.NODE_CLASS_MAPPINGS)
 
+        if self.config()["modules"]["ImageBatch"]:
+            from .modules import ImageBatch
+            modules.update(ImageBatch.NODE_CLASS_MAPPINGS)
+
         if self.config()["modules"]["ImageComposite"]:
             from .modules import ImageComposite
             modules.update(ImageComposite.NODE_CLASS_MAPPINGS)

README.MD

@@ -905,6 +905,53 @@ You also can change the fonts folder in config.
 
 </details>
 
+---
+### Image Batch
+![image_batch](images/image_batch.png)
+> Allows you to manage batch of images.
+
+<details>
+  <summary>Nodes:</summary>
+
+### Get
+![image_batch_get](images/image_batch_get.png)
+> Get image from images batch.
+
+<details>
+  <summary>Params:</summary>
+
+* index `[1 - *]` - Number of image.
+
+</details>
+
+### Remove
+![image_batch_remove](images/image_batch_remove.png)
+> Remove image from images batch.
+
+<details>
+  <summary>Params:</summary>
+
+* index `[1 - *]` - Number of image.
+
+</details>
+
+### Fork
+![image_batch_fork](images/image_batch_fork.png)
+> Split images from batch.
+
+<details>
+  <summary>Params:</summary>
+
+* priority `[first, second]` - If there are an odd number of images, choose which stream to send the image from the center to.
+
+</details>
+
+### Join
+![image_batch_join](images/image_batch_join.png)
+> Combine images together into batch.
+
+</details>
+
 </details>
 
 ## Configuration:

config.json

@@ -3,6 +3,7 @@
   "modules": {
     "AlphaChanel": true,
     "Clamp": true,
+    "ImageBatch": true,
     "ImageComposite": true,
     "ImageContainer": true,
     "ImageDraw": true,

modules/ImageBatch.py

@@ -0,0 +1,136 @@
+import torch
+
+
+class ImageBatchGet:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "index": ("INT", {
+                    "default": 1,
+                    "min": 1,
+                    "step": 1
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "image_batch_get"
+    CATEGORY = "image/batch"
+
+    def image_batch_get(self, images, index):
+        batch = images.shape[0]
+        index = min(batch, index - 1)
+
+        return (images[index].unsqueeze(0),)
+
+
+class ImageBatchRemove:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "index": ("INT", {
+                    "default": 1,
+                    "min": 1,
+                    "step": 1
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "image_batch_get"
+    CATEGORY = "image/batch"
+
+    def image_batch_get(self, images, index):
+        batch = images.shape[0]
+        index = min(batch, index - 1)
+
+        return (torch.cat((images[:index], images[index + 1:]), dim=0),)
+
+
+class ImageBatchFork:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "priority": (["first", "second"],),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE", "IMAGE")
+    FUNCTION = "image_batch_join"
+    CATEGORY = "image/batch"
+
+    def image_batch_join(self, images, priority):
+        batch = images.shape[0]
+
+        if batch == 1:
+            return images, images
+        elif batch % 2 == 0:
+            first = batch // 2
+            second = batch // 2
+        else:
+            if priority == "first":
+                first = batch // 2 + 1
+                second = batch // 2
+            elif priority == "second":
+                first = batch // 2
+                second = batch // 2 + 1
+            else:
+                raise ValueError("Not existing priority.")
+
+        return images[:first], images[-second:]
+
+
+class ImageBatchJoin:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images_a": ("IMAGE",),
+                "images_b": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "image_batch_join"
+    CATEGORY = "image/batch"
+
+    def image_batch_join(self, images_a, images_b):
+        height_a, width_a, channels_a = images_a[0].shape
+        height_b, width_b, channels_b = images_b[0].shape
+
+        if height_a != height_b:
+            raise ValueError("Height of images_a not equals of images_b. You can use ImageTransformResize for fix it.")
+
+        if width_a != width_b:
+            raise ValueError("Width of images_a not equals of images_b. You can use ImageTransformResize for fix it.")
+
+        if channels_a != channels_b:
+            raise ValueError("Channels of images_a not equals of images_b. Your can add or delete alpha channels with AlphaChanel module.")
+
+        return (torch.cat((images_a, images_b)),)
+
+
+NODE_CLASS_MAPPINGS = {
+    "ImageBatchGet": ImageBatchGet,
+    "ImageBatchRemove": ImageBatchRemove,
+    "ImageBatchFork": ImageBatchFork,
+    "ImageBatchJoin": ImageBatchJoin
+}