Upload Batch_6.py

Batch_6.py

@@ -0,0 +1,124 @@
+# Save as: ComfyUI/custom_nodes/batch_merge_6_any.py
+# Restart ComfyUI after saving.
+
+import torch
+
+
+class Batch_6:
+    """
+    Takes up to 6 IMAGE inputs (each can be a single image [H,W,C] or a batch [B,H,W,C]),
+    RGB or RGBA, and concatenates them into one batch.
+
+    - IMAGES_1 is required (so the node can always run).
+    - IMAGES_2..IMAGES_6 are optional (can be left unconnected).
+
+    Channel handling:
+      - If ANY input is RGBA (C=4), output will be RGBA.
+      - RGB inputs (C=3) will be upgraded to RGBA by adding alpha=1.
+      - If all inputs are RGB, output stays RGB.
+
+    Requirements:
+      - All images must share the same H and W (no resizing/cropping is done).
+      - Channels must be 3 or 4.
+    """
+
+    CATEGORY = "image/batch"
+    FUNCTION = "merge"
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("IMAGES_OUT",)
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "IMAGES_1": ("IMAGE",),
+            },
+            "optional": {
+                "IMAGES_2": ("IMAGE",),
+                "IMAGES_3": ("IMAGE",),
+                "IMAGES_4": ("IMAGE",),
+                "IMAGES_5": ("IMAGE",),
+                "IMAGES_6": ("IMAGE",),
+            },
+        }
+
+    @staticmethod
+    def _normalize_to_batch(t: torch.Tensor) -> torch.Tensor:
+        # Accept [H,W,C] as single image and convert to [1,H,W,C]
+        if t.dim() == 3:
+            return t.unsqueeze(0)
+        if t.dim() == 4:
+            return t
+        raise ValueError(f"Expected IMAGE tensor with 3 or 4 dims, got shape {tuple(t.shape)}")
+
+    @staticmethod
+    def _ensure_channels(t: torch.Tensor) -> int:
+        if t.dim() != 4:
+            raise ValueError(f"Expected [B,H,W,C], got shape {tuple(t.shape)}")
+        c = int(t.shape[-1])
+        if c not in (3, 4):
+            raise ValueError(f"Expected RGB/RGBA (C=3 or 4), got C={c}")
+        return c
+
+    def merge(self, IMAGES_1, IMAGES_2=None, IMAGES_3=None, IMAGES_4=None, IMAGES_5=None, IMAGES_6=None):
+        inputs = [IMAGES_1, IMAGES_2, IMAGES_3, IMAGES_4, IMAGES_5, IMAGES_6]
+
+        tensors = []
+        for idx, x in enumerate(inputs, start=1):
+            if x is None:
+                continue
+            if not isinstance(x, torch.Tensor):
+                raise TypeError(f"IMAGES_{idx} is not a torch.Tensor (got {type(x)})")
+
+            x = self._normalize_to_batch(x)
+            self._ensure_channels(x)
+            tensors.append(x)
+
+        if len(tensors) == 0:
+            # Shouldn't happen because IMAGES_1 is required, but keep it safe.
+            raise ValueError("No images provided.")
+
+        # Use first input as reference for device/dtype/size
+        ref = tensors[0]
+        device = ref.device
+        dtype = ref.dtype
+        H = int(ref.shape[1])
+        W = int(ref.shape[2])
+
+        # Decide output channels: RGBA if any input is RGBA
+        target_c = 4 if any(int(t.shape[-1]) == 4 for t in tensors) else 3
+
+        prepared = []
+        for i, t in enumerate(tensors):
+            # Align device/dtype
+            if t.device != device or t.dtype != dtype:
+                t = t.to(device=device, dtype=dtype)
+
+            # Validate size
+            if int(t.shape[1]) != H or int(t.shape[2]) != W:
+                raise ValueError(
+                    f"Size mismatch: input #{i+1} has [H,W]=[{int(t.shape[1])},{int(t.shape[2])}] "
+                    f"but expected [{H},{W}]."
+                )
+
+            c = int(t.shape[-1])
+
+            # Upgrade RGB -> RGBA if needed
+            if target_c == 4 and c == 3:
+                alpha = torch.ones((int(t.shape[0]), H, W, 1), device=device, dtype=dtype)
+                t = torch.cat([t, alpha], dim=-1)
+
+            # (No need to drop alpha because target_c is 3 only if all are 3)
+            prepared.append(t)
+
+        out = torch.cat(prepared, dim=0)
+        return (out,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "Batch_6": Batch_6,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Batch_6": "Batch 6",
+}