ltx2

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)</code></pre></details>
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&#91;:, None&#93;.float() * emb&#91;None, :&#93;
emb = scale * emb
emb = torch.cat(&#91;torch.sin(emb), torch.cos(emb)&#93;, dim=-1)

if flip_sin_to_cos:
    emb = torch.cat(&#91;emb&#91;:, half_dim:&#93;, emb&#91;:, :half_dim&#93;&#93;, dim=-1)

if embedding_dim % 2 == 1:
    emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
return emb</code></pre></details>
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) -&gt; 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</code></pre></details>
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) -&gt; 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,
    )</code></pre></details>
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) -&gt; torch.Tensor:
    timesteps_proj = self.time_proj(timestep)
    timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype))
    return timesteps_emb</code></pre></details>
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&#91;torch.dtype&#93; = None,
) -&gt; Tuple&#91;torch.Tensor, torch.Tensor&#93;:
    embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype)
    return self.linear(self.silu(embedded_timestep)), embedded_timestep</code></pre></details>
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 = &quot;gelu_tanh&quot;):
    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 == &quot;gelu_tanh&quot;:
        self.act_1 = torch.nn.GELU(approximate=&quot;tanh&quot;)
    elif act_fn == &quot;silu&quot;:
        self.act_1 = torch.nn.SiLU()
    else:
        raise ValueError(f&quot;Unknown activation function: &#123;act_fn&#125;&quot;)
    self.linear_2 = torch.nn.Linear(in_features=hidden_size, out_features=out_features, bias=True)

def forward(self, caption: torch.Tensor) -&gt; 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</code></pre></details>
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 = &quot;&quot;,
):
    super().__init__()
    self.proj = ReplicatedLinear(
        in_features,
        out_features,
        bias=bias,
        quant_config=quant_config,
        prefix=f&quot;&#123;prefix&#125;.fc_in&quot;,
    )
    self.act = nn.GELU(approximate=&quot;tanh&quot;)

def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
    return self.act(self.proj(x)&#91;0&#93;)</code></pre></details>
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 = &quot;&quot;,
) -&gt; None:
    super().__init__()
    inner_dim = int(dim * mult)
    project_in = GELUApprox(
        dim,
        inner_dim,
        bias=bias,
        quant_config=quant_config,
        prefix=f&quot;&#123;prefix&#125;.ffn&quot;,
    )
    project_out = ReplicatedLinear(
        inner_dim,
        dim_out,
        bias=bias,
        quant_config=quant_config,
        prefix=f&quot;&#123;prefix&#125;.ffn.fc_out&quot;,
    )
    self.net = nn.ModuleList(&#91;project_in, nn.Identity(), project_out&#93;)

def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
    x = self.net&#91;0&#93;(x)
    x = self.net&#91;1&#93;(x)
    x = self.net&#91;2&#93;(x)&#91;0&#93;
    return x</code></pre></details>
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) -&gt; int:
    return self&#91;0&#93;

@property
def channels(self) -&gt; int:
    return self&#91;1&#93;

@property
def frames(self) -&gt; int:
    return self&#91;2&#93;

@property
def height(self) -&gt; int:
    return self&#91;3&#93;

@property
def width(self) -&gt; int:
    return self&#91;4&#93;

@staticmethod
def from_torch_shape(shape: torch.Size) -&gt; &quot;VideoLatentShape&quot;:
    return VideoLatentShape(shape)

def to_torch_shape(self) -&gt; torch.Size:
    return torch.Size(self)</code></pre></details>
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) -&gt; int:
    return self&#91;0&#93;

@property
def channels(self) -&gt; int:
    return self&#91;1&#93;

@property
def frames(self) -&gt; int:
    return self&#91;2&#93;

@property
def mel_bins(self) -&gt; int:
    return self&#91;3&#93;

def to_torch_shape(self) -&gt; torch.Size:
    return torch.Size(self)

@staticmethod
def from_torch_shape(shape: torch.Size) -&gt; &quot;AudioLatentShape&quot;:
    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,
) -&gt; &quot;AudioLatentShape&quot;:
    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))</code></pre></details>
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) -&gt; Tuple&#91;int, int, int&#93;:
    return self._patch_size

