ltx2.py

# SPDX-License-Identifier: Apache-2.0
"""
LTX-2 transformer implementation
"""

from dataclasses import dataclass, replace
from enum import Enum
import functools
import math
import os
from pathlib import Path
from typing import Any, Optional, Tuple, Callable

import numpy as np

import torch
import torch.nn as nn
from einops import rearrange, repeat

from fastvideo.attention.backends.sdpa import SDPAMetadata
from fastvideo.attention.layer import DistributedAttention, LocalAttention
from fastvideo.configs.models.dits import LTX2VideoConfig
from fastvideo.distributed.communication_op import (
    sequence_model_parallel_all_gather,
    sequence_model_parallel_all_gather_with_unpad,
    sequence_model_parallel_all_to_all_4D,
    sequence_model_parallel_shard,
)
from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size
from fastvideo.forward_context import ForwardContext, get_forward_context, set_forward_context
from fastvideo.layers.linear import ReplicatedLinear
from fastvideo.layers.quantization.base_config import QuantizationConfig
from fastvideo.logger import init_logger
from fastvideo.models.dits.base import BaseDiT
from fastvideo.platforms import AttentionBackendEnum

logger = init_logger(__name__)

# QuACK provides a fast NVFP4 RMSNorm used only for the LTX-2.3 refine stage.
# Import defensively: upstream installs without QuACK must still work, simply
# falling back to ``torch.nn.functional.rms_norm`` (the non-refine path).
try:
    from quack import rmsnorm as _quack_rmsnorm
except Exception:  # pragma: no cover - optional dependency
    _quack_rmsnorm = None


def _is_ltx2_refine_stage() -> bool:
    """Return True when running LTX-2 stage-2 refinement denoising.

    Only the refine stage uses the QuACK fast path; everything else (and any
    install lacking QuACK) uses the standard Torch RMSNorm, which reproduces
    LTX-2.0 numerics exactly.
    """
    if _quack_rmsnorm is None:
        return False
    try:
        forward_ctx = get_forward_context()
    except AssertionError:
        return False
    forward_batch = forward_ctx.forward_batch
    if forward_batch is None:
        return False
    return forward_batch.extra.get("ltx2_fp4_stage_profile") == "refine"


def _rms_norm_dispatch(
    x: torch.Tensor,
    eps: float,
    weight: torch.Tensor | None = None,
) -> torch.Tensor:
    """Use QuACK RMSNorm only for LTX-2 refine stage, else Torch RMSNorm."""
    if _is_ltx2_refine_stage():
        return _quack_rmsnorm(x, weight=weight, eps=eps)
    return torch.nn.functional.rms_norm(
        x, (x.shape[-1], ), weight=weight, eps=eps)


class StageAwareRMSNorm(nn.RMSNorm):
    """Use torch.nn.RMSNorm for base stage and QuACK for refine stage.

    When QuACK is unavailable or outside the refine stage this behaves
    identically to ``torch.nn.RMSNorm`` (LTX-2.0 q_norm/k_norm behavior).
    """

    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
        super().__init__(hidden_size, eps=eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if _is_ltx2_refine_stage():
            return _quack_rmsnorm(
                x,
                weight=self.weight,
                eps=self.eps,
            )
        return super().forward(x)


def _supports_prequantized_input(linear: ReplicatedLinear) -> bool:
    """Whether ``linear``'s quant method is willing to accept a
    pre-quantized ``(x_fp4, x_scale, x_global_sf)`` input tuple.

    Used by the LTX-2 attention forward path so that a single input
    tensor can be quantized once and reused across q/k/v projections
    when they share the same source (self-attention).
    """
    quant_method = getattr(linear, "quant_method", None)
    if not callable(getattr(quant_method, "quantize_input", None)):
        return False
    wants_prequant = getattr(quant_method, "wants_prequantized_input", None)
    if callable(wants_prequant):
        try:
            return bool(wants_prequant())
        except Exception:
            return False
    return True


def _linear_project_with_optional_prequant(
    linear: ReplicatedLinear,
    x: torch.Tensor,
    pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None,
) -> torch.Tensor:
    """Project ``x`` through ``linear``, optionally bypassing the
    in-method quantize step when ``pre_quantized`` is supplied.

    When the linear's quant method does not support pre-quantized
    inputs (e.g. ``UnquantizedLinearMethod``), falls back to the
    standard ``linear(x)`` call and discards the bias-pass-through
    tuple element.
    """
    if pre_quantized is None or not _supports_prequantized_input(linear):
        return linear(x)[0]
    bias = linear.bias if not linear.skip_bias_add else None
    return linear.quant_method.apply(  # type: ignore[union-attr]
        linear,
        x,
        bias=bias,
        pre_quantized=pre_quantized,
    )


def get_timestep_embedding(
    timesteps: torch.Tensor,
    embedding_dim: int,
    flip_sin_to_cos: bool = False,
    downscale_freq_shift: float = 1,
    scale: float = 1,
    max_period: int = 10000,
) -> torch.Tensor:
    """Sinusoidal timestep embedding used by LTX-2 AdaLN."""
    if len(timesteps.shape) != 1:
        raise ValueError("Timesteps should be a 1d-array")

    half_dim = embedding_dim // 2
    exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device)
    exponent = exponent / (half_dim - downscale_freq_shift)

    emb = torch.exp(exponent)
    emb = timesteps[:, None].float() * emb[None, :]
    emb = scale * emb
    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)

    if flip_sin_to_cos:
        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)

    if embedding_dim % 2 == 1:
        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
    return emb


class TimestepEmbedding(torch.nn.Module):
    """Two-layer MLP to project timestep embeddings."""
    def __init__(
        self,
        in_channels: int,
        time_embed_dim: int,
        out_dim: int | None = None,
        post_act_fn: str | None = None,
        cond_proj_dim: int | None = None,
        sample_proj_bias: bool = True,
    ):
        super().__init__()
        self.linear_1 = torch.nn.Linear(in_channels, time_embed_dim, sample_proj_bias)
        self.cond_proj = torch.nn.Linear(cond_proj_dim, in_channels, bias=False) if cond_proj_dim is not None else None
        self.act = torch.nn.SiLU()
        time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim
        self.linear_2 = torch.nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias)
        self.post_act = None if post_act_fn is None else None

    def forward(self, sample: torch.Tensor, condition: torch.Tensor | None = None) -> torch.Tensor:
        if condition is not None:
            sample = sample + self.cond_proj(condition)
        sample = self.linear_1(sample)
        if self.act is not None:
            sample = self.act(sample)
        sample = self.linear_2(sample)
        if self.post_act is not None:
            sample = self.post_act(sample)
        return sample


class Timesteps(torch.nn.Module):
    """Sinusoidal timestep embedding wrapper with scaling knobs."""
    def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1):
        super().__init__()
        self.num_channels = num_channels
        self.flip_sin_to_cos = flip_sin_to_cos
        self.downscale_freq_shift = downscale_freq_shift
        self.scale = scale

    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
        return get_timestep_embedding(
            timesteps,
            self.num_channels,
            flip_sin_to_cos=self.flip_sin_to_cos,
            downscale_freq_shift=self.downscale_freq_shift,
            scale=self.scale,
        )


class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module):
    """PixArt-Alpha timestep embedding used by LTX-2 AdaLN."""
    def __init__(self, embedding_dim: int, size_emb_dim: int):
        super().__init__()
        self.outdim = size_emb_dim
        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)

    def forward(self, timestep: torch.Tensor, hidden_dtype: torch.dtype) -> torch.Tensor:
        timesteps_proj = self.time_proj(timestep)
        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype))
        return timesteps_emb


class AdaLayerNormSingle(torch.nn.Module):
    """AdaLN-single modulation that emits scale/shift/gates for LTX-2."""
    def __init__(self, embedding_dim: int, embedding_coefficient: int = 6):
        super().__init__()
        self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings(
            embedding_dim,
            size_emb_dim=embedding_dim // 3,
        )
        self.silu = torch.nn.SiLU()
        self.linear = torch.nn.Linear(embedding_dim, embedding_coefficient * embedding_dim, bias=True)

    def forward(
        self,
        timestep: torch.Tensor,
        hidden_dtype: Optional[torch.dtype] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype)
        return self.linear(self.silu(embedded_timestep)), embedded_timestep


# Number of AdaLN scale/shift/gate rows. LTX-2.0 emits 6 (shift/scale/gate for
# self-attn and FFN). LTX-2.3 with cross_attention_adaln emits 3 extra rows for
# the text cross-attention shift/scale/gate.
ADALN_NUM_BASE_PARAMS = 6
ADALN_NUM_CROSS_ATTN_PARAMS = 3


def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int:
    return ADALN_NUM_BASE_PARAMS + (ADALN_NUM_CROSS_ATTN_PARAMS
                                    if cross_attention_adaln else 0)


class PixArtAlphaTextProjection(torch.nn.Module):
    """Caption projection MLP used by LTX-2."""
    def __init__(self, in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = "gelu_tanh"):
        super().__init__()
        if out_features is None:
            out_features = hidden_size
        self.linear_1 = torch.nn.Linear(in_features=in_features, out_features=hidden_size, bias=True)
        if act_fn == "gelu_tanh":
            self.act_1 = torch.nn.GELU(approximate="tanh")
        elif act_fn == "silu":
            self.act_1 = torch.nn.SiLU()
        else:
            raise ValueError(f"Unknown activation function: {act_fn}")
        self.linear_2 = torch.nn.Linear(in_features=hidden_size, out_features=out_features, bias=True)

    def forward(self, caption: torch.Tensor) -> torch.Tensor:
        hidden_states = self.linear_1(caption)
        hidden_states = self.act_1(hidden_states)
        hidden_states = self.linear_2(hidden_states)
        return hidden_states


