--- title: ltx2 ---

LTX-2 transformer implementation

fastvideo.models.dits.ltx2.SDPAMetadata

SDPAMetadata(current_timestep: int, attn_mask: torch.Tensor | None = None, is_causal: bool = False, *, VSA_sparsity: float = 0.0)

fastvideo.models.dits.ltx2.DistributedAttention

DistributedAttention(num_heads: int, head_size: int, 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)

Distributed attention layer.

fastvideo.models.dits.ltx2.DistributedAttention.forward

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, original_seq_len: int | None = None, replicated_q: torch.Tensor | None = None, replicated_k: torch.Tensor | None = None, replicated_v: torch.Tensor | None = None, freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, attn_metadata: AttentionMetadata | None = None) -> tuple[torch.Tensor, torch.Tensor | None]

Forward pass for distributed attention.

Parameters

q: torch.Tensor

Query tensor [batch_size, seq_len, num_heads, head_dim]

k: torch.Tensor

Key tensor [batch_size, seq_len, num_heads, head_dim]

v: torch.Tensor

Value tensor [batch_size, seq_len, num_heads, head_dim]

original_seq_len: int

Original (unpadded) full sequence length

replicated_q: Optional[torch.Tensor]

Replicated query tensor, typically for text tokens

replicated_k: Optional[torch.Tensor]

Replicated key tensor

replicated_v: Optional[torch.Tensor]

Replicated value tensor

Returns

tuple[torch.Tensor, torch.Tensor | None]

Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing: - o (torch.Tensor): Output tensor after attention for the main sequence - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided

fastvideo.models.dits.ltx2.LocalAttention

LocalAttention(num_heads: int, head_size: int, num_kv_heads: int | None = None, softmax_scale: float | None = None, causal: bool = False, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, default_backend: AttentionBackendEnum | None = None, **extra_impl_args)

Attention layer.

fastvideo.models.dits.ltx2.LocalAttention.forward

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor

Apply local attention between query, key and value tensors.

Parameters

q: torch.Tensor

Query tensor of shape [batch_size, seq_len, num_heads, head_dim]

k: torch.Tensor

Key tensor of shape [batch_size, seq_len, num_heads, head_dim]

v: torch.Tensor

Value tensor of shape [batch_size, seq_len, num_heads, head_dim]

Returns

torch.Tensor

torch.Tensor: Output tensor after local attention

fastvideo.models.dits.ltx2.effective_attention_backend

effective_attention_backend(config: object) -> AttentionBackendEnum | None

Return the attention backend that a component built from config follows.

Some components choose part of their own structure by backend: Wan picks its VSA transformer block, LTX-2 its distributed-attention path, and the Gemma and T5-Gemma text encoders switch Hugging Face attention to SDPA. That choice must match the attention layers inside the component, and those layers resolve through :func:get_attn_backend. This function applies the same order:

  1. The decision a loader recorded on config with :func:record_resolved_attention_backend. Every component built by a loader has one when a backend was requested.
  2. While a loader is building a component (a _component_attention_backend_scope is active), that scope's rule: its backend, and the environment variable only when the scope consults it. A scope that deliberately ignores the environment variable, such as the dense DMD teacher and critic, keeps ignoring it here.
  3. With no loader involved, the FASTVIDEO_ATTENTION_BACKEND environment variable. This covers components constructed directly, such as MagiHuman's lazily built T5-Gemma encoder or a script that builds a transformer itself.

Reading only the recorded decision is not enough: a directly constructed component has no recorded decision (None), so it would ignore the environment variable while its attention layers still follow it, and the component's structure would not match its layers.

fastvideo.models.dits.ltx2.LTX2VideoConfig

LTX2VideoConfig(arch_config: DiTArchConfig = LTX2VideoArchConfig(), prefix: str = 'ltx2', quant_config: QuantizationConfig | None = None, *, _resolved_attention_backend: AttentionBackendEnum | None = None)

Main configuration for LTX-2 transformer.

fastvideo.models.dits.ltx2.sequence_model_parallel_all_gather

sequence_model_parallel_all_gather(input_: torch.Tensor, dim: int = -1) -> torch.Tensor

All-gather the input tensor across model parallel group.

fastvideo.models.dits.ltx2.sequence_model_parallel_all_gather_with_unpad

sequence_model_parallel_all_gather_with_unpad(input_: torch.Tensor, original_seq_len: int, dim: int = -1) -> torch.Tensor

All-gather the input tensor and remove padding.

Parameters

input_: torch.Tensor

Sharded (and possibly padded) tensor to gather

original_seq_len: int

Original sequence length before padding

dim: int

Dimension to gather along (default: -1)

Returns

Tensor: torch.Tensor

Gathered and unpadded tensor