def get_token_count(self, tgt_shape: VideoLatentShape) -&gt; int:
    return math.prod(tgt_shape.to_torch_shape()&#91;2:&#93;) // math.prod(self._patch_size)

def patchify(self, latents: torch.Tensor) -&gt; torch.Tensor:
    return rearrange(
        latents,
        &quot;b c (f p1) (h p2) (w p3) -&gt; b (f h w) (c p1 p2 p3)&quot;,
        p1=self._patch_size&#91;0&#93;,
        p2=self._patch_size&#91;1&#93;,
        p3=self._patch_size&#91;2&#93;,
    )

def unpatchify(self, latents: torch.Tensor, output_shape: VideoLatentShape) -&gt; torch.Tensor:
    patch_grid_frames = output_shape.frames // self._patch_size&#91;0&#93;
    patch_grid_height = output_shape.height // self._patch_size&#91;1&#93;
    patch_grid_width = output_shape.width // self._patch_size&#91;2&#93;
    return rearrange(
        latents,
        &quot;b (f h w) (c p q) -&gt; b c f (h p) (w q)&quot;,
        f=patch_grid_frames,
        h=patch_grid_height,
        w=patch_grid_width,
        p=self._patch_size&#91;1&#93;,
        q=self._patch_size&#91;2&#93;,
    )

def get_patch_grid_bounds(
    self,
    output_shape: VideoLatentShape,
    device: Optional&#91;torch.device&#93; = None,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;Get patch grid bounds for RoPE computation.

    Args:
        output_shape: Shape of the video latent tensor
        device: Device to create tensors on
    &quot;&quot;&quot;
    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&#91;0&#93;, device=device),
        torch.arange(start=0, end=height, step=self._patch_size&#91;1&#93;, device=device),
        torch.arange(start=0, end=width, step=self._patch_size&#91;2&#93;, device=device),
        indexing=&quot;ij&quot;,
    )

    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,
        &quot;c f h w bounds -&gt; b c (f h w) bounds&quot;,
        b=batch_size,
        bounds=2,
    )
    return latent_coords</code></pre></details>
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
&quot;&quot;&quot;
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&#91;0&#93;, device=device),
    torch.arange(start=0, end=height, step=self._patch_size&#91;1&#93;, device=device),
    torch.arange(start=0, end=width, step=self._patch_size&#91;2&#93;, device=device),
    indexing=&quot;ij&quot;,
)

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,
    &quot;c f h w bounds -&gt; b c (f h w) bounds&quot;,
    b=batch_size,
    bounds=2,
)
return latent_coords</code></pre></details>
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,
) -&gt; 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) -&gt; int:
    return tgt_shape.frames

def patchify(self, latents: torch.Tensor) -&gt; torch.Tensor:
    return rearrange(latents, &quot;b c t f -&gt; b t (c f)&quot;)

def unpatchify(
    self,
    latents: torch.Tensor,
    output_shape: AudioLatentShape,
) -&gt; torch.Tensor:
    return rearrange(
        latents,
        &quot;b t (c f) -&gt; b c t f&quot;,
        c=output_shape.channels,
        f=output_shape.mel_bins,
    )

