import re import torch # Placeholder fill values (0..1) for unoccupied frames in the output IMAGE batch. # Alpha output is independent of this choice; see `invert_alpha` for convention. _PLACEHOLDER_FILL = { "Black": 0.0, "Gray": 0.5, "White": 1.0, } class easy_ImageBatch: """ Place up to N keyframes on a 1-based VFX timeline, then output a *cut window* of that timeline as a fixed-length IMAGE batch with placeholder fill (Black / 50% Gray / White) for unoccupied frames. Returns four outputs: - `image_batch` — IMAGE batch of length `total_frames`, representing the cut window starting at VFX frame `cut_start_frame` and lasting `total_frames` frames. Selected positions carry the picked images; everything else is the placeholder. - `alpha_batch` — MASK at the same length. Default convention: selected = 0.0 (black), empty = 1.0 (white). Toggle `invert_alpha` flips to compositing-style (selected = 1.0). - `selected_image_batch` — IMAGE containing only the keyframes that actually landed *inside the cut*, packed back-to-back in ascending order (no placeholders). Empty if nothing was placed. - `selected_frames` — STRING, comma-separated VFX frame numbers of the placed-inside-cut keyframes (e.g. `"41, 63, 78, 88, 98, 121"`). Same format as the `source_frames` input → round-trippable into another `easy_ImageBatch`. Frame-numbering conventions --------------------------- All frame numbers (`source_frames`, `imageN_frame`, `cut_start_frame`) are **VFX-numbered, 1-based**. Internally: source pick index = vfx_frame - 1 # always output (cut) index = vfx_frame - cut_start_frame # cut window `source_batch[0]` is treated as VFX frame 1. There is no exposed knob to renumber `source_batch` — if your stack represents a different VFX range, renumber externally before connecting. Cut window ---------- The output represents VFX frames `[cut_start_frame .. cut_start_frame + total_frames - 1]`. Frames outside this window are placed but silently *dropped from the cut* (a single summary line is logged at the end listing them). Frames inside the window are placed at output index `vfx - cut_start_frame`. The default `cut_start_frame = 1` means the cut starts at VFX frame 1, so the full timeline shows up — equivalent to "no cut". Three input modes, freely mixable --------------------------------- 1. Per-slot images (up to 4): connect `imageN` and set `imageN_frame` to its VFX timeline position. The image lands at output index `imageN_frame - cut_start_frame` if that's inside the cut. 2. `source_batch` (optional pre-batched IMAGE, e.g. straight from Load Video) + the 4 slots: when a slot's `imageN` is *not* connected, the node pulls VFX frame `imageN_frame` from `source_batch` (i.e. `source_batch[imageN_frame - 1]`) and places it at output index `imageN_frame - cut_start_frame`. Collapses the "Load Video → Get Images From Batch In Range × N" helper chain into a single node. 3. `source_batch` + `source_frames` (optional comma/newline/whitespace separated string like `"1, 27, 41, 63, 85, 120"`): for keyframe counts beyond 4. Each number picks `source_batch[N - 1]` and places it at output index `N - cut_start_frame` (if inside the cut). Bad tokens warn and are skipped; frames not present in `source_batch` warn; frames outside the cut window are dropped silently and summarised at the end. Important: when `source_frames` is non-empty, the 4 manual slots contribute ONLY where `imageN` is explicitly connected — the `imageN_frame` defaults do NOT pick from `source_batch` in this mode. The list fully controls the selection, with explicit `imageN` connections used solely to override individual positions. Priority order (later wins): `source_frames` list → 4 manual slots. Explicit `imageN` inputs always override list entries at the same cut-output position. Output dedup is automatic (set-based). At least one of `image1` or `source_batch` must be provided. All keyframes (and `source_batch` frames) must share the same H/W/C. Useful for preparing sparse control sequences for video models like Wan 2.2, where placeholder frames indicate no latent update. Example A — cut_start_frame=1 (no cut, full timeline) ----------------------------------------------------- Inputs: source_batch = Load Video (121 frames) source_frames = "1, 27, 41, 63, 78, 88, 98, 121" cut_start_frame = 1 total_frames = 121 placeholder_color = "Gray" Outputs: image_batch — 121 frames; the 8 picked frames sit at their natural positions (1→idx 0, 27→idx 26, …, 121→idx 120). Rest is gray. selected_image_batch — 8 frames, packed back-to-back. selected_frames — "1, 27, 41, 63, 78, 88, 98, 121" Example B — cut_start_frame=41, total_frames=81 (cut) ----------------------------------------------------- Same inputs as Example A, but cut_start_frame=41 / total_frames=81. Outputs: image_batch — 81 frames representing VFX 41..121. Frames 41/63/78/88/98/121 land at output indices 0/22/37/47/57/80. Frames 1, 27 are *outside the cut* and dropped. selected_image_batch — 6 frames (the in-cut picks). selected_frames — "41, 63, 78, 88, 98, 121" Console summary: — "cut window: frames 41..121 (81 frames). 6 placed; 2 outside cut: 1, 27." """ @classmethod def INPUT_TYPES(cls): return { "required": { "total_frames": ("INT", { "default": 81, "min": 1, "max": 1024, "step": 1, "display": "number" }), "cut_start_frame": ("INT", { "default": 1, "min": 1, "max": 999999, "step": 1, "display": "number" }), "placeholder_color": (list(_PLACEHOLDER_FILL.keys()), { "default": "Black" }), "invert_alpha": ("BOOLEAN", { "default": False, "label_on": "compositing", "label_off": "inpaint", }), "source_frames": ("STRING", { "default": "", "multiline": True, "placeholder": "extra frames from source_batch, e.g. 1, 27, 41, 63, 85 (commas and/or newlines)", }), "image1_frame": ("INT", { "default": 5, "min": 1, "max": 999999, "step": 1, "display": "number" }), }, "optional": { "source_batch": ("IMAGE", ), "image1": ("IMAGE", ), "image2": ("IMAGE", ), "image2_frame": ("INT", { "default": 9, "min": 1, "max": 999999, "step": 1, "display": "number" }), "image3": ("IMAGE", ), "image3_frame": ("INT", { "default": 13, "min": 1, "max": 999999, "step": 1, "display": "number" }), "image4": ("IMAGE", ), "image4_frame": ("INT", { "default": 17, "min": 1, "max": 999999, "step": 1, "display": "number" }), } } RETURN_TYPES = ("IMAGE", "MASK", "IMAGE", "STRING", ) RETURN_NAMES = ("image_batch", "alpha_batch", "selected_image_batch", "selected_frames", ) FUNCTION = "create_batch" CATEGORY = "Koolook/Image" OUTPUT_NODE = False # Not necessarily an output node, but can be chained def create_batch( self, total_frames, cut_start_frame, placeholder_color, invert_alpha, source_frames, image1_frame, source_batch=None, image1=None, image2=None, image2_frame=None, image3=None, image3_frame=None, image4=None, image4_frame=None, ): # Validate per-slot image inputs (must be single-frame tensors, # not pre-batched). source_batch is the only accepted multi-frame input. for slot_name, slot_img in ( ("image1", image1), ("image2", image2), ("image3", image3), ("image4", image4), ): if slot_img is not None and slot_img.shape[0] != 1: raise ValueError( f"{slot_name} should be a single IMAGE (batch size 1), not a pre-batched " "tensor. Use the 'source_batch' input for pre-batched stacks." ) # Determine reference H/W/C/device/dtype from the first available source. # Priority: image1 → source_batch → image2 → image3 → image4. reference = None for candidate in (image1, source_batch, image2, image3, image4): if candidate is not None: reference = candidate break if reference is None: raise ValueError( "easy_ImageBatch needs at least one of: image1, source_batch, image2, image3, image4." ) h, w, c = reference.shape[1:] device = reference.device dtype = reference.dtype # Cross-check that all provided inputs share H/W/C. for label, tensor in ( ("image1", image1), ("source_batch", source_batch), ("image2", image2), ("image3", image3), ("image4", image4), ): if tensor is not None and tensor.shape[1:] != (h, w, c): raise ValueError( f"{label} shape {tuple(tensor.shape[1:])} does not match reference " f"({h}, {w}, {c}). All inputs must share the same H/W/C." ) fill_value = _PLACEHOLDER_FILL[placeholder_color] image_batch = torch.full((total_frames, h, w, c), fill_value, device=device, dtype=dtype) # Alpha (default convention): 1.0 (white) = empty, set to 0.0 on occupied frames below. # `invert_alpha` flips the result to compositing-style at the end. alpha_batch = torch.ones((total_frames, h, w), device=device, dtype=torch.float32) # Track timeline indices that actually got a keyframe placed. Used at the # end to build `selected_image_batch` (just the placed frames, packed) # and `selected_frames` (their VFX numbers as a comma-separated string). placed_indices: set[int] = set() # VFX frame numbers that were referenced but landed outside the cut # window. Collected once and printed as a single summary line at the # end so each frame doesn't generate its own console warning. outside_cut: list[int] = [] source_len = source_batch.shape[0] if source_batch is not None else 0 # 1) source_frames list (optional, additive). Each token is a VFX-numbered # frame: it picks source_batch[N - 1] (VFX 1-based, hardcoded) AND # places it at output index N - cut_start_frame. The 4 manual slots # run after this and override per-position when needed. # # The field is multiline: tokens may be separated by commas, newlines, # spaces, tabs, or any mix of them ("1, 27\n41 63" → [1, 27, 41, 63]). source_frames_str = (source_frames or "").strip() if source_frames_str: if source_batch is None: print( "[easy_ImageBatch] source_frames is set but source_batch is not connected; " "ignoring the list (it only picks frames from source_batch)." ) else: tokens = [t for t in re.split(r"[,\s]+", source_frames_str) if t] for tok in tokens: try: f = int(tok) except ValueError: print( f"[easy_ImageBatch] source_frames: ignoring non-integer token '{tok}'." ) continue source_index = f - 1 output_index = f - cut_start_frame if not (0 <= source_index < source_len): print( f"[easy_ImageBatch] source_frames: VFX frame {f} doesn't exist in " f"source_batch (size {source_len}, covers frames 1..{source_len}); " "skipping." ) continue if not (0 <= output_index < total_frames): outside_cut.append(f) continue image_batch[output_index] = source_batch[source_index] alpha_batch[output_index] = 0.0 placed_indices.add(output_index) # 2) Resolve each manual slot to (frame_image, timeline_frame). A slot # contributes iff (a) its imageN is connected, or (b) source_batch # is connected (in which case imageN_frame doubles as the source # pick index). Manual slots run AFTER the source_frames list so an # explicit imageN always wins at its target position. slots = ( (image1, image1_frame, "image1"), (image2, image2_frame, "image2"), (image3, image3_frame, "image3"), (image4, image4_frame, "image4"), ) # When the user is driving the node from `source_frames`, the 4 manual # `imageN_frame` defaults (0, 4, 8, 12) shouldn't sneak in extra picks # from source_batch — only the explicit list (and any explicitly # connected imageN inputs) should contribute. So the source_batch # fallback for an UN-connected slot is only active when the list is # empty (the original 4-keyframe workflow). slot_source_fallback = source_batch is not None and not source_frames_str for slot_img, slot_frame, slot_label in slots: if slot_frame is None: continue # Optional slot with no frame value — skip. output_index = slot_frame - cut_start_frame if slot_img is not None: frame_tensor = slot_img[0] elif slot_source_fallback: # Pull VFX frame `slot_frame` from source_batch (1-based). source_index = slot_frame - 1 if not (0 <= source_index < source_len): print( f"[easy_ImageBatch] {slot_label}: VFX frame {slot_frame} doesn't " f"exist in source_batch (size {source_len}, covers frames " f"1..{source_len}); skipping." ) continue frame_tensor = source_batch[source_index] else: # Either source_batch isn't connected, or source_frames is # being used and no explicit imageN was wired — leave this # slot empty so it doesn't pollute the selection. continue if 0 <= output_index < total_frames: image_batch[output_index] = frame_tensor alpha_batch[output_index] = 0.0 placed_indices.add(output_index) else: outside_cut.append(slot_frame) if invert_alpha: alpha_batch = 1.0 - alpha_batch # Build the "selected only" outputs: keyframes packed back-to-back in # ascending timeline order, plus their VFX frame numbers as a string # that round-trips cleanly into the `source_frames` input. sorted_indices = sorted(placed_indices) if sorted_indices: selected_image_batch = image_batch[sorted_indices] else: # No keyframes placed — emit a 0-length batch with the same H/W/C. selected_image_batch = image_batch[:0] selected_frames_out = ", ".join( str(idx + cut_start_frame) for idx in sorted_indices ) # One-line summary of the cut window. Always logged so it's easy to # confirm what the node produced; mentions outside-cut drops when any. cut_end = cut_start_frame + total_frames - 1 if outside_cut: unique_outside = sorted(set(outside_cut)) outside_str = ", ".join(str(f) for f in unique_outside) print( f"[easy_ImageBatch] cut window: frames {cut_start_frame}..{cut_end} " f"({total_frames} frames). {len(sorted_indices)} placed; " f"{len(unique_outside)} outside cut: {outside_str}." ) else: print( f"[easy_ImageBatch] cut window: frames {cut_start_frame}..{cut_end} " f"({total_frames} frames). {len(sorted_indices)} placed." ) return (image_batch, alpha_batch, selected_image_batch, selected_frames_out, ) # Node mappings for ComfyUI registration NODE_CLASS_MAPPINGS = { "easy_ImageBatch": easy_ImageBatch } NODE_DISPLAY_NAME_MAPPINGS = { "easy_ImageBatch": "Easy Image Batch (Koolook)" }