class GELUApprox(nn.Module):
    """Linear + tanh-approximate GELU used by LTX-2 FFN."""
    def __init__(
        self,
        in_features: int,
        out_features: int,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        self.proj = ReplicatedLinear(
            in_features,
            out_features,
            quant_config=quant_config,
            prefix=f"{prefix}.fc_in",
        )
        self.act = nn.GELU(approximate="tanh")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.act(self.proj(x)[0])


class FeedForward(nn.Module):
    """LTX-2 FFN: GELUApprox -> Identity -> Linear."""
    def __init__(
        self,
        dim: int,
        dim_out: int,
        mult: int = 4,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        inner_dim = int(dim * mult)
        project_in = GELUApprox(
            dim,
            inner_dim,
            quant_config=quant_config,
            prefix=f"{prefix}.ffn",
        )
        project_out = ReplicatedLinear(
            inner_dim,
            dim_out,
            quant_config=quant_config,
            prefix=f"{prefix}.ffn.fc_out",
        )
        self.net = nn.ModuleList([project_in, nn.Identity(), project_out])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.net[0](x)
        x = self.net[1](x)
        x = self.net[2](x)[0]
        return x


class VideoLatentShape(tuple):
    """Helper for (B, C, T, H, W) latent shapes."""
    @property
    def batch(self) -> int:
        return self[0]

    @property
    def channels(self) -> int:
        return self[1]

    @property
    def frames(self) -> int:
        return self[2]

    @property
    def height(self) -> int:
        return self[3]

    @property
    def width(self) -> int:
        return self[4]

    @staticmethod
    def from_torch_shape(shape: torch.Size) -> "VideoLatentShape":
        return VideoLatentShape(shape)

    def to_torch_shape(self) -> torch.Size:
        return torch.Size(self)


class AudioLatentShape(tuple):
    """Helper for (B, C, T, F) audio latent shapes."""

    @property
    def batch(self) -> int:
        return self[0]

    @property
    def channels(self) -> int:
        return self[1]

    @property
    def frames(self) -> int:
        return self[2]

    @property
    def mel_bins(self) -> int:
        return self[3]

    def to_torch_shape(self) -> torch.Size:
        return torch.Size(self)

    @staticmethod
    def from_torch_shape(shape: torch.Size) -> "AudioLatentShape":
        return AudioLatentShape(shape)

    @staticmethod
    def from_duration(
        batch: int,
        duration: float,
        channels: int,
        mel_bins: int,
        sample_rate: int,
        hop_length: int,
        audio_latent_downsample_factor: int,
    ) -> "AudioLatentShape":
        latents_per_second = float(sample_rate) / float(
            hop_length) / float(audio_latent_downsample_factor)
        return AudioLatentShape(
            (batch, channels, round(duration * latents_per_second), mel_bins))


class VideoLatentPatchifier:
    """Patchify/unpatchify latent tokens for LTX-2."""
    def __init__(self, patch_size: int):
        self._patch_size = (1, patch_size, patch_size)

    @property
    def patch_size(self) -> Tuple[int, int, int]:
        return self._patch_size

    def get_token_count(self, tgt_shape: VideoLatentShape) -> int:
        return math.prod(tgt_shape.to_torch_shape()[2:]) // math.prod(self._patch_size)

    def patchify(self, latents: torch.Tensor) -> torch.Tensor:
        return rearrange(
            latents,
            "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)",
            p1=self._patch_size[0],
            p2=self._patch_size[1],
            p3=self._patch_size[2],
        )

    def unpatchify(self, latents: torch.Tensor, output_shape: VideoLatentShape) -> torch.Tensor:
        patch_grid_frames = output_shape.frames // self._patch_size[0]
        patch_grid_height = output_shape.height // self._patch_size[1]
        patch_grid_width = output_shape.width // self._patch_size[2]
        return rearrange(
            latents,
            "b (f h w) (c p q) -> b c f (h p) (w q)",
            f=patch_grid_frames,
            h=patch_grid_height,
            w=patch_grid_width,
            p=self._patch_size[1],
            q=self._patch_size[2],
        )

    def get_patch_grid_bounds(
        self,
        output_shape: VideoLatentShape,
        device: Optional[torch.device] = None,
    ) -> torch.Tensor:
        """Get patch grid bounds for RoPE computation.

        Args:
            output_shape: Shape of the video latent tensor
            device: Device to create tensors on
        """
        frames = output_shape.frames
        height = output_shape.height
        width = output_shape.width
        batch_size = output_shape.batch

        grid_coords = torch.meshgrid(
            torch.arange(start=0, end=frames, step=self._patch_size[0], device=device),
            torch.arange(start=0, end=height, step=self._patch_size[1], device=device),
            torch.arange(start=0, end=width, step=self._patch_size[2], device=device),
            indexing="ij",
        )

        patch_starts = torch.stack(grid_coords, dim=0)
        patch_size_delta = torch.tensor(
            self._patch_size,
            device=patch_starts.device,
            dtype=patch_starts.dtype,
        ).view(3, 1, 1, 1)
        patch_ends = patch_starts + patch_size_delta
        latent_coords = torch.stack((patch_starts, patch_ends), dim=-1)
        latent_coords = repeat(
            latent_coords,
            "c f h w bounds -> b c (f h w) bounds",
            b=batch_size,
            bounds=2,
        )
        return latent_coords


class AudioLatentPatchifier:
    """Patchify/unpatchify audio latents and compute timing bounds."""

    def __init__(
        self,
        patch_size: int,
        sample_rate: int,
        hop_length: int,
        audio_latent_downsample_factor: int,
        is_causal: bool = True,
        shift: int = 0,
    ) -> None:
        self.hop_length = hop_length
        self.sample_rate = sample_rate
        self.audio_latent_downsample_factor = audio_latent_downsample_factor
        self.is_causal = is_causal
        self.shift = shift
        self._patch_size = (1, patch_size, patch_size)

    def get_token_count(self, tgt_shape: AudioLatentShape) -> int:
        return tgt_shape.frames

    def patchify(self, latents: torch.Tensor) -> torch.Tensor:
        return rearrange(latents, "b c t f -> b t (c f)")

    def unpatchify(
        self,
        latents: torch.Tensor,
        output_shape: AudioLatentShape,
    ) -> torch.Tensor:
        return rearrange(
            latents,
            "b t (c f) -> b c t f",
            c=output_shape.channels,
            f=output_shape.mel_bins,
        )

    def get_patch_grid_bounds(
        self,
        output_shape: AudioLatentShape,
        device: Optional[torch.device] = None,
    ) -> torch.Tensor:
        """Get patch grid bounds for audio RoPE computation.

        Args:
            output_shape: Shape of the audio latent tensor
            device: Device to create tensors on
        """
        start_timings = self._get_audio_latent_time_in_sec(
            self.shift,
            output_shape.frames + self.shift,
            torch.float32,
            device,
        )
        start_timings = start_timings.unsqueeze(0).expand(output_shape.batch,
                                                          -1).unsqueeze(1)

        end_timings = self._get_audio_latent_time_in_sec(
            self.shift + 1,
            output_shape.frames + self.shift + 1,
            torch.float32,
            device,
        )
        end_timings = end_timings.unsqueeze(0).expand(output_shape.batch,
                                                      -1).unsqueeze(1)

        return torch.stack([start_timings, end_timings], dim=-1)

    def _get_audio_latent_time_in_sec(
        self,
        start_latent: int,
        end_latent: int,
        dtype: torch.dtype,
        device: Optional[torch.device],
    ) -> torch.Tensor:
        resolved_device = device or torch.device("cpu")
        audio_latent_frame = torch.arange(
            start_latent, end_latent, dtype=dtype, device=resolved_device)
        audio_mel_frame = audio_latent_frame * self.audio_latent_downsample_factor
        if self.is_causal:
            causal_offset = 1
            audio_mel_frame = (
                audio_mel_frame + causal_offset -
                self.audio_latent_downsample_factor).clamp(min=0)
        return audio_mel_frame * self.hop_length / self.sample_rate


def _get_pixel_coords(
    latent_coords: torch.Tensor,
    scale_factors: tuple[int, int, int],
    fps: float | None,
    causal_fix: bool = True,
) -> torch.Tensor:
    broadcast_shape = [1] * latent_coords.ndim
    broadcast_shape[1] = -1
    scale_tensor = torch.tensor(
        scale_factors,
        device=latent_coords.device,
        dtype=torch.float32,
    ).view(*broadcast_shape)
    pixel_coords = latent_coords.to(torch.float32) * scale_tensor
    if causal_fix:
        pixel_coords[:, 0, ...] = (
            pixel_coords[:, 0, ...] + 1 - scale_factors[0]).clamp(min=0)
    if fps:
        pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps
    return pixel_coords


def _to_denoised(
    sample: torch.Tensor,
    velocity: torch.Tensor,
    sigma: torch.Tensor,
    calc_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
    if isinstance(sigma, torch.Tensor):
        sigma = sigma.to(calc_dtype)
    while sigma.ndim < sample.ndim:
        sigma = sigma.unsqueeze(-1)
    return (sample.to(calc_dtype) - velocity.to(calc_dtype) * sigma).to(sample.dtype)


def _debug_block_log_line(message: str) -> None:
    if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") != "1":
        return
    log_path = os.getenv("LTX2_PIPELINE_DEBUG_PATH", "")
    if not log_path:
        return
    log_dir = os.path.dirname(log_path)
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)
    with open(log_path, "a", encoding="utf-8") as f:
        f.write(message + "\n")


def _debug_transformer_args(prefix: str, args: "TransformerArgs | None") -> None:
    if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") != "1" or args is None:
        return
    pe_cos, pe_sin = args.positional_embeddings
    cross_cos = None
    cross_sin = None
    if args.cross_positional_embeddings is not None:
        cross_cos, cross_sin = args.cross_positional_embeddings
    mask = args.context_mask
    if mask is None:
        mask_summary = "mask=None"
    else:
        finite = torch.isfinite(mask)
        finite_sum = mask[finite].sum().item() if finite.any() else 0.0
        mask_summary = (
            f"mask_min={mask.min().item():.6f} "
            f"mask_max={mask.max().item():.6f} "
            f"mask_finite_sum={finite_sum:.6f} "
            f"mask_finite_count={finite.sum().item()}"
        )
    _debug_block_log_line(
        f"{prefix}:x_sum={args.x.float().sum().item():.6f} "
        f"context_sum={args.context.float().sum().item():.6f} "
        f"t_sum={args.timesteps.float().sum().item():.6f} "
        f"emb_sum={args.embedded_timestep.float().sum().item():.6f} "
        f"pe_cos_sum={pe_cos.float().sum().item():.6f} "
        f"pe_sin_sum={pe_sin.float().sum().item():.6f} "
        f"cross_pe_cos_sum={(cross_cos.float().sum().item() if cross_cos is not None else 0.0):.6f} "
        f"cross_pe_sin_sum={(cross_sin.float().sum().item() if cross_sin is not None else 0.0):.6f} "
        f"{mask_summary}"
    )


class LTXRopeType(Enum):
    """LTX-2 rotary variants (interleaved vs split)."""
    INTERLEAVED = "interleaved"
    SPLIT = "split"


DEFAULT_LTX2_SCALE_FACTORS = (8, 32, 32)
DEFAULT_LTX2_AUDIO_CHANNELS = 8
DEFAULT_LTX2_AUDIO_MEL_BINS = 16
DEFAULT_LTX2_AUDIO_SAMPLE_RATE = 16000
DEFAULT_LTX2_AUDIO_HOP_LENGTH = 160
DEFAULT_LTX2_AUDIO_DOWNSAMPLE = 4
# The vocoder upsamples from mel spectrograms (16kHz sample rate) to 24kHz audio
DEFAULT_LTX2_VOCODER_OUTPUT_SAMPLE_RATE = 24000


def apply_ltx_rotary_emb(
    input_tensor: torch.Tensor,
    freqs_cis: tuple[torch.Tensor, torch.Tensor],
    rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
) -> torch.Tensor:
    """Apply LTX-2 rotary embeddings to a tensor."""
    if rope_type == LTXRopeType.INTERLEAVED:
        return _apply_ltx_interleaved_rotary_emb(input_tensor, *freqs_cis)
    if rope_type == LTXRopeType.SPLIT:
        return _apply_ltx_split_rotary_emb(input_tensor, *freqs_cis)
    raise ValueError(f"Invalid rope type: {rope_type}")


def apply_ltx_rotary_emb_4d(
    input_tensor: torch.Tensor,
    freqs_cis: tuple[torch.Tensor, torch.Tensor],
    rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
) -> torch.Tensor:
    """Apply LTX-2 rotary embeddings to a 4D tensor [B, H, T, D].
    
    This is used for applying RoPE after all-to-all in distributed attention,
    where the tensor is already in [B, H, T, D] format.
    """
    cos_freqs, sin_freqs = freqs_cis
    if rope_type == LTXRopeType.INTERLEAVED:
        # For interleaved, cos/sin have shape [B, T, inner_dim]
        # Need to reshape to [B, H, T, D] format
        # Actually interleaved doesn't have per-head rotations, so we broadcast
        t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2)
        t1, t2 = t_dup.unbind(dim=-1)
        t_dup = torch.stack((-t2, t1), dim=-1)
        input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)")
        return input_tensor * cos_freqs + input_tensor_rot * sin_freqs
    if rope_type == LTXRopeType.SPLIT:
        # For split, cos/sin already have shape [B, H, T, D/2]
        # input_tensor is [B, H, T, D]
        split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2)
        first_half_input = split_input[..., :1, :]
        second_half_input = split_input[..., 1:, :]

        output = split_input * cos_freqs.unsqueeze(-2)
        first_half_output = output[..., :1, :]
        second_half_output = output[..., 1:, :]

        first_half_output.addcmul_(-sin_freqs.unsqueeze(-2), second_half_input)
        second_half_output.addcmul_(sin_freqs.unsqueeze(-2), first_half_input)

        return rearrange(output, "... d r -> ... (d r)")
    raise ValueError(f"Invalid rope type: {rope_type}")


def _apply_ltx_interleaved_rotary_emb(
    input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor
) -> torch.Tensor:
    t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2)
    t1, t2 = t_dup.unbind(dim=-1)
    t_dup = torch.stack((-t2, t1), dim=-1)
    input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)")
    return input_tensor * cos_freqs + input_tensor_rot * sin_freqs


def _apply_ltx_split_rotary_emb(
    input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor
) -> torch.Tensor:
    needs_reshape = False
    if input_tensor.ndim != 4 and cos_freqs.ndim == 4:
        b, h, t, _ = cos_freqs.shape
        input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2)
        needs_reshape = True

    split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2)
    first_half_input = split_input[..., :1, :]
    second_half_input = split_input[..., 1:, :]

    output = split_input * cos_freqs.unsqueeze(-2)
    first_half_output = output[..., :1, :]
    second_half_output = output[..., 1:, :]

    first_half_output.addcmul_(-sin_freqs.unsqueeze(-2), second_half_input)
    second_half_output.addcmul_(sin_freqs.unsqueeze(-2), first_half_input)

    output = rearrange(output, "... d r -> ... (d r)")
    if needs_reshape:
        output = output.swapaxes(1, 2).reshape(b, t, -1)
    return output


@functools.lru_cache(maxsize=5)
def generate_ltx_freq_grid_np(
    positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
) -> torch.Tensor:
    """Generate LTX-2 rotary frequencies with high-precision numpy."""
    theta = positional_embedding_theta
    start = 1
    end = theta
    n_elem = 2 * positional_embedding_max_pos_count
    pow_indices = np.power(
        theta,
        np.linspace(
            np.log(start) / np.log(theta),
            np.log(end) / np.log(theta),
            inner_dim // n_elem,
            dtype=np.float64,
        ),
    )
    return torch.tensor(pow_indices * math.pi / 2, dtype=torch.float32)


@functools.lru_cache(maxsize=5)
def generate_ltx_freq_grid_pytorch(
    positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int
) -> torch.Tensor:
    """Generate LTX-2 rotary frequencies in torch for speed."""
    theta = positional_embedding_theta
    start = 1
    end = theta
    n_elem = 2 * positional_embedding_max_pos_count
    indices = theta ** (
        torch.linspace(
            math.log(start, theta),
            math.log(end, theta),
            inner_dim // n_elem,
            dtype=torch.float32,
        )
    )
    indices = indices.to(dtype=torch.float32)
    return indices * math.pi / 2


def _ltx_get_fractional_positions(indices_grid: torch.Tensor, max_pos: list[int]) -> torch.Tensor:
    n_pos_dims = indices_grid.shape[1]
    if n_pos_dims != len(max_pos):
        raise ValueError(
            f"Number of position dimensions ({n_pos_dims}) must match max_pos length ({len(max_pos)})"
        )
    return torch.stack(
        [indices_grid[:, i] / max_pos[i] for i in range(n_pos_dims)],
        dim=-1,
    )


def _ltx_generate_freqs(
    indices: torch.Tensor, indices_grid: torch.Tensor, max_pos: list[int], use_middle_indices_grid: bool
) -> torch.Tensor:
    if use_middle_indices_grid:
        indices_grid_start, indices_grid_end = indices_grid[..., 0], indices_grid[..., 1]
        indices_grid = (indices_grid_start + indices_grid_end) / 2.0
    elif len(indices_grid.shape) == 4:
        indices_grid = indices_grid[..., 0]

    fractional_positions = _ltx_get_fractional_positions(indices_grid, max_pos)
    indices = indices.to(device=fractional_positions.device)
    freqs = (indices * (fractional_positions.unsqueeze(-1) * 2 - 1)).transpose(-1, -2).flatten(2)
    return freqs


def _ltx_split_freqs_cis(
    freqs: torch.Tensor, pad_size: int, num_attention_heads: int
) -> tuple[torch.Tensor, torch.Tensor]:
    cos_freq = freqs.cos()
    sin_freq = freqs.sin()

    if pad_size != 0:
        cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
        sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size])
        cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1)
        sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1)

    b = cos_freq.shape[0]
    t = cos_freq.shape[1]
    cos_freq = cos_freq.reshape(b, t, num_attention_heads, -1)
    sin_freq = sin_freq.reshape(b, t, num_attention_heads, -1)
    cos_freq = torch.swapaxes(cos_freq, 1, 2)
    sin_freq = torch.swapaxes(sin_freq, 1, 2)
    return cos_freq, sin_freq


def _ltx_interleaved_freqs_cis(freqs: torch.Tensor, pad_size: int) -> tuple[torch.Tensor, torch.Tensor]:
    cos_freq = freqs.cos().repeat_interleave(2, dim=-1)
    sin_freq = freqs.sin().repeat_interleave(2, dim=-1)
    if pad_size != 0:
        cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
        sin_padding = torch.zeros_like(cos_freq[:, :, :pad_size])
        cos_freq = torch.cat([cos_padding, cos_freq], dim=-1)
        sin_freq = torch.cat([sin_padding, sin_freq], dim=-1)
    return cos_freq, sin_freq