def get_patch_grid_bounds(
    self,
    output_shape: AudioLatentShape,
    device: Optional&#91;torch.device&#93; = None,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;Get patch grid bounds for audio RoPE computation.

    Args:
        output_shape: Shape of the audio latent tensor
        device: Device to create tensors on
    &quot;&quot;&quot;
    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(&#91;start_timings, end_timings&#93;, dim=-1)

def _get_audio_latent_time_in_sec(
    self,
    start_latent: int,
    end_latent: int,
    dtype: torch.dtype,
    device: Optional&#91;torch.device&#93;,
) -&gt; torch.Tensor:
    resolved_device = device or torch.device(&quot;cpu&quot;)
    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</code></pre></details>
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
&quot;&quot;&quot;
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(&#91;start_timings, end_timings&#93;, dim=-1)</code></pre></details>
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 &#91;B, H, T, D&#93; format.
&quot;&quot;&quot;
cos_freqs, sin_freqs = freqs_cis
if rope_type == LTXRopeType.INTERLEAVED:
    # For interleaved, cos/sin have shape &#91;B, T, inner_dim&#93;
    # Need to reshape to &#91;B, H, T, D&#93; format
    # Actually interleaved doesn&#x27;t have per-head rotations, so we broadcast
    t_dup = rearrange(input_tensor, &quot;... (d r) -&gt; ... d r&quot;, r=2)
    t1, t2 = t_dup.unbind(dim=-1)
    t_dup = torch.stack((-t2, t1), dim=-1)
    input_tensor_rot = rearrange(t_dup, &quot;... d r -&gt; ... (d r)&quot;)
    return input_tensor * cos_freqs + input_tensor_rot * sin_freqs
if rope_type == LTXRopeType.SPLIT:
    # For split, cos/sin already have shape &#91;B, H, T, D/2&#93;
    # input_tensor is &#91;B, H, T, D&#93;
    split_input = rearrange(input_tensor, &quot;... (d r) -&gt; ... d r&quot;, d=2)
    first_half_input = split_input&#91;..., :1, :&#93;
    second_half_input = split_input&#91;..., 1:, :&#93;

    output = split_input * cos_freqs.unsqueeze(-2)
    first_half_output = output&#91;..., :1, :&#93;
    second_half_output = output&#91;..., 1:, :&#93;

    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, &quot;... d r -&gt; ... (d r)&quot;)
raise ValueError(f&quot;Invalid rope type: &#123;rope_type&#125;&quot;)</code></pre></details>
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&#91;1&#93;, 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&#91;-1&#93;
    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&#91;1&#93;
    cos_freq, sin_freq = _ltx_interleaved_freqs_cis(freqs, dim % n_elem)
return cos_freq.to(out_dtype), sin_freq.to(out_dtype)</code></pre></details>
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&#91;int&#93;,
    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,
) -&gt; 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,
) -&gt; tuple&#91;torch.Tensor, torch.Tensor&#93;:
    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&#91;-1&#93;)
    embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.shape&#91;-1&#93;)
    return timestep, embedded_timestep

def _prepare_context(
    self,
    context: torch.Tensor,
    x: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
) -&gt; tuple&#91;torch.Tensor, torch.Tensor | None&#93;:
    batch_size = x.shape&#91;0&#93;
    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&#91;-1&#93;)
    return context, attention_mask

def _prepare_attention_mask(self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype) -&gt; 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&#91;0&#93;, 1, -1, attention_mask.shape&#91;-1&#93;)) * torch.finfo(x_dtype).max

def _prepare_positional_embeddings(
    self,
    positions: torch.Tensor,
    inner_dim: int,
    max_pos: list&#91;int&#93;,
    use_middle_indices_grid: bool,
    num_attention_heads: int,
    x_dtype: torch.dtype,
) -&gt; 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) -&gt; 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&#91;0&#93;, 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&#91;0&#93;,
            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,
    )</code></pre></details>
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&#91;int&#93;,
    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,
) -&gt; 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) -&gt; TransformerArgs:
    transformer_args = self.simple_preprocessor.prepare(modality)
    cross_pe = self.simple_preprocessor._prepare_positional_embeddings(
        positions=modality.positions&#91;:, 0:1, :&#93;,
        inner_dim=self.audio_cross_attention_dim,
        max_pos=&#91;self.cross_pe_max_pos&#93;,
        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&#91;0&#93;,
        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,
) -&gt; tuple&#91;torch.Tensor, torch.Tensor&#93;:
    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&#91;-1&#93;)
    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&#91;-1&#93;)

    return scale_shift_timestep, gate_noise_timestep</code></pre></details>
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&#91;0&#93;,
    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,
)</code></pre></details>
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&#91;AttentionBackendEnum, ...&#93; | None = None,
    prefix: str = &quot;&quot;,
    **extra_impl_args,
) -&gt; 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&#91;torch.Tensor, torch.Tensor&#93; | None = None,
) -&gt; tuple&#91;torch.Tensor, torch.Tensor | None&#93;:
    &quot;&quot;&quot;Forward pass with LTX-2 style RoPE application.

    Args:
        q: Query tensor &#91;batch_size, seq_len, num_heads, head_dim&#93;
        k: Key tensor &#91;batch_size, seq_len, num_heads, head_dim&#93;
        v: Value tensor &#91;batch_size, seq_len, num_heads, head_dim&#93;
        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 &#91;B, H, T, D&#93;

    Returns:
        Tuple of output tensor and optional replicated output
    &quot;&quot;&quot;
    assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, &quot;Expected 4D tensors&quot;
    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), &quot;Replicated QKV is not supported for VSA now&quot;
        if gate_compress is None:
            raise ValueError(&quot;gate_compress must be provided when using VIDEO_SPARSE_ATTN&quot;)

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

        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&#91;:, head_start:head_end, :original_seq_len, :&#93;
            sin_local = sin&#91;:, head_start:head_end, :original_seq_len, :&#93;
            qk_part = qkvg&#91;:batch_size * 2&#93;.transpose(1, 2)
            qk_part = apply_ltx_rotary_emb_4d(qk_part, (cos_local, sin_local), self.rope_type)
            qkvg&#91;:batch_size * 2&#93; = 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&#91;call-arg&#93;
            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(&#91;q, k, v&#93;, dim=0)  # &#91;3*batch, seq_len, num_heads, head_dim&#93;

    # 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&#91;1&#93; - original_seq_len
    qkv = qkv&#91;:, :original_seq_len, :, :&#93;

    # 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&#x27;s heads: &#91;B, H/SP, T, D&#93;
        cos_local = cos&#91;:, head_start:head_end, :original_seq_len, :&#93;
        sin_local = sin&#91;:, head_start:head_end, :original_seq_len, :&#93;

        # Apply RoPE to Q and K together (first 2*batch_size in dim 0)
        qk_part = qkv&#91;:batch_size * 2&#93;
        # Transpose to &#91;2*B, H/SP, T, D&#93; 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 &#91;2*B, T, H/SP, D&#93;
        qkv&#91;:batch_size * 2&#93; = 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(&#91;replicated_q, replicated_k, replicated_v&#93;,
                                   dim=0)  # &#91;3, seq_len, num_heads, head_dim&#93;
        heads_per_rank = num_heads // world_size
        replicated_qkv = replicated_qkv&#91;:, :, local_rank * heads_per_rank:(local_rank + 1) * heads_per_rank&#93;
        qkv = torch.cat(&#91;qkv, replicated_qkv&#93;, 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&#91;:, split_idx:&#93;
        output = output&#91;:, :split_idx&#93;
        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</code></pre></details>
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 &#91;batch_size, seq_len, num_heads, head_dim&#93;
    k: Key tensor &#91;batch_size, seq_len, num_heads, head_dim&#93;
    v: Value tensor &#91;batch_size, seq_len, num_heads, head_dim&#93;
    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 &#91;B, H, T, D&#93;

Returns:
    Tuple of output tensor and optional replicated output
&quot;&quot;&quot;
assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, &quot;Expected 4D tensors&quot;
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), &quot;Replicated QKV is not supported for VSA now&quot;
    if gate_compress is None:
        raise ValueError(&quot;gate_compress must be provided when using VIDEO_SPARSE_ATTN&quot;)

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

    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&#91;:, head_start:head_end, :original_seq_len, :&#93;
        sin_local = sin&#91;:, head_start:head_end, :original_seq_len, :&#93;
        qk_part = qkvg&#91;:batch_size * 2&#93;.transpose(1, 2)
        qk_part = apply_ltx_rotary_emb_4d(qk_part, (cos_local, sin_local), self.rope_type)
        qkvg&#91;:batch_size * 2&#93; = 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&#91;call-arg&#93;
        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(&#91;q, k, v&#93;, dim=0)  # &#91;3*batch, seq_len, num_heads, head_dim&#93;

# 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&#91;1&#93; - original_seq_len
qkv = qkv&#91;:, :original_seq_len, :, :&#93;

# 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&#x27;s heads: &#91;B, H/SP, T, D&#93;
    cos_local = cos&#91;:, head_start:head_end, :original_seq_len, :&#93;
    sin_local = sin&#91;:, head_start:head_end, :original_seq_len, :&#93;

    # Apply RoPE to Q and K together (first 2*batch_size in dim 0)
    qk_part = qkv&#91;:batch_size * 2&#93;
    # Transpose to &#91;2*B, H/SP, T, D&#93; 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 &#91;2*B, T, H/SP, D&#93;
    qkv&#91;:batch_size * 2&#93; = 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(&#91;replicated_q, replicated_k, replicated_v&#93;,
                               dim=0)  # &#91;3, seq_len, num_heads, head_dim&#93;
    heads_per_rank = num_heads // world_size
    replicated_qkv = replicated_qkv&#91;:, :, local_rank * heads_per_rank:(local_rank + 1) * heads_per_rank&#93;
    qkv = torch.cat(&#91;qkv, replicated_qkv&#93;, 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&#91;:, split_idx:&#93;
    output = output&#91;:, :split_idx&#93;
    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</code></pre></details>
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&#91;AttentionBackendEnum, ...&#93; | None = None,
    **extra_impl_args,
) -&gt; 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&#91;torch.Tensor, torch.Tensor&#93; | None = None,
    ltx_k_freqs_cis: tuple&#91;torch.Tensor, torch.Tensor&#93; | None = None,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;Forward pass with LTX-2 style RoPE application.

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

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

    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 &#91;B, H, T, D&#93; 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 &#91;B, T, H, D&#93;
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)

    if self.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
        if gate_compress is None:
            raise ValueError(&quot;gate_compress must be provided when using VIDEO_SPARSE_ATTN&quot;)
        output = self.attn_impl.forward(  # type: ignore&#91;call-arg&#93;
            q, k, v, gate_compress, ctx_attn_metadata)
    else:
        output = self.attn_impl.forward(q, k, v, ctx_attn_metadata)
    return output</code></pre></details>
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 &#91;batch_size, seq_len, num_heads, head_dim&#93;
    k: Key tensor &#91;batch_size, seq_len, num_heads, head_dim&#93;
    v: Value tensor &#91;batch_size, seq_len, num_heads, head_dim&#93;
    ltx_freqs_cis: LTX-2 style RoPE (cos, sin) for Q with shape &#91;B, H, T, D&#93;
    ltx_k_freqs_cis: LTX-2 style RoPE (cos, sin) for K (if different from Q)

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

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 &#91;B, H, T, D&#93; 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 &#91;B, T, H, D&#93;
    q = q.transpose(1, 2)
    k = k.transpose(1, 2)

