From 179c462e9b702e816c6ce17436622490a5910ada Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 5 Nov 2023 11:37:37 +0200
Subject: [PATCH] Add image concanate for custom grids

---
 nodes.py | 37 +++++++++++++++++++++++++++++++------
 1 file changed, 31 insertions(+), 6 deletions(-)

diff --git a/nodes.py b/nodes.py
index 09cccd3..c793c42 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1100,7 +1100,32 @@ class SaveImageWithAlpha:
 
         return { "ui": { "images": results } }
 
+class ImageConcanate:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "image1": ("IMAGE",),
+            "image2": ("IMAGE",),
+            "direction": (
+            [   'right',
+                'down',
+            ],
+            {
+            "default": 'right'
+             }),
+        }}
 
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "concanate"
+    CATEGORY = "KJNodes"
+
+    def concanate(self, image1, image2, direction):
+        if direction == 'right':
+            row = torch.cat((image1, image2), dim=2)
+        elif direction == 'down':
+            row = torch.cat((image1, image2), dim=1)
+        return (row,)
+    
 class ImageGridComposite2x2:
     @classmethod
     def INPUT_TYPES(s):
@@ -1108,8 +1133,7 @@ class ImageGridComposite2x2:
             "image1": ("IMAGE",),
             "image2": ("IMAGE",),
             "image3": ("IMAGE",),
-            "image4": ("IMAGE",),
-            
+            "image4": ("IMAGE",),   
         }}
 
     RETURN_TYPES = ("IMAGE",)
@@ -1134,8 +1158,7 @@ class ImageGridComposite3x3:
             "image6": ("IMAGE",),
             "image7": ("IMAGE",),
             "image8": ("IMAGE",),
-            "image9": ("IMAGE",),
-            
+            "image9": ("IMAGE",),     
         }}
 
     RETURN_TYPES = ("IMAGE",)
@@ -1173,7 +1196,8 @@ NODE_CLASS_MAPPINGS = {
     "SaveImageWithAlpha": SaveImageWithAlpha,
     "ReverseImageBatch": ReverseImageBatch,
     "ImageGridComposite2x2": ImageGridComposite2x2,
-    "ImageGridComposite3x3": ImageGridComposite3x3
+    "ImageGridComposite3x3": ImageGridComposite3x3,
+    "ImageConcanate": ImageConcanate
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "INTConstant": "INT Constant",
@@ -1198,5 +1222,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "SaveImageWithAlpha": "SaveImageWithAlpha",
     "ReverseImageBatch": "ReverseImageBatch",
     "ImageGridComposite2x2": "ImageGridComposite2x2",
-    "ImageGridComposite3x3": "ImageGridComposite3x3"
+    "ImageGridComposite3x3": "ImageGridComposite3x3",
+    "ImageConcanate": "ImageConcanate"
 }
\ No newline at end of file