def precompute_ltx_freqs_cis(
    indices_grid: torch.Tensor,
    dim: int,
    out_dtype: torch.dtype,
    theta: float = 10000.0,
    max_pos: list[int] | None = None,
    use_middle_indices_grid: bool = False,
    num_attention_heads: int = 32,
    rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
    freq_grid_generator: Callable[[float, int, int], torch.Tensor] = generate_ltx_freq_grid_pytorch,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Precompute LTX-2 rotary cos/sin grids for (t, x, y) positions."""
    if max_pos is None:
        max_pos = [20, 2048, 2048]

    indices = freq_grid_generator(theta, indices_grid.shape[1], dim)
    freqs = _ltx_generate_freqs(indices, indices_grid, max_pos, use_middle_indices_grid)

    if rope_type == LTXRopeType.SPLIT:
        expected_freqs = dim // 2
        current_freqs = freqs.shape[-1]
        pad_size = expected_freqs - current_freqs
        cos_freq, sin_freq = _ltx_split_freqs_cis(freqs, pad_size, num_attention_heads)
    else:
        n_elem = 2 * indices_grid.shape[1]
        cos_freq, sin_freq = _ltx_interleaved_freqs_cis(freqs, dim % n_elem)
    return cos_freq.to(out_dtype), sin_freq.to(out_dtype)


@dataclass(frozen=True)
class TransformerArgs:
    """Pack transformer inputs for LTX-2 blocks."""
    x: torch.Tensor
    context: torch.Tensor
    context_mask: torch.Tensor | None
    timesteps: torch.Tensor
    embedded_timestep: torch.Tensor
    positional_embeddings: torch.Tensor
    cross_positional_embeddings: torch.Tensor | None
    cross_scale_shift_timestep: torch.Tensor | None
    cross_gate_timestep: torch.Tensor | None
    enabled: bool
    # LTX-2.3 cross-attention AdaLN prompt timestep embedding (None for 2.0).
    prompt_timestep: torch.Tensor | None = None


@dataclass(frozen=True)
class Modality:
    """Lightweight modality container for LTX-2 inputs."""
    enabled: bool
    latent: torch.Tensor
    timesteps: torch.Tensor
    positions: torch.Tensor
    context: torch.Tensor
    context_mask: torch.Tensor | None = None
    # LTX-2.3 cross-attention AdaLN sigma timestep (None for 2.0).
    sigma: torch.Tensor | None = None


class TransformerArgsPreprocessor:
    """Prepare LTX-2 transformer inputs (patchify, AdaLN, rope)."""
    def __init__(  # noqa: PLR0913
        self,
        patchify_proj: torch.nn.Linear,
        adaln: AdaLayerNormSingle,
        caption_projection: PixArtAlphaTextProjection,
        inner_dim: int,
        max_pos: list[int],
        num_attention_heads: int,
        use_middle_indices_grid: bool,
        timestep_scale_multiplier: int,
        double_precision_rope: bool,
        positional_embedding_theta: float,
        rope_type: LTXRopeType,
        prompt_adaln: AdaLayerNormSingle | None = None,
    ) -> None:
        self.patchify_proj = patchify_proj
        self.adaln = adaln
        self.caption_projection = caption_projection
        self.inner_dim = inner_dim
        self.max_pos = max_pos
        self.num_attention_heads = num_attention_heads
        self.use_middle_indices_grid = use_middle_indices_grid
        self.timestep_scale_multiplier = timestep_scale_multiplier
        self.double_precision_rope = double_precision_rope
        self.positional_embedding_theta = positional_embedding_theta
        self.rope_type = rope_type
        # LTX-2.3 cross-attention AdaLN prompt timestep embedder (None for 2.0).
        self.prompt_adaln = prompt_adaln

    def _prepare_timestep(
        self,
        timestep: torch.Tensor,
        batch_size: int,
        hidden_dtype: torch.dtype,
        adaln: AdaLayerNormSingle | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        adaln = self.adaln if adaln is None else adaln
        timestep = timestep * self.timestep_scale_multiplier
        timestep, embedded_timestep = adaln(timestep.flatten(), hidden_dtype=hidden_dtype)
        timestep = timestep.view(batch_size, -1, timestep.shape[-1])
        embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.shape[-1])
        return timestep, embedded_timestep

    def _prepare_context(
        self,
        context: torch.Tensor,
        x: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        batch_size = x.shape[0]
        if context.device != x.device:
            context = context.to(x.device)
        if context.dtype != x.dtype:
            context = context.to(x.dtype)
        if attention_mask is not None and attention_mask.device != x.device:
            attention_mask = attention_mask.to(x.device)
        # When ``caption_proj_before_connector`` is set the caption projection
        # lives in the text encoder, so it is None here and we skip it.
        if self.caption_projection is not None:
            context = self.caption_projection(context)
            context = context.view(batch_size, -1, x.shape[-1])
        return context, attention_mask

    def _prepare_attention_mask(self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype) -> torch.Tensor | None:
        if attention_mask is None or torch.is_floating_point(attention_mask):
            return attention_mask
        return (attention_mask - 1).to(x_dtype).reshape(
            (attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
        ) * torch.finfo(x_dtype).max

    def _prepare_positional_embeddings(
        self,
        positions: torch.Tensor,
        inner_dim: int,
        max_pos: list[int],
        use_middle_indices_grid: bool,
        num_attention_heads: int,
        x_dtype: torch.dtype,
    ) -> torch.Tensor:
        if self.double_precision_rope:
            freq_grid_generator = generate_ltx_freq_grid_np
        else:
            freq_grid_generator = generate_ltx_freq_grid_pytorch
        return precompute_ltx_freqs_cis(
            positions,
            dim=inner_dim,
            out_dtype=x_dtype,
            theta=self.positional_embedding_theta,
            max_pos=max_pos,
            use_middle_indices_grid=use_middle_indices_grid,
            num_attention_heads=num_attention_heads,
            rope_type=self.rope_type,
            freq_grid_generator=freq_grid_generator,
        )

    def prepare(self, modality: Modality) -> TransformerArgs:
        # `rearrange`-produced latent tokens can be non-contiguous with large
        # leading strides. Materialize a contiguous view before Linear to avoid
        # invalid BF16 GEMM descriptors on some CUDA/cuBLAS stacks.
        latent = modality.latent
        if not latent.is_contiguous():
            latent = latent.contiguous()
        x = self.patchify_proj(latent)
        timestep, embedded_timestep = self._prepare_timestep(
            modality.timesteps, x.shape[0], modality.latent.dtype, self.adaln)
        prompt_timestep = None
        if self.prompt_adaln is not None and modality.sigma is not None:
            prompt_timestep, _ = self._prepare_timestep(
                modality.sigma,
                x.shape[0],
                modality.latent.dtype,
                self.prompt_adaln,
            )
        context, attention_mask = self._prepare_context(modality.context, x, modality.context_mask)
        attention_mask = self._prepare_attention_mask(attention_mask, modality.latent.dtype)
        pe = self._prepare_positional_embeddings(
            positions=modality.positions,
            inner_dim=self.inner_dim,
            max_pos=self.max_pos,
            use_middle_indices_grid=self.use_middle_indices_grid,
            num_attention_heads=self.num_attention_heads,
            x_dtype=modality.latent.dtype,
        )
        return TransformerArgs(
            x=x,
            context=context,
            context_mask=attention_mask,
            timesteps=timestep,
            embedded_timestep=embedded_timestep,
            positional_embeddings=pe,
            cross_positional_embeddings=None,
            cross_scale_shift_timestep=None,
            cross_gate_timestep=None,
            enabled=modality.enabled,
            prompt_timestep=prompt_timestep,
        )


class MultiModalTransformerArgsPreprocessor:
    """Prepare transformer args for audio/video cross-attention paths."""

    def __init__(  # noqa: PLR0913
        self,
        patchify_proj: torch.nn.Linear,
        adaln: AdaLayerNormSingle,
        caption_projection: PixArtAlphaTextProjection,
        cross_scale_shift_adaln: AdaLayerNormSingle,
        cross_gate_adaln: AdaLayerNormSingle,
        inner_dim: int,
        max_pos: list[int],
        num_attention_heads: int,
        cross_pe_max_pos: int,
        use_middle_indices_grid: bool,
        audio_cross_attention_dim: int,
        timestep_scale_multiplier: int,
        double_precision_rope: bool,
        positional_embedding_theta: float,
        rope_type: LTXRopeType,
        av_ca_timestep_scale_multiplier: int,
        prompt_adaln: AdaLayerNormSingle | None = None,
    ) -> None:
        self.simple_preprocessor = TransformerArgsPreprocessor(
            patchify_proj=patchify_proj,
            adaln=adaln,
            caption_projection=caption_projection,
            inner_dim=inner_dim,
            max_pos=max_pos,
            num_attention_heads=num_attention_heads,
            use_middle_indices_grid=use_middle_indices_grid,
            timestep_scale_multiplier=timestep_scale_multiplier,
            double_precision_rope=double_precision_rope,
            positional_embedding_theta=positional_embedding_theta,
            rope_type=rope_type,
            prompt_adaln=prompt_adaln,
        )
        self.cross_scale_shift_adaln = cross_scale_shift_adaln
        self.cross_gate_adaln = cross_gate_adaln
        self.cross_pe_max_pos = cross_pe_max_pos
        self.audio_cross_attention_dim = audio_cross_attention_dim
        self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier

    def prepare(self, modality: Modality) -> TransformerArgs:
        transformer_args = self.simple_preprocessor.prepare(modality)
        cross_pe = self.simple_preprocessor._prepare_positional_embeddings(
            positions=modality.positions[:, 0:1, :],
            inner_dim=self.audio_cross_attention_dim,
            max_pos=[self.cross_pe_max_pos],
            use_middle_indices_grid=True,
            num_attention_heads=self.simple_preprocessor.num_attention_heads,
            x_dtype=modality.latent.dtype,
        )

        cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep(
            timestep=modality.timesteps,
            timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier,
            batch_size=transformer_args.x.shape[0],
            hidden_dtype=modality.latent.dtype,
        )

        return replace(
            transformer_args,
            cross_positional_embeddings=cross_pe,
            cross_scale_shift_timestep=cross_scale_shift_timestep,
            cross_gate_timestep=cross_gate_timestep,
        )

    def _prepare_cross_attention_timestep(
        self,
        timestep: torch.Tensor,
        timestep_scale_multiplier: int,
        batch_size: int,
        hidden_dtype: torch.dtype,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        timestep = timestep * timestep_scale_multiplier

        av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier

        scale_shift_timestep, _ = self.cross_scale_shift_adaln(
            timestep.flatten(),
            hidden_dtype=hidden_dtype,
        )
        scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1])
        gate_noise_timestep, _ = self.cross_gate_adaln(
            timestep.flatten() * av_ca_factor,
            hidden_dtype=hidden_dtype,
        )
        gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1])

        return scale_shift_timestep, gate_noise_timestep


@dataclass
class TransformerConfig:
    """Attention/FFN dims for LTX-2 transformer blocks."""
    dim: int
    heads: int
    d_head: int
    context_dim: int
    # LTX-2.3 gated extensions (default OFF == LTX-2.0 behavior).
    apply_gated_attention: bool = False
    cross_attention_adaln: bool = False


class LTXDistributedAttention(DistributedAttention):
    """LTX-2 specialized DistributedAttention that handles LTX-style RoPE internally."""

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        rope_type: LTXRopeType,
        num_kv_heads: int | None = None,
        softmax_scale: float | None = None,
        causal: bool = False,
        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
        super().__init__(
            num_heads=num_heads,
            head_size=head_size,
            num_kv_heads=num_kv_heads,
            softmax_scale=softmax_scale,
            causal=causal,
            supported_attention_backends=supported_attention_backends,
            prefix=prefix,
            **extra_impl_args,
        )
        self.rope_type = rope_type

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        original_seq_len: int,
        replicated_q: torch.Tensor | None = None,
        replicated_k: torch.Tensor | None = None,
        replicated_v: torch.Tensor | None = None,
        gate_compress: torch.Tensor | None = None,
        ltx_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Forward pass with LTX-2 style RoPE application.

        Args:
            q: Query tensor [batch_size, seq_len, num_heads, head_dim]
            k: Key tensor [batch_size, seq_len, num_heads, head_dim]
            v: Value tensor [batch_size, seq_len, num_heads, head_dim]
            replicated_q: Replicated query tensor for text tokens
            replicated_k: Replicated key tensor
            replicated_v: Replicated value tensor
            original_seq_len: Original (unpadded) full sequence length
            ltx_freqs_cis: LTX-2 style RoPE (cos, sin) with shape [B, H, T, D]

        Returns:
            Tuple of output tensor and optional replicated output
        """
        assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors"
        batch_size, _, num_heads, _ = q.shape
        local_rank = get_sp_parallel_rank()
        world_size = get_sp_world_size()

        forward_context: ForwardContext = get_forward_context()
        ctx_attn_metadata = forward_context.attn_metadata

        use_vsa = self.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN

        if use_vsa:
            assert (
                replicated_q is None and replicated_k is None
                and replicated_v is None
            ), "Replicated QKV is not supported for VSA now"
            if gate_compress is None:
                raise ValueError(
                    "gate_compress must be provided when using VIDEO_SPARSE_ATTN"
                )

            qkvg = torch.cat([q, k, v, gate_compress],
                             dim=0)  # [4*batch, seq_len, heads, head_dim]
            qkvg = sequence_model_parallel_all_to_all_4D(qkvg,
                                                         scatter_dim=2,
                                                         gather_dim=1)
            pad_seq_len = qkvg.shape[1] - original_seq_len
            qkvg = qkvg[:, :original_seq_len, :, :]

            if ltx_freqs_cis is not None:
                cos, sin = ltx_freqs_cis
                heads_per_rank = num_heads // world_size
                head_start = local_rank * heads_per_rank
                head_end = head_start + heads_per_rank
                cos_local = cos[:, head_start:head_end, :original_seq_len, :]
                sin_local = sin[:, head_start:head_end, :original_seq_len, :]
                qk_part = qkvg[:batch_size * 2].transpose(1, 2)
                qk_part = apply_ltx_rotary_emb_4d(
                    qk_part, (cos_local, sin_local), self.rope_type)
                qkvg[:batch_size * 2] = qk_part.transpose(1, 2)

            qkvg = self.attn_impl.preprocess_qkv(qkvg, ctx_attn_metadata)
            q, k, v, gate_compress = qkvg.chunk(4, dim=0)
            output = self.attn_impl.forward(  # type: ignore[call-arg]
                q, k, v, gate_compress, ctx_attn_metadata)
            output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)
            output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len))

            output = sequence_model_parallel_all_to_all_4D(output,
                                                           scatter_dim=1,
                                                           gather_dim=2)
            return output, None

        # Stack QKV
        qkv = torch.cat([q, k, v], dim=0)  # [3*batch, seq_len, num_heads, head_dim]

        # Redistribute heads across sequence dimension
        qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1)

        # After all-to-all, each rank has the full sequence but only a subset of heads
        pad_seq_len = qkv.shape[1] - original_seq_len
        qkv = qkv[:, :original_seq_len, :, :]

        # Apply LTX-2 style RoPE after all-to-all (when we have full sequence)
        if ltx_freqs_cis is not None:
            cos, sin = ltx_freqs_cis
            heads_per_rank = num_heads // world_size
            head_start = local_rank * heads_per_rank
            head_end = head_start + heads_per_rank
            # Slice to this rank's heads: [B, H/SP, T, D]
            cos_local = cos[:, head_start:head_end, :original_seq_len, :]
            sin_local = sin[:, head_start:head_end, :original_seq_len, :]

            # Apply RoPE to Q and K together (first 2*batch_size in dim 0)
            qk_part = qkv[:batch_size * 2]
            # Transpose to [2*B, H/SP, T, D] for LTX RoPE application
            qk_part = qk_part.transpose(1, 2)
            qk_part = apply_ltx_rotary_emb_4d(qk_part, (cos_local, sin_local), self.rope_type)
            # Transpose back to [2*B, T, H/SP, D]
            qkv[:batch_size * 2] = qk_part.transpose(1, 2)

        # Apply backend-specific preprocess_qkv
        qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata)

        # Concatenate with replicated QKV if provided
        if replicated_q is not None:
            assert replicated_k is not None and replicated_v is not None
            replicated_qkv = torch.cat(
                [replicated_q, replicated_k, replicated_v],
                dim=0)  # [3, seq_len, num_heads, head_dim]
            heads_per_rank = num_heads // world_size
            replicated_qkv = replicated_qkv[:, :, local_rank *
                                            heads_per_rank:(local_rank + 1) *
                                            heads_per_rank]
            qkv = torch.cat([qkv, replicated_qkv], dim=1)

        q, k, v = qkv.chunk(3, dim=0)

        output = self.attn_impl.forward(q, k, v, ctx_attn_metadata)

        # Redistribute back if using sequence parallelism
        replicated_output = None
        if replicated_q is not None:
            split_idx = original_seq_len
            replicated_output = output[:, split_idx:]
            output = output[:, :split_idx]
            replicated_output = sequence_model_parallel_all_gather(
                replicated_output.contiguous(), dim=2)

        # Apply backend-specific postprocess_output
        output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)

        output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len))

        output = sequence_model_parallel_all_to_all_4D(output, scatter_dim=1, gather_dim=2)

        return output, replicated_output