if self.backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
    if gate_compress is None:
        raise ValueError(&quot;gate_compress must be provided when using VIDEO_SPARSE_ATTN&quot;)
    output = self.attn_impl.forward(  # type: ignore&#91;call-arg&#93;
        q, k, v, gate_compress, ctx_attn_metadata)
else:
    output = self.attn_impl.forward(q, k, v, ctx_attn_metadata)
return output</code></pre></details>
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&#91;AttentionBackendEnum, ...&#93;,
    apply_gated_attention: bool = False,
    quant_config: QuantizationConfig | None = None,
    prefix: str = &quot;&quot;,
) -&gt; 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&quot;&#123;prefix&#125;.to_q&quot;,
    )
    self.to_k = ReplicatedLinear(
        context_dim,
        inner_dim,
        bias=True,
        quant_config=quant_config,
        prefix=f&quot;&#123;prefix&#125;.to_k&quot;,
    )
    self.to_v = ReplicatedLinear(
        context_dim,
        inner_dim,
        bias=True,
        quant_config=quant_config,
        prefix=f&quot;&#123;prefix&#125;.to_v&quot;,
    )
    # 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&quot;&#123;prefix&#125;.to_gate_logits&quot;,
        )
    self.to_out = nn.ModuleList(&#91;
        ReplicatedLinear(
            inner_dim,
            query_dim,
            bias=True,
            quant_config=quant_config,
            prefix=f&quot;&#123;prefix&#125;.to_out&quot;,
        ),
        nn.Identity(),
    &#93;)

    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&quot;&#123;prefix&#125;.to_gate_compress&quot;,
        )
    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&#91;torch.Tensor, torch.Tensor&#93; | None = None,
    k_pe: tuple&#91;torch.Tensor, torch.Tensor&#93; | None = None,
) -&gt; torch.Tensor:
    context = x if context is None else context

    q_pre_quantized = (
        self.to_q.quant_method.quantize_input(x)  # type: ignore&#91;union-attr&#93;
        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&#91;union-attr&#93;
                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)&#91;0&#93; if self.to_gate_logits is not None else None)
    gate_compress = (self.to_gate_compress(context)&#91;0&#93; 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&#91;1&#93;
    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&#91;0&#93;(out)&#91;0&#93;
    return self.to_out&#91;1&#93;(out)</code></pre></details>
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&#91;union-attr&#93;
    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&#91;union-attr&#93;
            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)&#91;0&#93; if self.to_gate_logits is not None else None)
