Upload Custom_Batch_Output.py

Custom_Batch_Output.py

@@ -0,0 +1,86 @@
+# Save as: ComfyUI/custom_nodes/special_batch_split.py
+# Restart ComfyUI after saving.
+
+import torch
+
+
+class SpecialBatchSplit:
+    """
+    Input:
+      - images (IMAGE batch, typically torch.Tensor [B, H, W, C])
+
+    Outputs:
+      - Batch_Up:  [ ID 7 ] + [ IDs 9..25 ] + [ IDs 27..31 ] + [ IDs 33..36 ]
+      - Rife_x3:   [ ID 4 ] + [ ID 37 ]  (2-image batch)
+
+    Indexing is 0-based and ranges are inclusive (e.g., 9..25 includes both 9 and 25).
+
+    Safety behavior:
+      - If the input batch is too small (needs at least indices up to 37 => B >= 38),
+        or input is not a proper IMAGE tensor, the node returns the original input batch
+        for BOTH outputs.
+    """
+
+    CATEGORY = "image/batch"
+    FUNCTION = "make_special_batches"
+
+    RETURN_TYPES = ("IMAGE", "IMAGE")
+    RETURN_NAMES = ("Batch_Up", "Rife_x3")
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {"images": ("IMAGE",)}}
+
+    @staticmethod
+    def _normalize_to_batch(images: torch.Tensor) -> torch.Tensor:
+        # Accept single image [H,W,C] and convert to [1,H,W,C]
+        if images.dim() == 3:
+            return images.unsqueeze(0)
+        return images
+
+    def make_special_batches(self, images):
+        # Basic validation + safe fallback
+        if not isinstance(images, torch.Tensor):
+            return (images, images)
+
+        images = self._normalize_to_batch(images)
+
+        # Expect [B,H,W,C]
+        if images.dim() != 4:
+            return (images, images)
+
+        b = int(images.shape[0])
+
+        # Need indices up to 37 => batch size at least 38
+        if b < 38:
+            return (images, images)
+
+        # Build Batch_Up indices (inclusive ranges)
+        batch_up_indices = (
+            [7]
+            + list(range(9, 26))   # 9..25
+            + list(range(27, 32))  # 27..31
+            + list(range(33, 37))  # 33..36
+        )
+
+        # Build Rife_x3 indices
+        rife_x3_indices = [4, 37]
+
+        # Gather using index_select (works on GPU/CPU, preserves dtype/device)
+        device = images.device
+        idx_up = torch.tensor(batch_up_indices, dtype=torch.long, device=device)
+        idx_rife = torch.tensor(rife_x3_indices, dtype=torch.long, device=device)
+
+        batch_up = torch.index_select(images, 0, idx_up).clone()
+        rife_x3 = torch.index_select(images, 0, idx_rife).clone()
+
+        return (batch_up, rife_x3)
+
+
+NODE_CLASS_MAPPINGS = {
+    "Custom_Batch_Output": Custom_Batch_Output,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Custom_Batch_Output": "Custom_Batch_Output",
+}