class LTXLocalAttention(LocalAttention):
    """LTX-2 specialized LocalAttention that handles LTX-style RoPE internally."""

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        rope_type: LTXRopeType,
        num_kv_heads: int | None = None,
        softmax_scale: float | None = None,
        causal: bool = False,
        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
        **extra_impl_args,
    ) -> None:
        super().__init__(
            num_heads=num_heads,
            head_size=head_size,
            num_kv_heads=num_kv_heads,
            softmax_scale=softmax_scale,
            causal=causal,
            supported_attention_backends=supported_attention_backends,
            **extra_impl_args,
        )
        self.rope_type = rope_type

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        gate_compress: torch.Tensor | None = None,
        ltx_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
        ltx_k_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> torch.Tensor:
        """Forward pass with LTX-2 style RoPE application.

        Args:
            q: Query tensor [batch_size, seq_len, num_heads, head_dim]
            k: Key tensor [batch_size, seq_len, num_heads, head_dim]
            v: Value tensor [batch_size, seq_len, num_heads, head_dim]
            ltx_freqs_cis: LTX-2 style RoPE (cos, sin) for Q with shape [B, H, T, D]
            ltx_k_freqs_cis: LTX-2 style RoPE (cos, sin) for K (if different from Q)

        Returns:
            Output tensor after local attention
        """
        assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors"

        forward_context: ForwardContext = get_forward_context()
        ctx_attn_metadata = forward_context.attn_metadata

        # Apply LTX-2 style RoPE
        if ltx_freqs_cis is not None:
            # Use separate K RoPE if provided (for cross-attention), otherwise use same as Q
            k_freqs = ltx_k_freqs_cis if ltx_k_freqs_cis is not None else ltx_freqs_cis
            # Transpose to [B, H, T, D] for RoPE application
            q = q.transpose(1, 2)
            k = k.transpose(1, 2)
            q = apply_ltx_rotary_emb_4d(q, ltx_freqs_cis, self.rope_type)
            k = apply_ltx_rotary_emb_4d(k, k_freqs, self.rope_type)
            # Transpose back to [B, T, H, D]
            q = q.transpose(1, 2)
            k = k.transpose(1, 2)

        if self.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
            if gate_compress is None:
                raise ValueError(
                    "gate_compress must be provided when using VIDEO_SPARSE_ATTN"
                )
            output = self.attn_impl.forward(  # type: ignore[call-arg]
                q, k, v, gate_compress, ctx_attn_metadata)
        else:
            output = self.attn_impl.forward(q, k, v, ctx_attn_metadata)
        return output


