feat: 🚨 Image Stack node (horizontal and vertical stack)

with dynamic inputs

nodes/image_utils.py

@@ -0,0 +1,41 @@
+import torch
+
+from ..log import log
+
+
+class StackImages:
+    """Stack the input images horizontally or vertically"""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {"vertical": ("BOOLEAN", {"default": False})}}
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "stack"
+    CATEGORY = "mtb/image utils"
+
+    def stack(self, vertical, **kwargs):
+        if not kwargs:
+            raise ValueError("At least one tensor must be provided.")
+
+        tensors = list(kwargs.values())
+        log.debug(
+            f"Stacking {len(tensors)} tensors {'vertically' if vertical else 'horizontally'}"
+        )
+        log.debug(list(kwargs.keys()))
+
+        ref_shape = tensors[0].shape
+        for tensor in tensors[1:]:
+            if tensor.shape[1:] != ref_shape[1:]:
+                raise ValueError(
+                    "All tensors must have the same dimensions except for the stacking dimension."
+                )
+
+        dim = 1 if vertical else 2
+
+        stacked_tensor = torch.cat(tensors, dim=dim)
+
+        return (stacked_tensor,)
+
+
+__nodes__ = [StackImages]

web/comfy_shared.js

@@ -115,6 +115,7 @@ export const dynamic_connection = (
     // make inputs sequential again
     for (let i = 0; i < node.inputs.length; i++) {
       node.inputs[i].label = `${connectionPrefix}${i + 1}`
+      node.inputs[i].name = `${connectionPrefix}${i + 1}`
     }
   }
 

web/mtb_widgets.js

@@ -877,6 +877,30 @@ const mtb_widgets = {
 
         break
       }
+      case 'Stack Images (mtb)': {
+        const onNodeCreated = nodeType.prototype.onNodeCreated
+        nodeType.prototype.onNodeCreated = function () {
+          const r = onNodeCreated
+            ? onNodeCreated.apply(this, arguments)
+            : undefined
+          this.addInput(`image_1`, 'IMAGE')
+          return r
+        }
+
+        const onConnectionsChange = nodeType.prototype.onConnectionsChange
+        nodeType.prototype.onConnectionsChange = function (
+          type,
+          index,
+          connected,
+          link_info
+        ) {
+          const r = onConnectionsChange
+            ? onConnectionsChange.apply(this, arguments)
+            : undefined
+          shared.dynamic_connection(this, index, connected, 'image_', 'IMAGE')
+        }
+        break
+      }
       case 'Save Tensors (mtb)': {
         const onDrawBackground = nodeType.prototype.onDrawBackground
         nodeType.prototype.onDrawBackground = function (ctx, canvas) {