scripts/ltx_av_guide.py

"""AV-aware guide injection, which is what image- and video-to-video need.

ComfyUI already ships the guide mechanism as `LTXVAddGuide`: it encodes a still
or a clip, appends the resulting tokens to the end of the latent sequence, tells
the transformer where those tokens belong through `keyframe_idxs`, and masks
them out of the noise. `LTXVCropGuides` strips them back off after sampling.

Every piece of that already works for LTX-2.5's audio-video model:

* `LTXAVModel._process_input` forwards `keyframe_idxs` and the denoise mask to
  the video branch alone, and patchifies audio separately.
* `_process_output` restores the video branch to `orig_shape` - guides still
  attached - which is exactly what `LTXVCropGuides` expects to trim.
* `model_base` unpacks a packed denoise mask into separate video and audio
  masks, so a video-only mask is a supported input and not a workaround.

The one thing that does not work is the node. `append_keyframe` ends with

    latent_image = torch.cat([latent_image, guiding_latent], dim=2)

and an LTX-2.5 latent is a `comfy.nested_tensor.NestedTensor` holding the video
and audio latents as a pair. `NestedTensor` is not a `torch.Tensor` and defines
no `__torch_function__`, so `torch.cat` raises. The same module defines
`cat_nested`, which would do the right thing - nothing anywhere calls it.

The guard one line earlier,

    if latent_image.shape[1] != in_channels or guiding_latent.shape[1] != in_channels:
        raise ValueError("Adding guide to a combined AV latent is not supported.")

never fires, because `NestedTensor.shape` proxies to `tensors[0].shape` - the
video latent, whose channel count is the 128 the check is looking for. So the
real failure is a `TypeError` from `torch.cat`, and that error message is stale:
guides on a combined AV latent are supported by everything downstream of it.

Nothing here reimplements the guide maths. It unwraps the pair, hands the video
half to the stock node, and re-wraps the result, so the arithmetic stays the
vendor's and stays in one place. The audio branch is left untouched on purpose:
it carries its own coordinates and its own length, `CFGGuider.sample` pads a
video-only denoise mask with ones to cover it, and `model_base` splits the
packed mask apart again.
"""

from __future__ import annotations

from pathlib import Path


def unwrap(node_output):
    """ComfyUI nodes return `io.NodeOutput`; the values live behind it."""
    for attr in ("result", "_result", "values"):
        if hasattr(node_output, attr):
            value = getattr(node_output, attr)
            if value is not None:
                return value
    return node_output


def split_av(latent):
    """`(video, extras)`, where `extras` is the audio half when there is one.

    Kept tolerant of a plain tensor so the same code path serves a video-only
    checkpoint without a branch at every call site.
    """
    samples = latent["samples"]
    if getattr(samples, "is_nested", False):
        tensors = list(samples.unbind())
        return tensors[0], tensors[1:]
    return samples, []


def video_noise_mask(latent):
    """The video half of a noise mask, or None.

    A mask may arrive nested (from `LTXVConcatAVLatent`) or already video-only
    (from a guide). Both reduce to the same thing here.
    """
    mask = latent.get("noise_mask")
    if mask is None:
        return None
    if getattr(mask, "is_nested", False):
        return mask.unbind()[0]
    return mask


def join_av(latent, video, extras, noise_mask=None):
    """Rebuild the latent dict, re-nesting the audio half if there was one."""
    import comfy.nested_tensor

    out = dict(latent)
    if extras:
        out["samples"] = comfy.nested_tensor.NestedTensor((video, *extras))
    else:
        out["samples"] = video
    if noise_mask is None:
        out.pop("noise_mask", None)
    else:
        out["noise_mask"] = noise_mask
    return out


def add_av_guide(positive, negative, video_vae, latent, image, frame_idx=0,
                 strength=1.0):
    """`LTXVAddGuide` against the video branch of an AV latent.

    `negative` is required by the node because it writes `keyframe_idxs` into
    both conditionings. At CFG 1.0 there is no negative to speak of, so callers
    pass the positive twice and discard the second return - the guide tokens are
    identical either way, and the distilled schedule never reads it.
    """
    from comfy_extras.nodes_lt import LTXVAddGuide

    video, extras = split_av(latent)
    video_latent = {"samples": video}
    mask = video_noise_mask(latent)
    if mask is not None:
        video_latent["noise_mask"] = mask

    positive, negative, guided = unwrap(LTXVAddGuide.execute(
        positive, negative, video_vae, video_latent, image, frame_idx, strength))
    return positive, negative, join_av(
        latent, guided["samples"], extras, guided.get("noise_mask"))