class LTXSelfAttention(nn.Module):
    """LTX-2 attention block with RMSNorm + FastVideo LocalAttention."""
    def __init__(
        self,
        query_dim: int,
        context_dim: int | None,
        heads: int,
        dim_head: int,
        norm_eps: float,
        rope_type: LTXRopeType,
        supported_attention_backends: tuple[AttentionBackendEnum, ...],
        apply_gated_attention: bool = False,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        inner_dim = dim_head * heads
        context_dim = query_dim if context_dim is None else context_dim

        self.heads = heads
        self.dim_head = dim_head
        self.rope_type = rope_type

        # StageAwareRMSNorm is identical to torch.nn.RMSNorm outside the
        # LTX-2.3 refine stage, so default-off behavior matches LTX-2.0.
        self.q_norm = StageAwareRMSNorm(inner_dim, eps=norm_eps)
        self.k_norm = StageAwareRMSNorm(inner_dim, eps=norm_eps)
        self.to_q = ReplicatedLinear(
            query_dim,
            inner_dim,
            bias=True,
            quant_config=quant_config,
            prefix=f"{prefix}.to_q",
        )
        self.to_k = ReplicatedLinear(
            context_dim,
            inner_dim,
            bias=True,
            quant_config=quant_config,
            prefix=f"{prefix}.to_k",
        )
        self.to_v = ReplicatedLinear(
            context_dim,
            inner_dim,
            bias=True,
            quant_config=quant_config,
            prefix=f"{prefix}.to_v",
        )
        # LTX-2.3 gated attention. DISTINCT from the VSA-QAT to_gate_compress
        # gate below: this gate multiplies the *self-attention output* by
        # 2*sigmoid(to_gate_logits(x)). None (default) == LTX-2.0 behavior.
        self.to_gate_logits: ReplicatedLinear | None = None
        if apply_gated_attention:
            self.to_gate_logits = ReplicatedLinear(
                query_dim,
                heads,
                bias=True,
                quant_config=quant_config,
                prefix=f"{prefix}.to_gate_logits",
            )
        self.to_out = nn.ModuleList([
            ReplicatedLinear(
                inner_dim,
                query_dim,
                bias=True,
                quant_config=quant_config,
                prefix=f"{prefix}.to_out",
            ),
            nn.Identity(),
        ])

        self.attn = LTXLocalAttention(
            num_heads=heads,
            head_size=dim_head,
            rope_type=rope_type,
            dropout_rate=0,
            softmax_scale=None,
            causal=False,
            supported_attention_backends=supported_attention_backends,
        )
        self.to_gate_compress: ReplicatedLinear | None = None
        if self.attn.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
            self.to_gate_compress = ReplicatedLinear(
                context_dim,
                inner_dim,
                bias=True,
                quant_config=quant_config,
                prefix=f"{prefix}.to_gate_compress",
            )
        self.attn_masked = LTXLocalAttention(
            num_heads=heads,
            head_size=dim_head,
            rope_type=rope_type,
            dropout_rate=0,
            softmax_scale=None,
            causal=False,
            supported_attention_backends=(AttentionBackendEnum.TORCH_SDPA,),
        )

    def forward(
        self,
        x: torch.Tensor,
        context: torch.Tensor | None = None,
        mask: torch.Tensor | None = None,
        pe: tuple[torch.Tensor, torch.Tensor] | None = None,
        k_pe: tuple[torch.Tensor, torch.Tensor] | None = None,
    ) -> torch.Tensor:
        context = x if context is None else context

        q_pre_quantized = (
            self.to_q.quant_method.quantize_input(x)  # type: ignore[union-attr]
            if _supports_prequantized_input(self.to_q) else None)
        kv_pre_quantized = None
        if (_supports_prequantized_input(self.to_k)
                and _supports_prequantized_input(self.to_v)):
            if context is x and q_pre_quantized is not None:
                kv_pre_quantized = q_pre_quantized
            else:
                kv_pre_quantized = self.to_k.quant_method.quantize_input(  # type: ignore[union-attr]
                    context)

        q = _linear_project_with_optional_prequant(self.to_q, x, q_pre_quantized)
        k = _linear_project_with_optional_prequant(self.to_k, context, kv_pre_quantized)
        v = _linear_project_with_optional_prequant(self.to_v, context, kv_pre_quantized)
        gate_logits = (self.to_gate_logits(x)[0]
                       if self.to_gate_logits is not None else None)
        gate_compress = (self.to_gate_compress(context)[0]
                         if self.to_gate_compress is not None else None)

        q = self.q_norm(q)
        k = self.k_norm(k)

        # RoPE is applied inside LTXLocalAttention
        b, q_len, _ = q.shape
        k_len = k.shape[1]
        q = q.view(b, q_len, self.heads, self.dim_head)
        k = k.view(b, k_len, self.heads, self.dim_head)
        v = v.view(b, k_len, self.heads, self.dim_head)
        if gate_compress is not None:
            gate_compress = gate_compress.view(b, k_len, self.heads,
                                               self.dim_head)

        if mask is not None:
            if mask.ndim == 2:
                mask = mask.unsqueeze(0)
            if mask.ndim == 3:
                mask = mask.unsqueeze(1)
            try:
                forward_ctx = get_forward_context()
                current_timestep = forward_ctx.current_timestep
                forward_batch = forward_ctx.forward_batch
            except AssertionError:
                current_timestep = 0
                forward_batch = None
            attn_metadata = SDPAMetadata(
                current_timestep=current_timestep,
                attn_mask=mask,
            )
            with set_forward_context(
                current_timestep=current_timestep,
                attn_metadata=attn_metadata,
                forward_batch=forward_batch,
            ):
                out = self.attn_masked(q,
                                       k,
                                       v,
                                       ltx_freqs_cis=pe,
                                       ltx_k_freqs_cis=k_pe)
        else:
            out = self.attn(q,
                            k,
                            v,
                            gate_compress=gate_compress,
                            ltx_freqs_cis=pe,
                            ltx_k_freqs_cis=k_pe)
        # LTX-2.3 gated attention: scale the attention output per-head by
        # 2*sigmoid(gate_logits). No-op for LTX-2.0 (gate_logits is None).
        if gate_logits is not None:
            out = out.view(b, q_len, self.heads, self.dim_head)
            gates = 2.0 * torch.sigmoid(gate_logits).to(out.dtype)
            out = out * gates.unsqueeze(-1)
        out = out.reshape(b, q_len, -1)
        out = self.to_out[0](out)[0]
        return self.to_out[1](out)


class LTXDistributedSelfAttention(nn.Module):
    """LTX-2 attention block with RMSNorm + LTXDistributedAttention for SP."""

    def __init__(
        self,
        query_dim: int,
        context_dim: int | None,
        heads: int,
        dim_head: int,
        norm_eps: float,
        rope_type: LTXRopeType,
        supported_attention_backends: tuple[AttentionBackendEnum, ...],
        apply_gated_attention: bool = False,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        inner_dim = dim_head * heads
        context_dim = query_dim if context_dim is None else context_dim

        self.heads = heads
        self.dim_head = dim_head
        self.rope_type = rope_type

        # StageAwareRMSNorm == torch.nn.RMSNorm outside the LTX-2.3 refine
        # stage, preserving LTX-2.0 numerics by default.
        self.q_norm = StageAwareRMSNorm(inner_dim, eps=norm_eps)
        self.k_norm = StageAwareRMSNorm(inner_dim, eps=norm_eps)
        self.to_q = ReplicatedLinear(
            query_dim,
            inner_dim,
            bias=True,
            quant_config=quant_config,
            prefix=f"{prefix}.to_q",
        )
        self.to_k = ReplicatedLinear(
            context_dim,
            inner_dim,
            bias=True,
            quant_config=quant_config,
            prefix=f"{prefix}.to_k",
        )
        self.to_v = ReplicatedLinear(
            context_dim,
            inner_dim,
            bias=True,
            quant_config=quant_config,
            prefix=f"{prefix}.to_v",
        )
        # LTX-2.3 gated attention (distinct from VSA-QAT to_gate_compress).
        self.to_gate_logits: ReplicatedLinear | None = None
        if apply_gated_attention:
            self.to_gate_logits = ReplicatedLinear(
                query_dim,
                heads,
                bias=True,
                quant_config=quant_config,
                prefix=f"{prefix}.to_gate_logits",
            )
        self.to_out = nn.ModuleList([
            ReplicatedLinear(
                inner_dim,
                query_dim,
                bias=True,
                quant_config=quant_config,
                prefix=f"{prefix}.to_out",
            ),
            nn.Identity(),
        ])

        self.attn = LTXDistributedAttention(
            num_heads=heads,
            head_size=dim_head,
            rope_type=rope_type,
            causal=False,
            supported_attention_backends=supported_attention_backends,
            prefix=f"{prefix}.attn",
        )
        self.to_gate_compress: ReplicatedLinear | None = None
        if self.attn.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
            self.to_gate_compress = ReplicatedLinear(
                context_dim,
                inner_dim,
                bias=True,
                quant_config=quant_config,
                prefix=f"{prefix}.to_gate_compress",
            )

    def forward(
        self,
        x: torch.Tensor,
        context: torch.Tensor | None = None,
        pe: tuple[torch.Tensor, torch.Tensor] | None = None,
        k_pe: tuple[torch.Tensor, torch.Tensor] | None = None,
        original_seq_len: int | None = None,
    ) -> torch.Tensor:
        """Forward pass for distributed self-attention.

        Args:
            x: Query tensor [B, L, C]
            context: Key/Value tensor [B, L, C], or None for self-attention
            pe: Rotary position embeddings for Q (cos, sin) with shape [B, H, T_full, D]
                NOTE: For SP, this must be the FULL sequence RoPE, not sharded!
            k_pe: Rotary position embeddings for K (cos, sin), or None to use pe
            original_seq_len: Original (unpadded) full sequence length
        """
        context = x if context is None else context

        q_pre_quantized = (
            self.to_q.quant_method.quantize_input(x)  # type: ignore[union-attr]
            if _supports_prequantized_input(self.to_q) else None)
        kv_pre_quantized = None
        if (_supports_prequantized_input(self.to_k)
                and _supports_prequantized_input(self.to_v)):
            if context is x and q_pre_quantized is not None:
                kv_pre_quantized = q_pre_quantized
            else:
                kv_pre_quantized = self.to_k.quant_method.quantize_input(  # type: ignore[union-attr]
                    context)

        q = _linear_project_with_optional_prequant(self.to_q, x, q_pre_quantized)
        k = _linear_project_with_optional_prequant(self.to_k, context, kv_pre_quantized)
        v = _linear_project_with_optional_prequant(self.to_v, context, kv_pre_quantized)
        gate_logits = (self.to_gate_logits(x)[0]
                       if self.to_gate_logits is not None else None)
        gate_compress = (self.to_gate_compress(context)[0]
                         if self.to_gate_compress is not None else None)

        q = self.q_norm(q)
        k = self.k_norm(k)

        # RoPE is applied inside LTXDistributedAttention AFTER the all-to-all,
        # when each rank has the full sequence (but subset of heads).
        b, q_len, _ = q.shape
        k_len = k.shape[1]
        q = q.view(b, q_len, self.heads, self.dim_head)
        k = k.view(b, k_len, self.heads, self.dim_head)
        v = v.view(b, k_len, self.heads, self.dim_head)
        if gate_compress is not None:
            gate_compress = gate_compress.view(b, k_len, self.heads,
                                               self.dim_head)

        # Pass full RoPE to distributed attention - it will apply after all-to-all
        # and slice to this rank's heads
        if original_seq_len is None:
            raise ValueError("original_seq_len must be provided for distributed attention")
        out, _ = self.attn(
            q,
            k,
            v,
            original_seq_len,
            gate_compress=gate_compress,
            ltx_freqs_cis=pe,
        )

        # LTX-2.3 gated attention (no-op for LTX-2.0).
        if gate_logits is not None:
            out = out.view(b, q_len, self.heads, self.dim_head)
            gates = 2.0 * torch.sigmoid(gate_logits).to(out.dtype)
            out = out * gates.unsqueeze(-1)
        out = out.reshape(b, q_len, -1)
        out = self.to_out[0](out)[0]
        return self.to_out[1](out)


class BasicAVTransformerBlock(torch.nn.Module):
    """LTX-2 transformer block (audio/video + cross attention + FFN)."""

    def __init__(
        self,
        idx: int,
        video: TransformerConfig | None = None,
        audio: TransformerConfig | None = None,
        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
        norm_eps: float = 1e-6,
        use_distributed_attention: bool = False,
        cross_attention_adaln: bool = False,
        stg_block_idx: int = 29,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        self.idx = idx
        self.use_distributed_attention = use_distributed_attention
        # LTX-2.3 cross-attention AdaLN + per-sample STG (defaults reproduce 2.0).
        self.cross_attention_adaln = cross_attention_adaln
        self.stg_block_idx = stg_block_idx

        # Choose attention class based on SP mode
        # Self-attention and audio-video cross-attention use DistributedAttention when SP > 1
        # Text cross-attention uses LocalAttention (text embeddings are replicated)
        SelfAttnCls = LTXDistributedSelfAttention if use_distributed_attention else LTXSelfAttention
        CrossAttnCls = LTXSelfAttention  # Text cross-attention is always local
        video_self_attn_backends = (
            AttentionBackendEnum.FLASH_ATTN,
            AttentionBackendEnum.TORCH_SDPA,
            AttentionBackendEnum.VIDEO_SPARSE_ATTN,
        )
        dense_attn_backends = (
            AttentionBackendEnum.FLASH_ATTN,
            AttentionBackendEnum.TORCH_SDPA,
        )

        if video is not None:
            # Video self-attention - use distributed when SP > 1
            self.attn1 = SelfAttnCls(
                query_dim=video.dim,
                context_dim=None,
                heads=video.heads,
                dim_head=video.d_head,
                norm_eps=norm_eps,
                rope_type=rope_type,
                supported_attention_backends=video_self_attn_backends,
                apply_gated_attention=video.apply_gated_attention,
                quant_config=quant_config,
                prefix=f"{prefix}.blocks.{idx}.attn1",
            )
            # Text cross-attention - always local (text is replicated)
            self.attn2 = CrossAttnCls(
                query_dim=video.dim,
                context_dim=video.context_dim,
                heads=video.heads,
                dim_head=video.d_head,
                norm_eps=norm_eps,
                rope_type=rope_type,
                supported_attention_backends=dense_attn_backends,
                apply_gated_attention=video.apply_gated_attention,
                quant_config=quant_config,
                prefix=f"{prefix}.blocks.{idx}.attn2",
            )
            self.ff = FeedForward(
                video.dim,
                dim_out=video.dim,
                quant_config=quant_config,
                prefix=f"{prefix}.blocks.{idx}",
            )
            # 6 rows for LTX-2.0; 9 rows when cross_attention_adaln adds the
            # text cross-attention shift/scale/gate.
            video_sst_size = adaln_embedding_coefficient(cross_attention_adaln)
            self.scale_shift_table = torch.nn.Parameter(
                torch.empty(video_sst_size, video.dim))

        if audio is not None:
            # Audio self-attention - use distributed when SP > 1
            self.audio_attn1 = SelfAttnCls(
                query_dim=audio.dim,
                context_dim=None,
                heads=audio.heads,
                dim_head=audio.d_head,
                norm_eps=norm_eps,
                rope_type=rope_type,
                supported_attention_backends=dense_attn_backends,
                apply_gated_attention=audio.apply_gated_attention,
                quant_config=quant_config,
                prefix=f"{prefix}.blocks.{idx}.audio_attn1",
            )
            # Text cross-attention - always local (text is replicated)
            self.audio_attn2 = CrossAttnCls(
                query_dim=audio.dim,
                context_dim=audio.context_dim,
                heads=audio.heads,
                dim_head=audio.d_head,
                norm_eps=norm_eps,
                rope_type=rope_type,
                supported_attention_backends=dense_attn_backends,
                apply_gated_attention=audio.apply_gated_attention,
                quant_config=quant_config,
                prefix=f"{prefix}.blocks.{idx}.audio_attn2",
            )
            self.audio_ff = FeedForward(
                audio.dim,
                dim_out=audio.dim,
                quant_config=quant_config,
                prefix=f"{prefix}.blocks.{idx}.audio",
            )
            audio_sst_size = adaln_embedding_coefficient(cross_attention_adaln)
            self.audio_scale_shift_table = torch.nn.Parameter(
                torch.empty(audio_sst_size, audio.dim))

        if audio is not None and video is not None:
            # Audio-to-video cross-attention
            # Uses local attention - context is gathered from all SP ranks in forward()
            self.audio_to_video_attn = LTXSelfAttention(
                query_dim=video.dim,
                context_dim=audio.dim,
                heads=audio.heads,
                dim_head=audio.d_head,
                norm_eps=norm_eps,
                rope_type=rope_type,
                supported_attention_backends=dense_attn_backends,
                apply_gated_attention=video.apply_gated_attention,
                quant_config=quant_config,
                prefix=f"{prefix}.blocks.{idx}.audio_to_video_attn",
            )
            # Video-to-audio cross-attention
            # Uses local attention - context is gathered from all SP ranks in forward()
            self.video_to_audio_attn = LTXSelfAttention(
                query_dim=audio.dim,
                context_dim=video.dim,
                heads=audio.heads,
                dim_head=audio.d_head,
                norm_eps=norm_eps,
                rope_type=rope_type,
                supported_attention_backends=dense_attn_backends,
                apply_gated_attention=audio.apply_gated_attention,
                quant_config=quant_config,
                prefix=f"{prefix}.blocks.{idx}.video_to_audio_attn",
            )
            self.scale_shift_table_a2v_ca_audio = torch.nn.Parameter(torch.empty(5, audio.dim))
            self.scale_shift_table_a2v_ca_video = torch.nn.Parameter(torch.empty(5, video.dim))

        # LTX-2.3 cross-attention AdaLN prompt scale/shift tables (only created
        # when cross_attention_adaln is enabled, so absent for LTX-2.0).
        if self.cross_attention_adaln and video is not None:
            self.prompt_scale_shift_table = torch.nn.Parameter(
                torch.empty(2, video.dim))
        if self.cross_attention_adaln and audio is not None:
            self.audio_prompt_scale_shift_table = torch.nn.Parameter(
                torch.empty(2, audio.dim))

        self.norm_eps = norm_eps

    def _register_fsdp_backward_hooks_on_output(self, vx, ax):
        """Register backward hooks on output tensors to trigger FSDP2 unshard.
        
        FSDP2's module-level backward hooks don't fire when the module returns 
        dataclass outputs. We must register hooks directly on the output tensors.
        """
        if not hasattr(self, 'unshard'):
            return  # Not wrapped by FSDP2
            
        def make_unshard_hook():
            def hook(grad):
                self.unshard()
                return grad
            return hook
        
        if vx is not None and vx.requires_grad:
            vx.register_hook(make_unshard_hook())
        if ax is not None and ax.requires_grad:
            ax.register_hook(make_unshard_hook())

    def get_ada_values(
        self, scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, indices: slice
    ) -> tuple[torch.Tensor, ...]:
        num_ada_params = scale_shift_table.shape[0]
        ada_values = (
            scale_shift_table[indices].unsqueeze(0).unsqueeze(0).to(device=timestep.device, dtype=timestep.dtype)
            + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[:, :, indices, :]
        ).unbind(dim=2)
        return ada_values

    def get_av_ca_ada_values(
        self,
        scale_shift_table: torch.Tensor,
        batch_size: int,
        scale_shift_timestep: torch.Tensor,
        gate_timestep: torch.Tensor,
        num_scale_shift_values: int = 4,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        scale_shift_ada_values = self.get_ada_values(
            scale_shift_table[:num_scale_shift_values, :], batch_size, scale_shift_timestep, slice(None, None)
        )
        gate_ada_values = self.get_ada_values(
            scale_shift_table[num_scale_shift_values:, :], batch_size, gate_timestep, slice(None, None)
        )

        scale_shift_chunks = [t.squeeze(2) for t in scale_shift_ada_values]
        gate_ada_values = [t.squeeze(2) for t in gate_ada_values]

        return (*scale_shift_chunks, *gate_ada_values)

    def _apply_text_cross_attention(
        self,
        x: torch.Tensor,
        context: torch.Tensor,
        attn: Callable,
        scale_shift_table: torch.Tensor,
        prompt_scale_shift_table: torch.Tensor | None,
        timestep: torch.Tensor,
        prompt_timestep: torch.Tensor | None,
        context_mask: torch.Tensor | None,
    ) -> torch.Tensor:
        """Text cross-attention with optional LTX-2.3 cross-attention AdaLN.

        When ``cross_attention_adaln`` is False this is exactly the LTX-2.0
        path: ``attn(rms_norm(x), context=context, mask=context_mask)``.
        """
        if self.cross_attention_adaln:
            # The 3 extra AdaLN rows (6..9) carry text cross-attn shift/scale/gate.
            shift_q, scale_q, gate = self.get_ada_values(
                scale_shift_table, x.shape[0], timestep, slice(6, 9))
            return apply_cross_attention_adaln(
                x=x,
                context=context,
                attn=attn,
                q_shift=shift_q,
                q_scale=scale_q,
                q_gate=gate,
                prompt_scale_shift_table=prompt_scale_shift_table,
                prompt_timestep=prompt_timestep,
                context_mask=context_mask,
                norm_eps=self.norm_eps,
            )
        return attn(
            _rms_norm_dispatch(x, eps=self.norm_eps),
            context=context,
            mask=context_mask,
        )

    def forward(
        self,
        video: TransformerArgs | None,
        audio: TransformerArgs | None,
        video_original_seq_len: int | None = None,
        audio_original_seq_len: int | None = None,
        skip_cross_modal_attn: bool = False,
        skip_video_self_attn: bool | torch.Tensor = False,
        skip_audio_self_attn: bool | torch.Tensor = False,
    ) -> tuple[TransformerArgs | None, TransformerArgs | None]:
        """Forward pass for transformer block.

        Args:
            video: Video transformer args
            audio: Audio transformer args
            video_original_seq_len: Original unpadded video sequence length
            audio_original_seq_len: Original unpadded audio sequence length
            skip_cross_modal_attn: If True, skip A2V and V2A cross-modal
                attention (used for the modality-isolated CFG pass).
            skip_video_self_attn: If True, skip video self-attention
                in this block (used for STG perturbed pass).
            skip_audio_self_attn: If True, skip audio self-attention
                in this block (used for STG perturbed pass).
        """
        vx = video.x if video is not None else None
        ax = audio.x if audio is not None else None

        run_vx = video is not None and video.enabled and vx.numel() > 0
        run_ax = audio is not None and audio.enabled and ax.numel() > 0
        run_a2v = run_vx and run_ax
        run_v2a = run_ax and run_vx

        def _build_attn_keep_mask(
            values: torch.Tensor,
            skip_flag: bool | torch.Tensor,
        ) -> torch.Tensor:
            """STG keep-mask for self-attention.

            Returns a per-sample multiplier (1 keep, 0 perturb) broadcastable
            over ``values``. Only the configured ``stg_block_idx`` is perturbed.
            Supports both the LTX-2.0 bool flag and an LTX-2.3 per-sample
            tensor ``[B]`` (1/True == perturb). When skip_flag is False/0 the
            multiplier is all ones, reproducing the un-skipped LTX-2.0 path.
            """
            bsz = values.shape[0]
            keep = torch.ones((bsz, ), device=values.device, dtype=values.dtype)
            if self.idx == self.stg_block_idx:
                if torch.is_tensor(skip_flag):
                    if skip_flag.ndim == 0:
                        perturb = skip_flag.reshape(1).expand(bsz)
                    else:
                        if skip_flag.shape[0] != bsz:
                            raise ValueError(
                                "Per-sample STG mask batch size mismatch: "
                                f"got {skip_flag.shape[0]}, expected {bsz}")
                        perturb = skip_flag.reshape(bsz)
                    keep = 1.0 - perturb.to(device=values.device,
                                            dtype=values.dtype)
                elif bool(skip_flag):
                    keep.zero_()
            return keep.view(bsz, *([1] * (values.ndim - 1)))

        if run_vx:
            vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(
                self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3)
            )
            norm_vx = _rms_norm_dispatch(vx, eps=self.norm_eps) * (
                1 + vscale_msa) + vshift_msa
            video_self_attn_mask = _build_attn_keep_mask(vx, skip_video_self_attn)
            if self.use_distributed_attention:
                vx_attn1_out = self.attn1(
                    norm_vx,
                    pe=video.positional_embeddings,
                    original_seq_len=video_original_seq_len,
                )
            else:
                vx_attn1_out = self.attn1(norm_vx, pe=video.positional_embeddings)
            vx = vx + vx_attn1_out * vgate_msa * video_self_attn_mask
            # Text cross-attention: no SP mask needed (text is replicated).
            vx = vx + self._apply_text_cross_attention(
                x=vx,
                context=video.context,
                attn=self.attn2,
                scale_shift_table=self.scale_shift_table,
                prompt_scale_shift_table=getattr(
                    self, "prompt_scale_shift_table", None),
                timestep=video.timesteps,
                prompt_timestep=video.prompt_timestep,
                context_mask=video.context_mask,
            )

        if run_ax:
            ashift_msa, ascale_msa, agate_msa = self.get_ada_values(
                self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3)
            )
            norm_ax = _rms_norm_dispatch(ax, eps=self.norm_eps) * (
                1 + ascale_msa) + ashift_msa
            audio_self_attn_mask = _build_attn_keep_mask(ax, skip_audio_self_attn)
            if self.use_distributed_attention:
                ax_attn1_out = self.audio_attn1(
                    norm_ax,
                    pe=audio.positional_embeddings,
                    original_seq_len=audio_original_seq_len,
                )
            else:
                ax_attn1_out = self.audio_attn1(norm_ax, pe=audio.positional_embeddings)
            ax = ax + ax_attn1_out * agate_msa * audio_self_attn_mask
            # Text cross-attention: no SP mask needed (text is replicated).
            ax = ax + self._apply_text_cross_attention(
                x=ax,
                context=audio.context,
                attn=self.audio_attn2,
                scale_shift_table=self.audio_scale_shift_table,
                prompt_scale_shift_table=getattr(
                    self, "audio_prompt_scale_shift_table", None),
                timestep=audio.timesteps,
                prompt_timestep=audio.prompt_timestep,
                context_mask=audio.context_mask,
            )

        if (run_a2v or run_v2a) and not skip_cross_modal_attn:
            vx_norm3 = _rms_norm_dispatch(vx, eps=self.norm_eps)
            ax_norm3 = _rms_norm_dispatch(ax, eps=self.norm_eps)

            (
                scale_ca_audio_hidden_states_a2v,
                shift_ca_audio_hidden_states_a2v,
                scale_ca_audio_hidden_states_v2a,
                shift_ca_audio_hidden_states_v2a,
                gate_out_v2a,
            ) = self.get_av_ca_ada_values(
                self.scale_shift_table_a2v_ca_audio,
                ax.shape[0],
                audio.cross_scale_shift_timestep,
                audio.cross_gate_timestep,
            )

            (
                scale_ca_video_hidden_states_a2v,
                shift_ca_video_hidden_states_a2v,
                scale_ca_video_hidden_states_v2a,
                shift_ca_video_hidden_states_v2a,
                gate_out_a2v,
            ) = self.get_av_ca_ada_values(
                self.scale_shift_table_a2v_ca_video,
                vx.shape[0],
                video.cross_scale_shift_timestep,
                video.cross_gate_timestep,
            )

            if run_a2v:
                vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v) + shift_ca_video_hidden_states_a2v
                ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v
                # Audio-to-video cross-attention: Q from video (sharded), K/V from audio
                # For SP, gather audio context from all ranks before cross-attention
                if self.use_distributed_attention:
                    # Gather audio context from all SP ranks
                    ax_context = sequence_model_parallel_all_gather(ax_scaled, dim=1)
                    # Video Q: video tokens divide evenly (adjusted upstream), so simple slicing works
                    video_q_pe = None
                    if video.cross_positional_embeddings is not None:
                        sp_rank = get_sp_parallel_rank()
                        local_seq_len = vx_scaled.shape[1]
                        pe_tuple = video.cross_positional_embeddings
                        start_idx = sp_rank * local_seq_len
                        end_idx = start_idx + local_seq_len
                        video_q_pe = tuple(pe[:, :, start_idx:end_idx, :] for pe in pe_tuple)
                    # Audio K: gathered context may have padding, trim to original length
                    audio_k_pe = audio.cross_positional_embeddings
                    if audio_k_pe is not None:
                        original_audio_len = audio_k_pe[0].shape[2]
                        ax_context = ax_context[:, :original_audio_len, :]
                    vx = vx + (
                        self.audio_to_video_attn(
                            vx_scaled,
                            context=ax_context,
                            pe=video_q_pe,
                            k_pe=audio_k_pe,
                        )
                        * gate_out_a2v
                    )
                else:
                    ax_context = ax_scaled
                    video_q_pe = video.cross_positional_embeddings
                    audio_k_pe = audio.cross_positional_embeddings
                    vx = vx + (
                        self.audio_to_video_attn(
                            vx_scaled,
                            context=ax_context,
                            pe=video_q_pe,
                            k_pe=audio_k_pe,
                        )
                        * gate_out_a2v
                    )

            if run_v2a:
                ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a
                vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a
                # Video-to-audio cross-attention: Q from audio, K/V from video
                # For SP, gather both modalities and run cross-attention on full sequences
                # (audio is short, so no need to keep it sharded for cross-attention)
                if self.use_distributed_attention:
                    # Gather both modalities to full sequences
                    ax_full = sequence_model_parallel_all_gather(ax_scaled, dim=1)
                    vx_context = sequence_model_parallel_all_gather(vx_scaled, dim=1)
                    # Use full positional embeddings directly
                    audio_q_pe = audio.cross_positional_embeddings
                    video_k_pe = video.cross_positional_embeddings
                    # Trim gathered tensors to original lengths (may have padding from sharding)
                    if audio_q_pe is not None:
                        original_audio_len = audio_q_pe[0].shape[2]
                        ax_full = ax_full[:, :original_audio_len, :]
                    if video_k_pe is not None:
                        original_video_len = video_k_pe[0].shape[2]
                        vx_context = vx_context[:, :original_video_len, :]
                    # Run cross-attention on full audio
                    v2a_out_full = self.video_to_audio_attn(
                        ax_full,
                        context=vx_context,
                        pe=audio_q_pe,
                        k_pe=video_k_pe,
                    )
                    # Shard the output back to match ax's sharding
                    v2a_out, _ = sequence_model_parallel_shard(v2a_out_full, dim=1)
                    ax = ax + v2a_out * gate_out_v2a
                else:
                    vx_context = vx_scaled
                    audio_q_pe = audio.cross_positional_embeddings
                    video_k_pe = video.cross_positional_embeddings
                    ax = ax + (
                        self.video_to_audio_attn(
                            ax_scaled,
                            context=vx_context,
                            pe=audio_q_pe,
                            k_pe=video_k_pe,
                        )
                        * gate_out_v2a
                    )

        if run_vx:
            # slice(3, 6) selects the FFN shift/scale/gate. Identical to
            # slice(3, None) for the 6-row LTX-2.0 table, but correct for the
            # 9-row cross_attention_adaln table (rows 6..9 are the cross-attn).
            vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
                self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, 6)
            )
            vx_scaled = _rms_norm_dispatch(vx, eps=self.norm_eps) * (
                1 + vscale_mlp) + vshift_mlp
            vx = vx + self.ff(vx_scaled) * vgate_mlp

        if run_ax:
            ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(
                self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, 6)
            )
            ax_scaled = _rms_norm_dispatch(ax, eps=self.norm_eps) * (
                1 + ascale_mlp) + ashift_mlp
            ax = ax + self.audio_ff(ax_scaled) * agate_mlp

        if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1":
            video_sum = vx.float().sum().item() if vx is not None else 0.0
            audio_sum = ax.float().sum().item() if ax is not None else 0.0
            _debug_block_log_line(
                f"fastvideo:block={self.idx}:video_sum={video_sum:.6f} "
                f"audio_sum={audio_sum:.6f}"
            )

        # Register FSDP2 backward hooks on output tensors (module-level hooks don't
        # fire for dataclass outputs, so we must hook the tensors directly)
        self._register_fsdp_backward_hooks_on_output(vx, ax)

        return (
            replace(video, x=vx) if video is not None else None,
            replace(audio, x=ax) if audio is not None else None,
        )


