Upload salia_Get_First_and_Last_Batch.py

salia_Get_First_and_Last_Batch.py

@@ -0,0 +1,53 @@
+# Save this as: ComfyUI/custom_nodes/first_last_image.py
+# Restart ComfyUI after saving.
+
+import torch
+
+
+class FirstLastImageFromBatch:
+    """
+    Takes a batch of ComfyUI IMAGE tensors and outputs:
+      - first image (as a batch of size 1)
+      - last image  (as a batch of size 1)
+
+    ComfyUI IMAGE is typically a torch.Tensor with shape: [B, H, W, C]
+    """
+
+    CATEGORY = "image/batch"
+    FUNCTION = "first_last"
+    RETURN_TYPES = ("IMAGE", "IMAGE")
+    RETURN_NAMES = ("first", "last")
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {"images": ("IMAGE",)}}
+
+    def first_last(self, images):
+        if not isinstance(images, torch.Tensor):
+            raise TypeError(f"Expected IMAGE as torch.Tensor, got {type(images)}")
+
+        # Accept single image [H, W, C] as batch size 1
+        if images.dim() == 3:
+            images = images.unsqueeze(0)
+        elif images.dim() != 4:
+            raise ValueError(f"Expected images with 3 or 4 dims, got shape {tuple(images.shape)}")
+
+        b = images.shape[0]
+        if b < 1:
+            raise ValueError("Received an empty image batch (B=0).")
+
+        # Slice to keep batch dimension, return batches of size 1
+        first = images[0:1]
+        last = images[-1:]
+
+        # Optional: clone to avoid any accidental in-place edits elsewhere affecting outputs
+        return (first.clone(), last.clone())
+
+
+NODE_CLASS_MAPPINGS = {
+    "FirstLastImageFromBatch": FirstLastImageFromBatch,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "FirstLastImageFromBatch": "First + Last Image (from batch)",
+}