saliacoel commited on
Commit
10a8e09
·
verified ·
1 Parent(s): 11855b6

Upload Batch_Remove_Img_At.py

Browse files
Files changed (1) hide show
  1. 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
+ }