def apply_cross_attention_adaln(
    x: torch.Tensor,
    context: torch.Tensor,
    attn: Callable,
    q_shift: torch.Tensor,
    q_scale: torch.Tensor,
    q_gate: torch.Tensor,
    prompt_scale_shift_table: torch.Tensor | None,
    prompt_timestep: torch.Tensor | None,
    context_mask: torch.Tensor | None = None,
    norm_eps: float = 1e-6,
) -> torch.Tensor:
    """LTX-2.3 cross-attention AdaLN: modulate query (x) and key/value
    (context) by AdaLN scale/shift and gate the output.

    Only reached when ``cross_attention_adaln`` is enabled, so it never
    affects the LTX-2.0 path.
    """
    if prompt_scale_shift_table is None or prompt_timestep is None:
        raise ValueError(
            "cross_attention_adaln requires prompt scale/shift tables and prompt_timestep."
        )
    batch_size = x.shape[0]
    shift_kv, scale_kv = (
        prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
        + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1)
    ).unbind(dim=2)
    attn_input = _rms_norm_dispatch(x, eps=norm_eps) * (1 + q_scale) + q_shift
    encoder_hidden_states = context * (1 + scale_kv) + shift_kv
    return attn(attn_input, context=encoder_hidden_states,
                mask=context_mask) * q_gate


class LTXModelType(Enum):
    """Model type flags for LTX-2."""
    AudioVideo = "ltx av model"
    VideoOnly = "ltx video only model"
    AudioOnly = "ltx audio only model"

    def is_video_enabled(self) -> bool:
        return self in (LTXModelType.AudioVideo, LTXModelType.VideoOnly)

    def is_audio_enabled(self) -> bool:
        return self in (LTXModelType.AudioVideo, LTXModelType.AudioOnly)


