File size: 1,611 Bytes
e409f0f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | # Save as: ComfyUI/custom_nodes/batch_remove_first_last.py
# Restart ComfyUI after saving.
import torch
class BatchRemoveFirstLast:
"""
Takes an IMAGE batch and returns the same batch, except:
- removes the FIRST image (index 0)
- removes the LAST image (index B-1)
Output = images[1:-1]
Notes:
- If the batch has fewer than 3 images (B < 3), removing both ends would
produce an empty/invalid batch, so this node returns the original batch.
- If a single image comes in as [H, W, C], it is treated as a batch of 1.
"""
CATEGORY = "image/batch"
FUNCTION = "remove_first_last"
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("images",)
@classmethod
def INPUT_TYPES(cls):
return {"required": {"images": ("IMAGE",)}}
def remove_first_last(self, images):
if not isinstance(images, torch.Tensor):
# Best-effort fallback
return (images,)
# Normalize to [B, H, W, C]
if images.dim() == 3:
images = images.unsqueeze(0)
elif images.dim() != 4:
# Unknown shape; return unchanged
return (images,)
b = int(images.shape[0])
# If too small to safely remove both ends, return original batch
if b < 3:
return (images,)
out = images[1:-1].clone()
return (out,)
NODE_CLASS_MAPPINGS = {
"BatchRemoveFirstLast": BatchRemoveFirstLast,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"BatchRemoveFirstLast": "Batch Remove First + Last",
}
|