Spaces:

diffusers-modular
/

flux3-image-klein-bbox

Running on Zero

flux3-image-klein-bbox / diffusers /modular_pipelines /ltx2 /before_denoise.py
FLUX.3 layout-control canvas (custom HTML) — initial
3db1cbc verified 22 days ago
79 kB
# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
import inspect

import numpy as np
import torch

from ...models import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, LTX2VideoTransformer3DModel

# NOTE (modular.md gotcha #1): `LTX2ReferenceCondition` is a plain dataclass under `diffusers.pipelines.ltx2.*`, and
# modular blocks must not import from `diffusers.pipelines.*`. It belongs in the same neutral-module relocation as
# the other shared LTX-2 data/utilities enumerated in `encoders.py`. Imported from the pipelines path here only so
# the draft is runnable.
from ...pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
from ...schedulers import FlowMatchEulerDiscreteScheduler
from ...utils import logging
from ...utils.torch_utils import randn_tensor
from ..modular_pipeline import ModularPipelineBlocks, PipelineState
from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam


logger = logging.get_logger(__name__)


def calculate_shift(
    image_seq_len,
    base_seq_len: int = 256,
    max_seq_len: int = 4096,
    base_shift: float = 0.5,
    max_shift: float = 1.15,
):
    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
    b = base_shift - m * base_seq_len
    mu = image_seq_len * m + b
    return mu


# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
def retrieve_timesteps(
    scheduler,
    num_inference_steps: int | None = None,
    device: str | torch.device | None = None,
    timesteps: list[int] | None = None,
    sigmas: list[float] | None = None,
    **kwargs,
):
    r"""
    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
    Args:
        scheduler (`SchedulerMixin`):
            The scheduler to get timesteps from.
        num_inference_steps (`int`):
            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
            must be `None`.
        device (`str` or `torch.device`, *optional*):
            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
        timesteps (`list[int]`, *optional*):
            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
            `num_inference_steps` and `sigmas` must be `None`.
        sigmas (`list[float]`, *optional*):
            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
            `num_inference_steps` and `timesteps` must be `None`.
    Returns:
        `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
        second element is the number of inference steps.
    """
    if timesteps is not None and sigmas is not None:
        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
    if timesteps is not None:
        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
        if not accepts_timesteps:
            raise ValueError(
                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
                f" timestep schedules. Please check whether you are using the correct scheduler."
            )
        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
        timesteps = scheduler.timesteps
        num_inference_steps = len(timesteps)
    elif sigmas is not None:
        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
        if not accept_sigmas:
            raise ValueError(
                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
                f" sigmas schedules. Please check whether you are using the correct scheduler."
            )
        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
        timesteps = scheduler.timesteps
        num_inference_steps = len(timesteps)
    else:
        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
        timesteps = scheduler.timesteps
    return timesteps, num_inference_steps


# The pack/normalize helpers below mirror the static methods on `diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline`.
# Redefined here (not imported) because modular blocks must not import from `diffusers.pipelines.*` (gotcha #1).
def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor:
    batch_size, num_channels, num_frames, height, width = latents.shape
    post_patch_num_frames = num_frames // patch_size_t
    post_patch_height = height // patch_size
    post_patch_width = width // patch_size
    latents = latents.reshape(
        batch_size,
        -1,
        post_patch_num_frames,
        patch_size_t,
        post_patch_height,
        patch_size,
        post_patch_width,
        patch_size,
    )
    latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3)
    return latents


def _unpack_latents(
    latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1
) -> torch.Tensor:
    batch_size = latents.size(0)
    latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size)
    latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
    return latents


def _pack_audio_latents(
    latents: torch.Tensor, patch_size: int | None = None, patch_size_t: int | None = None
) -> torch.Tensor:
    # Audio latents of shape [B, C, L, M] (L = latent audio length, M = mel bins). With no patch sizes this packs to
    # [B, L, C * M] (implicit mel patch_size of M, temporal patch_size of 1).
    if patch_size is not None and patch_size_t is not None:
        batch_size, num_channels, latent_length, latent_mel_bins = latents.shape
        post_patch_latent_length = latent_length / patch_size_t
        post_patch_mel_bins = latent_mel_bins / patch_size
        latents = latents.reshape(
            batch_size, -1, post_patch_latent_length, patch_size_t, post_patch_mel_bins, patch_size
        )
        latents = latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2)
    else:
        latents = latents.transpose(1, 2).flatten(2, 3)  # [B, C, L, M] -> [B, L, C * M]
    return latents


def _normalize_latents(
    latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
) -> torch.Tensor:
    latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
    latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
    latents = (latents - latents_mean) * scaling_factor / latents_std
    return latents


def _normalize_audio_latents(
    latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor
) -> torch.Tensor:
    latents_mean = latents_mean.to(latents.device, latents.dtype)
    latents_std = latents_std.to(latents.device, latents.dtype)
    return (latents - latents_mean) / latents_std


def _create_noised_state(
    latents: torch.Tensor, noise_scale: float | torch.Tensor, generator: torch.Generator | None = None
) -> torch.Tensor:
    noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype)
    return noise_scale * noise + (1 - noise_scale) * latents


def _prepare_keyframe_coords(
    keyframe_latent_num_frames: int,
    keyframe_latent_height: int,
    keyframe_latent_width: int,
    pixel_frame_idx: int,
    num_pixel_frames: int,
    fps: float,
    patch_size: int,
    patch_size_t: int,
    scale_factors: tuple[int, int, int],
    device: torch.device,
) -> torch.Tensor:
    """
    Positional coordinates for a keyframe condition appended as extra tokens, mirroring
    `VideoConditionByKeyframeIndex.apply_to` in the reference implementation:
      - latent coords scaled to pixel space *without* the causal fix (non-zero-index keyframes don't need the
        first-frame causal adjustment),
      - temporal axis offset by `pixel_frame_idx` (the pixel-space index the keyframe appears at),
      - for single-pixel-frame keyframes the per-patch temporal extent is clamped to `[idx, idx + 1)` so the keyframe
        occupies one pixel timestep rather than the VAE-scaled range,
      - temporal coords divided by `fps` to give seconds.
    """
    grid_f = torch.arange(
        start=0, end=keyframe_latent_num_frames, step=patch_size_t, dtype=torch.float32, device=device
    )
    grid_h = torch.arange(start=0, end=keyframe_latent_height, step=patch_size, dtype=torch.float32, device=device)
    grid_w = torch.arange(start=0, end=keyframe_latent_width, step=patch_size, dtype=torch.float32, device=device)
    grid = torch.stack(torch.meshgrid(grid_f, grid_h, grid_w, indexing="ij"), dim=0)

    patch_size_delta = torch.tensor((patch_size_t, patch_size, patch_size), dtype=grid.dtype, device=device)
    patch_ends = grid + patch_size_delta.view(3, 1, 1, 1)

    latent_coords = torch.stack([grid, patch_ends], dim=-1)  # [3, N_F, N_H, N_W, 2]
    latent_coords = latent_coords.flatten(1, 3).unsqueeze(0)  # [1, 3, num_patches, 2]

    scale_tensor = torch.tensor(scale_factors, device=device, dtype=latent_coords.dtype)
    broadcast_shape = [1] * latent_coords.ndim
    broadcast_shape[1] = -1
    pixel_coords = latent_coords * scale_tensor.view(*broadcast_shape)

    pixel_coords[:, 0, :, :] = pixel_coords[:, 0, :, :] + pixel_frame_idx
    if num_pixel_frames == 1:
        pixel_coords[:, 0, :, 1:] = pixel_coords[:, 0, :, :1] + 1
    pixel_coords[:, 0, :, :] = pixel_coords[:, 0, :, :] / fps

    return pixel_coords