def crop_av_guides(positive, negative, latent):
    """`LTXVCropGuides` against the video branch, after sampling."""
    from comfy_extras.nodes_lt import LTXVCropGuides

    video, extras = split_av(latent)
    video_latent = {"samples": video}
    mask = video_noise_mask(latent)
    if mask is not None:
        video_latent["noise_mask"] = mask

    positive, negative, cropped = unwrap(
        LTXVCropGuides.execute(positive, negative, video_latent))
    return positive, negative, join_av(
        latent, cropped["samples"], extras, cropped.get("noise_mask"))


def load_image(path: Path):
    """A still in ComfyUI's IMAGE layout: `[1, H, W, C]`, float, 0..1."""
    import numpy as np
    import torch
    from PIL import Image

    with Image.open(path) as handle:
        array = np.asarray(handle.convert("RGB"), dtype=np.float32) / 255.0
    return torch.from_numpy(array).unsqueeze(0)


def load_video(path: Path, max_frames: int | None = None):
    """Decoded frames as `[T, H, W, C]`, float, 0..1.

    `av` is already a hard dependency of ComfyUI's video nodes, so this adds
    nothing that a working ComfyUI does not already have.
    """
    import av
    import numpy as np
    import torch

    frames = []
    with av.open(str(path)) as container:
        for frame in container.decode(video=0):
            frames.append(frame.to_ndarray(format="rgb24"))
            if max_frames is not None and len(frames) >= max_frames:
                break
    if not frames:
        raise SystemExit(f"no video frames decoded from {path}")
    array = np.stack(frames).astype(np.float32) / 255.0
    return torch.from_numpy(array)


def trim_to_guide_length(frames, time_scale: int = 8):
    """Crop to `8*n + 1` frames.

    `LTXVAddGuide` does this itself and says so in its tooltip, but silently.
    Doing it here means the count that actually reached the model is the count
    the log prints, which matters when a guide comes out shorter than asked.
    """
    count = int(frames.shape[0])
    if count <= 1:
        return frames[:1]
    usable = (count - 1) // time_scale * time_scale + 1
    return frames[:usable]