gate_compress = (self.to_gate_compress(context)&#91;0&#93; 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&#91;1&#93;
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&#91;0&#93;(out)&#91;0&#93;
return self.to_out&#91;1&#93;(out)</code></pre></details>
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&#91;AttentionBackendEnum, ...&#93;,
    apply_gated_attention: bool = False,
    quant_config: QuantizationConfig | None = None,
    prefix: str = &quot;&quot;,
) -&gt; 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&quot;&#123;prefix&#125;.to_q&quot;,
    )
    self.to_k = ReplicatedLinear(
        context_dim,
        inner_dim,
        bias=True,
        quant_config=quant_config,
        prefix=f&quot;&#123;prefix&#125;.to_k&quot;,
    )
    self.to_v = ReplicatedLinear(
        context_dim,
        inner_dim,
        bias=True,
        quant_config=quant_config,
        prefix=f&quot;&#123;prefix&#125;.to_v&quot;,
    )
    # 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&quot;&#123;prefix&#125;.to_gate_logits&quot;,
        )
    self.to_out = nn.ModuleList(&#91;
        ReplicatedLinear(
            inner_dim,
            query_dim,
            bias=True,
            quant_config=quant_config,
            prefix=f&quot;&#123;prefix&#125;.to_out&quot;,
        ),
        nn.Identity(),
    &#93;)

    self.attn = LTXDistributedAttention(
        num_heads=heads,
        head_size=dim_head,
        rope_type=rope_type,
        causal=False,
        supported_attention_backends=supported_attention_backends,
        prefix=f&quot;&#123;prefix&#125;.attn&quot;,
    )
    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&quot;&#123;prefix&#125;.to_gate_compress&quot;,
        )

