iclora.py

import logging

import comfy
import comfy.samplers
import comfy.sd
import comfy.utils
import comfy_extras.nodes_lt as nodes_lt
import folder_paths
import node_helpers
import torch
from comfy_api.latest import io

from .latents import LTXVDilateLatent
from .nodes_registry import NODES_DISPLAY_NAME_PREFIX, comfy_node
from .tiled_fusion_plan import snap_frames

FRAME_IDX_TOOLTIP = (
    "RoPE position of the reference, in pixel frames.\n"
    "0 aligns the reference with the generated clip. Use for video references.\n"
    ">0 places the reference at a specific pixel-frame position in the clip.\n"
    "-1 places the reference outside the generated frame range, while still using it "
    "for attention conditioning. Use for single-image references.\n"
    "For multi-frame references, >0 values are rounded down to 1 mod 8."
)


@comfy_node(name="LTXAddVideoICLoRAGuide")
class LTXAddVideoICLoRAGuide(io.ComfyNode):
    PATCHIFIER = nodes_lt.SymmetricPatchifier(1, start_end=True)

    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="LTXAddVideoICLoRAGuide",
            display_name=NODES_DISPLAY_NAME_PREFIX + " Add Video IC-LoRA Guide",
            category="Lightricks/IC-LoRA",
            description=(
                "Adds one or more conditioning frames starting at the specified frame index. "
                "Supports both single images and multi-frame videos. "
                "The latent_downscale_factor resizes input to a fraction of the target size "
                "(1 = original, 2 = half, 3 = third, etc.) for IC-LoRA on small grids. "
                "For LTXVTiledFusionSampler temporal windows, turn on use_streaming and "
                "set tile_frames to the same 8n+1 value as the sampler (97 for LTX)."
            ),
            inputs=[
                io.Conditioning.Input("positive"),
                io.Conditioning.Input("negative"),
                io.Vae.Input("vae"),
                io.Latent.Input(
                    "latent",
                    tooltip="Video-only latent to condition. Must be a 5D video latent (batch, channels, frames, height, width).",
                ),
                io.Image.Input("image"),
                io.Int.Input(
                    "frame_idx",
                    default=0,
                    min=-9999,
                    max=9999,
                    tooltip=FRAME_IDX_TOOLTIP,
                ),
                io.Float.Input(
                    "strength",
                    default=1.0,
                    min=0.0,
                    max=1.0,
                    step=0.01,
                ),
                io.Float.Input(
                    "latent_downscale_factor",
                    default=1.0,
                    min=1.0,
                    max=10.0,
                    step=1.0,
                    tooltip="For IC-LoRA on small grid. 1 means original size, 2 means half size, 3 means third, etc.",
                ),
                io.Combo.Input(
                    "crop",
                    options=["disabled", "center"],
                    default="disabled",
                    tooltip="Crop mode when resizing. 'center' crops to fit, 'disabled' stretches to fit.",
                ),
                io.Boolean.Input(
                    "use_tiled_encode",
                    default=False,
                    tooltip="Enable tiled VAE encoding for large resolutions/long videos to reduce memory usage.",
                ),
                io.Int.Input(
                    "tile_size",
                    default=256,
                    min=64,
                    max=512,
                    step=32,
                    tooltip="Spatial tile size for tiled encoding. Only used when use_tiled_encode is enabled.",
                ),
                io.Int.Input(
                    "tile_overlap",
                    default=64,
                    min=16,
                    max=256,
                    step=16,
                    tooltip="Overlap between tiles for tiled encoding. Only used when use_tiled_encode is enabled.",
                ),
                io.Boolean.Input(
                    "use_streaming",
                    optional=True,
                    default=False,
                    tooltip=(
                        "Off (default): encode the whole clip as one IC-LoRA guide. "
                        "On: re-encode one fresh guide per LTXVTiledFusionSampler "
                        "temporal window and write the plan the sampler reads. "
                        "Set tile_frames on this node and the sampler to the same "
                        "value (97 for LTX)."
                    ),
                ),
                io.Int.Input(
                    "tile_frames",
                    optional=True,
                    default=97,
                    min=0,
                    max=1000,
                    step=1,
                    tooltip=(
                        "Window length in pixel frames, used only when use_streaming "
                        "is on. Snapped to 8n+1. Must match LTXVTiledFusionSampler. "
                        "Ignored when use_streaming is off."
                    ),
                ),
            ],
            outputs=[
                io.Conditioning.Output("positive"),
                io.Conditioning.Output("negative"),
                io.Latent.Output("latent"),
            ],
        )

    @classmethod
    def resolve_frame_idx(
        cls,
        positive,
        frame_idx,
        latent_length,
        image_length,
        guide_latent,
        scale_factors,
    ):
        """Map frame_idx to the position append_keyframe should use.

        A negative value is a RoPE position before the clip, not an index counted back
        from its end, so it bypasses get_latent_index. That is what lets a one-frame
        reference sit at -1: adjacent to the first generated frame without landing in
        its slot, which is the only position a one-frame guide can occupy that makes it
        indistinguishable from the frame being generated. The length assert is
        meaningless there -- the guide is outside the clip by construction.
        """
        if frame_idx < 0:
            return frame_idx
        frame_idx, latent_idx = nodes_lt.LTXVAddGuide.get_latent_index(
            positive, latent_length, image_length, frame_idx, scale_factors
        )
        assert (
            latent_idx + guide_latent.shape[2] <= latent_length
        ), "Conditioning frames exceed the length of the latent sequence."
        return frame_idx

    @classmethod
    def encode(
        cls,
        vae,
        latent_width,
        latent_height,
        images,
        scale_factors,
        latent_downscale_factor,
        crop,
        use_tiled_encode,
        tile_size,
        tile_overlap,
    ):
        time_scale_factor, width_scale_factor, height_scale_factor = scale_factors
        num_frames_to_keep = (
            (images.shape[0] - 1) // time_scale_factor
        ) * time_scale_factor + 1
        images = images[:num_frames_to_keep]
        # Divide target size by latent_downscale_factor
        target_width = int(latent_width * width_scale_factor / latent_downscale_factor)
        target_height = int(
            latent_height * height_scale_factor / latent_downscale_factor
        )
        pixels = comfy.utils.common_upscale(
            images.movedim(-1, 1),
            target_width,
            target_height,
            "bilinear",
            crop=crop,
        ).movedim(1, -1)
        encode_pixels = pixels[:, :, :, :3]
        if use_tiled_encode:
            guide_latent = vae.encode_tiled(
                encode_pixels,
                tile_x=tile_size,
                tile_y=tile_size,
                overlap=tile_overlap,
            )
        else:
            guide_latent = vae.encode(encode_pixels)
        return encode_pixels, guide_latent

    @classmethod
    def encode_windowed_guides(
        cls,
        positive,
        negative,
        vae,
        latent_image,
        noise_mask,
        image,
        latent_length,
        latent_width,
        latent_height,
        scale_factors,
        strength,
        crop,
        tile_frames,
        attention_strength=1.0,
        attention_mask=None,
        frame_idx=0,
        use_tiled_encode=False,
    ):
        """Re-encode each fusion window as its own clip, or return None.

        Mid-clip causal latents cannot be sliced into windows, so each window's
        pixels are encoded as a clip starting at 0. The shared plan is written
        onto the conditioning so LTXVTiledFusionSampler pairs window w with
        guide block w. Returns None when the clip fits one window.
        """
        from .iclora_attention import append_guide_attention_entry
        from .tiled_fusion_plan import attach_temporal_plan, temporal_plan

        t_scale = scale_factors[0]
        tf_lat = max(1, (tile_frames - 1) // t_scale + 1)
        if not (0 < tf_lat < latent_length):
            return None
        if int(frame_idx) != 0:
            raise ValueError(
                "use_streaming encodes each window as a fresh clip at frame 0; "
                f"frame_idx={frame_idx} is not supported. Use the whole-clip path "
                "(use_streaming off) for stills or mid-clip placement."
            )
        if use_tiled_encode:
            raise ValueError(
                "use_streaming cannot use tiled VAE encode; tiled guide encodes "
                "imprint a grid into the conditioning. Leave use_tiled_encode off."
            )
        keep = ((image.shape[0] - 1) // t_scale) * t_scale + 1
        image = image[:keep]
        ts = temporal_plan(latent_length, tf_lat)
        span_px = (tf_lat - 1) * t_scale + 1
        need = ts[-1] * t_scale + span_px
        if image.shape[0] < need:
            raise ValueError(
                f"windowed guide needs {need} frames for {latent_length} "
                f"content latents, found {image.shape[0]}"
            )
        for w, t0 in enumerate(ts):
            span = image[t0 * t_scale : t0 * t_scale + span_px]
            _, guide_latent = cls.encode(
                vae,
                latent_width,
                latent_height,
                span,
                scale_factors,
                1.0,
                crop,
                False,
                256,
                64,
            )
            shape = list(guide_latent.shape[2:])
            toks = shape[0] * shape[1] * shape[2]
            positive, negative, latent_image, noise_mask = (
                nodes_lt.LTXVAddGuide.append_keyframe(
                    positive,
                    negative,
                    0,
                    latent_image,
                    noise_mask,
                    guide_latent,
                    strength,
                    scale_factors,
                    guide_mask=None,
                    latent_downscale_factor=1.0,
                    causal_fix=True,
                )
            )
            # Pixel mask is (1, 1, F, H, W) after normalize_mask. F==1 broadcasts;
            # otherwise it must match this window's encoded span, not the full clip.
            window_mask = attention_mask
            if window_mask is not None and window_mask.shape[2] != 1:
                start = t0 * t_scale
                window_mask = window_mask[:, :, start : start + span_px]
            positive = append_guide_attention_entry(
                positive,
                toks,
                shape,
                attention_strength=attention_strength,
                attention_mask=window_mask,
            )
            negative = append_guide_attention_entry(
                negative,
                toks,
                shape,
                attention_strength=attention_strength,
                attention_mask=window_mask,
            )
            logging.info(
                "window %d/%d: guide frames %d-%d -> %s",
                w + 1,
                len(ts),
                t0 * t_scale,
                t0 * t_scale + span_px - 1,
                tuple(shape),
            )
        plan = {
            "tile_frames": int(tile_frames),
            "tf_lat": int(tf_lat),
            "n_content": int(latent_length),
            "ts": [int(v) for v in ts],
        }
        positive = attach_temporal_plan(positive, plan)
        negative = attach_temporal_plan(negative, plan)
        logging.info("%d per-window guides (t0 = %s)", len(ts), plan["ts"])
        return io.NodeOutput(
            positive, negative, {"samples": latent_image, "noise_mask": noise_mask}
        )

    @classmethod
    def maybe_encode_windowed(
        cls,
        positive,
        negative,
        vae,
        latent_image,
        noise_mask,
        image,
        latent_length,
        latent_width,
        latent_height,
        scale_factors,
        strength,
        crop,
        use_streaming,
        tile_frames,
        latent_downscale_factor=1.0,
        attention_strength=1.0,
        attention_mask=None,
        frame_idx=0,
        use_tiled_encode=False,
    ):
        """Run the windowed encode only when use_streaming is on. Else None."""
        if not use_streaming:
            return None
        if latent_downscale_factor != 1:
            raise ValueError(
                "use_streaming requires latent_downscale_factor=1; "
                "windowed fusion guides cannot be dilated."
            )
        tf = snap_frames(int(tile_frames) if tile_frames and tile_frames > 0 else 97)
        return cls.encode_windowed_guides(
            positive,
            negative,
            vae,
            latent_image,
            noise_mask,
            image,
            latent_length,
            latent_width,
            latent_height,
            scale_factors,
            strength,
            crop,
            tf,
            attention_strength=attention_strength,
            attention_mask=attention_mask,
            frame_idx=frame_idx,
            use_tiled_encode=use_tiled_encode,
        )

    @classmethod
    def execute(
        cls,
        positive,
        negative,
        vae,
        latent,
        image,
        frame_idx,
        strength,
        latent_downscale_factor,
        crop,
        use_tiled_encode,
        tile_size,
        tile_overlap,
        use_streaming=False,
        tile_frames=0,
    ) -> io.NodeOutput:
        scale_factors = vae.downscale_index_formula
        latent_image = latent["samples"]
        noise_mask = nodes_lt.get_noise_mask(latent)

        _, _, latent_length, latent_height, latent_width = latent_image.shape

        windowed = cls.maybe_encode_windowed(
            positive,
            negative,
            vae,
            latent_image,
            noise_mask,
            image,
            latent_length,
            latent_width,
            latent_height,
            scale_factors,
            strength,
            crop,
            use_streaming,
            tile_frames,
            latent_downscale_factor=latent_downscale_factor,
            frame_idx=frame_idx,
            use_tiled_encode=use_tiled_encode,
        )
        if windowed is not None:
            return windowed

        time_scale_factor = scale_factors[0]
        num_frames_to_keep = (
            (image.shape[0] - 1) // time_scale_factor
        ) * time_scale_factor + 1
        causal_fix = frame_idx == 0 or num_frames_to_keep == 1
        if not causal_fix:
            image = torch.cat([image[:1], image], dim=0)

        image, guide_latent = cls.encode(
            vae,
            latent_width,
            latent_height,
            image,
            scale_factors,
            latent_downscale_factor,
            crop,
            use_tiled_encode,
            tile_size,
            tile_overlap,
        )

        if not causal_fix:
            guide_latent = guide_latent[:, :, 1:, :, :]
            image = image[1:]

        # Record original (pre-dilation) guide latent shape for spatial mask downsampling
        guide_orig_shape = list(guide_latent.shape[2:])  # [F, H_small, W_small]

        guide_mask = None

        # Dilate the latent if latent_downscale_factor > 1
        if latent_downscale_factor > 1:
            if (
                latent_width % latent_downscale_factor != 0
                or latent_height % latent_downscale_factor != 0
            ):
                raise ValueError(
                    f"Latent spatial size {latent_width}x{latent_height} must be divisible by latent_downscale_factor {latent_downscale_factor}"
                )

            dilated = LTXVDilateLatent().dilate_latent(
                {"samples": guide_latent},
                horizontal_scale=int(latent_downscale_factor),
                vertical_scale=int(latent_downscale_factor),
            )[0]
            guide_mask = dilated["noise_mask"]
            guide_latent = dilated["samples"]

        # Pre-filter token count = product of (potentially dilated) spatial dims
        iclora_tokens_added = (
            guide_latent.shape[2] * guide_latent.shape[3] * guide_latent.shape[4]
        )

        frame_idx = cls.resolve_frame_idx(
            positive, frame_idx, latent_length, len(image), guide_latent, scale_factors
        )

        positive, negative, latent_image, noise_mask = (
            nodes_lt.LTXVAddGuide.append_keyframe(
                positive,
                negative,
                frame_idx,
                latent_image,
                noise_mask,
                guide_latent,
                strength,
                scale_factors,
                guide_mask=guide_mask,
                latent_downscale_factor=latent_downscale_factor,
                causal_fix=causal_fix,
            )
        )

        # Track this guide in guide_attention_entries for per-reference attention control.
        from .iclora_attention import append_guide_attention_entry

        positive = append_guide_attention_entry(
            positive, iclora_tokens_added, guide_orig_shape
        )
        negative = append_guide_attention_entry(
            negative, iclora_tokens_added, guide_orig_shape
        )

        return io.NodeOutput(
            positive, negative, {"samples": latent_image, "noise_mask": noise_mask}
        )


@comfy_node(name="LTXAddVideoICLoRAGuideAdvanced")
class LTXAddVideoICLoRAGuideAdvanced(LTXAddVideoICLoRAGuide):
    """Extended IC-LoRA guide node with per-guide attention strength control.

    Same as LTXAddVideoICLoRAGuide, but adds attention_strength and
    attention_mask inputs to control how strongly this guide's conditioning
    influences generation via self-attention.
    """

    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="LTXAddVideoICLoRAGuideAdvanced",
            display_name=NODES_DISPLAY_NAME_PREFIX
            + " Add Video IC-LoRA Guide Advanced",
            category="Lightricks/IC-LoRA",
            description=(
                "Adds IC-LoRA guide conditioning with per-guide attention strength control. "
                "Same as LTXAddVideoICLoRAGuide, but allows controlling how strongly this "
                "guide influences generation via self-attention, optionally with a spatial mask. "
                "use_streaming + matching tile_frames re-encodes one guide per Tiled Fusion window."
            ),
            inputs=[
                io.Conditioning.Input("positive"),
                io.Conditioning.Input("negative"),
                io.Vae.Input("vae"),
                io.Latent.Input(
                    "latent",
                    tooltip="Video-only latent to condition. Must be a 5D video latent.",
                ),
                io.Image.Input("image"),
                io.Int.Input(
                    "frame_idx",
                    default=0,
                    min=-9999,
                    max=9999,
                    tooltip=FRAME_IDX_TOOLTIP,
                ),
                io.Float.Input("strength", default=1.0, min=0.0, max=1.0, step=0.01),
                io.Float.Input(
                    "latent_downscale_factor",
                    default=1.0,
                    min=1.0,
                    max=10.0,
                    step=1.0,
                    tooltip="For IC-LoRA on small grid. 1 = original size, 2 = half, etc.",
                ),
                io.Combo.Input(
                    "crop",
                    options=["disabled", "center"],
                    default="disabled",
                ),
                io.Boolean.Input("use_tiled_encode", default=False),
                io.Int.Input("tile_size", default=256, min=64, max=512, step=32),
                io.Int.Input("tile_overlap", default=64, min=16, max=256, step=16),
                io.Float.Input(
                    "attention_strength",
                    default=1.0,
                    min=0.0,
                    max=1.0,
                    step=0.01,
                    tooltip=(
                        "Controls how strongly this guide influences generation via "
                        "self-attention. 1.0 = full conditioning (default), 0.0 = ignore. "
                        "When an attention_mask is also provided, this multiplies the mask values."
                    ),
                ),
                io.Mask.Input(
                    "attention_mask",
                    optional=True,
                    tooltip=(
                        "Optional pixel-space spatial mask. Shape (F, H, W) or (H, W). "
                        "Values in [0, 1]. Controls per-region conditioning influence. "
                        "Multiplied by attention_strength."
                    ),
                ),
                io.Boolean.Input(
                    "use_streaming",
                    optional=True,
                    default=False,
                    tooltip=(
                        "Off (default): whole-clip guide. On: one fresh encode per "
                        "fusion temporal window. Must match the sampler's tile_frames."
                    ),
                ),
                io.Int.Input(
                    "tile_frames",
                    optional=True,
                    default=97,
                    min=0,
                    max=1000,
                    step=1,
                    tooltip=(
                        "Used only when use_streaming is on. Snapped to 8n+1. "
                        "97 for LTX. Must match the sampler."
                    ),
                ),
            ],
            outputs=[
                io.Conditioning.Output("positive"),
                io.Conditioning.Output("negative"),
                io.Latent.Output("latent"),
            ],
        )

    @classmethod
    def execute(
        cls,
        positive,
        negative,
        vae,
        latent,
        image,
        frame_idx,
        strength,
        latent_downscale_factor,
        crop,
        use_tiled_encode,
        tile_size,
        tile_overlap,
        attention_strength=1.0,
        attention_mask=None,
        use_streaming=False,
        tile_frames=0,
    ) -> io.NodeOutput:
        from .iclora_attention import normalize_mask
        from .latents import LTXVDilateLatent

        scale_factors = vae.downscale_index_formula
        latent_image = latent["samples"]
        noise_mask = nodes_lt.get_noise_mask(latent)

        _, _, latent_length, latent_height, latent_width = latent_image.shape

        windowed = cls.maybe_encode_windowed(
            positive,
            negative,
            vae,
            latent_image,
            noise_mask,
            image,
            latent_length,
            latent_width,
            latent_height,
            scale_factors,
            strength,
            crop,
            use_streaming,
            tile_frames,
            latent_downscale_factor=latent_downscale_factor,
            attention_strength=attention_strength,
            attention_mask=normalize_mask(attention_mask),
            frame_idx=frame_idx,
            use_tiled_encode=use_tiled_encode,
        )
        if windowed is not None:
            return windowed

        time_scale_factor = scale_factors[0]
        num_frames_to_keep = (
            (image.shape[0] - 1) // time_scale_factor
        ) * time_scale_factor + 1
        causal_fix = frame_idx == 0 or num_frames_to_keep == 1
        if not causal_fix:
            image = torch.cat([image[:1], image], dim=0)

        image, guide_latent = cls.encode(
            vae,
            latent_width,
            latent_height,
            image,
            scale_factors,
            latent_downscale_factor,
            crop,
            use_tiled_encode,
            tile_size,
            tile_overlap,
        )

        if not causal_fix:
            guide_latent = guide_latent[:, :, 1:, :, :]
            image = image[1:]

        guide_orig_shape = list(guide_latent.shape[2:])
        guide_mask = None

        if latent_downscale_factor > 1:
            if (
                latent_width % latent_downscale_factor != 0
                or latent_height % latent_downscale_factor != 0
            ):
                raise ValueError(
                    f"Latent spatial size {latent_width}x{latent_height} must be "
                    f"divisible by latent_downscale_factor {latent_downscale_factor}"
                )
            dilated = LTXVDilateLatent().dilate_latent(
                {"samples": guide_latent},
                horizontal_scale=int(latent_downscale_factor),
                vertical_scale=int(latent_downscale_factor),
            )[0]
            guide_mask = dilated["noise_mask"]
            guide_latent = dilated["samples"]

        iclora_tokens_added = (
            guide_latent.shape[2] * guide_latent.shape[3] * guide_latent.shape[4]
        )

        frame_idx = cls.resolve_frame_idx(
            positive, frame_idx, latent_length, len(image), guide_latent, scale_factors
        )

        positive, negative, latent_image, noise_mask = (
            nodes_lt.LTXVAddGuide.append_keyframe(
                positive,
                negative,
                frame_idx,
                latent_image,
                noise_mask,
                guide_latent,
                strength,
                scale_factors,
                guide_mask=guide_mask,
                latent_downscale_factor=latent_downscale_factor,
                causal_fix=causal_fix,
            )
        )

        from .iclora_attention import append_guide_attention_entry

        norm_mask = normalize_mask(attention_mask)
        positive = append_guide_attention_entry(
            positive,
            iclora_tokens_added,
            guide_orig_shape,
            attention_strength=attention_strength,
            attention_mask=norm_mask,
        )
        negative = append_guide_attention_entry(
            negative,
            iclora_tokens_added,
            guide_orig_shape,
            attention_strength=attention_strength,
            attention_mask=norm_mask,
        )

        return io.NodeOutput(
            positive, negative, {"samples": latent_image, "noise_mask": noise_mask}
        )


@comfy_node(name="LTXICLoRALoaderModelOnly")
class LTXICLoRALoaderModelOnly(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="LTXICLoRALoaderModelOnly",
            display_name=NODES_DISPLAY_NAME_PREFIX + " IC-LoRA Loader Model Only",
            category="Lightricks/IC-LoRA",
            description="Loads a LoRA model and extracts the latent_downscale_factor from the safetensors metadata.",
            inputs=[
                io.Model.Input("model"),
                io.Combo.Input(
                    "lora_name",
                    options=folder_paths.get_filename_list("loras"),
                ),
                io.Float.Input(
                    "strength_model",
                    default=1.0,
                    min=-100.0,
                    max=100.0,
                    step=0.01,
                ),
            ],
            outputs=[
                io.Model.Output("model"),
                io.Float.Output("latent_downscale_factor"),
            ],
        )

    @classmethod
    def execute(cls, model, lora_name, strength_model) -> io.NodeOutput:
        lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)

        # Load lora and extract metadata
        lora, metadata = comfy.utils.load_torch_file(
            lora_path, safe_load=True, return_metadata=True
        )

        # Extract latent_downscale_factor from metadata
        try:
            latent_downscale_factor = float(metadata["reference_downscale_factor"])
        except (KeyError, ValueError, TypeError):
            latent_downscale_factor = 1.0
            logging.warning(
                "Failed to extract reference_downscale_factor from metadata for %s, using 1.0",
                lora_path,
            )

        if strength_model == 0:
            return io.NodeOutput(model, latent_downscale_factor)

        model_lora, _ = comfy.sd.load_lora_for_models(
            model, None, lora, strength_model, 0
        )
        return io.NodeOutput(model_lora, latent_downscale_factor)