def encode_video_latent(video_vae, frames, width: int, height: int, length: int,
                        tile: int = 256, temporal: int = 16):
    """Encode pixels into the video branch at exactly `length` frames.

    Short input is held on its last frame rather than looped or zero-padded:
    a freeze reads as a still, while a loop invents motion the source never had
    and zeros inject a black flash the sampler then has to explain away.

    The encode is tiled and the VAE is walked on and off the device by hand, for
    the same reason `ltx_render_clips.decode` does it: calling `VAE.encode`
    directly bypasses ComfyUI's memory management, so nothing knows to evict the
    10.6 GB DiT first. Untiled, a 25-frame 512x320 clip is enough to end a 16 GB
    card three seconds into the run - measured, not guessed.

    `inference_mode` is required rather than merely tidy. `VAE.process_input`
    writes in place, and outside inference mode those tiles come back as
    ordinary tensors whose in-place write then raises.
    """
    import comfy.model_management
    import comfy.utils
    import torch

    pixels = comfy.utils.common_upscale(
        frames[:, :, :, :3].movedim(-1, 1), width, height,
        "bilinear", "disabled").movedim(1, -1)
    if pixels.shape[0] < length:
        tail = pixels[-1:].expand(length - pixels.shape[0], -1, -1, -1)
        pixels = torch.cat([pixels, tail], dim=0)
    else:
        pixels = pixels[:length]

    comfy.model_management.unload_all_models()
    comfy.model_management.soft_empty_cache()
    # The move stays OUTSIDE inference mode. `Module.to` rebinds every
    # parameter, and parameters rebound inside inference mode stay inference
    # tensors for the life of the process - so a later `VAE.encode` dies with
    # "Inference tensors cannot be saved for backward", nowhere near here.
    video_vae.first_stage_model.to(comfy.model_management.get_torch_device())
    with torch.inference_mode():
        latent = video_vae.encode_tiled(
            pixels, tile_x=tile, tile_y=tile, overlap=tile // 8,
            tile_t=temporal, overlap_t=temporal // 8)
    video_vae.first_stage_model.to("cpu")
    comfy.model_management.soft_empty_cache()
    return latent.clone()


def denoise_sigmas(sigmas, denoise: float):
    """The tail of a schedule, which is how a partial denoise is expressed.

    Sampling from sigma 1.0 discards the input entirely; starting partway down
    the curve is what makes video-to-video keep its source. The count is over
    steps, not sigmas, so `denoise=0.5` on eight steps runs four of them.
    """
    if denoise >= 1.0:
        return sigmas
    if not 0.0 < denoise < 1.0:
        raise SystemExit("denoise must be in (0, 1]")
    steps = max(1, int(round((len(sigmas) - 1) * denoise)))
    return sigmas[-(steps + 1):]


def upscale_av_latent(latent, video_vae, model_path: Path):
    """The vendor's 2x latent upscale, on the video branch only.

    `ltx_render_clips` skips this pass and says why: it doubles each spatial
    dimension, so the second sampling runs at four times the token count, and
    that was judged unaffordable on a 16 GB card. The long-form ladder later
    measured 353 frames at 6.26 GiB, which says the judgement deserved
    rechecking rather than inheriting.

    `LTXVLatentUpsampler` reads `samples["samples"]` and hands it to a conv
    stack, so it has the same NestedTensor problem `LTXVAddGuide` has, and takes
    the same treatment. The audio branch is left alone: its coordinates are
    temporal and a spatial upscale means nothing to it.
    """
    import comfy.model_management
    import folder_paths
    from comfy_extras.nodes_hunyuan import LatentUpscaleModelLoader
    from comfy_extras.nodes_lt_upsampler import LTXVLatentUpsampler

    model_path = Path(model_path)
    folder_paths.add_model_folder_path("latent_upscale_models", str(model_path.parent))
    upscaler = unwrap(LatentUpscaleModelLoader.execute(model_path.name))
    if isinstance(upscaler, (tuple, list)):
        upscaler = upscaler[0]

    video, extras = split_av(latent)
    out = LTXVLatentUpsampler().upsample_latent(
        {"samples": video}, upscaler, video_vae)
    if isinstance(out, (tuple, list)):
        out = out[0]
    comfy.model_management.soft_empty_cache()
    # The upsampler drops any noise mask; a guide would have to be re-applied
    # after this, and nothing here does that yet.
    return join_av(latent, out["samples"], extras)


#: Seconds of video the DiT's temporal RoPE was built for. `LTXAVModel` carries
#: `positional_embedding_max_pos=[20, 2048, 2048]`, and both
#: `_prepare_positional_embeddings` and the AV cross-attention path divide the
#: temporal coordinate by the frame rate first - so the 20 is seconds, not
#: frames.
TEMPORAL_MAX_SECONDS = 20.0


def fit_temporal_positions(model, seconds: float):
    """Keep the temporal RoPE inside the range it was trained on.

    `get_fractional_positions` divides each coordinate by its `max_pos`, so a
    clip longer than 20 s pushes the temporal fraction past 1.0 and the rotary
    frequencies extrapolate off the end of what the model ever saw. Raising
    `max_pos[0]` to the clip's own duration maps it back onto [0, 1] instead -
    the same trick position interpolation plays on a language model's context,
    and what "Train Short, Inference Long" (arXiv 2602.14027) argues is the
    actual failure behind long-horizon degradation.

    Returns the value it set, or None when the clip already fits and nothing was
    touched. Patching is safe because `generate_freq_grid_np`'s cache is keyed on
    theta, dimension count and device - never on `max_pos`.

    Interpolation is not free: it compresses temporal detail, because a second
    of video now occupies less of the rotary circle than it did in training.
    Below 20 s this does nothing at all, and that is deliberate.
    """
    inner = getattr(model, "model", model)
    diffusion = getattr(inner, "diffusion_model", inner)
    max_pos = getattr(diffusion, "positional_embedding_max_pos", None)
    if max_pos is None:
        raise SystemExit("this model exposes no positional_embedding_max_pos")
    if seconds <= TEMPORAL_MAX_SECONDS:
        return None
    max_pos[0] = float(seconds)
    return float(seconds)


def reset_temporal_positions(model):
    """Put `max_pos[0]` back to what the checkpoint shipped with.

    Needed because the patch is a mutation on a live model that outlives one
    request: without this, a clip rendered after a long one would silently
    inherit its interpolation and no longer be the render its settings describe.
    """
    inner = getattr(model, "model", model)
    diffusion = getattr(inner, "diffusion_model", inner)
    max_pos = getattr(diffusion, "positional_embedding_max_pos", None)
    if max_pos is not None:
        max_pos[0] = TEMPORAL_MAX_SECONDS
    return None