diff --git a/comfy_extras/nodes_lt_keyframes.py b/comfy_extras/nodes_lt_keyframes.py new file mode 100644 index 000000000..3cf07d8fb --- /dev/null +++ b/comfy_extras/nodes_lt_keyframes.py @@ -0,0 +1,1051 @@ +import math +import re + +import node_helpers +import torch +from comfy.ldm.lightricks.symmetric_patchifier import SymmetricPatchifier, latent_to_pixel_coords +from comfy_api.latest import ComfyExtension, io +from comfy_extras.nodes_lt import ( + LTXVAddGuide, + _append_guide_attention_entry, + conditioning_get_any_value, + get_keyframe_idxs, + get_noise_mask, +) +from typing_extensions import override + +DEFAULT_TEMPORAL_SCALE = 8 +_OCCUPIED_MASK_MAX = 1.0 - 1e-4 + + +def get_generated_keyframes(cond): + return conditioning_get_any_value(cond, "generated_keyframes", None) + + +def _parse_frame_index_list(value, field, expected_count, first, last, expected_desc, empty_hint): + """Parse a manual frame index override and validate count, uniqueness and range. + + Order is not significant — only the set of positions matters — so the list + is returned as written rather than sorted. ``expected_count`` may be None + when the list itself defines how many keyframes to add. + """ + parts = [part for part in re.split(r"[,\s]+", value.strip()) if part] + if not parts: + raise ValueError( + f"{field} is empty. Provide at least one index, or leave {field} empty {empty_hint}." + ) + try: + indices = [int(part) for part in parts] + except ValueError: + bad = ", ".join(repr(part) for part in parts if not re.fullmatch(r"-?\d+", part)) + raise ValueError( + f"{field} must be a comma-separated list of integers, but could not parse {bad}." + ) from None + + if expected_count is not None and len(indices) != expected_count: + raise ValueError( + f"{field} lists {len(indices)} index/indices but {expected_desc}. Provide exactly " + f"{expected_count}, or leave {field} empty {empty_hint}." + ) + + duplicates = sorted({index for index in indices if indices.count(index) > 1}) + if duplicates: + raise ValueError( + f"{field} must not place two keyframes on the same pixel frame, but " + f"{', '.join(str(index) for index in duplicates)} appear(s) more than once." + ) + + out_of_range = [index for index in indices if not first <= index <= last] + if out_of_range: + raise ValueError( + f"{field} must lie between {first} and {last}, but got " + f"{', '.join(str(index) for index in out_of_range)}." + ) + return indices + + +def _grow_guide_attention_entry(positive, negative, index, extra_pre_filter_count, extra_frames): + """Grow an existing guide_attention_entry when extending generated keyframes.""" + results = [] + for cond in (positive, negative): + existing = [] + for t in cond: + found = t[1].get("guide_attention_entries", None) + if found is not None: + existing = found + break + if index >= len(existing): + raise ValueError( + f"The generated keyframes recorded guide entry {index} but the conditioning only has " + f"{len(existing)}. The conditioning was rebuilt after they were added." + ) + entries = list(existing) + grown = dict(entries[index]) + grown["pre_filter_count"] = grown["pre_filter_count"] + extra_pre_filter_count + shape = list(grown["latent_shape"]) + shape[0] = shape[0] + extra_frames + grown["latent_shape"] = shape + entries[index] = grown + results.append(node_helpers.conditioning_set_values(cond, {"guide_attention_entries": entries})) + return results[0], results[1] + + +def _spaced_positions_keep_last(num_keyframes: int, num_frames: int) -> list[int]: + return ( + torch.linspace(0, num_frames - 1, num_keyframes + 1) + .round() + .to(torch.int64) + .tolist()[1:] + ) + + +def detailing_positions(num_frames: int, interval_frames: float) -> list[int]: + """About one detailing keyframe every ``interval_frames`` pixel frames. + + Skips frame 0 (already a standalone token) and keeps the last frame. + ``interval_frames`` 24 is about 1/s at 24 fps. + """ + if interval_frames <= 0: + raise ValueError(f"interval_frames must be > 0, got {interval_frames}") + if num_frames <= 1: + raise ValueError( + f"A {num_frames}-frame target has no pixel frames to place keyframes on." + ) + count = max(1, round((num_frames - 1) / interval_frames)) + positions = [index for index in _spaced_positions_keep_last(count, num_frames) if index != 0] + if not positions: + raise ValueError( + f"A {num_frames}-frame target has no pixel frames to place keyframes on." + ) + return positions + + +def free_detailing_slots(num_frames: int, interval_frames: float, occupied: set[int]) -> list[int]: + """Density candidates that are not already I2V / guides / detailing KFs.""" + taken = set(occupied) + positions = [ + index + for index in detailing_positions(num_frames, interval_frames) + if index not in taken + ] + if not positions: + raise ValueError( + "Every candidate detailing-keyframe pixel already has an image keyframe " + "or a guide. Leave at least one unoccupied frame, or pass frame_indices." + ) + return positions + + +def scale_frame_indices(indices: list[int], old_num_frames: int, new_num_frames: int) -> list[int]: + """Map pixel indices from one canvas length onto another (e.g. temporal x2).""" + if old_num_frames <= 1: + raise ValueError(f"Cannot scale keyframe indices from a {old_num_frames}-frame canvas.") + if new_num_frames <= 1: + raise ValueError(f"Cannot scale keyframe indices onto a {new_num_frames}-frame canvas.") + scale = (new_num_frames - 1) / (old_num_frames - 1) + remapped = [int(round(index * scale)) for index in indices] + duplicates = sorted({index for index in remapped if remapped.count(index) > 1}) + if duplicates: + raise ValueError( + "Scaling keyframe indices onto the new canvas collapsed " + f"{', '.join(str(index) for index in duplicates)} onto the same pixel frame." + ) + return remapped + + +def _as_int_set(value) -> set[int]: + if value is None: + return set() + if isinstance(value, (int, float)): + return {int(round(value))} + if isinstance(value, (set, frozenset)): + return {int(round(item)) for item in value} + if isinstance(value, (list, tuple)): + out = set() + for item in value: + out |= _as_int_set(item) + return out + if hasattr(value, "detach"): + value = value.detach() + if hasattr(value, "cpu"): + value = value.cpu() + if hasattr(value, "tolist"): + return _as_int_set(value.tolist()) + return {int(round(float(value)))} + + +def pixel_frames_from_keyframe_idxs(keyframe_idxs) -> set[int]: + """Unique RoPE *start* pixel times of extra guide / keyframe tokens. + + ``keyframe_idxs`` is ``(B, 3, tokens, 2)`` (t/h/w × start/end). A keyframe + at pixel 24 spans the half-open interval ``[24, 25)``, so only 24 is + occupied — the exclusive end is not a slot. + """ + if keyframe_idxs is None: + return set() + if keyframe_idxs.ndim >= 4: + starts = keyframe_idxs[:, 0, :, 0] + else: + starts = keyframe_idxs[:, 0] + return _as_int_set(starts) + + +def occupied_pixel_frames(latent, temporal_scale: int, num_frames: int, video_latent_frames=None) -> set[int]: + """Pixel frames that already hold an in-place image keyframe. + + Prefers ``noise_mask`` (0 means frozen / guided). Falls back to + non-zero latent frames when no mask is present. Each occupied latent + index maps to its representative pixel: 0, ``t * scale``, or last frame. + + ``video_latent_frames`` limits the scan to the video portion of T so + appended guide / detailing-keyframe tokens are not treated as video frames. + Extra-token pixel times come from ``pixel_frames_from_keyframe_idxs``. + """ + samples = latent["samples"] + latent_frames = samples.shape[2] + if video_latent_frames is None: + scan_frames = latent_frames + else: + scan_frames = min(max(int(video_latent_frames), 0), latent_frames) + taken = set() + + def add_latent_index(index: int) -> None: + if index <= 0: + pixel = 0 + elif index >= scan_frames - 1: + pixel = num_frames - 1 + else: + pixel = index * temporal_scale + taken.add(min(max(pixel, 0), num_frames - 1)) + + mask = latent.get("noise_mask") + if mask is not None and getattr(mask, "ndim", 0) >= 3: + for index in range(min(mask.shape[2], scan_frames)): + if torch.any(mask[:, :, index : index + 1] < _OCCUPIED_MASK_MAX): + add_latent_index(index) + return taken + + for index in range(scan_frames): + if torch.any(samples[:, :, index : index + 1] != 0): + add_latent_index(index) + return taken + + +def nearest_latent_index(pixel_frame: int, temporal_scale: int, num_latent_frames: int) -> int: + return min(max(round(pixel_frame / temporal_scale), 0), num_latent_frames - 1) + + +def should_copy_nearest_video_frames(keyframe_t, requested_count, has_recorded_indices, batched_singles): + """True when ``keyframes`` is a longer video to sample, not stacked keyframes.""" + return ( + not has_recorded_indices + and requested_count is not None + and not batched_singles + and keyframe_t > requested_count + ) + + +def keyframes_from_video(samples, indices, temporal_scale: int): + """Stack the nearest video latent frame at each pixel index.""" + if not torch.is_tensor(samples) or samples.ndim != 5: + raise ValueError( + "Initializing keyframes from a video needs a plain 5D video latent. " + "Split audio with Separate AV Latent first, and peel generated " + "keyframes before copying from the video." + ) + temporal_scale = int(temporal_scale) + if temporal_scale < 1: + raise ValueError(f"temporal_scale must be >= 1, got {temporal_scale}") + num_latent_frames = samples.shape[2] + if num_latent_frames < 1: + raise ValueError("The video latent has no frames to copy from.") + frames = [] + for pixel_frame in indices: + idx = nearest_latent_index(pixel_frame, temporal_scale, num_latent_frames) + frames.append(samples[:, :, idx : idx + 1]) + return torch.cat(frames, dim=2) + + +def _fit_keyframe_samples(keyframe_samples, samples, num_slots): + """Pad stacked keyframe tokens with zeros up to ``num_slots``; reject extras.""" + expected = ( + samples.shape[0], + samples.shape[1], + num_slots, + samples.shape[3], + samples.shape[4], + ) + have = keyframe_samples.shape[2] + if ( + keyframe_samples.shape[0] != expected[0] + or keyframe_samples.shape[1] != expected[1] + or keyframe_samples.shape[3:] != expected[3:] + ): + raise ValueError( + "The keyframes latent must hold whole latent frames at this latent's shape, expected " + f"{list(expected)} but got {list(keyframe_samples.shape)}. Resize it to this stage's " + "resolution first." + ) + if have > num_slots: + raise ValueError( + f"The keyframes latent holds {have} frame(s) but only {num_slots} free slot(s) " + "are available on this canvas. Pass fewer keyframes, or set frame_indices." + ) + if have < num_slots: + pad = torch.zeros( + (expected[0], expected[1], num_slots - have, expected[3], expected[4]), + dtype=keyframe_samples.dtype, + device=keyframe_samples.device, + ) + keyframe_samples = torch.cat([keyframe_samples, pad], dim=2) + return keyframe_samples + + +class LTXVAddGeneratedKeyframes(io.ComfyNode): + PATCHIFIER = SymmetricPatchifier(1, start_end=True) + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LTXVAddGeneratedKeyframes", + display_name="LTXV Add Generated Keyframes", + category="model/conditioning/ltxv", + search_aliases=["detailing", "dfr", "generated keyframes"], + description=( + "Append detailing keyframes to a video latent. Each keyframe is one latent " + "frame of tokens whose RoPE position spans a single pixel frame; they are " + "denoised with the video and are not part of the decoded output. Placement " + "is one slot every interval_frames pixels, skipping I2V frames, existing " + "guides, and detailing keyframes already on the cond. Connected keyframes " + "are content only and are re-placed on this canvas unless frame_indices is " + "set. Pull them back out with LTXV Separate Generated Keyframes. Requires a " + "checkpoint trained for generated keyframes (one carrying " + "keyframes_abs_pos_embedding)." + ), + inputs=[ + io.Conditioning.Input( + "positive", + tooltip="Positive conditioning the keyframes are attached to.", + ), + io.Conditioning.Input( + "negative", + tooltip="Negative conditioning the keyframes are attached to.", + ), + io.Vae.Input( + "vae", tooltip="Only used to read the latent scale factors." + ), + io.Latent.Input( + "latent", + tooltip=( + "Plain 5D video latent to generate keyframes alongside. Add them " + "before Concat AV Latent." + ), + ), + io.Int.Input( + "interval_frames", + optional=True, + default=24, + min=1, + max=1024, + tooltip=( + "Pixel-frame stride for auto placement. Default 24 is about one " + "keyframe per second at 24 fps. Occupied pixels are skipped. " + "Ignored when frame_indices is set." + ), + ), + io.Latent.Input( + "keyframes", + optional=True, + tooltip=( + "Optional content to initialize the new keyframes with. Connect " + "keyframes from an earlier Separate (same spatial size), or a " + "plain video latent to copy the nearest frame at each new slot " + "(e.g. after temporal upscale). These are still denoised, not " + "pinned as guides. Recorded indices on a keyframes latent are " + "ignored unless frame_indices is set. Only has an effect when " + "sampling starts below sigma 1." + ), + ), + io.String.Input( + "frame_indices", + optional=True, + default="", + tooltip=( + "Optional pixel-frame indices. Leave empty to place from " + "interval_frames on the current canvas. When set, this list is " + "the placement (connected keyframes are matched in order). The " + "last frame is allowed; frame 0 is not (it is already a " + "standalone token)." + ), + ), + ], + outputs=[ + io.Conditioning.Output( + display_name="positive", + tooltip="Positive conditioning with generated-keyframe attention attached.", + ), + io.Conditioning.Output( + display_name="negative", + tooltip="Negative conditioning with generated-keyframe attention attached.", + ), + io.Latent.Output( + display_name="latent", + tooltip="Video latent with generated keyframes appended on T.", + ), + ], + ) + + @classmethod + def parse_frame_indices(cls, frame_indices, num_pixel_frames, expected_count=None): + """Pixel frame positions from a manual override. + + Frame 0 is excluded (already a standalone token). The terminal frame is + allowed so a DFR segment grid that includes N-1 is legal. + """ + first, last = 1, num_pixel_frames - 1 + if last < first: + raise ValueError( + f"A {num_pixel_frames}-frame target has no pixel frames to place keyframes on." + ) + return _parse_frame_index_list( + frame_indices, + "frame_indices", + expected_count, + first, + last, + expected_desc=( + f"{expected_count} keyframe(s)" + if expected_count is not None + else "the list sets the count" + ), + empty_hint="to place them from interval_frames", + ) + + @classmethod + def keyframe_coords(cls, latent, frame_index, scale_factors): + """Pixel coordinates of one keyframe: the full spatial grid over [t, t + 1).""" + _, latent_coords = cls.PATCHIFIER.patchify(latent[:, :, :1]) + pixel_coords = latent_to_pixel_coords(latent_coords, scale_factors, causal_fix=True) + pixel_coords[:, 0] += frame_index + return pixel_coords + + @classmethod + def execute( + cls, + positive, + negative, + vae, + latent, + interval_frames=24, + keyframes=None, + frame_indices="", + ) -> io.NodeOutput: + samples = latent["samples"] + if not torch.is_tensor(samples) or samples.ndim != 5: + raise ValueError( + "Generated keyframes must be added to a plain video latent. Add them before " + "merging the video and audio latents with Concat AV Latent." + ) + + existing_record = get_generated_keyframes(positive) + if existing_record is not None: + prev_tokens_per_frame = existing_record["tokens_per_frame"] + if prev_tokens_per_frame != samples.shape[3] * samples.shape[4]: + raise ValueError( + f"The existing generated keyframes were added at {prev_tokens_per_frame} tokens per latent " + f"frame but this latent has {samples.shape[3] * samples.shape[4]}. The latent was rescaled " + "after they were added, so more keyframes cannot be appended to them." + ) + prev_first = existing_record["first_latent_frame"] + prev_count = existing_record["num_keyframes"] + if prev_first + prev_count != samples.shape[2]: + raise ValueError( + f"The existing generated keyframes end at latent frame {prev_first + prev_count} but the " + f"latent has {samples.shape[2]}. Something was appended after them, so more keyframes would " + "not be contiguous with the existing block. Add all the keyframes before any guides." + ) + + scale_factors = vae.downscale_index_formula + time_scale_factor = scale_factors[0] + keyframe_idxs, num_guide_frames = get_keyframe_idxs(positive, samples.shape) + num_target_frames = samples.shape[2] - num_guide_frames + num_pixel_frames = (num_target_frames - 1) * time_scale_factor + 1 + occupied = occupied_pixel_frames( + latent, + time_scale_factor, + num_pixel_frames, + video_latent_frames=num_target_frames, + ) + occupied |= pixel_frames_from_keyframe_idxs(keyframe_idxs) + prev_indices = list(existing_record["frame_indices"]) if existing_record is not None else [] + occupied |= set(prev_indices) + + if frame_indices and str(frame_indices).strip(): + indices = cls.parse_frame_indices(frame_indices, num_pixel_frames) + else: + indices = free_detailing_slots(num_pixel_frames, float(interval_frames), occupied) + + clashes = sorted(set(indices) & occupied) + if clashes: + raise ValueError( + f"frame_indices reuses pixel frame(s) {', '.join(str(i) for i in clashes)}, which already hold " + "an image keyframe, a guide, or a generated keyframe. Each keyframe needs its own pixel frame." + ) + + num_keyframes = len(indices) + if keyframes is not None: + keyframe_samples = keyframes["samples"] + if keyframe_samples.ndim != 5: + raise ValueError( + f"The keyframes latent must be 5 dimensional, got {list(keyframe_samples.shape)}." + ) + recorded_positions = keyframes.get("generated_keyframe_indices") + batched_singles = ( + keyframe_samples.shape[2] == 1 + and keyframe_samples.shape[0] != samples.shape[0] + ) + if batched_singles: + if keyframe_samples.shape[0] % samples.shape[0] != 0: + raise ValueError( + f"The keyframes latent batch ({keyframe_samples.shape[0]}) is not a " + f"multiple of the video latent batch ({samples.shape[0]}), so it " + "cannot be reshaped into per-video keyframes." + ) + stacked = keyframe_samples.shape[0] // samples.shape[0] + keyframe_samples = keyframe_samples.reshape( + samples.shape[0], + stacked, + keyframe_samples.shape[1], + *keyframe_samples.shape[3:], + ).movedim(1, 2) + batched_singles = False + if should_copy_nearest_video_frames( + keyframe_samples.shape[2], + num_keyframes, + recorded_positions is not None, + batched_singles, + ): + keyframe_samples = keyframes_from_video( + keyframe_samples, + indices, + time_scale_factor, + ) + else: + keyframe_samples = _fit_keyframe_samples( + keyframe_samples.to(samples), samples, num_keyframes + ) + keyframe_samples = keyframe_samples.to(samples) + else: + keyframe_samples = torch.zeros( + ( + samples.shape[0], + samples.shape[1], + num_keyframes, + samples.shape[3], + samples.shape[4], + ), + dtype=samples.dtype, + device=samples.device, + ) + + generated_coords = torch.cat( + [cls.keyframe_coords(samples, index, scale_factors) for index in indices], + dim=2, + ) + if keyframe_idxs is not None: + generated_coords = torch.cat([keyframe_idxs, generated_coords.to(keyframe_idxs)], dim=2) + + existing_entries = conditioning_get_any_value(positive, "guide_attention_entries", None) or [] + generated_keyframes = { + "first_latent_frame": ( + existing_record["first_latent_frame"] if existing_record else samples.shape[2] + ), + "num_keyframes": ( + existing_record["num_keyframes"] + num_keyframes if existing_record else num_keyframes + ), + "frame_indices": prev_indices + list(indices), + "num_pixel_frames": num_pixel_frames, + "guide_entry_index": ( + existing_record["guide_entry_index"] if existing_record else len(existing_entries) + ), + "tokens_per_frame": samples.shape[3] * samples.shape[4], + } + + values = { + "keyframe_idxs": generated_coords, + "generated_keyframes": generated_keyframes, + } + positive = node_helpers.conditioning_set_values(positive, values) + negative = node_helpers.conditioning_set_values(negative, values) + + if existing_record is not None: + positive, negative = _grow_guide_attention_entry( + positive, + negative, + existing_record["guide_entry_index"], + extra_pre_filter_count=num_keyframes * samples.shape[3] * samples.shape[4], + extra_frames=num_keyframes, + ) + else: + positive, negative = _append_guide_attention_entry( + positive, + negative, + pre_filter_count=num_keyframes * samples.shape[3] * samples.shape[4], + latent_shape=[num_keyframes, samples.shape[3], samples.shape[4]], + strength=1.0, + ) + + noise_mask = get_noise_mask(latent) + keyframe_noise_mask = torch.ones( + ( + noise_mask.shape[0], + 1, + num_keyframes, + noise_mask.shape[3], + noise_mask.shape[4], + ), + dtype=noise_mask.dtype, + device=noise_mask.device, + ) + + out = latent.copy() + out["samples"] = torch.cat([samples, keyframe_samples], dim=2) + out["noise_mask"] = torch.cat([noise_mask, keyframe_noise_mask], dim=2) + return io.NodeOutput(positive, negative, out) + + +class LTXVSeparateGeneratedKeyframes(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LTXVSeparateGeneratedKeyframes", + display_name="LTXV Separate Generated Keyframes", + category="model/conditioning/ltxv", + search_aliases=["detailing", "dfr", "peel keyframes"], + description=( + "Split the generated keyframes added by LTXV Add Generated Keyframes " + "back out of a sampled latent, and remove them from the conditioning. " + "Separate them before spatially upscaling the video latent. Do not run " + "LTXV Crop Guides first — it treats generated keyframes as disposable " + "guides and drops them." + ), + inputs=[ + io.Conditioning.Input("positive"), + io.Conditioning.Input("negative"), + io.Latent.Input("latent"), + io.Boolean.Input( + "keyframes_to_batch", + default=False, + tooltip=( + "Return the keyframes as a batch of single-frame latents. Leave off " + "to get them as one multi-frame latent, which is what the latent " + "upsampler and a later Add Generated Keyframes expect." + ), + ), + ], + outputs=[ + io.Conditioning.Output( + display_name="positive", + tooltip="Positive conditioning with generated-keyframe metadata removed.", + ), + io.Conditioning.Output( + display_name="negative", + tooltip="Negative conditioning with generated-keyframe metadata removed.", + ), + io.Latent.Output( + display_name="latent", + tooltip="Video latent with the generated keyframes stripped.", + ), + io.Latent.Output( + display_name="keyframes", + tooltip=( + "The peeled keyframes, labeled with generated_keyframe_indices and " + "generated_keyframe_num_frames. Feed these to a later Add Generated " + "Keyframes to initialize new slots, or to Generated Keyframes To " + "Guides to pin them as frozen image guides (indices are remapped " + "if the canvas length changed)." + ), + ), + ], + ) + + @classmethod + def strip_keyframe_idxs(cls, cond, first_token, num_tokens): + keyframe_idxs = conditioning_get_any_value(cond, "keyframe_idxs", None) + if keyframe_idxs is None: + return None + remaining = torch.cat( + [ + keyframe_idxs[:, :, :first_token], + keyframe_idxs[:, :, first_token + num_tokens :], + ], + dim=2, + ) + return remaining if remaining.shape[2] > 0 else None + + @classmethod + def strip_guide_entry(cls, cond, entry_index): + entries = conditioning_get_any_value(cond, "guide_attention_entries", None) + if not entries: + return None + if entry_index >= len(entries): + raise ValueError( + f"The generated keyframes recorded guide entry {entry_index} but the conditioning only has " + f"{len(entries)}. The conditioning was rebuilt after they were added." + ) + remaining = entries[:entry_index] + entries[entry_index + 1 :] + return remaining or None + + @classmethod + def execute(cls, positive, negative, latent, keyframes_to_batch=False) -> io.NodeOutput: + generated_keyframes = get_generated_keyframes(positive) + if generated_keyframes is None: + raise ValueError( + "This latent has no generated keyframes. Add them with LTXV Add Generated Keyframes first." + ) + + samples = latent["samples"] + if not torch.is_tensor(samples) or samples.ndim != 5: + raise ValueError( + "Generated keyframes must be separated from a plain video latent. Split the video and " + "audio latents with Separate AV Latent first." + ) + + tokens_per_frame = samples.shape[3] * samples.shape[4] + if generated_keyframes["tokens_per_frame"] != tokens_per_frame: + raise ValueError( + f"The generated keyframes were added at {generated_keyframes['tokens_per_frame']} tokens per " + f"latent frame but this latent has {tokens_per_frame}. The latent was rescaled after they were " + "added, so the keyframes no longer line up. Separate them before upscaling the latent." + ) + + first_frame = generated_keyframes["first_latent_frame"] + num_keyframes = generated_keyframes["num_keyframes"] + end_frame = first_frame + num_keyframes + if end_frame > samples.shape[2]: + raise ValueError( + f"The generated keyframes span latent frames [{first_frame}, {end_frame}) but the latent " + f"only has {samples.shape[2]}. It was recorded against a different latent." + ) + + keyframe_samples = samples[:, :, first_frame:end_frame].clone() + if keyframes_to_batch: + batch, channels, _, height, width = keyframe_samples.shape + keyframe_samples = keyframe_samples.movedim(2, 1).reshape( + batch * num_keyframes, channels, 1, height, width + ) + + video_samples = torch.cat( + [samples[:, :, :first_frame], samples[:, :, end_frame:]], dim=2 + ) + noise_mask = get_noise_mask(latent) + video_noise_mask = torch.cat( + [noise_mask[:, :, :first_frame], noise_mask[:, :, end_frame:]], dim=2 + ) + + _, num_guide_frames = get_keyframe_idxs(positive, samples.shape) + first_token = (first_frame - (samples.shape[2] - num_guide_frames)) * tokens_per_frame + entry_index = generated_keyframes["guide_entry_index"] + + outputs = [] + for cond in (positive, negative): + outputs.append( + node_helpers.conditioning_set_values( + cond, + { + "keyframe_idxs": cls.strip_keyframe_idxs( + cond, first_token, num_keyframes * tokens_per_frame + ), + "guide_attention_entries": cls.strip_guide_entry(cond, entry_index), + "generated_keyframes": None, + }, + ) + ) + + out = latent.copy() + out["samples"] = video_samples + out["noise_mask"] = video_noise_mask + num_pixel_frames = generated_keyframes.get("num_pixel_frames") + if num_pixel_frames is None: + num_pixel_frames = (first_frame - 1) * DEFAULT_TEMPORAL_SCALE + 1 + keyframes_out = { + "samples": keyframe_samples, + "generated_keyframe_indices": list(generated_keyframes["frame_indices"]), + "generated_keyframe_num_frames": int(num_pixel_frames), + } + return io.NodeOutput(outputs[0], outputs[1], out, keyframes_out) + + +class LTXVGeneratedKeyframesToGuides(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LTXVGeneratedKeyframesToGuides", + display_name="LTXV Generated Keyframes to Guides", + category="model/conditioning/ltxv", + search_aliases=["detailing", "dfr", "keyframe guides"], + description=( + "Pin generated keyframes from an earlier stage as frozen image guides on " + "a later canvas. They are decoded as standalone frames, resized if needed, " + "and written with noise_mask=0 so they are not denoised again. After a " + "temporal upscale, recorded indices are scaled from the canvas they were " + "generated on onto this one (same moments). Use override_frame_indices to " + "set positions explicitly. To keep generating (and denoising) keyframes at " + "new positions, use Add Generated Keyframes instead." + ), + inputs=[ + io.Conditioning.Input("positive"), + io.Conditioning.Input("negative"), + io.Vae.Input("vae"), + io.Latent.Input( + "latent", + tooltip="The target video latent to add the guides to, e.g. the temporally upscaled one.", + ), + io.Latent.Input( + "keyframes", + tooltip=( + "The keyframes output of LTXV Separate Generated Keyframes, which " + "carries the pixel frame index each keyframe was generated at." + ), + ), + io.Float.Input( + "strength", + default=1.0, + min=0.0, + max=10.0, + step=0.01, + tooltip="Guide strength. 1.0 is a hard pin; lower values relax it.", + ), + io.String.Input( + "override_frame_indices", + optional=True, + default="", + tooltip=( + "Optional — pin at these pixel frames instead of the recorded " + "(or auto-scaled) positions. Provide one index per keyframe. " + "Leave empty to reuse recorded positions, or to scale them when " + "the target canvas is a different length (e.g. after temporal x2)." + ), + ), + ], + outputs=[ + io.Conditioning.Output( + display_name="positive", + tooltip="Positive conditioning with the keyframes pinned as image guides.", + ), + io.Conditioning.Output( + display_name="negative", + tooltip="Negative conditioning with the keyframes pinned as image guides.", + ), + io.Latent.Output( + display_name="latent", + tooltip="Target video latent with the keyframes added as frozen guides.", + ), + ], + ) + + @classmethod + def decode_single_frames(cls, vae, keyframe_samples): + """Decode each keyframe latent on its own, never as one clip.""" + if keyframe_samples.shape[2] != 1: + batch, channels, num_keyframes, height, width = keyframe_samples.shape + keyframe_samples = keyframe_samples.movedim(2, 1).reshape( + batch * num_keyframes, channels, 1, height, width + ) + images = vae.decode(keyframe_samples) + if images.ndim == 5: + images = images.reshape(-1, *images.shape[-3:]) + return images + + @classmethod + def execute( + cls, + positive, + negative, + vae, + latent, + keyframes, + strength, + override_frame_indices="", + ) -> io.NodeOutput: + indices = keyframes.get("generated_keyframe_indices", None) + if indices is None: + raise ValueError( + "This latent does not carry generated keyframe positions. Connect the keyframes output " + "of LTXV Separate Generated Keyframes." + ) + if get_generated_keyframes(positive) is not None: + raise ValueError( + "This conditioning still carries generated keyframes. Connect the positive and " + "negative outputs of LTXV Separate Generated Keyframes." + ) + + samples = latent["samples"] + if not torch.is_tensor(samples) or samples.ndim != 5: + raise ValueError( + "Generated keyframe guides must be added to a plain video latent. Add them before " + "merging the video and audio latents with Concat AV Latent." + ) + if samples.shape[0] != 1: + raise ValueError( + f"Only a batch size of 1 is supported, got {samples.shape[0]}. Each guide is encoded from " + "one image, so it cannot differ across batch elements." + ) + + kf_samples = keyframes["samples"] + if kf_samples.shape[2] != 1: + batch, channels, num_keyframes, height, width = kf_samples.shape + kf_samples = kf_samples.movedim(2, 1).reshape( + batch * num_keyframes, channels, 1, height, width + ) + + resize_needed = kf_samples.shape[3:] != samples.shape[3:] + guides = cls.decode_single_frames(vae, kf_samples) if resize_needed else kf_samples + if guides.shape[0] != len(indices): + raise ValueError( + f"Got {guides.shape[0]} keyframes for {len(indices)} recorded positions." + ) + + _, num_guide_frames = get_keyframe_idxs(positive, samples.shape) + time_scale_factor = vae.downscale_index_formula[0] + num_pixel_frames = (samples.shape[2] - num_guide_frames - 1) * time_scale_factor + 1 + if override_frame_indices and str(override_frame_indices).strip(): + indices = _parse_frame_index_list( + override_frame_indices, + "override_frame_indices", + len(indices), + 1, + num_pixel_frames - 1, + expected_desc=f"the keyframes latent carries {len(indices)}", + empty_hint="to reuse the recorded positions", + ) + else: + old_len = keyframes.get("generated_keyframe_num_frames") + if old_len is not None and int(old_len) != num_pixel_frames: + indices = scale_frame_indices(list(indices), int(old_len), num_pixel_frames) + if indices and max(indices) >= num_pixel_frames: + raise ValueError( + f"Keyframe position {max(indices)} is outside this latent's {num_pixel_frames} frames. The " + "target was resized temporally after the keyframes were generated." + ) + + for index, frame_idx in enumerate(indices): + if resize_needed: + added = LTXVAddGuide.execute( + positive, + negative, + vae, + latent, + guides[index].unsqueeze(0), + int(frame_idx), + strength, + ) + positive, negative, latent = added[0], added[1], added[2] + else: + positive, negative, latent = cls.append_latent_keyframe( + positive, + negative, + vae, + latent, + guides[index : index + 1], + int(frame_idx), + strength, + ) + + return io.NodeOutput(positive, negative, latent) + + @classmethod + def append_latent_keyframe(cls, positive, negative, vae, latent, guiding_latent, frame_idx, strength): + """Append an already encoded keyframe latent the same way LTXVAddGuide would after encoding.""" + noise_mask = get_noise_mask(latent) + positive, negative, latent_image, noise_mask = LTXVAddGuide.append_keyframe( + positive, + negative, + frame_idx, + latent["samples"], + noise_mask, + guiding_latent, + strength, + vae.downscale_index_formula, + ) + positive, negative = _append_guide_attention_entry( + positive, + negative, + pre_filter_count=math.prod(guiding_latent.shape[2:]), + latent_shape=list(guiding_latent.shape[2:]), + strength=strength, + ) + out = latent.copy() + out["samples"] = latent_image + out["noise_mask"] = noise_mask + return positive, negative, out + + +class LTXVFreezeLatent(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LTXVFreezeLatent", + display_name="LTXV Freeze Latent", + category="model/latent/ltxv", + search_aliases=["noise mask", "freeze audio", "freeze video"], + description=( + "Set noise_mask to 0 so this latent is kept clean during sampling. " + "Works on video or audio. Typical uses: freeze audio before Concat AV " + "so it only provides cross-attention, or freeze any latent that should " + "not be denoised." + ), + inputs=[ + io.Latent.Input( + "latent", + tooltip="Video or audio latent to freeze. Audio is 4D; video is 5D.", + ), + ], + outputs=[ + io.Latent.Output(display_name="latent"), + ], + ) + + @classmethod + def execute(cls, latent) -> io.NodeOutput: + samples = latent["samples"] + if not torch.is_tensor(samples): + raise ValueError( + "Freeze Latent expects a plain tensor, not a concatenated AV latent. " + "Split with Separate AV Latent first." + ) + out = latent.copy() + if samples.ndim == 5: + batch, _, frames, _, _ = samples.shape + out["noise_mask"] = torch.zeros( + (batch, 1, frames, 1, 1), + dtype=torch.float32, + device=samples.device, + ) + elif samples.ndim == 4: + batch, _, frames, _ = samples.shape + out["noise_mask"] = torch.zeros( + (batch, 1, frames, 1), + dtype=torch.float32, + device=samples.device, + ) + else: + raise ValueError( + f"Expected a 4D audio or 5D video latent, got shape {list(samples.shape)}." + ) + return io.NodeOutput(out) + + +class LTXVKeyframesExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [ + LTXVAddGeneratedKeyframes, + LTXVSeparateGeneratedKeyframes, + LTXVGeneratedKeyframesToGuides, + LTXVFreezeLatent, + ] + + +async def comfy_entrypoint() -> LTXVKeyframesExtension: + return LTXVKeyframesExtension() diff --git a/nodes.py b/nodes.py index e8f3d1e6e..038b88350 100644 --- a/nodes.py +++ b/nodes.py @@ -2460,6 +2460,7 @@ async def init_builtin_extra_nodes(): "nodes_minimax_music.py", "nodes_minimax_h3.py", "nodes_lt.py", + "nodes_lt_keyframes.py", "nodes_hooks.py", "nodes_multigpu.py", "nodes_load_3d.py", diff --git a/tests-unit/comfy_extras_test/nodes_lt_keyframes_test.py b/tests-unit/comfy_extras_test/nodes_lt_keyframes_test.py new file mode 100644 index 000000000..b8a49cf77 --- /dev/null +++ b/tests-unit/comfy_extras_test/nodes_lt_keyframes_test.py @@ -0,0 +1,911 @@ +"""Unit tests for native LTXV generated-keyframe nodes and Freeze Latent. + +They cover keyframe placement, conditioning metadata, guide conversion, and freeze-mask behavior. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +mock_nodes = MagicMock() +mock_nodes.MAX_RESOLUTION = 16384 +mock_server = MagicMock() + + +def _conditioning_get_any_value(conditioning, key, default=None): + for t in conditioning: + if key in t[1]: + return t[1][key] + return default + + +def _get_noise_mask(latent): + noise_mask = latent.get("noise_mask", None) + latent_image = latent["samples"] + if noise_mask is None: + batch_size, _, latent_length, _, _ = latent_image.shape + noise_mask = torch.ones( + (batch_size, 1, latent_length, 1, 1), + dtype=torch.float32, + device=latent_image.device, + ) + else: + noise_mask = noise_mask.clone() + return noise_mask + + +def _get_keyframe_idxs(cond, latent_shape=None): + keyframe_idxs = _conditioning_get_any_value(cond, "keyframe_idxs", None) + if keyframe_idxs is None: + return None, 0 + if latent_shape is not None and len(latent_shape) == 5: + tokens_per_frame = latent_shape[-2] * latent_shape[-1] + num_keyframes = keyframe_idxs.shape[2] // tokens_per_frame + return keyframe_idxs, num_keyframes + return keyframe_idxs, 0 + + +def _append_guide_attention_entry(positive, negative, pre_filter_count, latent_shape, strength=1.0, attention_mask=None): + import node_helpers + + new_entry = { + "pre_filter_count": pre_filter_count, + "strength": strength, + "pixel_mask": None, + "latent_shape": latent_shape, + } + results = [] + for cond in (positive, negative): + existing = [] + for t in cond: + found = t[1].get("guide_attention_entries", None) + if found is not None: + existing = found + break + results.append( + node_helpers.conditioning_set_values(cond, {"guide_attention_entries": [*existing, new_entry]}) + ) + return results[0], results[1] + + +class _StubAddGuide: + calls = [] + + @classmethod + def append_keyframe( + cls, + positive, + negative, + frame_idx, + latent_image, + noise_mask, + guiding_latent, + strength, + scale_factors, + **kwargs, + ): + cls.calls.append({"method": "append_keyframe", "frame_idx": int(frame_idx), "strength": strength}) + mask = torch.full( + (noise_mask.shape[0], 1, guiding_latent.shape[2], noise_mask.shape[3], noise_mask.shape[4]), + max(0.0, 1.0 - strength), + dtype=noise_mask.dtype, + device=noise_mask.device, + ) + return ( + positive, + negative, + torch.cat([latent_image, guiding_latent], dim=2), + torch.cat([noise_mask, mask], dim=2), + ) + + @classmethod + def execute(cls, positive, negative, vae, latent, image, frame_idx, strength, **kwargs): + cls.calls.append({"method": "execute", "frame_idx": int(frame_idx), "strength": strength, "image": image}) + samples = latent["samples"] + out = latent.copy() + extra = torch.zeros( + (samples.shape[0], samples.shape[1], 1, samples.shape[3], samples.shape[4]), + dtype=samples.dtype, + device=samples.device, + ) + out["samples"] = torch.cat([samples, extra], dim=2) + return _NodeOutput(positive, negative, out) + + +class _NodeOutput: + def __init__(self, *args): + self.args = args + + def __getitem__(self, index): + return self.args[index] + + +_nodes_lt_stub = MagicMock() +_nodes_lt_stub.conditioning_get_any_value = _conditioning_get_any_value +_nodes_lt_stub.get_noise_mask = _get_noise_mask +_nodes_lt_stub.get_keyframe_idxs = _get_keyframe_idxs +_nodes_lt_stub._append_guide_attention_entry = _append_guide_attention_entry +_nodes_lt_stub.LTXVAddGuide = _StubAddGuide + +with patch.dict( + "sys.modules", + { + "nodes": mock_nodes, + "server": mock_server, + "comfy_extras.nodes_lt": _nodes_lt_stub, + }, +): + import comfy_extras.nodes_lt_keyframes as keyframes + + +def _zeros(shape): + return torch.zeros(shape) + + +def _empty_121(): + return {"samples": _zeros((1, 2, 16, 2, 1))} + + +def _empty_241(): + return {"samples": _zeros((1, 2, 31, 2, 1))} + + +def _cond(**extra): + return [({}, dict(extra))] + + +def _vae(): + return SimpleNamespace(downscale_index_formula=(8, 32, 32)) + + +def _mask(shape, occupied): + tensor = torch.ones(shape) + for frame in occupied: + tensor[:, :, frame] = 0.0 + return tensor + + +def _keyframe_idxs_at(starts, tokens_per_frame=1): + times = [] + for start in starts: + times.extend([start] * tokens_per_frame) + n = len(times) + coords = torch.zeros((1, 3, n, 2)) + for i, start in enumerate(times): + coords[0, 0, i, 0] = float(start) + coords[0, 0, i, 1] = float(start + 1) + coords[0, 1, i, 1] = 1.0 + coords[0, 2, i, 1] = 1.0 + return coords + + +@contextmanager +def _stub_get_keyframe_idxs(idxs, num_guide_frames): + original = keyframes.get_keyframe_idxs + keyframes.get_keyframe_idxs = lambda cond, shape=None: (idxs, num_guide_frames) + try: + yield + finally: + keyframes.get_keyframe_idxs = original + + +@contextmanager +def _stub_keyframe_coords(): + original = keyframes.LTXVAddGeneratedKeyframes.keyframe_coords + + def _fake(cls, latent, frame_index, scale_factors): + return torch.zeros((latent.shape[0], 3, latent.shape[3] * latent.shape[4], 2)) + + keyframes.LTXVAddGeneratedKeyframes.keyframe_coords = classmethod(_fake) + try: + yield + finally: + keyframes.LTXVAddGeneratedKeyframes.keyframe_coords = original + + +class TestPlacementHelpers: + def test_detailing_positions_121_24(self): + assert keyframes.detailing_positions(121, 24) == [24, 48, 72, 96, 120] + assert keyframes.free_detailing_slots(121, 24, occupied=set()) == [24, 48, 72, 96, 120] + assert keyframes.free_detailing_slots(241, 24, occupied={0, 48, 96, 144, 192, 240}) == [ + 24, 72, 120, 168, 216 + ] + + def test_free_slots_skip_last_frame_when_occupied(self): + assert keyframes.free_detailing_slots(121, 24, occupied={120}) == [24, 48, 72, 96] + + def test_free_slots_rejects_when_every_candidate_is_occupied(self): + with pytest.raises(ValueError, match="already has an image keyframe"): + keyframes.free_detailing_slots(121, 24, occupied={24, 48, 72, 96, 120}) + + def test_scale_frame_indices_temporal_x2(self): + assert keyframes.scale_frame_indices([24, 48, 72, 96, 120], 121, 241) == [ + 48, 96, 144, 192, 240 + ] + with pytest.raises(ValueError, match="from a 1-frame"): + keyframes.scale_frame_indices([0], 1, 241) + with pytest.raises(ValueError, match="onto a 1-frame"): + keyframes.scale_frame_indices([24], 121, 1) + + def test_scale_frame_indices_rejects_collapsed_duplicates(self): + with pytest.raises(ValueError, match="collapsed"): + keyframes.scale_frame_indices([0, 1], 121, 3) + + def test_detailing_positions_keeps_last_skips_zero(self): + positions = keyframes.detailing_positions(121, 24.0) + assert positions[0] != 0 + assert positions[-1] == 120 + + def test_detailing_positions_rejects_nonpositive_interval(self): + with pytest.raises(ValueError, match="interval_frames"): + keyframes.detailing_positions(121, 0) + + def test_detailing_positions_rejects_one_frame_canvas(self): + with pytest.raises(ValueError, match="no pixel frames"): + keyframes.detailing_positions(1, 24) + with pytest.raises(ValueError, match="no pixel frames"): + keyframes.free_detailing_slots(1, 24, occupied=set()) + + def test_keyframes_from_video_stacking_shape(self): + samples = torch.arange(1 * 2 * 4 * 2 * 1, dtype=torch.float32).reshape(1, 2, 4, 2, 1) + stacked = keyframes.keyframes_from_video(samples, [8, 16, 24], temporal_scale=8) + assert stacked.shape == (1, 2, 3, 2, 1) + assert torch.equal(stacked[:, :, 0:1], samples[:, :, 1:2]) + assert torch.equal(stacked[:, :, 1:2], samples[:, :, 2:3]) + assert torch.equal(stacked[:, :, 2:3], samples[:, :, 3:4]) + + def test_keyframes_from_video_rejects_non_video_and_bad_scale(self): + with pytest.raises(ValueError, match="plain 5D video latent"): + keyframes.keyframes_from_video([0], [8], 8) + with pytest.raises(ValueError, match="temporal_scale"): + keyframes.keyframes_from_video(_zeros((1, 2, 4, 2, 1)), [8], 0) + with pytest.raises(ValueError, match="no frames to copy"): + keyframes.keyframes_from_video(_zeros((1, 2, 0, 2, 1)), [8], 8) + + def test_nearest_latent_index_clamps(self): + assert keyframes.nearest_latent_index(0, 8, 4) == 0 + assert keyframes.nearest_latent_index(8, 8, 4) == 1 + assert keyframes.nearest_latent_index(999, 8, 4) == 3 + + def test_should_copy_nearest_video_frames(self): + assert keyframes.should_copy_nearest_video_frames(31, 5, False, False) is True + assert keyframes.should_copy_nearest_video_frames(5, 5, False, False) is False + assert keyframes.should_copy_nearest_video_frames(4, 5, False, False) is False + assert keyframes.should_copy_nearest_video_frames(31, 5, True, False) is False + assert keyframes.should_copy_nearest_video_frames(31, None, False, False) is False + assert keyframes.should_copy_nearest_video_frames(1, 5, False, True) is False + + def test_parse_frame_index_list_validates_count_range_and_duplicates(self): + assert keyframes._parse_frame_index_list( + "24, 48", "frame_indices", 2, 1, 120, "num_keyframes is 2", "to space them" + ) == [24, 48] + assert keyframes._parse_frame_index_list( + "24 48", "frame_indices", 2, 1, 120, "num_keyframes is 2", "to space them" + ) == [24, 48] + with pytest.raises(ValueError, match="lists 1"): + keyframes._parse_frame_index_list( + "24", "frame_indices", 2, 1, 120, "num_keyframes is 2", "to space them" + ) + with pytest.raises(ValueError, match="same pixel frame"): + keyframes._parse_frame_index_list( + "24,24", "frame_indices", 2, 1, 120, "num_keyframes is 2", "to space them" + ) + with pytest.raises(ValueError, match="must lie between"): + keyframes._parse_frame_index_list( + "0,24", "frame_indices", 2, 1, 120, "num_keyframes is 2", "to space them" + ) + with pytest.raises(ValueError, match="could not parse"): + keyframes._parse_frame_index_list( + "24,abc", "frame_indices", 2, 1, 120, "num_keyframes is 2", "to space them" + ) + + def test_parse_frame_index_list_allows_omitted_count(self): + assert keyframes._parse_frame_index_list( + "24,48,72", "frame_indices", None, 1, 120, "unused", "to auto-place" + ) == [24, 48, 72] + + def test_parse_frame_index_list_rejects_empty_separator_only(self): + with pytest.raises(ValueError, match="is empty"): + keyframes._parse_frame_index_list( + ",", "frame_indices", None, 1, 120, "unused", "to place them from interval_frames" + ) + with pytest.raises(ValueError, match="is empty"): + keyframes._parse_frame_index_list( + " , , ", "frame_indices", None, 1, 120, "unused", "to place them from interval_frames" + ) + + def test_add_parse_frame_indices_allows_last_frame(self): + assert keyframes.LTXVAddGeneratedKeyframes.parse_frame_indices("24,120", 121) == [24, 120] + with pytest.raises(ValueError, match="no pixel frames"): + keyframes.LTXVAddGeneratedKeyframes.parse_frame_indices("1", 1) + + def test_occupied_from_nonzero_samples_without_mask(self): + samples = _zeros((1, 2, 16, 2, 1)) + samples[0, 0, 0, 0, 0] = 1.0 + taken = keyframes.occupied_pixel_frames({"samples": samples}, 8, 121) + assert 0 in taken + assert 120 not in taken + + def test_occupied_prefers_noise_mask_over_nonzero_samples(self): + samples = _zeros((1, 2, 16, 2, 1)) + samples[0, 0, 0, 0, 0] = 1.0 + latent = {"samples": samples, "noise_mask": _mask((1, 1, 16, 1, 1), occupied=set())} + assert keyframes.occupied_pixel_frames(latent, 8, 121) == set() + + def test_occupied_ignores_appended_guide_frames(self): + latent = { + "samples": _zeros((1, 2, 21, 2, 1)), + "noise_mask": _mask((1, 1, 21, 1, 1), occupied={0, 16, 17, 18, 19, 20}), + } + taken = keyframes.occupied_pixel_frames(latent, 8, 121, video_latent_frames=16) + assert taken == {0} + + def test_pixel_frames_from_keyframe_idxs_uses_start_not_exclusive_end(self): + idxs = _keyframe_idxs_at([24]) + assert idxs[0, 0, :, 0].tolist() == [24.0] + assert idxs[0, 0, :, 1].tolist() == [25.0] + assert keyframes.pixel_frames_from_keyframe_idxs(idxs) == {24} + assert keyframes.pixel_frames_from_keyframe_idxs(None) == set() + + def test_pixel_frames_from_keyframe_idxs_rejects_malformed(self): + with pytest.raises((TypeError, AttributeError, IndexError, ValueError)): + keyframes.pixel_frames_from_keyframe_idxs("not-a-tensor") + with pytest.raises((TypeError, ValueError)): + keyframes._as_int_set(object()) + + +class TestNativeSchemas: + def test_generated_keyframe_nodes_use_ltxv_conditioning_category(self): + for cls, node_id, display_name in ( + ( + keyframes.LTXVAddGeneratedKeyframes, + "LTXVAddGeneratedKeyframes", + "LTXV Add Generated Keyframes", + ), + ( + keyframes.LTXVSeparateGeneratedKeyframes, + "LTXVSeparateGeneratedKeyframes", + "LTXV Separate Generated Keyframes", + ), + ( + keyframes.LTXVGeneratedKeyframesToGuides, + "LTXVGeneratedKeyframesToGuides", + "LTXV Generated Keyframes to Guides", + ), + ): + schema = cls.define_schema() + assert schema.node_id == node_id + assert schema.display_name == display_name + assert schema.category == "model/conditioning/ltxv" + assert "dfr" in schema.search_aliases + + def test_freeze_latent_uses_ltxv_latent_category(self): + schema = keyframes.LTXVFreezeLatent.define_schema() + assert schema.node_id == "LTXVFreezeLatent" + assert schema.display_name == "LTXV Freeze Latent" + assert schema.category == "model/latent/ltxv" + + +class TestAddGeneratedKeyframes: + def test_rejects_non_video_latent(self): + with pytest.raises(ValueError, match="plain video latent"): + keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), _cond(), _vae(), {"samples": torch.zeros(1, 2, 16, 2)} + ) + + def test_execute_rejects_separator_only_frame_indices(self): + with pytest.raises(ValueError, match="is empty"): + keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), _cond(), _vae(), _empty_121(), frame_indices="," + ) + + def test_execute_rejects_one_frame_canvas(self): + with pytest.raises(ValueError, match="no pixel frames"): + keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), _cond(), _vae(), {"samples": _zeros((1, 2, 1, 2, 1))} + ) + + def test_rejects_rescaled_or_noncontiguous_existing_keyframes(self): + latent = _empty_121() + with pytest.raises(ValueError, match="rescaled"): + keyframes.LTXVAddGeneratedKeyframes.execute( + _cond( + generated_keyframes={ + "tokens_per_frame": 99, + "first_latent_frame": 16, + "num_keyframes": 0, + } + ), + _cond(), + _vae(), + latent, + ) + with pytest.raises(ValueError, match="contiguous"): + keyframes.LTXVAddGeneratedKeyframes.execute( + _cond( + generated_keyframes={ + "tokens_per_frame": 2, + "first_latent_frame": 10, + "num_keyframes": 3, + } + ), + _cond(), + _vae(), + latent, + ) + + def test_execute_appends_zero_keyframes_on_t(self): + with _stub_keyframe_coords(): + _positive, _negative, out = keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), _cond(), _vae(), _empty_121() + ) + assert out["samples"].shape == (1, 2, 21, 2, 1) + assert out["noise_mask"].shape[2] == 21 + assert torch.all(out["noise_mask"][:, :, 16:21] == 1.0) + + def test_execute_copies_nearest_frames_from_longer_video(self): + video = {"samples": torch.arange(1 * 2 * 16 * 2 * 1, dtype=torch.float32).reshape(1, 2, 16, 2, 1)} + with _stub_keyframe_coords(): + _positive, _negative, out = keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), + _cond(), + _vae(), + _empty_121(), + frame_indices="24,48,72,96,120", + keyframes=video, + ) + assert out["samples"].shape[2] == 21 + stacked = out["samples"][:, :, 16:21] + source = video["samples"] + assert torch.equal(stacked[:, :, 0:1], source[:, :, 3:4]) + assert torch.equal(stacked[:, :, 4:5], source[:, :, 15:16]) + + def test_execute_keeps_stacked_keyframes_when_t_equals_count(self): + stacked = {"samples": torch.arange(1 * 2 * 5 * 2 * 1, dtype=torch.float32).reshape(1, 2, 5, 2, 1)} + with _stub_keyframe_coords(): + _positive, _negative, out = keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), + _cond(), + _vae(), + _empty_121(), + frame_indices="24,48,72,96,120", + keyframes=stacked, + ) + assert out["samples"].shape[2] == 21 + assert torch.equal(out["samples"][:, :, 16:21], stacked["samples"]) + + def test_execute_reshapes_batched_single_frame_keyframes(self): + batched = {"samples": _zeros((5, 2, 1, 2, 1))} + with _stub_keyframe_coords(): + _positive, _negative, out = keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), + _cond(), + _vae(), + _empty_121(), + frame_indices="24,48,72,96,120", + keyframes=batched, + ) + assert out["samples"].shape[2] == 21 + + def test_execute_records_density_slots_and_canvas_length(self): + with _stub_keyframe_coords(): + positive, _negative, _out = keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), _cond(), _vae(), _empty_121() + ) + record = positive[0][1]["generated_keyframes"] + assert record["frame_indices"] == [24, 48, 72, 96, 120] + assert record["num_pixel_frames"] == 121 + assert record["num_keyframes"] == 5 + assert record["first_latent_frame"] == 16 + assert record["guide_entry_index"] == 0 + entries = positive[0][1]["guide_attention_entries"] + assert len(entries) == 1 + assert entries[0]["pre_filter_count"] == 5 * 2 * 1 + assert entries[0]["latent_shape"] == [5, 2, 1] + + def test_execute_copies_from_video_using_auto_slots(self): + video = {"samples": torch.arange(1 * 2 * 16 * 2 * 1, dtype=torch.float32).reshape(1, 2, 16, 2, 1)} + with _stub_keyframe_coords(): + _positive, _negative, out = keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), _cond(), _vae(), _empty_121(), keyframes=video + ) + stacked = out["samples"][:, :, 16:21] + source = video["samples"] + assert torch.equal(stacked[:, :, 0:1], source[:, :, 3:4]) + assert torch.equal(stacked[:, :, 4:5], source[:, :, 15:16]) + + def test_execute_replaces_stacked_tokens_on_current_canvas(self): + stacked = { + "samples": torch.arange(1 * 2 * 5 * 2 * 1, dtype=torch.float32).reshape(1, 2, 5, 2, 1), + "generated_keyframe_indices": [24, 48, 72, 96, 120], + "generated_keyframe_num_frames": 121, + } + with _stub_keyframe_coords(): + positive, _negative, out = keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), _cond(), _vae(), _empty_121(), keyframes=stacked + ) + assert torch.equal(out["samples"][:, :, 16:21], stacked["samples"]) + assert positive[0][1]["generated_keyframes"]["frame_indices"] == [24, 48, 72, 96, 120] + + def test_execute_skips_i2v_last_frame_noise_mask(self): + latent = { + "samples": _zeros((1, 2, 16, 2, 1)), + "noise_mask": _mask((1, 1, 16, 1, 1), occupied={15}), + } + with _stub_keyframe_coords(): + positive, _negative, out = keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), _cond(), _vae(), latent + ) + indices = positive[0][1]["generated_keyframes"]["frame_indices"] + assert 120 not in indices + assert indices == [24, 48, 72, 96] + assert out["samples"].shape[2] == 20 + + def test_execute_replaces_stacked_tokens_on_longer_canvas(self): + stacked = { + "samples": _zeros((1, 2, 5, 2, 1)), + "generated_keyframe_indices": [24, 48, 72, 96, 120], + "generated_keyframe_num_frames": 121, + } + latent = _empty_241() + latent["noise_mask"] = _mask((1, 1, 31, 1, 1), occupied={0}) + with _stub_keyframe_coords(): + positive, _negative, _out = keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), _cond(), _vae(), latent, keyframes=stacked + ) + indices = positive[0][1]["generated_keyframes"]["frame_indices"] + assert indices != [24, 48, 72, 96, 120] + assert indices == [24, 48, 72, 96, 120, 144, 168, 192, 216, 240] + + def test_execute_skips_existing_guide_keyframe_idxs(self): + latent = { + "samples": _zeros((1, 2, 36, 2, 1)), + "noise_mask": _mask((1, 1, 36, 1, 1), occupied={0}), + } + idxs = _keyframe_idxs_at([48, 96, 144, 192, 240]) + with _stub_keyframe_coords(), _stub_get_keyframe_idxs(idxs, 5): + positive, _negative, _out = keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), _cond(), _vae(), latent + ) + assert positive[0][1]["generated_keyframes"]["frame_indices"] == [24, 72, 120, 168, 216] + + def test_execute_rejects_occupied_manual_indices(self): + latent = { + "samples": _zeros((1, 2, 16, 2, 1)), + "noise_mask": _mask((1, 1, 16, 1, 1), occupied={15}), + } + with pytest.raises(ValueError, match="reuses pixel frame"): + keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), _cond(), _vae(), latent, frame_indices="24,120" + ) + + def test_execute_rejects_wrong_spatial_size_keyframes(self): + stacked = {"samples": _zeros((1, 2, 5, 4, 4))} + with pytest.raises(ValueError, match="whole latent frames"): + keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), + _cond(), + _vae(), + _empty_121(), + frame_indices="24,48,72,96,120", + keyframes=stacked, + ) + + def test_execute_rejects_too_many_stacked_keyframes(self): + stacked = { + "samples": _zeros((1, 2, 6, 2, 1)), + "generated_keyframe_indices": [24, 48, 72, 96, 120, 8], + } + with pytest.raises(ValueError, match="only 5 free slot"): + keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), + _cond(), + _vae(), + _empty_121(), + frame_indices="24,48,72,96,120", + keyframes=stacked, + ) + + def test_execute_rejects_non_5d_keyframes(self): + with pytest.raises(ValueError, match="5 dimensional"): + keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), + _cond(), + _vae(), + _empty_121(), + frame_indices="24", + keyframes={"samples": torch.zeros(1, 2, 1, 2)}, + ) + + def test_execute_grows_existing_generated_block(self): + latent = _empty_121() + with _stub_keyframe_coords(): + positive, negative, out = keyframes.LTXVAddGeneratedKeyframes.execute( + _cond(), _cond(), _vae(), latent, frame_indices="24,48" + ) + positive, negative, out = keyframes.LTXVAddGeneratedKeyframes.execute( + positive, negative, _vae(), out, frame_indices="72,96" + ) + record = positive[0][1]["generated_keyframes"] + assert record["frame_indices"] == [24, 48, 72, 96] + assert record["num_keyframes"] == 4 + assert record["first_latent_frame"] == 16 + assert out["samples"].shape[2] == 20 + entries = positive[0][1]["guide_attention_entries"] + assert len(entries) == 1 + assert entries[0]["pre_filter_count"] == 4 * 2 * 1 + assert entries[0]["latent_shape"] == [4, 2, 1] + + +class TestSeparateGeneratedKeyframes: + def test_requires_generated_keyframes(self): + with pytest.raises(ValueError, match="no generated keyframes"): + keyframes.LTXVSeparateGeneratedKeyframes.execute(_cond(), _cond(), _empty_121()) + + def test_execute_peels_keyframes_and_indices(self): + video = _zeros((1, 2, 16, 2, 1)) + keys = torch.arange(1, 1 + 1 * 2 * 5 * 2 * 1, dtype=torch.float32).reshape(1, 2, 5, 2, 1) + samples = torch.cat([video, keys], dim=2) + record = { + "first_latent_frame": 16, + "num_keyframes": 5, + "frame_indices": [24, 48, 72, 96, 120], + "num_pixel_frames": 121, + "guide_entry_index": 0, + "tokens_per_frame": 2, + } + positive, negative, latent, peeled = keyframes.LTXVSeparateGeneratedKeyframes.execute( + _cond( + generated_keyframes=record, + guide_attention_entries=[{"keep": False}, {"keep": True}], + ), + _cond(generated_keyframes=record), + {"samples": samples}, + ) + assert latent["samples"].shape == (1, 2, 16, 2, 1) + assert peeled["samples"].shape == (1, 2, 5, 2, 1) + assert peeled["generated_keyframe_indices"] == [24, 48, 72, 96, 120] + assert peeled["generated_keyframe_num_frames"] == 121 + assert torch.equal(peeled["samples"], keys) + assert positive[0][1]["generated_keyframes"] is None + assert positive[0][1]["guide_attention_entries"] == [{"keep": True}] + assert negative[0][1]["generated_keyframes"] is None + + def test_execute_keyframes_to_batch(self): + samples = _zeros((1, 2, 18, 2, 1)) + record = { + "first_latent_frame": 16, + "num_keyframes": 2, + "frame_indices": [24, 48], + "guide_entry_index": 0, + "tokens_per_frame": 2, + } + _p, _n, _latent, peeled = keyframes.LTXVSeparateGeneratedKeyframes.execute( + _cond(generated_keyframes=record), + _cond(generated_keyframes=record), + {"samples": samples}, + keyframes_to_batch=True, + ) + assert peeled["samples"].shape == (2, 2, 1, 2, 1) + + def test_rejects_token_mismatch_and_short_latent(self): + record = { + "first_latent_frame": 16, + "num_keyframes": 5, + "frame_indices": [24, 48, 72, 96, 120], + "guide_entry_index": 0, + "tokens_per_frame": 99, + } + with pytest.raises(ValueError, match="rescaled"): + keyframes.LTXVSeparateGeneratedKeyframes.execute( + _cond(generated_keyframes=record), + _cond(generated_keyframes=record), + _empty_121(), + ) + record = dict(record) + record["tokens_per_frame"] = 2 + with pytest.raises(ValueError, match="only has"): + keyframes.LTXVSeparateGeneratedKeyframes.execute( + _cond(generated_keyframes=record), + _cond(generated_keyframes=record), + _empty_121(), + ) + + def test_strip_guide_entry(self): + remaining = keyframes.LTXVSeparateGeneratedKeyframes.strip_guide_entry( + [({}, {"guide_attention_entries": [{"a": 1}, {"b": 2}]})], 0 + ) + assert remaining == [{"b": 2}] + empty = keyframes.LTXVSeparateGeneratedKeyframes.strip_guide_entry( + [({}, {"guide_attention_entries": [{"a": 1}]})], 0 + ) + assert empty is None + with pytest.raises(ValueError, match="recorded guide entry"): + keyframes.LTXVSeparateGeneratedKeyframes.strip_guide_entry( + [({}, {"guide_attention_entries": [{"a": 1}]})], 5 + ) + + def test_rejects_non_video_latent(self): + record = { + "first_latent_frame": 0, + "num_keyframes": 1, + "frame_indices": [24], + "guide_entry_index": 0, + "tokens_per_frame": 2, + } + with pytest.raises(ValueError, match="plain video latent"): + keyframes.LTXVSeparateGeneratedKeyframes.execute( + _cond(generated_keyframes=record), + _cond(generated_keyframes=record), + {"samples": torch.zeros(1, 2, 16, 2)}, + ) + + +class TestGeneratedKeyframesToGuides: + def test_requires_recorded_indices(self): + with pytest.raises(ValueError, match="does not carry generated keyframe positions"): + keyframes.LTXVGeneratedKeyframesToGuides.execute( + _cond(), _cond(), _vae(), _empty_121(), {"samples": _zeros((1, 2, 5, 2, 1))}, 1.0 + ) + + def test_rejects_unseparated_conditioning(self): + with pytest.raises(ValueError, match="still carries generated keyframes"): + keyframes.LTXVGeneratedKeyframesToGuides.execute( + _cond(generated_keyframes={"num_keyframes": 1}), + _cond(), + _vae(), + _empty_121(), + {"samples": _zeros((1, 2, 1, 2, 1)), "generated_keyframe_indices": [24]}, + 1.0, + ) + + def test_rejects_non_video_and_batched_canvas(self): + kf = {"samples": _zeros((1, 2, 1, 2, 1)), "generated_keyframe_indices": [24]} + with pytest.raises(ValueError, match="plain video latent"): + keyframes.LTXVGeneratedKeyframesToGuides.execute( + _cond(), _cond(), _vae(), {"samples": torch.zeros(1, 2, 16, 2)}, kf, 1.0 + ) + with pytest.raises(ValueError, match="batch size of 1"): + keyframes.LTXVGeneratedKeyframesToGuides.execute( + _cond(), _cond(), _vae(), {"samples": _zeros((2, 2, 16, 2, 1))}, kf, 1.0 + ) + + def test_pins_same_size_keyframes_via_append(self): + _StubAddGuide.calls.clear() + kf = { + "samples": _zeros((1, 2, 2, 2, 1)), + "generated_keyframe_indices": [24, 48], + "generated_keyframe_num_frames": 121, + } + positive, negative, out = keyframes.LTXVGeneratedKeyframesToGuides.execute( + _cond(), _cond(), _vae(), _empty_121(), kf, 1.0 + ) + assert out["samples"].shape[2] == 18 + assert torch.all(out["noise_mask"][:, :, 16:] == 0.0) + assert [call["frame_idx"] for call in _StubAddGuide.calls] == [24, 48] + assert all(call["method"] == "append_keyframe" for call in _StubAddGuide.calls) + entries = positive[0][1]["guide_attention_entries"] + assert len(entries) == 2 + + def test_scales_indices_after_temporal_x2(self): + _StubAddGuide.calls.clear() + kf = { + "samples": _zeros((1, 2, 2, 2, 1)), + "generated_keyframe_indices": [24, 120], + "generated_keyframe_num_frames": 121, + } + keyframes.LTXVGeneratedKeyframesToGuides.execute( + _cond(), _cond(), _vae(), _empty_241(), kf, 1.0 + ) + assert [call["frame_idx"] for call in _StubAddGuide.calls] == [48, 240] + + def test_override_frame_indices(self): + _StubAddGuide.calls.clear() + kf = { + "samples": _zeros((1, 2, 2, 2, 1)), + "generated_keyframe_indices": [24, 48], + "generated_keyframe_num_frames": 121, + } + keyframes.LTXVGeneratedKeyframesToGuides.execute( + _cond(), _cond(), _vae(), _empty_121(), kf, 0.5, override_frame_indices="32,96" + ) + assert [call["frame_idx"] for call in _StubAddGuide.calls] == [32, 96] + assert all(call["strength"] == 0.5 for call in _StubAddGuide.calls) + + def test_resize_path_decodes_and_calls_add_guide(self): + _StubAddGuide.calls.clear() + vae = _vae() + decoded = [] + + def decode(samples): + decoded.append(tuple(samples.shape)) + return torch.zeros((samples.shape[0], 8, 8, 3)) + + vae.decode = decode + kf = { + "samples": _zeros((1, 2, 2, 4, 4)), + "generated_keyframe_indices": [24, 48], + "generated_keyframe_num_frames": 121, + } + _p, _n, out = keyframes.LTXVGeneratedKeyframesToGuides.execute( + _cond(), _cond(), vae, _empty_121(), kf, 1.0 + ) + assert decoded == [(2, 2, 1, 4, 4)] + assert [call["method"] for call in _StubAddGuide.calls] == ["execute", "execute"] + assert out["samples"].shape[2] == 18 + + def test_rejects_count_mismatch(self): + kf = { + "samples": _zeros((1, 2, 2, 2, 1)), + "generated_keyframe_indices": [24], + "generated_keyframe_num_frames": 121, + } + with pytest.raises(ValueError, match="2 keyframes for 1 recorded"): + keyframes.LTXVGeneratedKeyframesToGuides.execute( + _cond(), _cond(), _vae(), _empty_121(), kf, 1.0 + ) + + def test_override_rejects_separator_only_indices(self): + kf = { + "samples": _zeros((1, 2, 2, 2, 1)), + "generated_keyframe_indices": [24, 48], + "generated_keyframe_num_frames": 121, + } + with pytest.raises(ValueError, match="is empty"): + keyframes.LTXVGeneratedKeyframesToGuides.execute( + _cond(), _cond(), _vae(), _empty_121(), kf, 1.0, override_frame_indices="," + ) + + +class TestFreezeLatent: + def test_video_and_audio_masks(self): + video = keyframes.LTXVFreezeLatent.execute({"samples": _zeros((2, 4, 8, 3, 5))})[0] + assert video["noise_mask"].shape == (2, 1, 8, 1, 1) + assert video["noise_mask"].device.type == "cpu" + assert torch.all(video["noise_mask"] == 0) + audio = keyframes.LTXVFreezeLatent.execute({"samples": _zeros((1, 8, 16, 4))})[0] + assert audio["noise_mask"].shape == (1, 1, 16, 1) + assert torch.all(audio["noise_mask"] == 0) + + def test_preserves_extra_latent_keys(self): + out = keyframes.LTXVFreezeLatent.execute( + {"samples": _zeros((1, 4, 8, 2, 2)), "downscale_ratio_spacial": 32} + )[0] + assert out["downscale_ratio_spacial"] == 32 + + def test_rejects_av_and_wrong_rank(self): + with pytest.raises(ValueError, match="plain tensor"): + keyframes.LTXVFreezeLatent.execute({"samples": [0.0]}) + with pytest.raises(ValueError, match="4D audio or 5D video"): + keyframes.LTXVFreezeLatent.execute({"samples": _zeros((1, 2, 3))}) + + +class TestKeyframeCoords: + def test_single_pixel_span_at_requested_index(self): + latent = torch.zeros((1, 4, 1, 2, 2)) + coords = keyframes.LTXVAddGeneratedKeyframes.keyframe_coords(latent, 24, (8, 32, 32)) + assert coords.shape[0] == 1 + assert coords.shape[1] == 3 + assert coords.shape[-1] == 2 + starts = coords[0, 0, :, 0] + ends = coords[0, 0, :, 1] + assert torch.all(starts == 24) + assert torch.all(ends == 25) + + +def test_extension_registers_all_four_nodes(): + import asyncio + + ext = asyncio.run(keyframes.comfy_entrypoint()) + names = [cls.__name__ for cls in asyncio.run(ext.get_node_list())] + assert names == [ + "LTXVAddGeneratedKeyframes", + "LTXVSeparateGeneratedKeyframes", + "LTXVGeneratedKeyframesToGuides", + "LTXVFreezeLatent", + ]