class LTXModel(torch.nn.Module):
    """Core LTX-2 transformer stack for audio/video latents."""

    def __init__(  # noqa: PLR0913
        self,
        *,
        model_type: LTXModelType = LTXModelType.AudioVideo,
        num_attention_heads: int = 32,
        attention_head_dim: int = 128,
        in_channels: int = 128,
        out_channels: int = 128,
        num_layers: int = 48,
        cross_attention_dim: int = 4096,
        norm_eps: float = 1e-06,
        caption_channels: int = 3840,
        positional_embedding_theta: float = 10000.0,
        positional_embedding_max_pos: list[int] | None = None,
        timestep_scale_multiplier: int = 1000,
        use_middle_indices_grid: bool = True,
        audio_num_attention_heads: int = 32,
        audio_attention_head_dim: int = 64,
        audio_in_channels: int = 128,
        audio_out_channels: int = 128,
        audio_cross_attention_dim: int = 2048,
        audio_positional_embedding_max_pos: list[int] | None = None,
        av_ca_timestep_scale_multiplier: int = 1,
        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
        double_precision_rope: bool = False,
        cross_attention_adaln: bool = False,
        caption_proj_before_connector: bool = False,
        apply_gated_attention: bool = False,
        stg_block_idx: int = 29,
        use_distributed_attention: bool = False,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        self._enable_gradient_checkpointing = False
        # LTX-2.3 gated extensions (all default OFF == LTX-2.0 behavior).
        self.cross_attention_adaln = cross_attention_adaln
        self.caption_proj_before_connector = caption_proj_before_connector
        self.apply_gated_attention = apply_gated_attention
        self.stg_block_idx = stg_block_idx
        self.use_middle_indices_grid = use_middle_indices_grid
        self.rope_type = rope_type
        self.double_precision_rope = double_precision_rope
        self.timestep_scale_multiplier = timestep_scale_multiplier
        self.positional_embedding_theta = positional_embedding_theta
        self.model_type = model_type
        cross_pe_max_pos = None

        if model_type.is_video_enabled():
            if positional_embedding_max_pos is None:
                positional_embedding_max_pos = [20, 2048, 2048]
            self.positional_embedding_max_pos = positional_embedding_max_pos
            self.num_attention_heads = num_attention_heads
            self.inner_dim = num_attention_heads * attention_head_dim
            self._init_video(
                in_channels=in_channels,
                out_channels=out_channels,
                caption_channels=caption_channels,
                norm_eps=norm_eps,
                create_caption_projection=not caption_proj_before_connector,
            )

        if model_type.is_audio_enabled():
            if audio_positional_embedding_max_pos is None:
                audio_positional_embedding_max_pos = [20]
            self.audio_positional_embedding_max_pos = audio_positional_embedding_max_pos
            self.audio_num_attention_heads = audio_num_attention_heads
            self.audio_inner_dim = self.audio_num_attention_heads * audio_attention_head_dim
            self._init_audio(
                in_channels=audio_in_channels,
                out_channels=audio_out_channels,
                caption_channels=caption_channels,
                norm_eps=norm_eps,
                create_caption_projection=not caption_proj_before_connector,
            )

        if model_type.is_video_enabled() and model_type.is_audio_enabled():
            cross_pe_max_pos = max(self.positional_embedding_max_pos[0], self.audio_positional_embedding_max_pos[0])
            self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier
            self.audio_cross_attention_dim = audio_cross_attention_dim
            self._init_audio_video(num_scale_shift_values=4)

        self._init_preprocessors(cross_pe_max_pos)
        self._init_transformer_blocks(
            num_layers=num_layers,
            attention_head_dim=attention_head_dim if model_type.is_video_enabled() else 0,
            cross_attention_dim=cross_attention_dim,
            audio_attention_head_dim=audio_attention_head_dim if model_type.is_audio_enabled() else 0,
            audio_cross_attention_dim=audio_cross_attention_dim,
            norm_eps=norm_eps,
            use_distributed_attention=use_distributed_attention,
            quant_config=quant_config,
            prefix=prefix,
        )

    def _init_video(
        self,
        in_channels: int,
        out_channels: int,
        caption_channels: int,
        norm_eps: float,
        create_caption_projection: bool = True,
    ) -> None:
        self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True)
        self.adaln_single = AdaLayerNormSingle(
            self.inner_dim,
            embedding_coefficient=adaln_embedding_coefficient(
                self.cross_attention_adaln),
        )
        # When caption_proj_before_connector is set the caption projection
        # lives in the text encoder, so it is omitted here (LTX-2.3 layout).
        if create_caption_projection:
            self.caption_projection = PixArtAlphaTextProjection(
                in_features=caption_channels,
                hidden_size=self.inner_dim,
            )
        if self.cross_attention_adaln:
            self.prompt_adaln_single = AdaLayerNormSingle(
                self.inner_dim,
                embedding_coefficient=2,
            )
        self.scale_shift_table = torch.nn.Parameter(torch.empty(2, self.inner_dim))
        self.norm_out = torch.nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=norm_eps)
        self.proj_out = torch.nn.Linear(self.inner_dim, out_channels)

    def _init_audio(
        self,
        in_channels: int,
        out_channels: int,
        caption_channels: int,
        norm_eps: float,
        create_caption_projection: bool = True,
    ) -> None:
        self.audio_patchify_proj = torch.nn.Linear(in_channels, self.audio_inner_dim, bias=True)
        self.audio_adaln_single = AdaLayerNormSingle(
            self.audio_inner_dim,
            embedding_coefficient=adaln_embedding_coefficient(
                self.cross_attention_adaln),
        )
        if create_caption_projection:
            self.audio_caption_projection = PixArtAlphaTextProjection(
                in_features=caption_channels,
                hidden_size=self.audio_inner_dim,
            )
        if self.cross_attention_adaln:
            self.audio_prompt_adaln_single = AdaLayerNormSingle(
                self.audio_inner_dim,
                embedding_coefficient=2,
            )
        self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(2, self.audio_inner_dim))
        self.audio_norm_out = torch.nn.LayerNorm(self.audio_inner_dim, elementwise_affine=False, eps=norm_eps)
        self.audio_proj_out = torch.nn.Linear(self.audio_inner_dim, out_channels)

    def _init_audio_video(self, num_scale_shift_values: int) -> None:
        self.av_ca_video_scale_shift_adaln_single = AdaLayerNormSingle(
            self.inner_dim,
            embedding_coefficient=num_scale_shift_values,
        )
        self.av_ca_audio_scale_shift_adaln_single = AdaLayerNormSingle(
            self.audio_inner_dim,
            embedding_coefficient=num_scale_shift_values,
        )
        self.av_ca_a2v_gate_adaln_single = AdaLayerNormSingle(
            self.inner_dim,
            embedding_coefficient=1,
        )
        self.av_ca_v2a_gate_adaln_single = AdaLayerNormSingle(
            self.audio_inner_dim,
            embedding_coefficient=1,
        )

    def _init_preprocessors(self, cross_pe_max_pos: int | None = None) -> None:
        if self.model_type.is_video_enabled() and self.model_type.is_audio_enabled():
            self.video_args_preprocessor = MultiModalTransformerArgsPreprocessor(
                patchify_proj=self.patchify_proj,
                adaln=self.adaln_single,
                caption_projection=getattr(self, "caption_projection", None),
                cross_scale_shift_adaln=self.av_ca_video_scale_shift_adaln_single,
                cross_gate_adaln=self.av_ca_a2v_gate_adaln_single,
                inner_dim=self.inner_dim,
                max_pos=self.positional_embedding_max_pos,
                num_attention_heads=self.num_attention_heads,
                cross_pe_max_pos=cross_pe_max_pos,
                use_middle_indices_grid=self.use_middle_indices_grid,
                audio_cross_attention_dim=self.audio_cross_attention_dim,
                timestep_scale_multiplier=self.timestep_scale_multiplier,
                double_precision_rope=self.double_precision_rope,
                positional_embedding_theta=self.positional_embedding_theta,
                rope_type=self.rope_type,
                av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier,
                prompt_adaln=getattr(self, "prompt_adaln_single", None),
            )
            self.audio_args_preprocessor = MultiModalTransformerArgsPreprocessor(
                patchify_proj=self.audio_patchify_proj,
                adaln=self.audio_adaln_single,
                caption_projection=getattr(self, "audio_caption_projection", None),
                cross_scale_shift_adaln=self.av_ca_audio_scale_shift_adaln_single,
                cross_gate_adaln=self.av_ca_v2a_gate_adaln_single,
                inner_dim=self.audio_inner_dim,
                max_pos=self.audio_positional_embedding_max_pos,
                num_attention_heads=self.audio_num_attention_heads,
                cross_pe_max_pos=cross_pe_max_pos,
                use_middle_indices_grid=self.use_middle_indices_grid,
                audio_cross_attention_dim=self.audio_cross_attention_dim,
                timestep_scale_multiplier=self.timestep_scale_multiplier,
                double_precision_rope=self.double_precision_rope,
                positional_embedding_theta=self.positional_embedding_theta,
                rope_type=self.rope_type,
                av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier,
                prompt_adaln=getattr(self, "audio_prompt_adaln_single", None),
            )
        elif self.model_type.is_video_enabled():
            self.video_args_preprocessor = TransformerArgsPreprocessor(
                patchify_proj=self.patchify_proj,
                adaln=self.adaln_single,
                caption_projection=getattr(self, "caption_projection", None),
                inner_dim=self.inner_dim,
                max_pos=self.positional_embedding_max_pos,
                num_attention_heads=self.num_attention_heads,
                use_middle_indices_grid=self.use_middle_indices_grid,
                timestep_scale_multiplier=self.timestep_scale_multiplier,
                double_precision_rope=self.double_precision_rope,
                positional_embedding_theta=self.positional_embedding_theta,
                rope_type=self.rope_type,
                prompt_adaln=getattr(self, "prompt_adaln_single", None),
            )
        elif self.model_type.is_audio_enabled():
            self.audio_args_preprocessor = TransformerArgsPreprocessor(
                patchify_proj=self.audio_patchify_proj,
                adaln=self.audio_adaln_single,
                caption_projection=getattr(self, "audio_caption_projection", None),
                inner_dim=self.audio_inner_dim,
                max_pos=self.audio_positional_embedding_max_pos,
                num_attention_heads=self.audio_num_attention_heads,
                use_middle_indices_grid=self.use_middle_indices_grid,
                timestep_scale_multiplier=self.timestep_scale_multiplier,
                double_precision_rope=self.double_precision_rope,
                positional_embedding_theta=self.positional_embedding_theta,
                rope_type=self.rope_type,
                prompt_adaln=getattr(self, "audio_prompt_adaln_single", None),
            )

    def _init_transformer_blocks(
        self,
        num_layers: int,
        attention_head_dim: int,
        cross_attention_dim: int,
        audio_attention_head_dim: int,
        audio_cross_attention_dim: int,
        norm_eps: float,
        use_distributed_attention: bool = False,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        video_config = (
            TransformerConfig(
                dim=self.inner_dim,
                heads=self.num_attention_heads,
                d_head=attention_head_dim,
                context_dim=cross_attention_dim,
                apply_gated_attention=self.apply_gated_attention,
                cross_attention_adaln=self.cross_attention_adaln,
            )
            if self.model_type.is_video_enabled()
            else None
        )
        audio_config = (
            TransformerConfig(
                dim=self.audio_inner_dim,
                heads=self.audio_num_attention_heads,
                d_head=audio_attention_head_dim,
                context_dim=audio_cross_attention_dim,
                apply_gated_attention=self.apply_gated_attention,
                cross_attention_adaln=self.cross_attention_adaln,
            )
            if self.model_type.is_audio_enabled()
            else None
        )
        self.use_distributed_attention = use_distributed_attention
        self.transformer_blocks = torch.nn.ModuleList(
            [
                BasicAVTransformerBlock(
                    idx=idx,
                    video=video_config,
                    audio=audio_config,
                    rope_type=self.rope_type,
                    norm_eps=norm_eps,
                    use_distributed_attention=use_distributed_attention,
                    cross_attention_adaln=self.cross_attention_adaln,
                    stg_block_idx=self.stg_block_idx,
                    quant_config=quant_config,
                    prefix=prefix,
                )
                for idx in range(num_layers)
            ]
        )

    def _process_transformer_blocks(
        self,
        video: TransformerArgs | None,
        audio: TransformerArgs | None,
        video_original_seq_len: int | None = None,
        audio_original_seq_len: int | None = None,
        skip_cross_modal_attn: bool = False,
        skip_video_self_attn_blocks: list[int] | None = None,
        skip_audio_self_attn_blocks: list[int] | None = None,
    ) -> tuple[TransformerArgs | None, TransformerArgs | None]:
        # Convert once so per-block membership checks stay O(1).
        skip_video_self_attn_block_set = set(skip_video_self_attn_blocks or [])
        skip_audio_self_attn_block_set = set(skip_audio_self_attn_blocks or [])

        for idx, block in enumerate(self.transformer_blocks):
            skip_v_sa = idx in skip_video_self_attn_block_set
            skip_a_sa = idx in skip_audio_self_attn_block_set
            video, audio = block(
                video=video,
                audio=audio,
                video_original_seq_len=video_original_seq_len,
                audio_original_seq_len=audio_original_seq_len,
                skip_cross_modal_attn=skip_cross_modal_attn,
                skip_video_self_attn=skip_v_sa,
                skip_audio_self_attn=skip_a_sa,
            )
        return video, audio

    def _process_output(
        self,
        scale_shift_table: torch.Tensor,
        norm_out: torch.nn.LayerNorm,
        proj_out: torch.nn.Linear,
        x: torch.Tensor,
        embedded_timestep: torch.Tensor,
    ) -> torch.Tensor:
        scale_shift_values = (
            scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + embedded_timestep[:, :, None]
        )
        shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
        x = norm_out(x)
        x = x * (1 + scale) + shift
        x = proj_out(x)
        return x

    def forward(
        self,
        video: Modality | None,
        audio: Modality | None,
        video_original_seq_len: int | None = None,
        audio_original_seq_len: int | None = None,
        skip_cross_modal_attn: bool = False,
        skip_video_self_attn_blocks: list[int] | None = None,
        skip_audio_self_attn_blocks: list[int] | None = None,
    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
        """Forward pass through the LTX model.

        Args:
            video: Video modality input
            audio: Audio modality input
            video_original_seq_len: Original unpadded video sequence length
            audio_original_seq_len: Original unpadded audio sequence length
            skip_cross_modal_attn: If True, skip A2V and V2A cross-modal
                attention in all transformer blocks (modality-isolated
                CFG pass).
            skip_video_self_attn_blocks: Block indices where video
                self-attention is skipped (STG perturbed pass).
            skip_audio_self_attn_blocks: Block indices where audio
                self-attention is skipped (STG perturbed pass).
        """
        if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1":
            _debug_block_log_line(
                "fastvideo:patchify_proj"
                f":video_w_sum={self.patchify_proj.weight.float().sum().item():.6f} "
                f"video_b_sum={self.patchify_proj.bias.float().sum().item():.6f} "
                f"audio_w_sum={self.audio_patchify_proj.weight.float().sum().item():.6f} "
                f"audio_b_sum={self.audio_patchify_proj.bias.float().sum().item():.6f}"
            )
        if not self.model_type.is_video_enabled() and video is not None:
            raise ValueError("Video is not enabled for this model")
        if not self.model_type.is_audio_enabled() and audio is not None:
            raise ValueError("Audio is not enabled for this model")

        video_args = self.video_args_preprocessor.prepare(video) if video is not None else None
        audio_args = self.audio_args_preprocessor.prepare(audio) if audio is not None else None
        _debug_transformer_args("fastvideo:prep_video", video_args)
        _debug_transformer_args("fastvideo:prep_audio", audio_args)
        video_out, audio_out = self._process_transformer_blocks(
            video_args,
            audio_args,
            video_original_seq_len=video_original_seq_len,
            audio_original_seq_len=audio_original_seq_len,
            skip_cross_modal_attn=skip_cross_modal_attn,
            skip_video_self_attn_blocks=skip_video_self_attn_blocks,
            skip_audio_self_attn_blocks=skip_audio_self_attn_blocks,
        )

        vx = (
            self._process_output(
                self.scale_shift_table, self.norm_out, self.proj_out, video_out.x, video_out.embedded_timestep
            )
            if video_out is not None
            else None
        )
        ax = (
            self._process_output(
                self.audio_scale_shift_table,
                self.audio_norm_out,
                self.audio_proj_out,
                audio_out.x,
                audio_out.embedded_timestep,
            )
            if audio_out is not None
            else None
        )
        return vx, ax


class LTX2Transformer3DModel(BaseDiT):
    """
    LTX-2 transformer using native FastVideo LTX-2 modules.
    """

    param_names_mapping = LTX2VideoConfig().param_names_mapping
    reverse_param_names_mapping = LTX2VideoConfig().reverse_param_names_mapping
    lora_param_names_mapping = LTX2VideoConfig().lora_param_names_mapping
    _fsdp_shard_conditions = LTX2VideoConfig()._fsdp_shard_conditions

    def __init__(self, config: LTX2VideoConfig, hf_config: dict[str, Any]):
        super().__init__(config=config, hf_config=hf_config)

        arch = config.arch_config

        # Get SP world size for distributed attention
        sp_world_size = get_sp_world_size()
        use_vsa_backend = os.getenv("FASTVIDEO_ATTENTION_BACKEND",
                                    "") == "VIDEO_SPARSE_ATTN"
        use_distributed_attention = sp_world_size > 1 or use_vsa_backend

        # Validate that attention heads are divisible by SP world size
        if sp_world_size > 1:
            assert arch.num_attention_heads % sp_world_size == 0, (
                f"The number of video attention heads ({arch.num_attention_heads}) "
                f"must be divisible by the sequence parallel size ({sp_world_size})"
            )
            assert arch.audio_num_attention_heads % sp_world_size == 0, (
                f"The number of audio attention heads ({arch.audio_num_attention_heads}) "
                f"must be divisible by the sequence parallel size ({sp_world_size})"
            )
            logger.info(
                f"LTX2 sequence parallelism enabled with SP world size {sp_world_size}"
            )
        elif use_vsa_backend:
            logger.info(
                "LTX2 VSA enabled with SP world size 1; using distributed "
                "attention path for VSA compatibility")

        model_type = LTXModelType.AudioVideo
        self.model = LTXModel(
            model_type=model_type,
            num_attention_heads=arch.num_attention_heads,
            attention_head_dim=arch.attention_head_dim,
            in_channels=arch.in_channels,
            out_channels=arch.out_channels,
            num_layers=arch.num_layers,
            cross_attention_dim=arch.cross_attention_dim,
            norm_eps=arch.norm_eps,
            caption_channels=arch.caption_channels,
            positional_embedding_theta=arch.positional_embedding_theta,
            positional_embedding_max_pos=arch.positional_embedding_max_pos,
            timestep_scale_multiplier=arch.timestep_scale_multiplier,
            use_middle_indices_grid=arch.use_middle_indices_grid,
            rope_type=LTXRopeType(arch.rope_type),
            double_precision_rope=arch.double_precision_rope,
            audio_num_attention_heads=arch.audio_num_attention_heads,
            audio_attention_head_dim=arch.audio_attention_head_dim,
            audio_in_channels=arch.audio_in_channels,
            audio_out_channels=arch.audio_out_channels,
            audio_cross_attention_dim=arch.audio_cross_attention_dim,
            audio_positional_embedding_max_pos=arch.audio_positional_embedding_max_pos,
            av_ca_timestep_scale_multiplier=arch.av_ca_timestep_scale_multiplier,
            cross_attention_adaln=arch.cross_attention_adaln,
            caption_proj_before_connector=arch.caption_proj_before_connector,
            apply_gated_attention=arch.apply_gated_attention,
            stg_block_idx=arch.stg_block_idx,
            use_distributed_attention=use_distributed_attention,
            quant_config=config.quant_config,
            prefix=config.prefix,
        )

        self.patchifier = VideoLatentPatchifier(patch_size=arch.patch_size[1])
        self.audio_patchifier = AudioLatentPatchifier(
            patch_size=DEFAULT_LTX2_AUDIO_MEL_BINS,
            sample_rate=DEFAULT_LTX2_AUDIO_SAMPLE_RATE,
            hop_length=DEFAULT_LTX2_AUDIO_HOP_LENGTH,
            audio_latent_downsample_factor=DEFAULT_LTX2_AUDIO_DOWNSAMPLE,
            is_causal=True,
            shift=0,
        )

        self.hidden_size = arch.num_attention_heads * arch.attention_head_dim
        self.num_attention_heads = arch.num_attention_heads
        self.num_channels_latents = arch.num_channels_latents

        if os.getenv("LTX2_DEBUG_DETAIL", "0") == "1":
            detail_path = os.getenv("LTX2_PIPELINE_DEBUG_DETAIL_PATH", "")
            if detail_path:
                self._attach_debug_detail_hooks(detail_path)

    def _attach_debug_detail_hooks(self, log_path: str) -> None:
        path = Path(log_path)
        path.parent.mkdir(parents=True, exist_ok=True)
        if path.exists():
            path.unlink()

        def _format_sum(tensor: torch.Tensor | None) -> str:
            if tensor is None:
                return "None"
            return f"{tensor.float().sum().item():.6f}"

        def _hook_factory(block_idx: int, name: str):
            def _hook(_module, _inputs, outputs):  # noqa: ANN001
                out = outputs[0] if isinstance(outputs, tuple) else outputs
                out_sum = _format_sum(out if torch.is_tensor(out) else None)
                with path.open("a", encoding="utf-8") as f:
                    f.write(f"fastvideo:{block_idx}:{name}:out_sum={out_sum}\n")
            return _hook

        for block in self.model.transformer_blocks:
            idx = block.idx
            for name in (
                "attn1",
                "attn2",
                "ff",
                "audio_attn1",
                "audio_attn2",
                "audio_ff",
                "audio_to_video_attn",
                "video_to_audio_attn",
            ):
                if hasattr(block, name):
                    getattr(block, name).register_forward_hook(
                        _hook_factory(idx, name))

        def _output_hook(label: str):
            def _hook(_module, _inputs, outputs):  # noqa: ANN001
                out = outputs[0] if isinstance(outputs, tuple) else outputs
                out_sum = _format_sum(out if torch.is_tensor(out) else None)
                with path.open("a", encoding="utf-8") as f:
                    f.write(f"fastvideo:output:{label}:out_sum={out_sum}\n")
            return _hook

        if hasattr(self.model, "proj_out"):
            self.model.proj_out.register_forward_hook(_output_hook("proj_out"))
        if hasattr(self.model, "audio_proj_out"):
            self.model.audio_proj_out.register_forward_hook(
                _output_hook("audio_proj_out"))

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
        timestep: torch.LongTensor,
        encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
        encoder_attention_mask: torch.Tensor | None = None,
        guidance=None,
        audio_hidden_states: torch.Tensor | None = None,
        audio_encoder_hidden_states: torch.Tensor | None = None,
        audio_timestep: torch.Tensor | None = None,
        audio_encoder_attention_mask: torch.Tensor | None = None,
        video_sigma: torch.Tensor | None = None,
        audio_sigma: torch.Tensor | None = None,
        skip_cross_modal_attn: bool = False,
        skip_video_self_attn_blocks: list[int] | None = None,
        skip_audio_self_attn_blocks: list[int] | None = None,
        video_position_offset_sec: float = 0.0,
        **kwargs,
    ) -> torch.Tensor:
        if isinstance(encoder_hidden_states, list):
            encoder_hidden_states = encoder_hidden_states[0]
        # Get SP parameters
        sp_world_size = get_sp_world_size()

        # Get fps for position computation
        fps = None
        try:
            forward_ctx = get_forward_context()
        except AssertionError:
            forward_ctx = None
        if forward_ctx is not None and forward_ctx.forward_batch is not None:
            fps_value = forward_ctx.forward_batch.fps
            if isinstance(fps_value, list):
                fps_value = fps_value[0] if fps_value else None
            if fps_value is not None:
                fps = float(fps_value)

        # Patchify video latents
        video_shape = VideoLatentShape.from_torch_shape(hidden_states.shape)
        latents = self.patchifier.patchify(hidden_states)
        # LTX-2.3 cross-attention AdaLN consumes a per-sample sigma timestep.
        # When the denoising stage does not supply one (LTX-2.0, or any caller
        # that omits it) derive it from the video timestep. The downstream
        # ``prompt_adaln`` is None for LTX-2.0 so this value is ignored there.
        if video_sigma is None:
            video_sigma = timestep[:, 0] if timestep.ndim > 1 else timestep

        # Shard video latents and timestep across SP ranks
        video_original_seq_len = latents.shape[1]
        video_padded_seq_len = video_original_seq_len
        video_timestep = timestep
        if sp_world_size > 1:
            latents, video_original_seq_len = sequence_model_parallel_shard(latents, dim=1)
            # Shard timestep along sequence dimension (timestep has shape [batch, seq_len])
            video_timestep, _ = sequence_model_parallel_shard(timestep, dim=1)
            current_seq_len = latents.shape[1]
            video_padded_seq_len = current_seq_len * sp_world_size
        # Compute RoPE positions for the FULL sequence (before sharding)
        # This is necessary because sequence sharding may not align to frame boundaries.
        # The full RoPE will be passed to DistributedAttention which applies it after
        # the all-to-all (when each rank has the full sequence but subset of heads).
        positions = self.patchifier.get_patch_grid_bounds(
            video_shape, device=hidden_states.device
        )
        positions = _get_pixel_coords(
            positions,
            DEFAULT_LTX2_SCALE_FACTORS,
            fps=fps,
            causal_fix=True,
        )
        if video_position_offset_sec:
            positions[:, 0, ...] = positions[:, 0, ...] + float(video_position_offset_sec)
        positions = positions.to(hidden_states.dtype)

        # Pad positions to match padded sequence length for SP cross-attention
        if sp_world_size > 1 and video_padded_seq_len > video_original_seq_len:
            padding_needed = video_padded_seq_len - video_original_seq_len
            last_pos = positions[:, :, -1:, :].expand(-1, -1, padding_needed, -1)
            positions = torch.cat([positions, last_pos], dim=2)

        video_modality = Modality(
            enabled=True,
            latent=latents,
            timesteps=video_timestep,
            positions=positions,
            context=encoder_hidden_states,
            context_mask=encoder_attention_mask,
            sigma=video_sigma,
        )
        if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1":
            video_head = latents.flatten()[:8].float().tolist()
            video_flat = latents.float().flatten()
            video_checksum = (video_flat * torch.arange(video_flat.numel(), device=video_flat.device)).sum().item()
            _debug_block_log_line(
                "fastvideo:modality_video"
                f":latent_sum={latents.float().sum().item():.6f} "
                f"latent_shape={tuple(latents.shape)} "
                f"positions_sum={positions.float().sum().item():.6f} "
                f"positions_shape={tuple(positions.shape)} "
                f"latent_head={video_head} "
                f"latent_checksum={video_checksum:.6f}"
            )

        # Process audio modality if provided
        audio_modality = None
        audio_shape = None
        audio_original_seq_len = 0

        if audio_hidden_states is not None and audio_encoder_hidden_states is not None and audio_timestep is not None:
            audio_shape = AudioLatentShape.from_torch_shape(audio_hidden_states.shape)
            audio_latents = self.audio_patchifier.patchify(audio_hidden_states)
            if audio_sigma is None:
                audio_sigma = (
                    audio_timestep[:, 0]
                    if audio_timestep.ndim > 1 else audio_timestep
                )

            # Shard audio latents and timestep across SP ranks
            audio_original_seq_len = audio_latents.shape[1]
            sharded_audio_timestep = audio_timestep
            if sp_world_size > 1:
                audio_latents, audio_original_seq_len = sequence_model_parallel_shard(audio_latents, dim=1)
                # Shard audio timestep along sequence dimension
                sharded_audio_timestep, _ = sequence_model_parallel_shard(audio_timestep, dim=1)

            # Compute audio RoPE positions for the FULL sequence (before sharding)
            # Same as video: full RoPE is applied after all-to-all in DistributedAttention
            audio_positions = self.audio_patchifier.get_patch_grid_bounds(
                audio_shape, device=audio_hidden_states.device
            )

            audio_modality = Modality(
                enabled=True,
                latent=audio_latents,
                timesteps=sharded_audio_timestep,
                positions=audio_positions,
                context=audio_encoder_hidden_states,
                context_mask=audio_encoder_attention_mask,
                sigma=audio_sigma,
            )
            if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1":
                audio_head = audio_latents.flatten()[:8].float().tolist()
                audio_flat = audio_latents.float().flatten()
                audio_checksum = (audio_flat * torch.arange(audio_flat.numel(), device=audio_flat.device)).sum().item()
                _debug_block_log_line(
                    "fastvideo:modality_audio"
                    f":latent_sum={audio_latents.float().sum().item():.6f} "
                    f"latent_shape={tuple(audio_latents.shape)} "
                    f"positions_sum={audio_positions.float().sum().item():.6f} "
                    f"positions_shape={tuple(audio_positions.shape)} "
                    f"latent_head={audio_head} "
                    f"latent_checksum={audio_checksum:.6f}"
                )

        # Run transformer with original sequence lengths
        video_out, audio_out = self.model(
            video=video_modality,
            audio=audio_modality,
            video_original_seq_len=video_original_seq_len,
            audio_original_seq_len=audio_original_seq_len,
            skip_cross_modal_attn=skip_cross_modal_attn,
            skip_video_self_attn_blocks=skip_video_self_attn_blocks,
            skip_audio_self_attn_blocks=skip_audio_self_attn_blocks,
        )

        # Denoised prediction
        if video_out is not None and video_modality is not None:
            video_out = _to_denoised(
                video_modality.latent,
                video_out,
                video_modality.timesteps,
            )
        if audio_out is not None and audio_modality is not None:
            audio_out = _to_denoised(
                audio_modality.latent,
                audio_out,
                audio_modality.timesteps,
            )

        # Gather and unpad video output
        if sp_world_size > 1 and video_out is not None:
            video_out = sequence_model_parallel_all_gather_with_unpad(
                video_out, video_original_seq_len, dim=1
            )

        # Gather and unpad audio output
        if sp_world_size > 1 and audio_out is not None:
            audio_out = sequence_model_parallel_all_gather_with_unpad(
                audio_out, audio_original_seq_len, dim=1
            )

        # Unpatchify
        video_out = self.patchifier.unpatchify(
            video_out, output_shape=video_shape)
        if audio_out is None or audio_shape is None:
            return video_out
        audio_out = self.audio_patchifier.unpatchify(
            audio_out, output_shape=audio_shape)
        return video_out, audio_out

# Entry point for model registry
EntryClass = LTX2Transformer3DModel