class LTX2TextInputStep(ModularPipelineBlocks):
    model_name = "ltx2"

    @property
    def description(self) -> str:
        return (
            "Input processing step that expands the connector text conditioning (cond and uncond) by "
            "`num_videos_per_prompt`, so it matches the `batch_size * num_videos_per_prompt` batch of the video and "
            "audio latents. Runs at the head of the denoise stage, which keeps the text-conditioning stage's outputs "
            "reusable across denoise runs with different `num_videos_per_prompt`."
        )

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"),
            InputParam(
                "connector_prompt_embeds",
                type_hint=torch.Tensor,
                required=True,
                description="Video-branch text conditioning (cond).",
            ),
            InputParam(
                "connector_audio_prompt_embeds",
                type_hint=torch.Tensor,
                required=True,
                description="Audio-branch text conditioning (cond).",
            ),
            InputParam(
                "connector_attention_mask",
                type_hint=torch.Tensor,
                required=True,
                description="Binary text attention mask (cond).",
            ),
            InputParam(
                "negative_connector_prompt_embeds",
                type_hint=torch.Tensor,
                required=True,
                description="Video-branch text conditioning (uncond).",
            ),
            InputParam(
                "negative_connector_audio_prompt_embeds",
                type_hint=torch.Tensor,
                required=True,
                description="Audio-branch text conditioning (uncond).",
            ),
            InputParam(
                "negative_connector_attention_mask",
                type_hint=torch.Tensor,
                required=True,
                description="Binary text attention mask (uncond).",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                "connector_prompt_embeds",
                type_hint=torch.Tensor,
                description="Video-branch text conditioning (cond), expanded per prompt.",
            ),
            OutputParam(
                "connector_audio_prompt_embeds",
                type_hint=torch.Tensor,
                description="Audio-branch text conditioning (cond), expanded per prompt.",
            ),
            OutputParam(
                "connector_attention_mask",
                type_hint=torch.Tensor,
                description="Binary text attention mask (cond), expanded per prompt.",
            ),
            OutputParam(
                "negative_connector_prompt_embeds",
                type_hint=torch.Tensor,
                description="Video-branch text conditioning (uncond), expanded per prompt.",
            ),
            OutputParam(
                "negative_connector_audio_prompt_embeds",
                type_hint=torch.Tensor,
                description="Audio-branch text conditioning (uncond), expanded per prompt.",
            ),
            OutputParam(
                "negative_connector_attention_mask",
                type_hint=torch.Tensor,
                description="Binary text attention mask (uncond), expanded per prompt.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        # `repeat_interleave` keeps each prompt's copies contiguous, matching how the latents are laid out
        # (`batch_size * num_videos_per_prompt`, prompt-major) and how `image_latents` are expanded downstream.
        num_videos = block_state.num_videos_per_prompt
        for name in self.intermediate_output_names:
            setattr(block_state, name, getattr(block_state, name).repeat_interleave(num_videos, dim=0))

        self.set_block_state(state, block_state)
        return components, state


class LTX2SetTimestepsStep(ModularPipelineBlocks):
    model_name = "ltx2"

    @property
    def description(self) -> str:
        return (
            "Sets the flow-match timesteps for the video `scheduler` and produces a deep-copied `audio_scheduler` "
            "(with the same schedule) so the audio latents are denoised on an independent scheduler state."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("num_inference_steps", default=30),
            InputParam.template("timesteps"),
            InputParam.template("sigmas"),
            InputParam.template("height", default=512),
            InputParam.template("width", default=704),
            InputParam(
                "num_frames",
                type_hint=int,
                default=None,
                description=(
                    "The number of frames in the generated video. Omit to auto-predict via the `duration_head` "
                    "(see `LTX2AutoDurationStep`)."
                ),
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam("timesteps", type_hint=torch.Tensor),
            OutputParam("num_inference_steps", type_hint=int),
            OutputParam(
                "audio_scheduler",
                description="Independent deep copy of `scheduler` used to update the audio latents in the loop.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        device = components._execution_device

        num_inference_steps = block_state.num_inference_steps
        sigmas = block_state.sigmas
        if sigmas is None:
            sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)

        # Resolution-aware (dynamic) shift: `mu` is computed from the actual packed video sequence length, matching
        # `latents.shape[1]` in the standard pipeline. The packed length divides the latent grid by the transformer
        # patch sizes (see `_pack_latents`). `num_frames` must already be a concrete int here -- when auto-duration is
        # used, the duration step resolves it upstream, before this block runs.
        latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1
        latent_height = block_state.height // components.vae_spatial_compression_ratio
        latent_width = block_state.width // components.vae_spatial_compression_ratio
        video_seq_len = (
            (latent_num_frames // components.transformer_temporal_patch_size)
            * (latent_height // components.transformer_spatial_patch_size)
            * (latent_width // components.transformer_spatial_patch_size)
        )
        mu = calculate_shift(
            video_seq_len,
            components.scheduler.config.get("base_image_seq_len", 1024),
            components.scheduler.config.get("max_image_seq_len", 4096),
            components.scheduler.config.get("base_shift", 0.95),
            components.scheduler.config.get("max_shift", 2.05),
        )

        block_state.audio_scheduler = copy.deepcopy(components.scheduler)
        retrieve_timesteps(
            block_state.audio_scheduler, num_inference_steps, device, block_state.timesteps, sigmas=sigmas, mu=mu
        )
        block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
            components.scheduler, num_inference_steps, device, block_state.timesteps, sigmas=sigmas, mu=mu
        )

        # Set begin index to skip the nonzero().item() call in scheduler init, which triggers a GPU sync.
        components.scheduler.set_begin_index(0)
        block_state.audio_scheduler.set_begin_index(0)

        self.set_block_state(state, block_state)
        return components, state


class LTX2PrepareLatentsStep(ModularPipelineBlocks):
    model_name = "ltx2"

    @property
    def description(self) -> str:
        return (
            "Prepares the packed video noise latents for text-to-video generation. `noise_scale` is declared with a "
            "`None` default and resolved to 0.0 here rather than declared as 0.0: a blockset keeps the first "
            "non-`None` default across its blocks, so a literal 0.0 would shadow the condition workflow's "
            "`None -> sigmas[0] or 1.0` resolution wherever the two share a blockset (`LTX2AutoBlocks`). The "
            "resolved value is written back to state for `LTX2PrepareAudioLatentsStep`."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [ComponentSpec("transformer", LTX2VideoTransformer3DModel)]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("height", default=512),
            InputParam.template("width", default=704),
            InputParam(
                "num_frames",
                type_hint=int,
                default=None,
                description=(
                    "The number of frames in the generated video. Omit to auto-predict via the `duration_head` "
                    "(see `LTX2AutoDurationStep`)."
                ),
            ),
            InputParam.template("latents"),
            InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"),
            InputParam(
                "noise_scale",
                type_hint=float,
                default=None,
                description=(
                    "Interpolation factor between random noise and any provided latents. `None` (default) resolves "
                    "to 0.0, which keeps the provided latents."
                ),
            ),
            InputParam.template("generator"),
            InputParam(
                "batch_size",
                type_hint=int,
                required=True,
                description="The number of prompts being denoised, used to expand conditioning per prompt.",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam("latents", type_hint=torch.Tensor, description="Packed noisy video latents."),
            OutputParam(
                "noise_scale",
                type_hint=float,
                description="The resolved interpolation factor, forwarded to the audio latents step.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        device = components._execution_device

        batch_size = block_state.batch_size * block_state.num_videos_per_prompt
        num_channels_latents = components.transformer.config.in_channels
        spatial_patch = components.transformer_spatial_patch_size
        temporal_patch = components.transformer_temporal_patch_size

        if block_state.noise_scale is None:
            block_state.noise_scale = 0.0

        if block_state.latents is not None:
            latents = block_state.latents
            if latents.ndim == 5:
                latents = _normalize_latents(
                    latents,
                    components.latents_mean,
                    components.latents_std,
                    components.vae_scaling_factor,
                )
                latents = _pack_latents(latents, spatial_patch, temporal_patch)
            latents = _create_noised_state(latents, block_state.noise_scale, block_state.generator)
            block_state.latents = latents.to(device=device, dtype=torch.float32)
        else:
            latent_height = block_state.height // components.vae_spatial_compression_ratio
            latent_width = block_state.width // components.vae_spatial_compression_ratio
            latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1

            shape = (batch_size, num_channels_latents, latent_num_frames, latent_height, latent_width)
            latents = randn_tensor(shape, generator=block_state.generator, device=device, dtype=torch.float32)
            block_state.latents = _pack_latents(latents, spatial_patch, temporal_patch)

        self.set_block_state(state, block_state)
        return components, state


class LTX2Image2VideoPrepareLatentsStep(ModularPipelineBlocks):
    model_name = "ltx2"

    @property
    def description(self) -> str:
        return (
            "Prepares image-to-video latents: blends the pre-encoded `image_latents` (kept clean on the first latent "
            "frame) with the noise `latents` via a conditioning mask. Expects pure-noise `latents` from "
            "`LTX2PrepareLatentsStep`."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [ComponentSpec("transformer", LTX2VideoTransformer3DModel)]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                "image_latents",
                type_hint=torch.Tensor,
                required=True,
                description="VAE-encoded reference-image latents used for image-to-video conditioning.",
            ),
            InputParam.template("latents", required=True),
            InputParam.template("height", default=512),
            InputParam.template("width", default=704),
            InputParam(
                "num_frames",
                type_hint=int,
                default=None,
                description=(
                    "The number of frames in the generated video. Omit to auto-predict via the `duration_head` "
                    "(see `LTX2AutoDurationStep`)."
                ),
            ),
            InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"),
            InputParam(
                "batch_size",
                type_hint=int,
                required=True,
                description="The number of prompts being denoised, used to expand conditioning per prompt.",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                "latents", type_hint=torch.Tensor, description="Packed noisy latents with image conditioning."
            ),
            OutputParam(
                "conditioning_mask",
                type_hint=torch.Tensor,
                description="Packed per-token mask marking the clean (image-conditioned) first latent frame.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        device = components._execution_device

        batch_size = block_state.batch_size * block_state.num_videos_per_prompt
        spatial_patch = components.transformer_spatial_patch_size
        temporal_patch = components.transformer_temporal_patch_size

        latent_height = block_state.height // components.vae_spatial_compression_ratio
        latent_width = block_state.width // components.vae_spatial_compression_ratio
        latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1

        init_latents = block_state.image_latents.to(device=device, dtype=torch.float32)
        if init_latents.shape[0] < batch_size:
            init_latents = init_latents.repeat_interleave(batch_size // init_latents.shape[0], dim=0)
        init_latents = init_latents.repeat(1, 1, latent_num_frames, 1, 1)

        conditioning_mask = torch.zeros(
            batch_size, 1, latent_num_frames, latent_height, latent_width, device=device, dtype=torch.float32
        )
        conditioning_mask[:, :, 0] = 1.0

        noise = _unpack_latents(
            block_state.latents, latent_num_frames, latent_height, latent_width, spatial_patch, temporal_patch
        )
        latents = init_latents * conditioning_mask + noise * (1 - conditioning_mask)

        block_state.conditioning_mask = _pack_latents(conditioning_mask, spatial_patch, temporal_patch).squeeze(-1)
        block_state.latents = _pack_latents(latents, spatial_patch, temporal_patch)

        self.set_block_state(state, block_state)
        return components, state


class LTX2PrepareAudioLatentsStep(ModularPipelineBlocks):
    model_name = "ltx2"

    @property
    def description(self) -> str:
        return (
            "Prepares the packed audio noise latents and derives the audio latent frame count. `noise_scale` is "
            "declared with a `None` default for the reason given on `LTX2PrepareLatentsStep`, which runs first and "
            "writes the resolved value back to state."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [ComponentSpec("audio_vae", AutoencoderKLLTX2Audio)]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                "num_frames",
                type_hint=int,
                default=None,
                description=(
                    "The number of frames in the generated video. Omit to auto-predict via the `duration_head` "
                    "(see `LTX2AutoDurationStep`)."
                ),
            ),
            InputParam(
                "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video."
            ),
            InputParam(
                "audio_latents",
                type_hint=torch.Tensor,
                default=None,
                description="Optional pre-encoded audio latents; random noise is used when not provided.",
            ),
            InputParam(
                "noise_scale",
                type_hint=float,
                default=None,
                description=(
                    "Interpolation factor between random noise and any provided `audio_latents`. Resolved upstream "
                    "by `LTX2PrepareLatentsStep` (0.0, which keeps the provided latents)."
                ),
            ),
            InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"),
            InputParam.template("generator"),
            InputParam(
                "batch_size",
                type_hint=int,
                required=True,
                description="The number of prompts being denoised, used to expand conditioning per prompt.",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam("audio_latents", type_hint=torch.Tensor, description="Packed noisy audio latents."),
            OutputParam(
                "audio_num_frames",
                type_hint=int,
                kwargs_type="denoiser_input_fields",
                description="Number of audio latent frames.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        device = components._execution_device

        batch_size = block_state.batch_size * block_state.num_videos_per_prompt
        num_channels_latents = components.audio_vae.config.latent_channels
        num_mel_bins = components.audio_vae.config.mel_bins

        duration_s = block_state.num_frames / block_state.frame_rate
        audio_latents_per_second = (
            components.audio_sampling_rate
            / components.audio_hop_length
            / float(components.audio_vae_temporal_compression_ratio)
        )
        audio_num_frames = round(duration_s * audio_latents_per_second)

        if block_state.audio_latents is not None:
            audio_latents = block_state.audio_latents
            if audio_latents.ndim == 4:
                audio_num_frames = audio_latents.shape[2]
                audio_latents = _pack_audio_latents(audio_latents)
            audio_latents = _normalize_audio_latents(
                audio_latents, components.audio_latents_mean, components.audio_latents_std
            )
            audio_latents = _create_noised_state(audio_latents, block_state.noise_scale, block_state.generator)
            block_state.audio_latents = audio_latents.to(device=device, dtype=torch.float32)
        else:
            latent_mel_bins = num_mel_bins // components.audio_vae_mel_compression_ratio
            shape = (batch_size, num_channels_latents, audio_num_frames, latent_mel_bins)
            audio_latents = randn_tensor(shape, generator=block_state.generator, device=device, dtype=torch.float32)
            block_state.audio_latents = _pack_audio_latents(audio_latents)

        block_state.audio_num_frames = audio_num_frames

        self.set_block_state(state, block_state)
        return components, state


class LTX2PrepareCoordsStep(ModularPipelineBlocks):
    model_name = "ltx2"

    @property
    def description(self) -> str:
        return (
            "Pre-computes the video and audio RoPE positional coordinates (constant across denoising steps). CFG/STG "
            "batch duplication is left to the denoise step."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [ComponentSpec("transformer", LTX2VideoTransformer3DModel)]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("height", default=512),
            InputParam.template("width", default=704),
            InputParam(
                "num_frames",
                type_hint=int,
                default=None,
                description=(
                    "The number of frames in the generated video. Omit to auto-predict via the `duration_head` "
                    "(see `LTX2AutoDurationStep`)."
                ),
            ),
            InputParam(
                "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video."
            ),
            InputParam("audio_num_frames", type_hint=int, required=True),
            InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"),
            InputParam(
                "batch_size",
                type_hint=int,
                required=True,
                description="The number of prompts being denoised, used to expand conditioning per prompt.",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                "video_coords",
                type_hint=torch.Tensor,
                kwargs_type="denoiser_input_fields",
                description="Video RoPE patch coordinates.",
            ),
            OutputParam(
                "audio_coords",
                type_hint=torch.Tensor,
                kwargs_type="denoiser_input_fields",
                description="Audio RoPE patch coordinates.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        device = components._execution_device

        batch_size = block_state.batch_size * block_state.num_videos_per_prompt
        latent_height = block_state.height // components.vae_spatial_compression_ratio
        latent_width = block_state.width // components.vae_spatial_compression_ratio
        latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1

        block_state.video_coords = components.transformer.rope.prepare_video_coords(
            batch_size, latent_num_frames, latent_height, latent_width, device, fps=block_state.frame_rate
        )
        block_state.audio_coords = components.transformer.audio_rope.prepare_audio_coords(
            batch_size, block_state.audio_num_frames, device
        )

        self.set_block_state(state, block_state)
        return components, state


class LTX2ConditionPrepareLatentsStep(ModularPipelineBlocks):
    model_name = "ltx2"

    @property
    def description(self) -> str:
        return (
            "Prepares the packed video latents for condition-based generation. Conditions at latent index 0 "
            "overwrite the tokens at the first-frame positions (`VideoConditionByLatentIndex` semantics); conditions "
            "at other latent indices are appended to the sequence as keyframe tokens carrying their own RoPE "
            "coordinates (`VideoConditionByKeyframeIndex` semantics). Emits the per-token `conditioning_mask` and "
            "`clean_latents` that drive the loop's masked timestep and x0 blend. Unlike `LTX2PrepareLatentsStep` "
            "this starts from zeros and samples the noise last, once the conditioning mask is known."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("transformer", LTX2VideoTransformer3DModel),
            ComponentSpec("vae", AutoencoderKLLTX2Video),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                "condition_latents",
                type_hint=list,
                required=True,
                description="Per-condition normalized VAE latents of shape [1, C, F, H, W].",
            ),
            InputParam(
                "condition_strengths",
                type_hint=list,
                required=True,
                description="Per-condition conditioning strengths.",
            ),
            InputParam(
                "condition_indices",
                type_hint=list,
                required=True,
                description="Per-condition latent frame index at which the condition is applied.",
            ),
            InputParam(
                "condition_pixel_frames",
                type_hint=list,
                required=True,
                description="Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords.",
            ),
            InputParam.template("latents"),
            InputParam.template("height", default=512),
            InputParam.template("width", default=704),
            InputParam(
                "num_frames",
                type_hint=int,
                default=None,
                description=(
                    "The number of frames in the generated video. Omit to auto-predict via the `duration_head` "
                    "(see `LTX2AutoDurationStep`)."
                ),
            ),
            InputParam(
                "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video."
            ),
            InputParam(
                "noise_scale",
                type_hint=float,
                default=None,
                description=(
                    "Initial noise level for the un-conditioned tokens. `None` (default) resolves to `sigmas[0]` "
                    "when custom `sigmas` are supplied, else 1.0."
                ),
            ),
            InputParam.template("sigmas"),
            InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"),
            InputParam(
                "batch_size",
                type_hint=int,
                required=True,
                description="The number of prompts being denoised, used to expand conditioning per prompt.",
            ),
            InputParam.template("generator"),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                "latents",
                type_hint=torch.Tensor,
                description="Packed noisy video latents, with any keyframe condition tokens appended.",
            ),
            OutputParam(
                "conditioning_mask",
                type_hint=torch.Tensor,
                description=(
                    "Packed per-token conditioning strengths of shape [B, S, 1] in [0, 1]: 1 at fully-conditioned "
                    "positions, 0 at free positions."
                ),
            ),
            OutputParam(
                "clean_latents",
                type_hint=torch.Tensor,
                description="Clean condition latents at conditioned positions, zeros elsewhere; same shape as `latents`.",
            ),
            OutputParam(
                "appended_coords",
                type_hint=torch.Tensor,
                description=(
                    "RoPE coordinates of shape [B, 3, num_keyframe_tokens, 2] for the appended keyframe tokens, "
                    "zero-width when there are none."
                ),
            ),
            OutputParam(
                "base_token_count",
                type_hint=int,
                description="Number of generated-video tokens, i.e. the sequence length before appended tokens.",
            ),
            OutputParam(
                "noise_scale",
                type_hint=float,
                description="The resolved initial noise level, forwarded to the audio latents step.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        device = components._execution_device

        batch_size = block_state.batch_size * block_state.num_videos_per_prompt
        spatial_patch = components.transformer_spatial_patch_size
        temporal_patch = components.transformer_temporal_patch_size
        frame_scale_factor = components.vae_temporal_compression_ratio

        latent_height = block_state.height // components.vae_spatial_compression_ratio
        latent_width = block_state.width // components.vae_spatial_compression_ratio
        latent_num_frames = (block_state.num_frames - 1) // frame_scale_factor + 1

        # Noise level the un-conditioned tokens start at: the first (largest) sigma when custom sigmas are supplied,
        # else 1.0. Matches `LTX2ConditionPipeline.__call__`.
        noise_scale = block_state.noise_scale
        if noise_scale is None:
            noise_scale = block_state.sigmas[0] if block_state.sigmas is not None else 1.0

        if isinstance(block_state.generator, list):
            logger.warning(
                f"{self.__class__.__name__} does not support using a list of generators. The first generator in the"
                f" list will be used for all (pseudo-)random operations."
            )

        if block_state.latents is not None:
            latents = _normalize_latents(
                block_state.latents,
                components.latents_mean,
                components.latents_std,
                components.vae_scaling_factor,
            )
        else:
            # Zeros rather than a Gaussian sample: the noise is mixed in at the end, once the mask is known.
            shape = (
                batch_size,
                components.transformer.config.in_channels,
                latent_num_frames,
                latent_height,
                latent_width,
            )
            latents = torch.zeros(shape, device=device, dtype=torch.float32)

        conditioning_mask = latents.new_zeros((batch_size, 1, latent_num_frames, latent_height, latent_width))
        latents = _pack_latents(latents, spatial_patch, temporal_patch)
        conditioning_mask = _pack_latents(conditioning_mask, spatial_patch, temporal_patch)  # [B, S, 1]

        if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape[:2]:
            raise ValueError(
                f"Provided `latents` tensor packs to shape {latents.shape}, but the expected packed shape is "
                f"{conditioning_mask.shape[:2] + (components.transformer.config.in_channels,)}."
            )

        base_token_count = latents.shape[1]
        condition_latents_packed = [
            _pack_latents(cond, spatial_patch, temporal_patch) for cond in block_state.condition_latents
        ]

        # First-frame conditions (latent index 0): overwrite the tokens at the first-frame positions. Condition
        # tensors carry batch 1 and broadcast across the generation batch.
        clean_latents = torch.zeros_like(latents)
        for cond, strength, latent_idx in zip(
            condition_latents_packed, block_state.condition_strengths, block_state.condition_indices
        ):
            if latent_idx != 0:
                continue
            num_cond_tokens = cond.size(1)
            latents[:, :num_cond_tokens] = cond
            conditioning_mask[:, :num_cond_tokens] = strength
            clean_latents[:, :num_cond_tokens] = cond

        # Non-first-frame ("keyframe") conditions (latent index > 0): append as extra tokens with an all-`strength`
        # conditioning mask and their own coords. At denoising step i they see an effective noise level of
        # (1 - strength) * sigma_i.
        scale_factors = (
            frame_scale_factor,
            components.vae_spatial_compression_ratio,
            components.vae_spatial_compression_ratio,
        )
        keyframe_tokens, keyframe_masks, keyframe_coords = [], [], []
        for cond_5d, cond_packed, strength, latent_idx, num_pixel_frames in zip(
            block_state.condition_latents,
            condition_latents_packed,
            block_state.condition_strengths,
            block_state.condition_indices,
            block_state.condition_pixel_frames,
        ):
            if latent_idx == 0:
                continue

            _, _, kf_latent_frames, kf_latent_height, kf_latent_width = cond_5d.shape
            coords = _prepare_keyframe_coords(
                keyframe_latent_num_frames=kf_latent_frames,
                keyframe_latent_height=kf_latent_height,
                keyframe_latent_width=kf_latent_width,
                pixel_frame_idx=(latent_idx - 1) * frame_scale_factor + 1,
                num_pixel_frames=num_pixel_frames,
                fps=block_state.frame_rate,
                patch_size=spatial_patch,
                patch_size_t=temporal_patch,
                scale_factors=scale_factors,
                device=device,
            )

            keyframe_tokens.append(cond_packed.expand(batch_size, -1, -1))
            keyframe_masks.append(
                torch.full(
                    (batch_size, cond_packed.shape[1], 1),
                    float(strength),
                    device=device,
                    dtype=conditioning_mask.dtype,
                )
            )
            keyframe_coords.append(coords.expand(batch_size, -1, -1, -1))

        if keyframe_tokens:
            keyframe_tokens = torch.cat(keyframe_tokens, dim=1)
            latents = torch.cat([latents, keyframe_tokens], dim=1)
            clean_latents = torch.cat([clean_latents, keyframe_tokens], dim=1)
            conditioning_mask = torch.cat([conditioning_mask, torch.cat(keyframe_masks, dim=1)], dim=1)
            appended_coords = torch.cat(keyframe_coords, dim=2)
        else:
            appended_coords = torch.zeros((batch_size, 3, 0, 2), device=device, dtype=torch.float32)

        # Mask semantics: 0 -> fully noised, 1 -> kept clean, in between -> noise level (1 - mask) * noise_scale.
        noise = randn_tensor(
            latents.shape, generator=block_state.generator, device=latents.device, dtype=latents.dtype
        )
        scaled_mask = (1.0 - conditioning_mask) * noise_scale
        block_state.latents = noise * scaled_mask + latents * (1 - scaled_mask)

        block_state.conditioning_mask = conditioning_mask
        block_state.clean_latents = clean_latents
        block_state.appended_coords = appended_coords
        block_state.base_token_count = base_token_count
        block_state.noise_scale = noise_scale

        self.set_block_state(state, block_state)
        return components, state


class LTX2InContextPrepareLatentsStep(ModularPipelineBlocks):
    model_name = "ltx2"

    @property
    def description(self) -> str:
        return (
            "Prepares the packed video latents for in-context (IC-LoRA) generation. Same frame-condition handling as "
            "`LTX2ConditionPrepareLatentsStep` (first-frame overwrite, keyframe token append), then appends the "
            "encoded reference tokens after the keyframes with a per-token `conditioning_mask` of their own "
            "strength, giving a single `[base | keyframe | reference]` sequence. Mirrors "
            "`LTX2InContextPipeline.prepare_latents`, which likewise re-implements the condition version rather "
            "than extending it -- the two are kept side by side so each reads top to bottom."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("transformer", LTX2VideoTransformer3DModel),
            ComponentSpec("vae", AutoencoderKLLTX2Video),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                "condition_latents",
                type_hint=list,
                required=True,
                description="Per-condition normalized VAE latents of shape [1, C, F, H, W].",
            ),
            InputParam(
                "condition_strengths",
                type_hint=list,
                required=True,
                description="Per-condition conditioning strengths.",
            ),
            InputParam(
                "condition_indices",
                type_hint=list,
                required=True,
                description="Per-condition latent frame index at which the condition is applied.",
            ),
            InputParam(
                "condition_pixel_frames",
                type_hint=list,
                required=True,
                description="Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords.",
            ),
            InputParam(
                "reference_conditions",
                type_hint=list,
                default=None,
                description=(
                    "`LTX2ReferenceCondition` (or list of them); only their `strength` is read here. Omit for "
                    "IC-LoRAs that carry their behavior in the adapter weights and take no reference video."
                ),
            ),
            InputParam(
                "reference_latents",
                type_hint=torch.Tensor,
                default=None,
                description=(
                    "Packed reference tokens of shape [1, total_reference_tokens, C], or `None` when no reference "
                    "conditions were supplied (`LTX2AutoReferenceEncoderStep` is skipped)."
                ),
            ),
            InputParam(
                "reference_coords",
                type_hint=torch.Tensor,
                default=None,
                description="RoPE coordinates for the reference tokens.",
            ),
            InputParam(
                "reference_token_counts",
                type_hint=list,
                default=None,
                description="Per-reference token counts, in `reference_conditions` order.",
            ),
            InputParam.template("latents"),
            InputParam.template("height", default=512),
            InputParam.template("width", default=704),
            InputParam(
                "num_frames",
                type_hint=int,
                default=None,
                description=(
                    "The number of frames in the generated video. Omit to auto-predict via the `duration_head` "
                    "(see `LTX2AutoDurationStep`)."
                ),
            ),
            InputParam(
                "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video."
            ),
            InputParam(
                "noise_scale",
                type_hint=float,
                default=None,
                description=(
                    "Initial noise level for the un-conditioned tokens. `None` (default) resolves to `sigmas[0]` "
                    "when custom `sigmas` are supplied, else 1.0."
                ),
            ),
            InputParam.template("sigmas"),
            InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"),
            InputParam(
                "batch_size",
                type_hint=int,
                required=True,
                description="The number of prompts being denoised, used to expand conditioning per prompt.",
            ),
            InputParam.template("generator"),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                "latents",
                type_hint=torch.Tensor,
                description="Packed noisy video latents, with keyframe and reference tokens appended.",
            ),
            OutputParam(
                "conditioning_mask",
                type_hint=torch.Tensor,
                description=(
                    "Packed per-token conditioning strengths of shape [B, S, 1] in [0, 1]: 1 at fully-conditioned "
                    "positions, 0 at free positions."
                ),
            ),
            OutputParam(
                "clean_latents",
                type_hint=torch.Tensor,
                description="Clean condition latents at conditioned positions, zeros elsewhere; same shape as `latents`.",
            ),
            OutputParam(
                "appended_coords",
                type_hint=torch.Tensor,
                description=(
                    "RoPE coordinates of shape [B, 3, num_keyframe_tokens + num_reference_tokens, 2] for the "
                    "appended tokens, zero-width when there are none."
                ),
            ),
            OutputParam(
                "base_token_count",
                type_hint=int,
                description="Number of generated-video tokens, i.e. the sequence length before appended tokens.",
            ),
            OutputParam(
                "num_ref_tokens",
                type_hint=int,
                description="Number of reference tokens, which sit at the very end of the sequence.",
            ),
            OutputParam(
                "noise_scale",
                type_hint=float,
                description="The resolved initial noise level, forwarded to the audio latents step.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        device = components._execution_device

        batch_size = block_state.batch_size * block_state.num_videos_per_prompt
        spatial_patch = components.transformer_spatial_patch_size
        temporal_patch = components.transformer_temporal_patch_size
        frame_scale_factor = components.vae_temporal_compression_ratio

        latent_height = block_state.height // components.vae_spatial_compression_ratio
        latent_width = block_state.width // components.vae_spatial_compression_ratio
        latent_num_frames = (block_state.num_frames - 1) // frame_scale_factor + 1

        noise_scale = block_state.noise_scale
        if noise_scale is None:
            noise_scale = block_state.sigmas[0] if block_state.sigmas is not None else 1.0

        if isinstance(block_state.generator, list):
            logger.warning(
                f"{self.__class__.__name__} does not support using a list of generators. The first generator in the"
                f" list will be used for all (pseudo-)random operations."
            )

        if block_state.latents is not None:
            latents = _normalize_latents(
                block_state.latents,
                components.latents_mean,
                components.latents_std,
                components.vae_scaling_factor,
            )
        else:
            shape = (
                batch_size,
                components.transformer.config.in_channels,
                latent_num_frames,
                latent_height,
                latent_width,
            )
            latents = torch.zeros(shape, device=device, dtype=torch.float32)

        conditioning_mask = latents.new_zeros((batch_size, 1, latent_num_frames, latent_height, latent_width))
        latents = _pack_latents(latents, spatial_patch, temporal_patch)
        conditioning_mask = _pack_latents(conditioning_mask, spatial_patch, temporal_patch)  # [B, S, 1]

        if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape[:2]:
            raise ValueError(
                f"Provided `latents` tensor packs to shape {latents.shape}, but the expected packed shape is "
                f"{conditioning_mask.shape[:2] + (components.transformer.config.in_channels,)}."
            )

        base_token_count = latents.shape[1]
        condition_latents_packed = [
            _pack_latents(cond, spatial_patch, temporal_patch) for cond in block_state.condition_latents
        ]

        # First-frame conditions (latent index 0): overwrite the tokens at the first-frame positions.
        clean_latents = torch.zeros_like(latents)
        for cond, strength, latent_idx in zip(
            condition_latents_packed, block_state.condition_strengths, block_state.condition_indices
        ):
            if latent_idx != 0:
                continue
            num_cond_tokens = cond.size(1)
            latents[:, :num_cond_tokens] = cond
            conditioning_mask[:, :num_cond_tokens] = strength
            clean_latents[:, :num_cond_tokens] = cond

        # Non-first-frame ("keyframe") conditions (latent index > 0): append as extra tokens with their own coords.
        scale_factors = (
            frame_scale_factor,
            components.vae_spatial_compression_ratio,
            components.vae_spatial_compression_ratio,
        )
        keyframe_tokens, keyframe_masks, keyframe_coords = [], [], []
        for cond_5d, cond_packed, strength, latent_idx, num_pixel_frames in zip(
            block_state.condition_latents,
            condition_latents_packed,
            block_state.condition_strengths,
            block_state.condition_indices,
            block_state.condition_pixel_frames,
        ):
            if latent_idx == 0:
                continue

            _, _, kf_latent_frames, kf_latent_height, kf_latent_width = cond_5d.shape
            coords = _prepare_keyframe_coords(
                keyframe_latent_num_frames=kf_latent_frames,
                keyframe_latent_height=kf_latent_height,
                keyframe_latent_width=kf_latent_width,
                pixel_frame_idx=(latent_idx - 1) * frame_scale_factor + 1,
                num_pixel_frames=num_pixel_frames,
                fps=block_state.frame_rate,
                patch_size=spatial_patch,
                patch_size_t=temporal_patch,
                scale_factors=scale_factors,
                device=device,
            )

            keyframe_tokens.append(cond_packed.expand(batch_size, -1, -1))
            keyframe_masks.append(
                torch.full(
                    (batch_size, cond_packed.shape[1], 1),
                    float(strength),
                    device=device,
                    dtype=conditioning_mask.dtype,
                )
            )
            keyframe_coords.append(coords.expand(batch_size, -1, -1, -1))

        # Seeded with a zero-width block so the concat below is unconditional whether or not anything is appended.
        appended_coords = [torch.zeros((batch_size, 3, 0, 2), device=device, dtype=torch.float32)]
        if keyframe_tokens:
            keyframe_tokens = torch.cat(keyframe_tokens, dim=1)
            latents = torch.cat([latents, keyframe_tokens], dim=1)
            clean_latents = torch.cat([clean_latents, keyframe_tokens], dim=1)
            conditioning_mask = torch.cat([conditioning_mask, torch.cat(keyframe_masks, dim=1)], dim=1)
            appended_coords.append(torch.cat(keyframe_coords, dim=2))

        # Reference (IC-LoRA) tokens, appended last so they sit at the very end of the sequence -- the video
        # self-attention mask is built off that placement. Same mechanism as the keyframes above: a per-token
        # `conditioning_mask` of the reference's strength, matching `VideoConditionByReferenceLatent` in the
        # reference implementation. Absent for IC-LoRAs that take no reference video (camera control, style, ...),
        # which `LTX2InContextPipeline` supports too -- `LTX2AutoReferenceEncoderStep` is then skipped.
        num_ref_tokens = 0
        if block_state.reference_latents is not None:
            reference_conditions = block_state.reference_conditions
            if isinstance(reference_conditions, LTX2ReferenceCondition):
                reference_conditions = [reference_conditions]
            reference_tokens = block_state.reference_latents.expand(batch_size, -1, -1)
            num_ref_tokens = reference_tokens.shape[1]
            # Per-reference token counts come from the encoder rather than an equal split of the total, so
            # references of differing lengths (a reference video shorter than `num_frames`) get the right strengths.
            reference_masks = [
                torch.full(
                    (batch_size, num_tokens, 1),
                    float(ref_cond.strength),
                    device=device,
                    dtype=conditioning_mask.dtype,
                )
                for ref_cond, num_tokens in zip(reference_conditions, block_state.reference_token_counts)
            ]
            latents = torch.cat([latents, reference_tokens], dim=1)
            clean_latents = torch.cat([clean_latents, reference_tokens], dim=1)
            conditioning_mask = torch.cat([conditioning_mask, torch.cat(reference_masks, dim=1)], dim=1)
            appended_coords.append(block_state.reference_coords.expand(batch_size, -1, -1, -1))

        # Mask semantics: 0 -> fully noised, 1 -> kept clean, in between -> noise level (1 - mask) * noise_scale.
        noise = randn_tensor(
            latents.shape, generator=block_state.generator, device=latents.device, dtype=latents.dtype
        )
        scaled_mask = (1.0 - conditioning_mask) * noise_scale
        block_state.latents = noise * scaled_mask + latents * (1 - scaled_mask)

        block_state.conditioning_mask = conditioning_mask
        block_state.clean_latents = clean_latents
        block_state.appended_coords = torch.cat(appended_coords, dim=2)
        block_state.base_token_count = base_token_count
        block_state.num_ref_tokens = num_ref_tokens
        block_state.noise_scale = noise_scale

        self.set_block_state(state, block_state)
        return components, state


class LTX2BuildVideoSelfAttentionMaskStep(ModularPipelineBlocks):
    model_name = "ltx2"

    @property
    def description(self) -> str:
        return (
            "Builds the video self-attention mask over the `[base | keyframe | reference]` token sequence for "
            "in-context generation, mirroring `build_attention_mask` in the reference implementation. Each "
            "`LTX2ReferenceCondition` is its own attention group:\n"
            "  - base <-> base, base <-> keyframe, keyframe <-> keyframe: 1.0 (full attention)\n"
            "  - base <-> reference group: that group's slice of `reference_cross_mask`, broadcast symmetrically "
            "across the base-token axis\n"
            "  - reference group <-> itself: 1.0 (a group fully attends to itself)\n"
            "  - reference group <-> any other appended group (keyframes, other references): 0.0\n"
            "Note the cross blocks span only the *base* tokens: keyframe tokens are appended conditioning like the "
            "references are, and the two are masked off from each other. Emitted tagged `denoiser_input_fields`, so "
            "it reaches the transformer's `video_self_attention_mask` argument without any change to "
            "`LTX2LoopDenoiser`. Only run when the references carry a per-token strength (see "
            "`LTX2AutoBuildVideoSelfAttentionMaskStep`); attention is otherwise left unmasked."
        )

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("latents", required=True),
            InputParam(
                "base_token_count",
                type_hint=int,
                required=True,
                description="Number of generated-video tokens, i.e. the sequence length before appended tokens.",
            ),
            InputParam(
                "num_ref_tokens",
                type_hint=int,
                required=True,
                description="Number of reference tokens, which sit at the very end of the sequence.",
            ),
            InputParam(
                "reference_cross_mask",
                type_hint=torch.Tensor,
                required=True,
                description="Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens].",
            ),
            InputParam(
                "reference_token_counts",
                type_hint=list,
                required=True,
                description="Per-reference token counts, used to split `reference_cross_mask` into attention groups.",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                "video_self_attention_mask",
                type_hint=torch.Tensor,
                kwargs_type="denoiser_input_fields",
                description="Multiplicative self-attention mask of shape [B, S, S] with values in [0, 1].",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        device = components._execution_device

        batch_size, total_tokens, _ = block_state.latents.shape
        # Cross-attention partners are the base tokens only. The prefix (base + keyframe tokens) attends fully to
        # itself, but keyframe tokens are appended conditioning just like the references, so the two groups are
        # masked off from each other.
        num_base_tokens = block_state.base_token_count
        num_prefix_tokens = total_tokens - block_state.num_ref_tokens
        cross = block_state.reference_cross_mask.to(device=device, dtype=torch.float32)

        # Start from zeros so the keyframe<->reference and reference<->reference blocks stay masked without explicit
        # assignment. Each guidance pass is its own single-batch forward, so this is built at the generation batch
        # size -- the standard pipeline's expand to 2B for the CFG batch has no counterpart here.
        attn_mask = torch.zeros((batch_size, total_tokens, total_tokens), device=device, dtype=torch.float32)
        attn_mask[:, :num_prefix_tokens, :num_prefix_tokens] = 1.0

        # One attention group per reference condition, in the order the encoder emitted their tokens.
        offset = num_prefix_tokens
        for group_cross in torch.split(cross, block_state.reference_token_counts, dim=1):
            n = group_cross.shape[1]
            attn_mask[:, :num_base_tokens, offset : offset + n] = group_cross.unsqueeze(1)
            attn_mask[:, offset : offset + n, :num_base_tokens] = group_cross.unsqueeze(2)
            attn_mask[:, offset : offset + n, offset : offset + n] = 1.0
            offset += n

        block_state.video_self_attention_mask = attn_mask

        self.set_block_state(state, block_state)
        return components, state


class LTX2ConditionSetTimestepsStep(ModularPipelineBlocks):
    model_name = "ltx2"

    @property
    def description(self) -> str:
        return (
            "Sets the flow-match timesteps for the video `scheduler` and produces a deep-copied `audio_scheduler` "
            "(with the same schedule) so the audio latents are denoised on an independent scheduler state. Unlike "
            "`LTX2SetTimestepsStep`, the resolution-aware shift `mu` is read off the packed `latents` sequence "
            "length, which for condition workflows includes the appended keyframe tokens -- so this block runs "
            "*after* `LTX2ConditionPrepareLatentsStep`, matching `LTX2ConditionPipeline`."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("num_inference_steps", default=30),
            InputParam.template("timesteps"),
            InputParam.template("sigmas"),
            InputParam.template("latents", required=True),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam("timesteps", type_hint=torch.Tensor),
            OutputParam("num_inference_steps", type_hint=int),
            OutputParam(
                "audio_scheduler",
                description="Independent deep copy of `scheduler` used to update the audio latents in the loop.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        device = components._execution_device

        num_inference_steps = block_state.num_inference_steps
        sigmas = block_state.sigmas
        if sigmas is None:
            sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)

        mu = calculate_shift(
            block_state.latents.shape[1],
            components.scheduler.config.get("base_image_seq_len", 1024),
            components.scheduler.config.get("max_image_seq_len", 4096),
            components.scheduler.config.get("base_shift", 0.95),
            components.scheduler.config.get("max_shift", 2.05),
        )

        block_state.audio_scheduler = copy.deepcopy(components.scheduler)
        retrieve_timesteps(
            block_state.audio_scheduler, num_inference_steps, device, block_state.timesteps, sigmas=sigmas, mu=mu
        )
        block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
            components.scheduler, num_inference_steps, device, block_state.timesteps, sigmas=sigmas, mu=mu
        )

        # Set begin index to skip the nonzero().item() call in scheduler init, which triggers a GPU sync.
        components.scheduler.set_begin_index(0)
        block_state.audio_scheduler.set_begin_index(0)

        self.set_block_state(state, block_state)
        return components, state


class LTX2ConditionPrepareAudioLatentsStep(ModularPipelineBlocks):
    model_name = "ltx2"

    @property
    def description(self) -> str:
        return (
            "Prepares the packed audio noise latents and derives the audio latent frame count, for condition-based "
            "generation. Two deliberate differences from `LTX2PrepareAudioLatentsStep`:\n"
            "  1. the noise is sampled directly in the packed shape [B, L, C * M], matching "
            "`LTX2ConditionPipeline.prepare_audio_latents`. The text-to-video/image-to-video pipelines sample "
            "unpacked [B, C, L, M] and pack afterwards; both draw the same number of values from the generator but "
            "lay them out differently, so sampling the wrong way silently desynchronizes the audio noise (and, "
            "through the joint attention, the video too).\n"
            "  2. `noise_scale` is declared with a `None` default: a blockset keeps the first non-`None` default "
            "across its blocks, so the text-to-video default of 0.0 would otherwise shadow the condition workflow's "
            "`None -> sigmas[0] or 1.0` resolution. `LTX2ConditionPrepareLatentsStep` runs first and writes the "
            "resolved value back to state."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [ComponentSpec("audio_vae", AutoencoderKLLTX2Audio)]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                "num_frames",
                type_hint=int,
                default=None,
                description=(
                    "The number of frames in the generated video. Omit to auto-predict via the `duration_head` "
                    "(see `LTX2AutoDurationStep`)."
                ),
            ),
            InputParam(
                "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video."
            ),
            InputParam(
                "audio_latents",
                type_hint=torch.Tensor,
                default=None,
                description="Optional pre-encoded audio latents; random noise is used when not provided.",
            ),
            InputParam(
                "noise_scale",
                type_hint=float,
                default=None,
                description=(
                    "Initial noise level applied to any provided `audio_latents`. Resolved upstream by "
                    "`LTX2ConditionPrepareLatentsStep` (`sigmas[0]` when custom sigmas are supplied, else 1.0)."
                ),
            ),
            InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"),
            InputParam.template("generator"),
            InputParam(
                "batch_size",
                type_hint=int,
                required=True,
                description="The number of prompts being denoised, used to expand conditioning per prompt.",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam("audio_latents", type_hint=torch.Tensor, description="Packed noisy audio latents."),
            OutputParam(
                "audio_num_frames",
                type_hint=int,
                kwargs_type="denoiser_input_fields",
                description="Number of audio latent frames.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        device = components._execution_device

        batch_size = block_state.batch_size * block_state.num_videos_per_prompt
        num_channels_latents = components.audio_vae.config.latent_channels
        latent_mel_bins = components.audio_vae.config.mel_bins // components.audio_vae_mel_compression_ratio

        duration_s = block_state.num_frames / block_state.frame_rate
        audio_latents_per_second = (
            components.audio_sampling_rate
            / components.audio_hop_length
            / float(components.audio_vae_temporal_compression_ratio)
        )
        audio_num_frames = round(duration_s * audio_latents_per_second)

        if block_state.audio_latents is not None:
            audio_latents = block_state.audio_latents
            if audio_latents.ndim == 4:
                audio_num_frames = audio_latents.shape[2]
                audio_latents = _pack_audio_latents(audio_latents)
            audio_latents = _normalize_audio_latents(
                audio_latents, components.audio_latents_mean, components.audio_latents_std
            )
            audio_latents = _create_noised_state(audio_latents, block_state.noise_scale, block_state.generator)
            block_state.audio_latents = audio_latents.to(device=device, dtype=torch.float32)
        else:
            # Sample directly in packed shape, following `LTX2ConditionPipeline.prepare_audio_latents` -- see the
            # block description for why the unpacked-then-pack order used by text-to-video is not interchangeable.
            packed_shape = (batch_size, audio_num_frames, num_channels_latents * latent_mel_bins)
            block_state.audio_latents = randn_tensor(
                packed_shape, generator=block_state.generator, device=device, dtype=torch.float32
            )

        block_state.audio_num_frames = audio_num_frames

        self.set_block_state(state, block_state)
        return components, state


class LTX2ConditionPrepareCoordsStep(ModularPipelineBlocks):
    model_name = "ltx2"

    @property
    def description(self) -> str:
        return (
            "Pre-computes the video and audio RoPE positional coordinates (constant across denoising steps) and "
            "appends the keyframe-condition coordinates onto the video coordinates, matching the order the "
            "keyframe tokens were appended to the latent sequence. CFG/STG batch duplication is left to the "
            "denoise step."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [ComponentSpec("transformer", LTX2VideoTransformer3DModel)]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("height", default=512),
            InputParam.template("width", default=704),
            InputParam(
                "num_frames",
                type_hint=int,
                default=None,
                description=(
                    "The number of frames in the generated video. Omit to auto-predict via the `duration_head` "
                    "(see `LTX2AutoDurationStep`)."
                ),
            ),
            InputParam(
                "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video."
            ),
            InputParam("audio_num_frames", type_hint=int, required=True),
            InputParam(
                "appended_coords",
                type_hint=torch.Tensor,
                required=True,
                description="RoPE coordinates for the appended keyframe tokens, zero-width when there are none.",
            ),
            InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"),
            InputParam(
                "batch_size",
                type_hint=int,
                required=True,
                description="The number of prompts being denoised, used to expand conditioning per prompt.",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                "video_coords",
                type_hint=torch.Tensor,
                kwargs_type="denoiser_input_fields",
                description="Video RoPE patch coordinates, with the keyframe-condition coordinates appended.",
            ),
            OutputParam(
                "audio_coords",
                type_hint=torch.Tensor,
                kwargs_type="denoiser_input_fields",
                description="Audio RoPE patch coordinates.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        device = components._execution_device

        batch_size = block_state.batch_size * block_state.num_videos_per_prompt
        latent_height = block_state.height // components.vae_spatial_compression_ratio
        latent_width = block_state.width // components.vae_spatial_compression_ratio
        latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1

        video_coords = components.transformer.rope.prepare_video_coords(
            batch_size, latent_num_frames, latent_height, latent_width, device, fps=block_state.frame_rate
        )
        block_state.video_coords = torch.cat([video_coords, block_state.appended_coords.to(video_coords.dtype)], dim=2)
        block_state.audio_coords = components.transformer.audio_rope.prepare_audio_coords(
            batch_size, block_state.audio_num_frames, device
        )

        self.set_block_state(state, block_state)
        return components, state