fastvideo.models.dits.ltx2.sequence_model_parallel_all_to_all_4D

sequence_model_parallel_all_to_all_4D(input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1) -> torch.Tensor

All-to-all communication of 4D tensors (e.g. QKV matrices) across sequence parallel group.

fastvideo.models.dits.ltx2.sequence_model_parallel_shard

sequence_model_parallel_shard(input_: torch.Tensor, dim: int = 1) -> tuple[torch.Tensor, int]

Shard the input tensor across model parallel group with optional padding.

Parameters

input_: torch.Tensor

Input tensor to shard

dim: int

Dimension to shard along (default: 1)

Returns

tuple: tuple[torch.Tensor, int]

(sharded_tensor, original_seq_len) - sharded_tensor: The sharded (and possibly padded) tensor - original_seq_len: Original sequence length before padding

fastvideo.models.dits.ltx2.get_sp_parallel_rank

get_sp_parallel_rank() -> int

Return my rank for the sequence model parallel group.

fastvideo.models.dits.ltx2.get_sp_world_size

get_sp_world_size() -> int

Return world size for the sequence model parallel group.

fastvideo.models.dits.ltx2.ForwardContext

ForwardContext(current_timestep: int, attn_metadata: AttentionMetadata, forward_batch: Optional[ForwardBatch] = None)

fastvideo.models.dits.ltx2.get_forward_context

get_forward_context() -> ForwardContext

Get the current forward context.

fastvideo.models.dits.ltx2.set_forward_context

set_forward_context(current_timestep, attn_metadata, forward_batch: Optional[ForwardBatch] = None)

A context manager that stores the current forward context, can be attention metadata, etc. Here we can inject common logic for every model forward pass.

fastvideo.models.dits.ltx2.ReplicatedLinear

ReplicatedLinear(input_size: int, output_size: int, bias: bool = True, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, prefix: str = '')

Replicated linear layer.

Parameters

input_size: int

input dimension of the linear layer.

output_size: int

output dimension of the linear layer.

bias: bool

If true, add bias.

skip_bias_add: bool

If true, skip adding bias but instead return it.

params_dtype: torch.dtype | None

Data type for the parameters.

quant_config: QuantizationConfig | None

Quantization configure.

prefix: str

The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj)

fastvideo.models.dits.ltx2.ReplicatedLinear.weight_loader

weight_loader(param: Parameter, loaded_weight: torch.Tensor) -> None

fastvideo.models.dits.ltx2.ReplicatedLinear.forward

forward(x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]

fastvideo.models.dits.ltx2.ReplicatedLinear.extra_repr

extra_repr() -> str

fastvideo.models.dits.ltx2.ReplicatedLinear.get_shape_mapping

get_shape_mapping() -> dict

Get the mapping from (input_shape, output_shape) to layer types.

fastvideo.models.dits.ltx2.ReplicatedLinear.reset_shape_tracking

reset_shape_tracking() -> None

Clear tracked shapes and layer type mappings.

fastvideo.models.dits.ltx2.ReplicatedLinear.print_shape_summary

print_shape_summary() -> None

Log a summary of all unique shapes and their layer types.

fastvideo.models.dits.ltx2.QuantizationConfig

QuantizationConfig()

Base class for quantization configs.

fastvideo.models.dits.ltx2.QuantizationConfig.get_name

get_name() -> QuantizationMethods

Name of the quantization method.

fastvideo.models.dits.ltx2.QuantizationConfig.get_supported_act_dtypes

get_supported_act_dtypes() -> list[torch.dtype]

List of supported activation dtypes.

fastvideo.models.dits.ltx2.QuantizationConfig.get_min_capability

get_min_capability() -> int

Minimum GPU capability to support the quantization method.

E.g., 70 for Volta, 75 for Turing, 80 for Ampere. This requirement is due to the custom CUDA kernels used by the quantization method.

fastvideo.models.dits.ltx2.QuantizationConfig.get_config_filenames

get_config_filenames() -> list[str]

List of filenames to search for in the model directory.

fastvideo.models.dits.ltx2.QuantizationConfig.from_config

from_config(config: dict[str, Any]) -> QuantizationConfig

Create a config class from the model's quantization config.

fastvideo.models.dits.ltx2.QuantizationConfig.override_quantization_method

override_quantization_method(hf_quant_cfg, user_quant) -> QuantizationMethods | None

Detects if this quantization method can support a given checkpoint format by overriding the user specified quantization method -- this method should only be overwritten by subclasses in exceptional circumstances

fastvideo.models.dits.ltx2.QuantizationConfig.get_from_keys

get_from_keys(config: dict[str, Any], keys: list[str]) -> Any

Get a value from the model's quantization config.

fastvideo.models.dits.ltx2.QuantizationConfig.get_from_keys_or

