from __future__ import annotations from dataclasses import dataclass 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, } # Upper bound on how many frames a single ``N-M`` range token may expand to, so # a typo like ``1-99999999`` can't balloon the frame list into millions of # entries. Comfortably above the node's max cut window (total_frames <= 1024). _MAX_RANGE_SPAN = 8192 @dataclass(frozen=True) class OffsetPlan: """Result of planning offset-mode placements (pure; no torch). - ``placements`` — ``(source_index, output_index, vfx_frame)`` for each incoming frame that lands inside the cut window, ascending by position. ``source_index`` is the frame's ordinal in ``keyframes_insert``. - ``outside_cut`` — VFX frame numbers paired to a frame but whose output index fell outside ``[0, total_frames)`` (dropped, like select mode). - ``unused_source`` — count of incoming frames with no position to fill (more frames than listed positions). - ``missing_positions`` — listed positions with no incoming frame left (more positions than frames). """ placements: list[tuple[int, int, int]] outside_cut: list[int] unused_source: int missing_positions: list[int] def plan_offset_placements( positions: list[int], source_count: int, total_frames: int, cut_start_frame: int, ) -> OffsetPlan: """Map a packed sequence of ``source_count`` frames onto keyframe positions. Offset mode (the inverse of selecting): the i-th incoming frame is assigned to the i-th position, with positions taken ascending and de-duplicated so the mapping round-trips cleanly with the ascending ``selected_frames`` output. A position is placed iff its output index ``vfx - cut_start_frame`` falls inside ``[0, total_frames)``; otherwise it is reported in ``outside_cut`` and dropped (matching select-mode cut behaviour). Pure Python — no torch — so the placement logic is unit-testable in CI without the heavyweight tensor dependency. """ ordered = sorted(dict.fromkeys(int(p) for p in positions)) paired = min(source_count, len(ordered)) placements: list[tuple[int, int, int]] = [] outside_cut: list[int] = [] for source_index in range(paired): vfx = ordered[source_index] output_index = vfx - cut_start_frame if 0 <= output_index < total_frames: placements.append((source_index, output_index, vfx)) else: outside_cut.append(vfx) unused_source = max(0, source_count - len(ordered)) missing_positions = ordered[source_count:] if source_count < len(ordered) else [] return OffsetPlan(placements, outside_cut, unused_source, missing_positions) def plan_source_base_fill( total_frames: int, cut_start_frame: int, base_len: int ) -> tuple[list[tuple[int, int]], int]: """Map each output index in the cut window onto a ``source_batch`` index. Used by insert-over-source mode to lay down the source video as the background before the inserts overwrite it. Output index ``i`` represents VFX frame ``cut_start_frame + i``, and ``source_batch[0]`` is VFX frame 1, so ``source_index = i + cut_start_frame - 1``. Returns ``(placements, fallback_count)`` where ``placements`` is the list of ``(output_index, source_index)`` the source actually covers and ``fallback_count`` is how many cut frames fell beyond the source (and stay placeholder). Pure Python — no torch — so it is unit-testable in CI. """ placements: list[tuple[int, int]] = [] fallback = 0 for output_index in range(total_frames): source_index = output_index + cut_start_frame - 1 if 0 <= source_index < base_len: placements.append((output_index, source_index)) else: fallback += 1 return placements, fallback def plan_slot_overwrites( slot_frames: list[tuple[int, int]], total_frames: int, cut_start_frame: int, ) -> tuple[list[tuple[int, int, int]], list[int]]: """Map connected manual slots (image1-4) onto output indices — Layer 3. ``slot_frames`` is ``(slot_index, vfx_frame)`` for each CONNECTED slot, in ascending slot order (image1 = 0 … image4 = 3). Manual slots are the top compositing layer: they overwrite whatever the sequence or background put at the same output index. On a collision between two slots the higher-numbered slot wins (image4 > image1). A slot lands iff its output index ``vfx - cut_start_frame`` is inside ``[0, total_frames)``; otherwise the frame is reported in ``outside_cut``. Returns ``(placements, outside_cut)`` where ``placements`` is ``(slot_index, output_index, vfx)`` ascending by output index, one winner per index. Pure Python — no torch — so it is unit-testable in CI. """ winner: dict[int, tuple[int, int]] = {} # output_index -> (slot_index, vfx) outside_cut: list[int] = [] for slot_index, vfx in slot_frames: output_index = vfx - cut_start_frame if 0 <= output_index < total_frames: # Higher slot wins on a collision (image4 > image1), independent of # the order slots are supplied in. prev = winner.get(output_index) if prev is None or slot_index > prev[0]: winner[output_index] = (slot_index, vfx) else: outside_cut.append(vfx) placements = [ (slot_index, output_index, vfx) for output_index, (slot_index, vfx) in sorted(winner.items()) ] return placements, outside_cut def parse_frame_tokens(source_frames: str) -> tuple[list[int], list[str]]: """Parse a frame-list string into 1-based VFX frame numbers. Tokens are separated by commas, spaces, tabs, or newlines. A token is either a single integer (``7``) or an inclusive ascending range (``14-17`` → 14, 15, 16, 17). Returns ``(values, bad_tokens)`` where ``bad_tokens`` are the tokens that could not be parsed (non-integer, a descending range like ``5-1``, or a range wider than ``_MAX_RANGE_SPAN``) so the caller can warn. Pure Python — no torch — so it is unit-testable in CI. Shared by select picks and insert positions. """ values: list[int] = [] bad_tokens: list[str] = [] for tok in re.split(r"[,\s]+", (source_frames or "").strip()): if not tok: continue span = re.fullmatch(r"(\d+)-(\d+)", tok) if span: lo, hi = int(span.group(1)), int(span.group(2)) # Reject descending or absurdly large ranges (a typo like # "1-99999999" would otherwise expand into millions of entries). if lo <= hi <= lo + _MAX_RANGE_SPAN - 1: values.extend(range(lo, hi + 1)) else: bad_tokens.append(tok) continue try: values.append(int(tok)) except ValueError: bad_tokens.append(tok) return values, bad_tokens 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". Select mode — the layers (lowest first; a later layer overwrites an earlier one at the same output index) ----------------------------------------------------------------------- 1. Background: with `source_batch` connected and an EMPTY `source_frames` list, the source cut window passes straight through (source in → source out) as kept content (`alpha` 0.0). Any tail beyond the source stays placeholder, and that uncovered gap becomes the `selected_*` output — the frames to inpaint/generate (the inverse of the usual "selected = picks/slots"). Otherwise the background is `placeholder_color` and is not counted as "placed". 2. `source_frames` (optional comma/newline/whitespace string like `"1, 27, 41, 63"`, with inclusive ranges `"14-17"` → 14,15,16,17): each number picks `source_batch[N - 1]` and places it at output index `N - cut_start_frame` (if inside the cut), onto the placeholder. Bad tokens warn and are skipped; frames not present in `source_batch` warn; frames outside the cut window are dropped and summarised at the end. 3. Manual slots (`image1`-`image4`): connect `imageN` and set `imageN_frame` to its VFX position; the image overwrites that output frame on TOP of everything. A higher-numbered slot wins a slot-vs-slot collision (`image4` > `image1`). An UNconnected slot does nothing — `imageN_frame` is only a placement target for a wired `imageN`, never a `source_batch` pick. Output dedup is automatic (set-based). All inputs must share the same H/W/C. When no image source is connected at all, the node emits a clean placeholder batch sized by the `width`/`height` widgets instead of erroring. Insert modes (the inverse) — connect a packed sequence to `keyframes_insert` and the node instead *scatters* those frames onto the `source_frames` positions (i-th frame → i-th position, ascending). The background depends on `source_batch`: - `source_batch` NOT connected → gaps filled with the placeholder (offset / reconstruct mode). Round-trips with `selected_image_batch` + `selected_frames`. - `source_batch` connected → the inserts are composited *over* the `source_batch` cut window (insert-over-source mode): the source video is the background and the listed positions are overwritten with the insert frames. Connected manual slots (`image1`-`image4`) still apply here — they composite on TOP of the inserts (a slot beats an insert at the same position; the higher slot wins a slot collision). With an empty `source_frames` list and no slots nothing is placed: a clean placeholder batch (no `source_batch`) or a clean `source_batch` cut-window passthrough. `alpha_batch` marks the inserted positions and slot overwrites as placed. See `create_batch_from_frames`. 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", "tooltip": "Length of the output cut window in frames.", }), "cut_start_frame": ("INT", { "default": 1, "min": 1, "max": 999999, "step": 1, "display": "number", "tooltip": "First VFX frame represented by output index 0. All frame numbers are 1-based.", }), "placeholder_color": (list(_PLACEHOLDER_FILL.keys()), { "default": "Black", "tooltip": "Fill color for empty frames in image_batch. Alpha output is controlled separately.", }), "invert_alpha": ("BOOLEAN", { "default": False, "label_on": "compositing", "label_off": "inpaint", "tooltip": "Off: selected frames are black/0 and empty frames are white/1. On: selected frames are white/1 for compositing.", }), "source_frames": ("STRING", { "default": "", "multiline": True, "placeholder": "frames from source_batch, e.g. 1, 27, 41-50, 63 (commas/newlines; ranges with -)", "tooltip": "Optional 1-based VFX frame list to pick from source_batch (in insert mode, the destination positions). Separators: commas, spaces, tabs, or newlines. Inclusive ranges like 14-17 expand to 14,15,16,17.", }), "image1_frame": ("INT", { "default": 5, "min": 1, "max": 999999, "step": 1, "display": "number", "tooltip": "1-based VFX frame number for image1; the connected image is placed here on top (no effect when the slot is unwired).", }), }, "optional": { "keyframes_insert": ("IMAGE", { "tooltip": ( "Connect a packed sequence of frames to INSERT them onto the " "timeline at the positions from the frame list (source_frames): " "1st frame -> 1st number, 2nd -> 2nd, ... Connected image1-4 " "still composite on top of the inserts. The background depends on " "source_batch: if source_batch is NOT connected, the gaps are " "filled with placeholder_color (offset/reconstruct mode); if " "source_batch IS connected, the inserts are composited over the " "source frames in the cut window (insert-over-source mode). With " "an empty frame list and no slots nothing is placed: a clean " "placeholder batch (no source_batch) or a clean source-batch " "passthrough (with source_batch)." ), }), "keyframe_batch": ("IMAGE", { "tooltip": ( "Deprecated alias for keyframes_insert (the input was " "renamed before release). Use keyframes_insert; this alias " "keeps pre-rename workflows loading. Ignored when " "keyframes_insert is also connected." ), }), "source_batch": ("IMAGE", ), "image1": ("IMAGE", ), "image2": ("IMAGE", ), "image2_frame": ("INT", { "default": 9, "min": 1, "max": 999999, "step": 1, "display": "number", "tooltip": "1-based VFX frame number for image2; the connected image is placed here on top (no effect when the slot is unwired).", }), "image3": ("IMAGE", ), "image3_frame": ("INT", { "default": 13, "min": 1, "max": 999999, "step": 1, "display": "number", "tooltip": "1-based VFX frame number for image3; the connected image is placed here on top (no effect when the slot is unwired).", }), "image4": ("IMAGE", ), "image4_frame": ("INT", { "default": 17, "min": 1, "max": 999999, "step": 1, "display": "number", "tooltip": "1-based VFX frame number for image4; the connected image is placed here on top (no effect when the slot is unwired).", }), # width/height are intentionally LAST so they append to the end # of the widget list — inserting them earlier would shift the # positional widgets_values of every previously-saved workflow. "width": ("INT", { "default": 512, "min": 1, "max": 8192, "step": 8, "display": "number", "tooltip": "Fallback output width. Used ONLY when no image input provides dimensions (e.g. a fully empty node producing a clean placeholder batch). Ignored whenever an image source is connected.", }), "height": ("INT", { "default": 512, "min": 1, "max": 8192, "step": 8, "display": "number", "tooltip": "Fallback output height. Used ONLY when no image input provides dimensions (e.g. a fully empty node producing a clean placeholder batch). Ignored whenever an image source is connected.", }), } } 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, keyframes_insert=None, keyframe_batch=None, width=512, height=512, ): # Deprecated alias: `keyframe_batch` was renamed to `keyframes_insert` # before release. Honour a pre-rename workflow's connection so it keeps # loading; an explicit `keyframes_insert` wins if both are wired. if keyframes_insert is None and keyframe_batch is not None: keyframes_insert = keyframe_batch print( "[easy_ImageBatch] 'keyframe_batch' is deprecated; use " "'keyframes_insert'. Routing the connected batch to insert mode." ) # Validate per-slot image inputs (must be single-frame tensors, not # pre-batched) for BOTH modes — manual slots are now composited as the # top layer in insert mode too, not just select. 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." ) # Manual slots (image1-4) are the top compositing layer (Layer 3) in # every mode: a connected imageN overwrites output frame imageN_frame on # top of the sequence/background. Higher slot wins on a collision. slots = ( (image1, image1_frame), (image2, image2_frame), (image3, image3_frame), (image4, image4_frame), ) # Insert modes: when keyframes_insert is connected, switch from "select # frames out of source_batch" to "scatter this packed sequence onto the # frame-list positions", then composite the slots over the result. The # background depends on source_batch: # - no source_batch -> placeholder gaps (offset / reconstruct mode) # - with source_batch -> the inserts sit over the source cut window # (insert-over-source mode) # See create_batch_from_frames. if keyframes_insert is not None: return self.create_batch_from_frames( keyframes_insert, source_frames, total_frames, cut_start_frame, placeholder_color, invert_alpha, base_batch=source_batch, slots=slots, ) # Determine reference H/W/C/device/dtype from the first available source. # Priority: image1 → source_batch → image2 → image3 → image4. With no # image source connected at all, fall back to the width/height widgets # and emit a clean placeholder batch instead of erroring. reference = None for candidate in (image1, source_batch, image2, image3, image4): if candidate is not None: reference = candidate break if reference is None: if (source_frames or "").strip(): print( "[easy_ImageBatch] source_frames is set but no image source " "(source_batch / image1-4) is connected; the list is ignored " "and a clean placeholder batch is returned." ) return self._clean_placeholder_batch( total_frames, height, width, 3, placeholder_color, invert_alpha ) 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." ) # Compute the frame list early — it decides the background. source_frames_str = (source_frames or "").strip() source_len = source_batch.shape[0] if source_batch is not None else 0 # Background (Layer 1): with source_batch connected AND an empty frame # list, pass the source cut window straight through (the same backdrop # insert-over-source lays down) — "source in, source out" with nothing # else to do. Otherwise fill with the placeholder; a non-empty list # means select mode (pull the named frames onto a neutral background). # The background is NOT counted as "placed": only list picks and slots # are, so alpha / selected_* describe just those. 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/to-fill; set to 0.0 on # frames that carry real content. `invert_alpha` flips at the end. alpha_batch = torch.ones((total_frames, h, w), device=device, dtype=torch.float32) # Timeline indices with a keyframe placed (picks + slots), used to build # the selected_* outputs. `outside_cut` collects referenced VFX frames # that fell outside the window (summarised once at the end). placed_indices: set[int] = set() outside_cut: list[int] = [] # Passthrough (empty list + source): the source cut window is the kept # background — covered frames carry real content (alpha 0.0), the # uncovered tail stays placeholder. `covered_indices` lets the selection # below resolve to the GAP (the to-inpaint frames) rather than to the # picks/slots that drive select mode. passthrough = source_batch is not None and not source_frames_str covered_indices: set[int] = set() if passthrough: base_placements, _ = plan_source_base_fill( total_frames, cut_start_frame, source_len ) for output_index, source_index in base_placements: image_batch[output_index] = source_batch[source_index] alpha_batch[output_index] = 0.0 covered_indices.add(output_index) # 1) source_frames list (optional). Each token is a VFX-numbered frame: # it picks source_batch[N - 1] (VFX 1-based) AND places it at output # index N - cut_start_frame, onto the placeholder. Runs only when the # list is non-empty (an empty list is the passthrough case above); # the manual slots run after and override per-position. # # 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]). 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: values, bad_tokens = parse_frame_tokens(source_frames_str) for tok in bad_tokens: print( f"[easy_ImageBatch] source_frames: ignoring invalid token '{tok}'." ) for f in values: 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) Manual slots (Layer 3, top priority): a connected image1-4 # overwrites its output frame on top of the list picks / passthrough # background. Higher slot wins on a collision (image4 > image1). # Unconnected slots contribute nothing — imageN_frame is only a # placement target for a wired imageN, never a source_batch pick. self._apply_slot_overwrites( image_batch, alpha_batch, placed_indices, slots, total_frames, cut_start_frame, outside_cut, (h, w, c), ) # Passthrough selection: the "selected" frames are the GAP — placeholder # frames with no source coverage and no slot, i.e. the ones to inpaint / # generate. This is the inverse of select mode (where selected = the # placed picks/slots); covered source and slot overwrites are kept # content and excluded from the selection. if passthrough: content = covered_indices | placed_indices placed_indices = { i for i in range(total_frames) if i not in content } 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; surfaces outside-cut drops, the # passthrough background, and any source-shorter-than-cut fill. cut_end = cut_start_frame + total_frames - 1 notes: list[str] = [] if outside_cut: unique_outside = sorted(set(outside_cut)) notes.append( f"{len(unique_outside)} outside cut: " + ", ".join(str(f) for f in unique_outside) ) if passthrough: # kept = everything not in the gap (covered source + any slot fills). notes.append( f"source passthrough: {total_frames - len(sorted_indices)} kept, " f"{len(sorted_indices)} gap frame(s) selected to inpaint" ) suffix = (" " + "; ".join(notes) + ".") if notes else "" print( f"[easy_ImageBatch] cut window: frames {cut_start_frame}..{cut_end} " f"({total_frames} frames). {len(sorted_indices)} placed.{suffix}" ) return (image_batch, alpha_batch, selected_image_batch, selected_frames_out, ) def _clean_placeholder_batch( self, total_frames, h, w, c, placeholder_color, invert_alpha, device=None, dtype=None ): """Build a fully-empty batch (no placed frames) of the placeholder color. Used when nothing drives placement: a totally empty node (sized by the width/height widgets, so ``device``/``dtype`` are absent → CPU/float). ``image_batch`` is uniform placeholder; ``alpha_batch`` is all-empty (1.0, flipped by ``invert_alpha``); ``selected_*`` outputs are zero-length. Returns the standard 4-tuple. """ fill_value = _PLACEHOLDER_FILL[placeholder_color] if device is None: image_batch = torch.full((total_frames, h, w, c), fill_value) alpha_batch = torch.ones((total_frames, h, w), dtype=torch.float32) else: image_batch = torch.full( (total_frames, h, w, c), fill_value, device=device, dtype=dtype ) alpha_batch = torch.ones( (total_frames, h, w), device=device, dtype=torch.float32 ) if invert_alpha: alpha_batch = 1.0 - alpha_batch print( f"[easy_ImageBatch] clean batch: {total_frames} frames of " f"{placeholder_color} placeholder; nothing placed." ) return (image_batch, alpha_batch, image_batch[:0], "") def _apply_slot_overwrites( self, image_batch, alpha_batch, placed_indices, slots, total_frames, cut_start_frame, outside_cut, ref_hwc, ): """Composite connected manual slots on top of the batch (Layer 3). ``slots`` is ``((image, frame), …)`` for image1..image4 in ascending order. Only connected slots with a frame value participate; each must match ``ref_hwc`` (H/W/C). Higher slot wins on a collision (see ``plan_slot_overwrites``). Mutates ``image_batch`` / ``alpha_batch`` / ``placed_indices`` and extends ``outside_cut`` in place. """ h, w, c = ref_hwc connected: list[tuple[int, int]] = [] for slot_index, (img, frame) in enumerate(slots): if img is None or frame is None: continue if tuple(img.shape[1:]) != (h, w, c): raise ValueError( f"image{slot_index + 1} shape {tuple(img.shape[1:])} does not match " f"reference ({h}, {w}, {c}). All inputs must share the same H/W/C." ) connected.append((slot_index, frame)) placements, slot_outside = plan_slot_overwrites( connected, total_frames, cut_start_frame ) for slot_index, output_index, _vfx in placements: image_batch[output_index] = slots[slot_index][0][0] alpha_batch[output_index] = 0.0 placed_indices.add(output_index) outside_cut.extend(slot_outside) def create_batch_from_frames( self, keyframes_insert, source_frames, total_frames, cut_start_frame, placeholder_color, invert_alpha, base_batch=None, slots=((None, None), (None, None), (None, None), (None, None)), ): """Insert mode — active when ``keyframes_insert`` is connected. Scatter the packed ``keyframes_insert`` sequence onto the timeline at the positions listed in ``source_frames`` (i-th frame → i-th position, ascending). The background depends on ``base_batch`` (the connected ``source_batch``, or ``None``): - **No ``base_batch`` (offset / reconstruct mode).** Gaps are filled with ``placeholder_color``. Feed back a processed copy of a previous ``selected_image_batch`` plus the same ``selected_frames`` list to rebuild a full-length sequence in the original spacing. - **With ``base_batch`` (insert-over-source mode).** The cut window of ``base_batch`` is laid down first (VFX frame ``cut_start_frame + i`` → ``base_batch[cut_start_frame + i - 1]``; frames beyond the source fall back to the placeholder), then the inserts overwrite at the listed positions. The inserts are composited *over* the source video. Connected manual slots (``image1``-``image4``) are then composited on top as the highest layer — a slot overwrites an insert at the same position, and the higher-numbered slot wins a slot-vs-slot collision. ``alpha_batch`` marks the inserted positions and any slot overwrites as placed (0.0), regardless of the background — so ``selected_image_batch`` / ``selected_frames`` describe just those. With an empty frame list and no slots nothing is placed: a clean placeholder batch (no ``base_batch``) or a clean ``base_batch`` cut-window passthrough. """ h, w, c = keyframes_insert.shape[1:] device = keyframes_insert.device dtype = keyframes_insert.dtype # Background: source cut window (insert-over-source) or placeholder. fill_value = _PLACEHOLDER_FILL[placeholder_color] image_batch = torch.full((total_frames, h, w, c), fill_value, device=device, dtype=dtype) base_short_fallback = 0 if base_batch is not None: if tuple(base_batch.shape[1:]) != (h, w, c): raise ValueError( f"source_batch shape {tuple(base_batch.shape[1:])} does not match " f"keyframes_insert ({h}, {w}, {c}). Both must share the same H/W/C." ) base_placements, base_short_fallback = plan_source_base_fill( total_frames, cut_start_frame, base_batch.shape[0] ) for output_index, source_index in base_placements: image_batch[output_index] = base_batch[source_index] # Convention: 1.0 (white/empty), 0.0 on inserted positions only. alpha_batch = torch.ones((total_frames, h, w), device=device, dtype=torch.float32) # Parse the position list (commas, spaces, tabs, newlines, or any mix; # inclusive ranges like "5-7" expand via parse_frame_tokens). source_frames_str = (source_frames or "").strip() positions, bad_tokens = parse_frame_tokens(source_frames_str) for tok in bad_tokens: print(f"[easy_ImageBatch] source_frames: ignoring invalid token '{tok}'.") source_count = keyframes_insert.shape[0] plan = plan_offset_placements(positions, source_count, total_frames, cut_start_frame) placed_indices: set[int] = set() for source_index, output_index, _vfx in plan.placements: image_batch[output_index] = keyframes_insert[source_index] alpha_batch[output_index] = 0.0 placed_indices.add(output_index) # Manual slots (Layer 3): composite connected image1-4 on top of the # inserts and the source background. A slot beats an insert at the same # position; higher slot wins a slot-vs-slot collision. slot_outside: list[int] = [] self._apply_slot_overwrites( image_batch, alpha_batch, placed_indices, slots, total_frames, cut_start_frame, slot_outside, (h, w, c), ) if invert_alpha: alpha_batch = 1.0 - alpha_batch sorted_indices = sorted(placed_indices) if sorted_indices: selected_image_batch = image_batch[sorted_indices] else: selected_image_batch = image_batch[:0] selected_frames_out = ", ".join( str(idx + cut_start_frame) for idx in sorted_indices ) # One-line summary (mirrors select mode) with insert-specific notes so # silent surprises (extra frames, short lists, outside-cut, empty list, # source shorter than the cut) are visible without per-frame logs. mode_label = "insert-over-source mode" if base_batch is not None else "offset mode" cut_end = cut_start_frame + total_frames - 1 notes: list[str] = [] if plan.outside_cut: notes.append( f"{len(plan.outside_cut)} outside cut: " + ", ".join(str(f) for f in plan.outside_cut) ) if plan.unused_source: notes.append( f"{plan.unused_source} extra insert frame(s) had no listed position" ) if plan.missing_positions: notes.append( f"{len(plan.missing_positions)} position(s) had no insert frame: " + ", ".join(str(f) for f in plan.missing_positions) ) if base_short_fallback: notes.append( f"{base_short_fallback} cut frame(s) beyond source_batch filled with placeholder" ) if slot_outside: notes.append( f"{len(set(slot_outside))} slot frame(s) outside cut: " + ", ".join(str(f) for f in sorted(set(slot_outside))) ) if not source_frames_str: notes.append("frame list is empty — nothing inserted") suffix = (" " + "; ".join(notes) + ".") if notes else "" print( f"[easy_ImageBatch] {mode_label}: cut window frames {cut_start_frame}.." f"{cut_end} ({total_frames} frames). {len(sorted_indices)} placed.{suffix}" ) 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)" }