def forward(
    self,
    x: torch.Tensor,
    context: torch.Tensor | None = None,
    pe: tuple&#91;torch.Tensor, torch.Tensor&#93; | None = None,
    k_pe: tuple&#91;torch.Tensor, torch.Tensor&#93; | None = None,
    original_seq_len: int | None = None,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;Forward pass for distributed self-attention.

    Args:
        x: Query tensor &#91;B, L, C&#93;
        context: Key/Value tensor &#91;B, L, C&#93;, or None for self-attention
        pe: Rotary position embeddings for Q (cos, sin) with shape &#91;B, H, T_full, D&#93;
            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
    &quot;&quot;&quot;
    context = x if context is None else context

    q_pre_quantized = (
        self.to_q.quant_method.quantize_input(x)  # type: ignore&#91;union-attr&#93;
        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&#91;union-attr&#93;
                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)&#91;0&#93; if self.to_gate_logits is not None else None)
    gate_compress = (self.to_gate_compress(context)&#91;0&#93; 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&#91;1&#93;
    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&#x27;s heads
    if original_seq_len is None:
        raise ValueError(&quot;original_seq_len must be provided for distributed attention&quot;)
    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&#91;0&#93;(out)&#91;0&#93;
    return self.to_out&#91;1&#93;(out)</code></pre></details>
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 &#91;B, L, C&#93;
    context: Key/Value tensor &#91;B, L, C&#93;, or None for self-attention
    pe: Rotary position embeddings for Q (cos, sin) with shape &#91;B, H, T_full, D&#93;
        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
&quot;&quot;&quot;
context = x if context is None else context

q_pre_quantized = (
    self.to_q.quant_method.quantize_input(x)  # type: ignore&#91;union-attr&#93;
    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&#91;union-attr&#93;
            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)&#91;0&#93; if self.to_gate_logits is not None else None)
gate_compress = (self.to_gate_compress(context)&#91;0&#93; 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&#91;1&#93;
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&#x27;s heads
if original_seq_len is None:
    raise ValueError(&quot;original_seq_len must be provided for distributed attention&quot;)
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&#91;0&#93;(out)&#91;0&#93;
return self.to_out&#91;1&#93;(out)</code></pre></details>
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 = &quot;&quot;,
):
    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 &gt; 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 &gt; 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&quot;&#123;prefix&#125;.blocks.&#123;idx&#125;.attn1&quot;,
        )
        # 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&quot;&#123;prefix&#125;.blocks.&#123;idx&#125;.attn2&quot;,
        )
        self.ff = FeedForward(
            video.dim,
            dim_out=video.dim,
            bias=video.ff_bias,
            quant_config=quant_config,
            prefix=f&quot;&#123;prefix&#125;.blocks.&#123;idx&#125;&quot;,
        )
        # 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 &gt; 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&quot;&#123;prefix&#125;.blocks.&#123;idx&#125;.audio_attn1&quot;,
        )
        # 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&quot;&#123;prefix&#125;.blocks.&#123;idx&#125;.audio_attn2&quot;,
        )
        self.audio_ff = FeedForward(
            audio.dim,
            dim_out=audio.dim,
            bias=audio.ff_bias,
            quant_config=quant_config,
            prefix=f&quot;&#123;prefix&#125;.blocks.&#123;idx&#125;.audio&quot;,
        )
        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&quot;&#123;prefix&#125;.blocks.&#123;idx&#125;.audio_to_video_attn&quot;,
        )
        # 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&quot;&#123;prefix&#125;.blocks.&#123;idx&#125;.video_to_audio_attn&quot;,
        )
        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):
    &quot;&quot;&quot;Register backward hooks on output tensors to trigger FSDP2 unshard.

    FSDP2&#x27;s module-level backward hooks don&#x27;t fire when the module returns 
    dataclass outputs. We must register hooks directly on the output tensors.
    &quot;&quot;&quot;
    if not hasattr(self, &#x27;unshard&#x27;):
        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) -&gt; tuple&#91;torch.Tensor, ...&#93;:
    num_ada_params = scale_shift_table.shape&#91;0&#93;
    ada_values = (
        scale_shift_table&#91;indices&#93;.unsqueeze(0).unsqueeze(0).to(device=timestep.device, dtype=timestep.dtype) +
        timestep.reshape(batch_size, timestep.shape&#91;1&#93;, num_ada_params, -1)&#91;:, :, indices, :&#93;).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,
) -&gt; tuple&#91;torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor&#93;:
    scale_shift_ada_values = self.get_ada_values(scale_shift_table&#91;:num_scale_shift_values, :&#93;, batch_size,
                                                 scale_shift_timestep, slice(None, None))
    gate_ada_values = self.get_ada_values(scale_shift_table&#91;num_scale_shift_values:, :&#93;, batch_size, gate_timestep,
                                          slice(None, None))

    scale_shift_chunks = &#91;t.squeeze(2) for t in scale_shift_ada_values&#93;
    gate_ada_values = &#91;t.squeeze(2) for t in gate_ada_values&#93;

    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,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;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)``.
    &quot;&quot;&quot;
    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&#91;0&#93;, 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,
) -&gt; tuple&#91;TransformerArgs | None, TransformerArgs | None&#93;:
    &quot;&quot;&quot;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).
    &quot;&quot;&quot;
    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() &gt; 0
    run_ax = audio is not None and audio.enabled and ax.numel() &gt; 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,
    ) -&gt; torch.Tensor:
        &quot;&quot;&quot;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 ``&#91;B&#93;`` (1/True == perturb). When skip_flag is False/0 the
        multiplier is all ones, reproducing the un-skipped LTX-2.0 path.
        &quot;&quot;&quot;
        bsz = values.shape&#91;0&#93;
        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&#91;0&#93; != bsz:
                    raise ValueError(&quot;Per-sample STG mask batch size mismatch: &quot;
                                     f&quot;got &#123;skip_flag.shape&#91;0&#93;&#125;, expected &#123;bsz&#125;&quot;)
                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, *(&#91;1&#93; * (values.ndim - 1)))

    if run_vx:
        vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(self.scale_shift_table, vx.shape&#91;0&#93;,
                                                                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, &quot;prompt_scale_shift_table&quot;, 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&#91;0&#93;,
                                                                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, &quot;audio_prompt_scale_shift_table&quot;, 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&#91;0&#93;,
            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&#91;0&#93;,
            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&#91;1&#93;
                    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&#91;:, :, start_idx:end_idx, :&#93; 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&#91;0&#93;.shape&#91;2&#93;
                    ax_context = ax_context&#91;:, :original_audio_len, :&#93;
                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&#91;0&#93;.shape&#91;2&#93;
                    ax_full = ax_full&#91;:, :original_audio_len, :&#93;
                if video_k_pe is not None:
                    original_video_len = video_k_pe&#91;0&#93;.shape&#91;2&#93;
                    vx_context = vx_context&#91;:, :original_video_len, :&#93;
                # 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&#x27;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&#91;0&#93;,
                                                                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&#91;0&#93;,
                                                                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&#x27;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,
    )</code></pre></details>
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 = &#91;t.squeeze(2) for t in scale_shift_ada_values&#93;
gate_ada_values = &#91;t.squeeze(2) for t in gate_ada_values&#93;

return (*scale_shift_chunks, *gate_ada_values)</code></pre></details>
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).
&quot;&quot;&quot;
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() &gt; 0
run_ax = audio is not None and audio.enabled and ax.numel() &gt; 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,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;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 ``&#91;B&#93;`` (1/True == perturb). When skip_flag is False/0 the
    multiplier is all ones, reproducing the un-skipped LTX-2.0 path.
    &quot;&quot;&quot;
    bsz = values.shape&#91;0&#93;
    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&#91;0&#93; != bsz:
                raise ValueError(&quot;Per-sample STG mask batch size mismatch: &quot;
                                 f&quot;got &#123;skip_flag.shape&#91;0&#93;&#125;, expected &#123;bsz&#125;&quot;)
            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, *(&#91;1&#93; * (values.ndim - 1)))

if run_vx:
    vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(self.scale_shift_table, vx.shape&#91;0&#93;,
                                                            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, &quot;prompt_scale_shift_table&quot;, 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&#91;0&#93;,
                                                            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, &quot;audio_prompt_scale_shift_table&quot;, 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&#91;0&#93;,
        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&#91;0&#93;,
        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&#91;1&#93;
                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&#91;:, :, start_idx:end_idx, :&#93; 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&#91;0&#93;.shape&#91;2&#93;
                ax_context = ax_context&#91;:, :original_audio_len, :&#93;
            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&#91;0&#93;.shape&#91;2&#93;
                ax_full = ax_full&#91;:, :original_audio_len, :&#93;
            if video_k_pe is not None:
                original_video_len = video_k_pe&#91;0&#93;.shape&#91;2&#93;
                vx_context = vx_context&#91;:, :original_video_len, :&#93;
            # 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&#x27;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&#91;0&#93;,
                                                            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&#91;0&#93;,
                                                            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&#x27;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,
)</code></pre></details>
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.
&quot;&quot;&quot;
if prompt_scale_shift_table is None:
    raise ValueError(&quot;cross_attention_adaln requires a prompt scale/shift table.&quot;)
batch_size = x.shape&#91;0&#93;
kv_modulation = prompt_scale_shift_table&#91;None, None&#93;.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&#91;1&#93;, 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</code></pre></details>
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) -&gt; bool:
    return self in (LTXModelType.AudioVideo, LTXModelType.VideoOnly)

def is_audio_enabled(self) -&gt; bool:
    return self in (LTXModelType.AudioVideo, LTXModelType.AudioOnly)</code></pre></details>
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&#91;int&#93; | 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&#91;int&#93; | 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 = &quot;&quot;,
):
    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 = &#91;20, 2048, 2048&#93;
        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 = &#91;20&#93;
        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&#91;0&#93;, self.audio_positional_embedding_max_pos&#91;0&#93;)
        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) -&gt; torch.Tensor | None:
    &quot;&quot;&quot;Resolve the optional parameter at call time, matching upstream.&quot;&quot;&quot;
    return getattr(self, &quot;keyframes_abs_pos_embedding&quot;, None)