get_from_keys_or(config: dict[str, Any], keys: list[str], default: Any) -> Any

Get a optional value from the model's quantization config.

fastvideo.models.dits.ltx2.QuantizationConfig.get_quant_method

get_quant_method(layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None

Get the quantize method to use for the quantized layer.

Parameters

layer: torch.nn.Module

The layer for the quant method.

prefix: str

The full name of the layer in the state dict

Returns: The quantize method. None if the given layer doesn't support quant method.

fastvideo.models.dits.ltx2.QuantizationConfig.validate_runtime

validate_runtime(device: torch.device) -> None

Fail fast on a device the scheme cannot execute on.

Called once by loaders that select a serialized checkpoint scheme, before any weight is read. The default accepts every device; schemes with kernel or capability requirements override it.

fastvideo.models.dits.ltx2.QuantizationConfig.get_cache_scale

get_cache_scale(name: str) -> str | None

fastvideo.models.dits.ltx2.init_logger

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

fastvideo.models.dits.ltx2.BaseDiT

BaseDiT(config: DiTConfig, hf_config: dict[str, Any], **kwargs)

fastvideo.models.dits.ltx2.BaseDiT.forward

forward(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, guidance = None, **kwargs) -> torch.Tensor

fastvideo.models.dits.ltx2.BaseDiT.preprocess_loaded_weight

preprocess_loaded_weight(name: str, tensor: torch.Tensor) -> torch.Tensor

Adapt a checkpoint tensor before sharding; default models are unchanged.

fastvideo.models.dits.ltx2.BaseDiT.inference_config_overrides

inference_config_overrides() -> dict[str, Any]

Extra config.json keys a deployable checkpoint of this model needs.

fastvideo.models.dits.ltx2.BaseDiT.inference_export_expected_shape

inference_export_expected_shape(output_key: str, base_shape: tuple[int, ...]) -> tuple[int, ...]

The exported shape of output_key given the base checkpoint's shape.

fastvideo.models.dits.ltx2.AttentionBackendEnum

AttentionBackendEnum()

fastvideo.models.dits.ltx2.StageAwareRMSNorm

StageAwareRMSNorm(hidden_size: int, eps: float = 1e-06)

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).

Source
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,
            )
        # torch 2.12 added rms_norm to autocast's float32 cast policy
        # (aten/src/ATen/autocast_mode.h), so under autocast(bfloat16) this
        # norm now returns float32 where torch <= 2.11 returned bfloat16.
        # Restore the input dtype so autocast-blind consumers downstream
        # (e.g. the flash-attention custom op) don't receive upcast q/k.
        return _functional_rms_norm(x, self.eps, self.weight)

fastvideo.models.dits.ltx2.StageAwareRMSNorm.forward

forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    if _is_ltx2_refine_stage():
        return _quack_rmsnorm(
            x,
            weight=self.weight,
            eps=self.eps,
        )
    # torch 2.12 added rms_norm to autocast's float32 cast policy
    # (aten/src/ATen/autocast_mode.h), so under autocast(bfloat16) this
    # norm now returns float32 where torch <= 2.11 returned bfloat16.
    # Restore the input dtype so autocast-blind consumers downstream
    # (e.g. the flash-attention custom op) don't receive upcast q/k.
    return _functional_rms_norm(x, self.eps, self.weight)

fastvideo.models.dits.ltx2.get_timestep_embedding

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.

Source
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

fastvideo.models.dits.ltx2.TimestepEmbedding

TimestepEmbedding(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)

Two-layer MLP to project timestep embeddings.

Source
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

fastvideo.models.dits.ltx2.TimestepEmbedding.forward

forward(sample: torch.Tensor, condition: torch.Tensor | None = None) -> torch.Tensor
Source
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

fastvideo.models.dits.ltx2.Timesteps

Timesteps(num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1)

Sinusoidal timestep embedding wrapper with scaling knobs.

Source
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,
        )

fastvideo.models.dits.ltx2.Timesteps.forward

forward(timesteps: torch.Tensor) -> torch.Tensor
Source
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,
    )

fastvideo.models.dits.ltx2.PixArtAlphaCombinedTimestepSizeEmbeddings

PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim: int, size_emb_dim: int)

PixArt-Alpha timestep embedding used by LTX-2 AdaLN.

Source
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

fastvideo.models.dits.ltx2.PixArtAlphaCombinedTimestepSizeEmbeddings.forward

forward(timestep: torch.Tensor, hidden_dtype: torch.dtype) -> torch.Tensor
Source
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

fastvideo.models.dits.ltx2.AdaLayerNormSingle

AdaLayerNormSingle(embedding_dim: int, embedding_coefficient: int = 6)

AdaLN-single modulation that emits scale/shift/gates for LTX-2.

Source
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