def _patchify_audio_latent(latent):
    """Reshape audio latent (b, c, t, f) → ref_audio token dict (b, t, c*f)."""
    b, c, t, f = latent.shape
    ref_tokens = latent.permute(0, 2, 1, 3).reshape(b, t, c * f)
    return {"tokens": ref_tokens}


@comfy_node(name="LTXVSetAudioRefTokens")
class LTXVSetAudioRefTokens(io.ComfyNode):
    """Provide speaker identity context for audio generation.

    Patchifies the audio latent and attaches it as reference tokens on both
    positive and negative conditioning. The model prepends these tokens with
    negative temporal positions so they serve as identity context without
    being part of the generated output.

    Also outputs a frozen copy of the audio latent (noise_mask=0) for
    direct use in stage 2 without re-encoding.
    """

    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="LTXVSetAudioRefTokens",
            display_name=NODES_DISPLAY_NAME_PREFIX + " Set Audio Ref Tokens",
            category="Lightricks/IC-LoRA",
            description=(
                "Provides speaker identity context for audio generation by attaching "
                "reference audio tokens to the conditioning. The tokens are prepended "
                "with negative temporal positions so the model treats them as context "
                "rather than generation targets."
            ),
            inputs=[
                io.Conditioning.Input(
                    "positive",
                    tooltip="Positive conditioning to attach the reference audio tokens to.",
                ),
                io.Conditioning.Input(
                    "negative",
                    tooltip="Negative conditioning to attach the reference audio tokens to.",
                ),
                io.Latent.Input(
                    "audio_latent",
                    tooltip="Encoded audio latent from LTXV Audio VAE Encode.",
                ),
            ],
            outputs=[
                io.Conditioning.Output(
                    display_name="positive",
                    tooltip="Positive conditioning with reference audio tokens attached.",
                ),
                io.Conditioning.Output(
                    display_name="negative",
                    tooltip="Negative conditioning with reference audio tokens attached.",
                ),
                io.Latent.Output(
                    display_name="frozen_audio",
                    tooltip="Audio latent with noise_mask=0, fully frozen during denoising.",
                ),
            ],
        )

    @classmethod
    def execute(cls, positive, negative, audio_latent) -> io.NodeOutput:
        latent = audio_latent["samples"]
        ref_audio = _patchify_audio_latent(latent)
        positive = node_helpers.conditioning_set_values(
            positive, {"ref_audio": ref_audio}
        )
        negative = node_helpers.conditioning_set_values(
            negative, {"ref_audio": ref_audio}
        )

        frozen = audio_latent.copy()
        b, c, t, f = latent.shape
        frozen["noise_mask"] = torch.zeros((b, 1, t, 1), dtype=torch.float32)

        return io.NodeOutput(positive, negative, frozen)