def _init_video(
    self,
    in_channels: int,
    out_channels: int,
    caption_channels: int,
    norm_eps: float,
    create_caption_projection: bool = True,
) -&gt; 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,
) -&gt; 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) -&gt; 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) -&gt; 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, &quot;caption_projection&quot;, 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, &quot;prompt_adaln_single&quot;, 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, &quot;audio_caption_projection&quot;, 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, &quot;audio_prompt_adaln_single&quot;, 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, &quot;caption_projection&quot;, 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, &quot;prompt_adaln_single&quot;, 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, &quot;audio_caption_projection&quot;, 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, &quot;audio_prompt_adaln_single&quot;, 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 = &quot;&quot;,
) -&gt; 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(&#91;
        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)
    &#93;)

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&#91;int&#93; | None = None,
    skip_audio_self_attn_blocks: list&#91;int&#93; | None = None,
) -&gt; tuple&#91;TransformerArgs | None, TransformerArgs | None&#93;:
    # Convert once so per-block membership checks stay O(1).
    skip_video_self_attn_block_set = set(skip_video_self_attn_blocks or &#91;&#93;)
    skip_audio_self_attn_block_set = set(skip_audio_self_attn_blocks or &#91;&#93;)

    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,
) -&gt; torch.Tensor:
    scale_shift_values = (scale_shift_table&#91;None, None&#93;.to(device=x.device, dtype=x.dtype) +
                          embedded_timestep&#91;:, :, None&#93;)
    shift, scale = scale_shift_values&#91;:, :, 0&#93;, scale_shift_values&#91;:, :, 1&#93;
    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&#91;int&#93; | None = None,
    skip_audio_self_attn_blocks: list&#91;int&#93; | None = None,
) -&gt; tuple&#91;torch.Tensor | None, torch.Tensor | None&#93;:
    &quot;&quot;&quot;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).
    &quot;&quot;&quot;
    if not self.model_type.is_video_enabled() and video is not None:
        raise ValueError(&quot;Video is not enabled for this model&quot;)
    if not self.model_type.is_audio_enabled() and audio is not None:
        raise ValueError(&quot;Audio is not enabled for this model&quot;)

    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</code></pre></details>
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).
&quot;&quot;&quot;
if not self.model_type.is_video_enabled() and video is not None:
    raise ValueError(&quot;Video is not enabled for this model&quot;)