fastvideo.models.dits.ltx2.AdaLayerNormSingle.forward

forward(timestep: torch.Tensor, hidden_dtype: Optional[torch.dtype] = None) -> Tuple[torch.Tensor, torch.Tensor]
Source
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

fastvideo.models.dits.ltx2.adaln_embedding_coefficient

adaln_embedding_coefficient(cross_attention_adaln: bool) -> int
Source
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)

fastvideo.models.dits.ltx2.PixArtAlphaTextProjection

PixArtAlphaTextProjection(in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = 'gelu_tanh')

Caption projection MLP used by LTX-2.

Source
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

fastvideo.models.dits.ltx2.PixArtAlphaTextProjection.forward

forward(caption: torch.Tensor) -> torch.Tensor
Source
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

fastvideo.models.dits.ltx2.GELUApprox

GELUApprox(in_features: int, out_features: int, bias: bool = True, quant_config: QuantizationConfig | None = None, prefix: str = '')

Linear + tanh-approximate GELU used by LTX-2 FFN.

Source
class GELUApprox(nn.Module):
    """Linear + tanh-approximate GELU used by LTX-2 FFN."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        self.proj = ReplicatedLinear(
            in_features,
            out_features,
            bias=bias,
            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])

fastvideo.models.dits.ltx2.GELUApprox.forward

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

fastvideo.models.dits.ltx2.FeedForward

FeedForward(dim: int, dim_out: int, mult: int = 4, bias: bool = True, quant_config: QuantizationConfig | None = None, prefix: str = '')

LTX-2 FFN: GELUApprox -> Identity -> Linear.

Source
class FeedForward(nn.Module):
    """LTX-2 FFN: GELUApprox -> Identity -> Linear."""

    def __init__(
        self,
        dim: int,
        dim_out: int,
        mult: int = 4,
        bias: bool = True,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        inner_dim = int(dim * mult)
        project_in = GELUApprox(
            dim,
            inner_dim,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.ffn",
        )
        project_out = ReplicatedLinear(
            inner_dim,
            dim_out,
            bias=bias,
            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

fastvideo.models.dits.ltx2.FeedForward.forward

forward(x: torch.Tensor) -> torch.Tensor
Source
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

fastvideo.models.dits.ltx2.VideoLatentShape

VideoLatentShape()

Helper for (B, C, T, H, W) latent shapes.

Source
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)

fastvideo.models.dits.ltx2.VideoLatentShape.from_torch_shape

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

fastvideo.models.dits.ltx2.VideoLatentShape.to_torch_shape

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

fastvideo.models.dits.ltx2.AudioLatentShape

AudioLatentShape()

Helper for (B, C, T, F) audio latent shapes.

Source
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))

fastvideo.models.dits.ltx2.AudioLatentShape.to_torch_shape

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

fastvideo.models.dits.ltx2.AudioLatentShape.from_torch_shape

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

fastvideo.models.dits.ltx2.AudioLatentShape.from_duration

from_duration(batch: int, duration: float, channels: int, mel_bins: int, sample_rate: int, hop_length: int, audio_latent_downsample_factor: int) -> AudioLatentShape
Source
@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))

fastvideo.models.dits.ltx2.VideoLatentPatchifier

VideoLatentPatchifier(patch_size: int)

Patchify/unpatchify latent tokens for LTX-2.

Source
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

fastvideo.models.dits.ltx2.VideoLatentPatchifier.get_token_count

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

fastvideo.models.dits.ltx2.VideoLatentPatchifier.patchify

patchify(latents: torch.Tensor) -> torch.Tensor
Source
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],
    )

fastvideo.models.dits.ltx2.VideoLatentPatchifier.unpatchify

unpatchify(latents: torch.Tensor, output_shape: VideoLatentShape) -> torch.Tensor
Source
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],
    )

fastvideo.models.dits.ltx2.VideoLatentPatchifier.get_patch_grid_bounds

get_patch_grid_bounds(output_shape: VideoLatentShape, device: Optional[torch.device] = None) -> torch.Tensor

Get patch grid bounds for RoPE computation.

Parameters

output_shape: VideoLatentShape

Shape of the video latent tensor

device: Optional[torch.device]

Device to create tensors on

Source
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

fastvideo.models.dits.ltx2.AudioLatentPatchifier

AudioLatentPatchifier(patch_size: int, sample_rate: int, hop_length: int, audio_latent_downsample_factor: int, is_causal: bool = True, shift: int = 0)

Patchify/unpatchify audio latents and compute timing bounds.

Source
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

fastvideo.models.dits.ltx2.AudioLatentPatchifier.get_token_count

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

fastvideo.models.dits.ltx2.AudioLatentPatchifier.patchify

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

fastvideo.models.dits.ltx2.AudioLatentPatchifier.unpatchify

unpatchify(latents: torch.Tensor, output_shape: AudioLatentShape) -> torch.Tensor
Source
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,
    )

fastvideo.models.dits.ltx2.AudioLatentPatchifier.get_patch_grid_bounds

get_patch_grid_bounds(output_shape: AudioLatentShape, device: Optional[torch.device] = None) -> torch.Tensor

Get patch grid bounds for audio RoPE computation.

Parameters

output_shape: AudioLatentShape

Shape of the audio latent tensor

device: Optional[torch.device]

Device to create tensors on

Source
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)

fastvideo.models.dits.ltx2.LTXRopeType

LTXRopeType()

LTX-2 rotary variants (interleaved vs split).

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

fastvideo.models.dits.ltx2.apply_ltx_rotary_emb

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.

Source
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}")

fastvideo.models.dits.ltx2.apply_ltx_rotary_emb_4d

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.

Source
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}")

fastvideo.models.dits.ltx2.generate_ltx_freq_grid_float64

generate_ltx_freq_grid_float64(positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int) -> torch.Tensor

Generate an eager-cached, compile-safe float64 LTX-2 rotary grid.

Source
def generate_ltx_freq_grid_float64(positional_embedding_theta: float, positional_embedding_max_pos_count: int,
                                   inner_dim: int) -> torch.Tensor:
    """Generate an eager-cached, compile-safe float64 LTX-2 rotary grid."""
    if torch.compiler.is_compiling():
        return _generate_ltx_freq_grid_float64(positional_embedding_theta, positional_embedding_max_pos_count,
                                               inner_dim)
    return _cached_generate_ltx_freq_grid_float64(positional_embedding_theta, positional_embedding_max_pos_count,
                                                  inner_dim)

fastvideo.models.dits.ltx2.generate_ltx_freq_grid_pytorch

generate_ltx_freq_grid_pytorch(positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int) -> torch.Tensor

Generate an eager-cached, compile-safe float32 LTX-2 rotary grid.

Source
def generate_ltx_freq_grid_pytorch(positional_embedding_theta: float, positional_embedding_max_pos_count: int,
                                   inner_dim: int) -> torch.Tensor:
    """Generate an eager-cached, compile-safe float32 LTX-2 rotary grid."""
    if torch.compiler.is_compiling():
        return _generate_ltx_freq_grid_pytorch(positional_embedding_theta, positional_embedding_max_pos_count,
                                               inner_dim)
    return _cached_generate_ltx_freq_grid_pytorch(positional_embedding_theta, positional_embedding_max_pos_count,
                                                  inner_dim)

fastvideo.models.dits.ltx2.precompute_ltx_freqs_cis

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.

Source
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)

fastvideo.models.dits.ltx2.TransformerArgs

TransformerArgs(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, prompt_timestep: torch.Tensor | None = None)

Pack transformer inputs for LTX-2 blocks.

Source
@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

fastvideo.models.dits.ltx2.Modality

Modality(enabled: bool, latent: torch.Tensor, timesteps: torch.Tensor, positions: torch.Tensor, context: torch.Tensor, context_mask: torch.Tensor | None = None, sigma: torch.Tensor | None = None, keyframes_mask: torch.Tensor | None = None)

Lightweight modality container for LTX-2 inputs.

Source
@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
    # LTX-2.5 per-token keyframe marker, shaped (B, T, 1).
    keyframes_mask: torch.Tensor | None = None

fastvideo.models.dits.ltx2.apply_keyframes_absolute_embedding

apply_keyframes_absolute_embedding(hidden_states: torch.Tensor, keyframes_mask: torch.Tensor | None, embedding_provider: KeyframesEmbeddingProvider | None) -> torch.Tensor

Apply the learned LTX-2.5 marker to projected keyframe tokens.

Source
def apply_keyframes_absolute_embedding(
    hidden_states: torch.Tensor,
    keyframes_mask: torch.Tensor | None,
    embedding_provider: KeyframesEmbeddingProvider | None,
) -> torch.Tensor:
    """Apply the learned LTX-2.5 marker to projected keyframe tokens."""
    if embedding_provider is None or keyframes_mask is None:
        return hidden_states
    embedding = embedding_provider()
    if embedding is None:
        return hidden_states
    mask = (keyframes_mask > 0).to(dtype=hidden_states.dtype)
    return hidden_states + mask * embedding.to(dtype=hidden_states.dtype)

fastvideo.models.dits.ltx2.TransformerArgsPreprocessor

TransformerArgsPreprocessor(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, keyframes_embedding_provider: KeyframesEmbeddingProvider | None = None)

Prepare LTX-2 transformer inputs (patchify, AdaLN, rope).

Source
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,
        keyframes_embedding_provider: KeyframesEmbeddingProvider | 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
        self.keyframes_embedding_provider = keyframes_embedding_provider

    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_float64
        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)
        x = apply_keyframes_absolute_embedding(
            x,
            modality.keyframes_mask,
            self.keyframes_embedding_provider,
        )
        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,
        )

fastvideo.models.dits.ltx2.TransformerArgsPreprocessor.prepare

prepare(modality: Modality) -> TransformerArgs
Source
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)
    x = apply_keyframes_absolute_embedding(
        x,
        modality.keyframes_mask,
        self.keyframes_embedding_provider,
    )
    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,
    )

fastvideo.models.dits.ltx2.MultiModalTransformerArgsPreprocessor

MultiModalTransformerArgsPreprocessor(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, keyframes_embedding_provider: KeyframesEmbeddingProvider | None = None)

Prepare transformer args for audio/video cross-attention paths.

Source
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,
        keyframes_embedding_provider: KeyframesEmbeddingProvider | 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,
            keyframes_embedding_provider=keyframes_embedding_provider,
        )
        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

fastvideo.models.dits.ltx2.MultiModalTransformerArgsPreprocessor.prepare

prepare(modality: Modality) -> TransformerArgs
Source
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,
    )

fastvideo.models.dits.ltx2.TransformerConfig

TransformerConfig(dim: int, heads: int, d_head: int, context_dim: int, apply_gated_attention: bool = False, cross_attention_adaln: bool = False, ff_bias: bool = True)

Attention/FFN dims for LTX-2 transformer blocks.

Source
@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
    ff_bias: bool = True

fastvideo.models.dits.ltx2.LTXDistributedAttention

LTXDistributedAttention(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)

LTX-2 specialized DistributedAttention that handles LTX-style RoPE internally.

Source
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

fastvideo.models.dits.ltx2.LTXDistributedAttention.forward

forward(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.

Parameters

q: torch.Tensor

Query tensor [batch_size, seq_len, num_heads, head_dim]

k: torch.Tensor

Key tensor [batch_size, seq_len, num_heads, head_dim]

v: torch.Tensor

Value tensor [batch_size, seq_len, num_heads, head_dim]

replicated_q: torch.Tensor | None

Replicated query tensor for text tokens

replicated_k: torch.Tensor | None

Replicated key tensor

replicated_v: torch.Tensor | None

Replicated value tensor

original_seq_len: int

Original (unpadded) full sequence length

ltx_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None

LTX-2 style RoPE (cos, sin) with shape [B, H, T, D]

Returns

tuple[torch.Tensor, torch.Tensor | None]

Tuple of output tensor and optional replicated output

Source
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

fastvideo.models.dits.ltx2.LTXLocalAttention

LTXLocalAttention(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)

LTX-2 specialized LocalAttention that handles LTX-style RoPE internally.

Source
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

fastvideo.models.dits.ltx2.LTXLocalAttention.forward

forward(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.

Parameters

q: torch.Tensor

Query tensor [batch_size, seq_len, num_heads, head_dim]

k: torch.Tensor

Key tensor [batch_size, seq_len, num_heads, head_dim]

v: torch.Tensor

Value tensor [batch_size, seq_len, num_heads, head_dim]

ltx_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None

LTX-2 style RoPE (cos, sin) for Q with shape [B, H, T, D]

ltx_k_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None

LTX-2 style RoPE (cos, sin) for K (if different from Q)

Returns

torch.Tensor

Output tensor after local attention

Source
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

fastvideo.models.dits.ltx2.LTXSelfAttention

LTXSelfAttention(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 = '')

LTX-2 attention block with RMSNorm + FastVideo LocalAttention.

Source
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, ),
            default_backend=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)

fastvideo.models.dits.ltx2.LTXSelfAttention.forward

forward(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
Source
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)

fastvideo.models.dits.ltx2.LTXDistributedSelfAttention

LTXDistributedSelfAttention(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 = '')

LTX-2 attention block with RMSNorm + LTXDistributedAttention for SP.

Source
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)

fastvideo.models.dits.ltx2.LTXDistributedSelfAttention.forward

forward(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.

Parameters

x: torch.Tensor

Query tensor [B, L, C]

context: torch.Tensor | None

Key/Value tensor [B, L, C], or None for self-attention

pe: tuple[torch.Tensor, torch.Tensor] | None

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: tuple[torch.Tensor, torch.Tensor] | None

Rotary position embeddings for K (cos, sin), or None to use pe

original_seq_len: int | None

Original (unpadded) full sequence length

Source
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)

fastvideo.models.dits.ltx2.BasicAVTransformerBlock

BasicAVTransformerBlock(idx: int, video: TransformerConfig | None = None, audio: TransformerConfig | None = None, rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, norm_eps: float = 1e-06, use_distributed_attention: bool = False, cross_attention_adaln: bool = False, stg_block_idx: int = 29, quant_config: QuantizationConfig | None = None, prefix: str = '')

LTX-2 transformer block (audio/video + cross attention + FFN).

Source
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_attn_backends = (
            AttentionBackendEnum.FLASH_ATTN,
            AttentionBackendEnum.TORCH_SDPA,
            AttentionBackendEnum.ATTN_QAT_INFER,
            AttentionBackendEnum.ATTN_QAT_TRAIN,
        )
        video_self_attn_backends = video_attn_backends + (AttentionBackendEnum.VIDEO_SPARSE_ATTN, )
        audio_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=video_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,
                bias=video.ff_bias,
                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=audio_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=audio_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,
                bias=audio.ff_bias,
                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=audio_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=audio_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``. Applied unconditionally: the eager caller
            (``_process_transformer_blocks``) enforces that only the
            configured ``stg_block_idx`` block receives a truthy flag.
            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)
            # The ``idx == stg_block_idx`` gate lives in the eager
            # ``_process_transformer_blocks`` caller, not here: reading the
            # per-instance ``self.idx`` inside this compiled forward forces
            # torch.compile to specialize a distinct graph per block (48x),
            # blowing the Dynamo recompile limit and defeating regional
            # compilation. The caller only passes a truthy ``skip_flag`` to the
            # configured STG block, so applying it unconditionally here is
            # equivalent while keeping the graph identical across all 48 blocks.
            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

        # 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,
        )

fastvideo.models.dits.ltx2.BasicAVTransformerBlock.get_ada_values

get_ada_values(scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, indices: slice) -> tuple[torch.Tensor, ...]
Source
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

fastvideo.models.dits.ltx2.BasicAVTransformerBlock.get_av_ca_ada_values

get_av_ca_ada_values(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]
Source
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)

fastvideo.models.dits.ltx2.BasicAVTransformerBlock.forward

forward(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.

Parameters

video: TransformerArgs | None

Video transformer args

audio: TransformerArgs | None

Audio transformer args

video_original_seq_len: int | None

Original unpadded video sequence length

audio_original_seq_len: int | None

Original unpadded audio sequence length

skip_cross_modal_attn: bool

If True, skip A2V and V2A cross-modal attention (used for the modality-isolated CFG pass).

skip_video_self_attn: bool | torch.Tensor

If True, skip video self-attention in this block (used for STG perturbed pass).

skip_audio_self_attn: bool | torch.Tensor

If True, skip audio self-attention in this block (used for STG perturbed pass).

Source
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``. Applied unconditionally: the eager caller
        (``_process_transformer_blocks``) enforces that only the
        configured ``stg_block_idx`` block receives a truthy flag.
        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)
        # The ``idx == stg_block_idx`` gate lives in the eager
        # ``_process_transformer_blocks`` caller, not here: reading the
        # per-instance ``self.idx`` inside this compiled forward forces
        # torch.compile to specialize a distinct graph per block (48x),
        # blowing the Dynamo recompile limit and defeating regional
        # compilation. The caller only passes a truthy ``skip_flag`` to the
        # configured STG block, so applying it unconditionally here is
        # equivalent while keeping the graph identical across all 48 blocks.
        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

    # 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,
    )

fastvideo.models.dits.ltx2.apply_cross_attention_adaln

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-06) -> 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.

Source
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:
        raise ValueError("cross_attention_adaln requires a prompt scale/shift table.")
    batch_size = x.shape[0]
    kv_modulation = prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
    # LTX-2.5 can disable the prompt-side AdaLN MLP. In that mode the static
    # per-block table is the entire K/V modulation, making it timestep
    # independent and cacheable. Older checkpoints add the dynamic term.
    if prompt_timestep is not None:
        kv_modulation = kv_modulation + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1)
    shift_kv, scale_kv = kv_modulation.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

fastvideo.models.dits.ltx2.LTXModelType

LTXModelType()

Model type flags for LTX-2.

Source
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)

fastvideo.models.dits.ltx2.LTXModelType.is_video_enabled

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

fastvideo.models.dits.ltx2.LTXModelType.is_audio_enabled

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

fastvideo.models.dits.ltx2.LTXModel

LTXModel(*, 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, use_prompt_adaln_single: bool = True, ff_bias: bool = True, audio_ff_bias: bool = True, use_keyframes_abs_pos_embedding: 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 = '')

Core LTX-2 transformer stack for audio/video latents.

Source
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,
        use_prompt_adaln_single: bool = True,
        ff_bias: bool = True,
        audio_ff_bias: bool = True,
        use_keyframes_abs_pos_embedding: 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.use_prompt_adaln_single = use_prompt_adaln_single
        self.ff_bias = ff_bias
        self.audio_ff_bias = audio_ff_bias
        self.use_keyframes_abs_pos_embedding = use_keyframes_abs_pos_embedding
        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 _keyframes_embedding(self) -> torch.Tensor | None:
        """Resolve the optional parameter at call time, matching upstream."""
        return getattr(self, "keyframes_abs_pos_embedding", None)

    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.keyframes_abs_pos_embedding = (
            torch.nn.Parameter(torch.zeros(1, self.inner_dim)) if self.use_keyframes_abs_pos_embedding else None
        )
        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,
            )
        self.prompt_adaln_single = (
            AdaLayerNormSingle(self.inner_dim, embedding_coefficient=2)
            if self.cross_attention_adaln and self.use_prompt_adaln_single else None
        )
        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,
            )
        self.audio_prompt_adaln_single = (
            AdaLayerNormSingle(self.audio_inner_dim, embedding_coefficient=2)
            if self.cross_attention_adaln and self.use_prompt_adaln_single else None
        )
        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),
                keyframes_embedding_provider=self._keyframes_embedding,
            )
            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),
                keyframes_embedding_provider=self._keyframes_embedding,
            )
        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,
            ff_bias=self.ff_bias,
        ) 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,
            ff_bias=self.audio_ff_bias,
        ) 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):
            # Preserve the original STG block condition in eager code so it does not introduce a block index guard in the compiled forward
            is_stg_block = idx == block.stg_block_idx
            skip_v_sa = is_stg_block and idx in skip_video_self_attn_block_set
            skip_a_sa = is_stg_block and 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 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
        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

fastvideo.models.dits.ltx2.LTXModel.forward

forward(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.

Parameters

video: Modality | None

Video modality input

audio: Modality | None

Audio modality input

video_original_seq_len: int | None

Original unpadded video sequence length

audio_original_seq_len: int | None

Original unpadded audio sequence length

skip_cross_modal_attn: bool

If True, skip A2V and V2A cross-modal attention in all transformer blocks (modality-isolated CFG pass).

skip_video_self_attn_blocks: list[int] | None

Block indices where video self-attention is skipped (STG perturbed pass).

skip_audio_self_attn_blocks: list[int] | None

Block indices where audio self-attention is skipped (STG perturbed pass).

Source
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 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
    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

fastvideo.models.dits.ltx2.LTX2Transformer3DModel

LTX2Transformer3DModel(config: LTX2VideoConfig, hf_config: dict[str, Any])

LTX-2 transformer using native FastVideo LTX-2 modules.

Source
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
    _compile_conditions = LTX2VideoConfig()._compile_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 = effective_attention_backend(config) == AttentionBackendEnum.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,
            use_prompt_adaln_single=arch.use_prompt_adaln_single,
            ff_bias=arch.ff_bias,
            audio_ff_bias=arch.audio_ff_bias,
            use_keyframes_abs_pos_embedding=arch.use_keyframes_abs_pos_embedding,
            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

    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,
        keyframes_mask: 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 keyframes_mask is not None:
            if keyframes_mask.ndim != 3 or keyframes_mask.shape[-1] != 1:
                raise ValueError(
                    "keyframes_mask must contain patchified per-token markers with shape (B, T, 1)"
                )
            if keyframes_mask.shape[:2] != latents.shape[:2]:
                raise ValueError(
                    "keyframes_mask batch/sequence dimensions must match the patchified video latents: "
                    f"mask={tuple(keyframes_mask.shape)}, latents={tuple(latents.shape)}"
                )
            keyframes_mask = keyframes_mask.to(device=latents.device)
        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)
            if keyframes_mask is not None:
                keyframes_mask, _ = sequence_model_parallel_shard(keyframes_mask, 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,
            keyframes_mask=keyframes_mask,
        )

        # 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,
            )

        # 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

fastvideo.models.dits.ltx2.LTX2Transformer3DModel.forward

forward(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, keyframes_mask: 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
Source
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,
    keyframes_mask: 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 keyframes_mask is not None:
        if keyframes_mask.ndim != 3 or keyframes_mask.shape[-1] != 1:
            raise ValueError(
                "keyframes_mask must contain patchified per-token markers with shape (B, T, 1)"
            )
        if keyframes_mask.shape[:2] != latents.shape[:2]:
            raise ValueError(
                "keyframes_mask batch/sequence dimensions must match the patchified video latents: "
                f"mask={tuple(keyframes_mask.shape)}, latents={tuple(latents.shape)}"
            )
        keyframes_mask = keyframes_mask.to(device=latents.device)
    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)
        if keyframes_mask is not None:
            keyframes_mask, _ = sequence_model_parallel_shard(keyframes_mask, 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,
        keyframes_mask=keyframes_mask,
    )

    # 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,
        )

    # 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