File size: 42,340 Bytes
e4ab0d4 | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 | 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)"
}
|