Upload Batch_Remove_Img_At.py
Browse files- Batch_Remove_Img_At.py +72 -0
Batch_Remove_Img_At.py
ADDED
|
@@ -0,0 +1,72 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ComfyUI Custom Node: Remove (Exclude) Image @ Index from a Batch
|
| 2 |
+
# Save as:
|
| 3 |
+
# ComfyUI/custom_nodes/batch_remove_index/__init__.py
|
| 4 |
+
# Restart ComfyUI, then find it under category: "Batch/Index"
|
| 5 |
+
|
| 6 |
+
import torch
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
def _clamp_index(index: int, batch_size: int) -> int:
|
| 10 |
+
"""Clamp index into [0, batch_size-1]."""
|
| 11 |
+
if batch_size <= 0:
|
| 12 |
+
raise ValueError("Input batch is empty (batch_size <= 0).")
|
| 13 |
+
if index < 0 or index >= batch_size:
|
| 14 |
+
print(
|
| 15 |
+
f"[BatchRemoveIndex] index {index} out of range for batch_size {batch_size}; "
|
| 16 |
+
f"clamping to valid range."
|
| 17 |
+
)
|
| 18 |
+
index = max(0, min(index, batch_size - 1))
|
| 19 |
+
return index
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
class BatchRemoveImageAtIndex:
|
| 23 |
+
"""
|
| 24 |
+
Takes an IMAGE batch and an integer index (0-based),
|
| 25 |
+
outputs the batch WITHOUT the image at that index.
|
| 26 |
+
|
| 27 |
+
Example:
|
| 28 |
+
images = [img0,img1,img2,img3], index=2
|
| 29 |
+
output = [img0,img1,img3]
|
| 30 |
+
"""
|
| 31 |
+
|
| 32 |
+
@classmethod
|
| 33 |
+
def INPUT_TYPES(cls):
|
| 34 |
+
return {
|
| 35 |
+
"required": {
|
| 36 |
+
"images": ("IMAGE",),
|
| 37 |
+
"index": ("INT", {"default": 0, "min": 0, "max": 10**9}),
|
| 38 |
+
}
|
| 39 |
+
}
|
| 40 |
+
|
| 41 |
+
RETURN_TYPES = ("IMAGE",)
|
| 42 |
+
RETURN_NAMES = ("images",)
|
| 43 |
+
FUNCTION = "remove"
|
| 44 |
+
CATEGORY = "Batch/Index"
|
| 45 |
+
|
| 46 |
+
def remove(self, images, index):
|
| 47 |
+
if not torch.is_tensor(images):
|
| 48 |
+
raise TypeError("Expected 'images' to be a torch Tensor (ComfyUI IMAGE type).")
|
| 49 |
+
if images.ndim != 4:
|
| 50 |
+
raise ValueError(f"Expected 'images' with shape [B,H,W,C], got ndim={images.ndim}.")
|
| 51 |
+
|
| 52 |
+
b = images.shape[0]
|
| 53 |
+
if b <= 1:
|
| 54 |
+
# Producing an empty IMAGE batch often breaks downstream nodes.
|
| 55 |
+
raise ValueError(
|
| 56 |
+
f"Cannot remove an item from a batch of size {b} (would produce an empty batch)."
|
| 57 |
+
)
|
| 58 |
+
|
| 59 |
+
idx = _clamp_index(int(index), b)
|
| 60 |
+
|
| 61 |
+
# Concatenate everything except the excluded index
|
| 62 |
+
out = torch.cat([images[:idx], images[idx + 1 :]], dim=0)
|
| 63 |
+
return (out,)
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
NODE_CLASS_MAPPINGS = {
|
| 67 |
+
"BatchRemoveImageAtIndex": BatchRemoveImageAtIndex,
|
| 68 |
+
}
|
| 69 |
+
|
| 70 |
+
NODE_DISPLAY_NAME_MAPPINGS = {
|
| 71 |
+
"BatchRemoveImageAtIndex": "Batch: Remove Image @ Index",
|
| 72 |
+
}
|