if not self.model_type.is_audio_enabled() and audio is not None:
    raise ValueError(&quot;Audio is not enabled for this model&quot;)

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</code></pre></details>
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&#91;str, Any&#93;):
    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 &gt; 1 or use_vsa_backend

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

    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&#91;1&#93;)
    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&#91;torch.Tensor&#93;,
    timestep: torch.LongTensor,
    encoder_hidden_states_image: torch.Tensor | list&#91;torch.Tensor&#93; | 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&#91;int&#93; | None = None,
    skip_audio_self_attn_blocks: list&#91;int&#93; | None = None,
    video_position_offset_sec: float = 0.0,
    **kwargs,
) -&gt; torch.Tensor:
    if isinstance(encoder_hidden_states, list):
        encoder_hidden_states = encoder_hidden_states&#91;0&#93;
    # 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&#91;0&#93; 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&#91;:, 0&#93; if timestep.ndim &gt; 1 else timestep

    # Shard video latents and timestep across SP ranks
    video_original_seq_len = latents.shape&#91;1&#93;
    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&#91;-1&#93; != 1:
            raise ValueError(
                &quot;keyframes_mask must contain patchified per-token markers with shape (B, T, 1)&quot;
            )
        if keyframes_mask.shape&#91;:2&#93; != latents.shape&#91;:2&#93;:
            raise ValueError(
                &quot;keyframes_mask batch/sequence dimensions must match the patchified video latents: &quot;
                f&quot;mask=&#123;tuple(keyframes_mask.shape)&#125;, latents=&#123;tuple(latents.shape)&#125;&quot;
            )
        keyframes_mask = keyframes_mask.to(device=latents.device)
    if sp_world_size &gt; 1:
        latents, video_original_seq_len = sequence_model_parallel_shard(latents, dim=1)
        # Shard timestep along sequence dimension (timestep has shape &#91;batch, seq_len&#93;)
        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&#91;1&#93;
        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&#91;:, 0, ...&#93; = positions&#91;:, 0, ...&#93; + 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 &gt; 1 and video_padded_seq_len &gt; video_original_seq_len:
        padding_needed = video_padded_seq_len - video_original_seq_len
        last_pos = positions&#91;:, :, -1:, :&#93;.expand(-1, -1, padding_needed, -1)
        positions = torch.cat(&#91;positions, last_pos&#93;, 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&#91;:, 0&#93; if audio_timestep.ndim &gt; 1 else audio_timestep)

        # Shard audio latents and timestep across SP ranks
        audio_original_seq_len = audio_latents.shape&#91;1&#93;
        sharded_audio_timestep = audio_timestep
        if sp_world_size &gt; 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 &gt; 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 &gt; 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</code></pre></details>
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&#91;0&#93; 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&#91;:, 0&#93; if timestep.ndim &gt; 1 else timestep

# Shard video latents and timestep across SP ranks
video_original_seq_len = latents.shape&#91;1&#93;
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&#91;-1&#93; != 1:
        raise ValueError(
            &quot;keyframes_mask must contain patchified per-token markers with shape (B, T, 1)&quot;
        )
    if keyframes_mask.shape&#91;:2&#93; != latents.shape&#91;:2&#93;:
        raise ValueError(
            &quot;keyframes_mask batch/sequence dimensions must match the patchified video latents: &quot;
            f&quot;mask=&#123;tuple(keyframes_mask.shape)&#125;, latents=&#123;tuple(latents.shape)&#125;&quot;
        )
    keyframes_mask = keyframes_mask.to(device=latents.device)
if sp_world_size &gt; 1:
    latents, video_original_seq_len = sequence_model_parallel_shard(latents, dim=1)
    # Shard timestep along sequence dimension (timestep has shape &#91;batch, seq_len&#93;)
    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&#91;1&#93;
    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&#91;:, 0, ...&#93; = positions&#91;:, 0, ...&#93; + 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 &gt; 1 and video_padded_seq_len &gt; video_original_seq_len:
    padding_needed = video_padded_seq_len - video_original_seq_len
    last_pos = positions&#91;:, :, -1:, :&#93;.expand(-1, -1, padding_needed, -1)
    positions = torch.cat(&#91;positions, last_pos&#93;, 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&#91;:, 0&#93; if audio_timestep.ndim &gt; 1 else audio_timestep)

    # Shard audio latents and timestep across SP ranks
    audio_original_seq_len = audio_latents.shape&#91;1&#93;
    sharded_audio_timestep = audio_timestep
    if sp_world_size &gt; 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 &gt; 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 &gt; 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</code></pre></details>