# Copyright 2025 The Lightricks team and The HuggingFace Team. # All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import inspect from collections.abc import Iterable from dataclasses import dataclass from types import SimpleNamespace from typing import TYPE_CHECKING, Any import torch import torch.distributed import torch.nn as nn import torch.nn.functional as F from diffusers.models.embeddings import PixArtAlphaCombinedTimestepSizeEmbeddings, PixArtAlphaTextProjection from diffusers.utils import ( BaseOutput, is_torch_version, ) from torch.utils.checkpoint import checkpoint from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, QKVParallelLinear, RowParallelLinear, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available logger = init_logger(__name__) _RMSNORM_INIT_PARAMS = inspect.signature(RMSNorm.__init__).parameters def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) -> nn.Module: """Bridge diffusers' RMSNorm API onto vLLM's `has_weight` variant.""" kwargs: dict[str, Any] = {"eps": eps} if "elementwise_affine" in _RMSNORM_INIT_PARAMS: kwargs["elementwise_affine"] = elementwise_affine elif "has_weight" in _RMSNORM_INIT_PARAMS: kwargs["has_weight"] = elementwise_affine elif not elementwise_affine: raise TypeError("RMSNorm backend does not support disabling affine weights.") return RMSNorm(hidden_size, **kwargs) def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: cos, sin = freqs # Concrete pair count instead of -1 keeps SDPA shape static under torch.compile. x_real, x_imag = x.unflatten(2, (x.shape[2] // 2, 2)).unbind(-1) # [B, S, C // 2] x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2) out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) return out def apply_split_rotary_emb( x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor], *, head_dim: int, ) -> torch.Tensor: # `head_dim` is plumbed in (not inferred via `-1`) so SDPA shape stays static under torch.compile. cos, sin = freqs x_dtype = x.dtype needs_reshape = False if x.ndim != 4 and cos.ndim == 4: # cos is (#b, h, t, r) -> reshape x to (b, h, t, dim_per_head) # The cos/sin batch dim may only be broadcastable, so take batch size from x b = x.shape[0] _, h, t, _ = cos.shape x = x.reshape(b, t, h, head_dim).swapaxes(1, 2) needs_reshape = True # Split last dim (2*r) into (d=2, r) last = x.shape[-1] if last % 2 != 0: raise ValueError(f"Expected x.shape[-1] to be even for split rotary, got {last}.") r = last // 2 # (..., 2, r) split_x = x.reshape(*x.shape[:-1], 2, r).float() # Explicitly upcast to float first_x = split_x[..., :1, :] # (..., 1, r) second_x = split_x[..., 1:, :] # (..., 1, r) cos_u = cos.unsqueeze(-2) # broadcast to (..., 1, r) against (..., 2, r) sin_u = sin.unsqueeze(-2) out = split_x * cos_u first_out = out[..., :1, :] second_out = out[..., 1:, :] first_out.addcmul_(-sin_u, second_x) second_out.addcmul_(sin_u, first_x) out = out.reshape(*out.shape[:-2], last) if needs_reshape: out = out.swapaxes(1, 2).reshape(b, t, h * head_dim) out = out.to(dtype=x_dtype) return out @dataclass class AudioVisualModelOutput(BaseOutput): r""" Holds the output of an audiovisual model which produces both visual (e.g. video) and audio outputs. Args: sample (`torch.Tensor` of shape `(batch_size, num_video_tokens, out_channels)`): The patchified visual output conditioned on the `encoder_hidden_states` input. This is the transformer output before the pipeline unpacks it back into video latent dimensions. audio_sample (`torch.Tensor` of shape `(batch_size, num_audio_tokens, audio_out_channels)`): The patchified audio output of the audiovisual model before the pipeline unpacks it back into audio latent dimensions. """ sample: "torch.Tensor" # noqa: F821 audio_sample: "torch.Tensor" # noqa: F821 class LTX2AdaLayerNormSingle(nn.Module): r""" Norm layer adaptive layer norm single (adaLN-single). As proposed in PixArt-Alpha (see: https://huggingface.co/papers/2310.00426; Section 2.3) and adapted by the LTX-2.0 model. In particular, the number of modulation parameters to be calculated is now configurable. Parameters: embedding_dim (`int`): The size of each embedding vector. num_mod_params (`int`, *optional*, defaults to `6`): The number of modulation parameters which will be calculated in the first return argument. The default of 6 is standard, but sometimes we may want to have a different (usually smaller) number of modulation parameters. use_additional_conditions (`bool`, *optional*, defaults to `False`): Whether to use additional conditions for normalization or not. """ def __init__(self, embedding_dim: int, num_mod_params: int = 6, use_additional_conditions: bool = False): super().__init__() self.num_mod_params = num_mod_params self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings( embedding_dim, size_emb_dim=embedding_dim // 3, use_additional_conditions=use_additional_conditions ) self.silu = nn.SiLU() self.linear = nn.Linear(embedding_dim, self.num_mod_params * embedding_dim, bias=True) def forward( self, timestep: torch.Tensor, added_cond_kwargs: dict[str, torch.Tensor] | None = None, batch_size: int | None = None, hidden_dtype: torch.dtype | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: # No modulation happening here. added_cond_kwargs = added_cond_kwargs or {"resolution": None, "aspect_ratio": None} embedded_timestep = self.emb(timestep, **added_cond_kwargs, batch_size=batch_size, hidden_dtype=hidden_dtype) return self.linear(self.silu(embedded_timestep)), embedded_timestep class ColumnParallelApproxGELU(nn.Module): def __init__( self, dim_in: int, dim_out: int, *, approximate: str, bias: bool = True, quant_config: "QuantizationConfig | None" = None, prefix: str = "", ): super().__init__() self.proj = ColumnParallelLinear( dim_in, dim_out, bias=bias, gather_output=False, return_bias=False, quant_config=quant_config, prefix=f"{prefix}.proj" if prefix else "proj", ) self.approximate = approximate def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.proj(x) return F.gelu(x, approximate=self.approximate) class LTX2FeedForward(nn.Module): def __init__( self, dim: int, dim_out: int | None = None, mult: int = 4, activation_fn: str = "gelu-approximate", inner_dim: int | None = None, bias: bool = True, dropout: float = 0.0, final_dropout: bool = False, quant_config: "QuantizationConfig | None" = None, prefix: str = "", ) -> None: super().__init__() assert activation_fn == "gelu-approximate", "Only gelu-approximate is supported." inner_dim = inner_dim or int(dim * mult) dim_out = dim_out or dim dropout_layer: nn.Module = nn.Dropout(dropout) if dropout > 0 else nn.Identity() layers: list[nn.Module] = [ ColumnParallelApproxGELU( dim, inner_dim, approximate="tanh", bias=bias, quant_config=quant_config, prefix=f"{prefix}.net.0" if prefix else "net.0", ), dropout_layer, RowParallelLinear( inner_dim, dim_out, input_is_parallel=True, return_bias=False, quant_config=quant_config, prefix=f"{prefix}.net.2" if prefix else "net.2", ), ] if final_dropout: layers.append(nn.Dropout(dropout) if dropout > 0 else nn.Identity()) self.net = nn.ModuleList(layers) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: for module in self.net: hidden_states = module(hidden_states) return hidden_states class TensorParallelRMSNorm(nn.Module): """RMSNorm that computes stats across TP shards for q/k norm. LTX2 uses qk_norm="rms_norm_across_heads" while Q/K are tensor-parallel sharded. A local RMSNorm would compute statistics on only the local shard, which changes the normalization when TP > 1. We all-reduce the squared sum to match the global RMS across all heads. """ def __init__(self, hidden_size: int, eps: float = 1e-6, elementwise_affine: bool = True, tp_size: int = 1): super().__init__() self.hidden_size = hidden_size self.global_hidden_size = hidden_size * max(tp_size, 1) self.eps = eps self.tp_size = tp_size if elementwise_affine: self.weight = nn.Parameter(torch.ones(hidden_size)) else: self.register_parameter("weight", None) def _all_reduce(self, tensor: torch.Tensor) -> None: if not torch.distributed.is_initialized(): return try: import vllm.distributed.parallel_state as vllm_parallel_state tp_group = getattr(vllm_parallel_state, "_TP", None) except Exception: tp_group = None if tp_group is not None and hasattr(tp_group, "all_reduce"): tp_group.all_reduce(tensor) return if tp_group is not None: torch.distributed.all_reduce(tensor, group=tp_group) else: torch.distributed.all_reduce(tensor) def forward(self, x: torch.Tensor) -> torch.Tensor: x_dtype = x.dtype x_float = x.float() local_sum = x_float.pow(2).sum(dim=-1, keepdim=True) if self.tp_size > 1: self._all_reduce(local_sum) inv_rms = torch.rsqrt(local_sum / self.global_hidden_size + self.eps) out = x_float * inv_rms if self.weight is not None: out = out * self.weight.float() return out.to(dtype=x_dtype) class LTX2AudioVideoAttnProcessor: r""" Processor for implementing attention (SDPA is used by default if you're using PyTorch 2.0) for the LTX-2.0 model. Compared to the LTX-1.0 model, we allow the RoPE embeddings for the queries and keys to be separate so that we can support audio-to-video (a2v) and video-to-audio (v2a) cross attention. """ _attention_backend = None _parallel_config = None def __init__(self): if is_torch_version("<", "2.0"): raise ValueError( "LTX attention processors require a minimum PyTorch version of 2.0. " "Please upgrade your PyTorch installation." ) @staticmethod def _to_padding_mask(attention_mask: torch.Tensor) -> torch.Tensor: # Convert additive/expanded masks into a 2D padding mask for flash-attn. if attention_mask.ndim > 2: if attention_mask.is_floating_point(): valid = attention_mask >= 0 else: valid = attention_mask.to(torch.bool) b = valid.shape[0] key_len = valid.shape[-1] valid = valid.reshape(b, -1, key_len).all(dim=1) attention_mask = valid if attention_mask.is_floating_point(): attention_mask = attention_mask >= 0 if attention_mask.dtype != torch.bool: attention_mask = attention_mask.to(torch.bool) return attention_mask @staticmethod def _is_sp_enabled() -> bool: if not is_forward_context_available(): return False try: od_config = get_forward_context().omni_diffusion_config parallel_config = getattr(od_config, "parallel_config", None) if od_config is not None else None return getattr(parallel_config, "sequence_parallel_size", 1) > 1 except Exception: return False def _prepare_attention_mask( self, attn: "LTX2Attention", hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None, attention_mask: torch.Tensor | None, batch_size: int, sequence_length: int, ) -> torch.Tensor | None: if attention_mask is None: return None if self._is_sp_enabled(): # In SP, Ulysses expects a 2D padding mask that matches query length. # For cross-attention, encoder sequence length != query length, so drop the mask. if encoder_hidden_states is not None and encoder_hidden_states.shape[1] != hidden_states.shape[1]: return None return self._to_padding_mask(attention_mask) attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) if attn.attn.attn_backend.get_name().upper() == "FLASH_ATTN": attention_mask = self._to_padding_mask(attention_mask) return attention_mask @staticmethod def _project_qkv( attn: "LTX2Attention", hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None, is_self_attention: bool, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if is_self_attention and attn.to_qkv is not None: qkv, _ = attn.to_qkv(hidden_states) q_heads = getattr(attn, "query_num_heads", attn.heads) kv_heads = getattr(attn, "kv_num_heads", attn.heads) q_size = q_heads * attn.head_dim kv_size = kv_heads * attn.head_dim query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1) return query, key, value query = attn.to_q(hidden_states) if isinstance(query, tuple): query = query[0] if encoder_hidden_states is None: raise ValueError("encoder_hidden_states is required for cross-attention projection.") key = attn.to_k(encoder_hidden_states) if isinstance(key, tuple): key = key[0] value = attn.to_v(encoder_hidden_states) if isinstance(value, tuple): value = value[0] return query, key, value @staticmethod def _slice_rope_for_tp( rope: tuple[torch.Tensor, torch.Tensor] | None, attn_module: "LTX2Attention", ) -> tuple[torch.Tensor, torch.Tensor] | None: if rope is None: return None cos, sin = rope tp_size = get_tensor_model_parallel_world_size() if tp_size <= 1: return rope tp_rank = get_tensor_model_parallel_rank() if cos.ndim == 4: if cos.shape[1] != attn_module.heads: local_heads = cos.shape[1] // tp_size if local_heads == attn_module.heads: start = tp_rank * local_heads end = start + local_heads cos = cos[:, start:end, :, :] sin = sin[:, start:end, :, :] elif cos.ndim == 3: local_dim = attn_module.heads * attn_module.head_dim if cos.shape[-1] != local_dim: if cos.shape[-1] == local_dim * tp_size: start = tp_rank * local_dim end = start + local_dim cos = cos[..., start:end] sin = sin[..., start:end] return cos, sin def __call__( self, attn: "LTX2Attention", hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: is_self_attention = encoder_hidden_states is None batch_size, sequence_length, _ = hidden_states.shape if is_self_attention else encoder_hidden_states.shape attention_mask = self._prepare_attention_mask( attn=attn, hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask, batch_size=batch_size, sequence_length=sequence_length, ) # Compute gate logits from original hidden_states (before attention) if attn.to_gate_logits is not None: gate_logits = attn.to_gate_logits(hidden_states) if is_self_attention: encoder_hidden_states = hidden_states query, key, value = self._project_qkv( attn=attn, hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, is_self_attention=is_self_attention, ) query = attn.norm_q(query) key = attn.norm_k(key) if query_rotary_emb is not None: query_rotary_emb = self._slice_rope_for_tp(query_rotary_emb, attn) if key_rotary_emb is not None: key_rotary_emb = self._slice_rope_for_tp(key_rotary_emb, attn) if attn.rope_type == "interleaved": query = apply_interleaved_rotary_emb(query, query_rotary_emb) key = apply_interleaved_rotary_emb( key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb ) elif attn.rope_type == "split": query = apply_split_rotary_emb(query, query_rotary_emb, head_dim=attn.head_dim) key = apply_split_rotary_emb( key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb, head_dim=attn.head_dim ) query = query.unflatten(2, (attn.heads, attn.head_dim)) key = key.unflatten(2, (attn.heads, attn.head_dim)) value = value.unflatten(2, (attn.heads, attn.head_dim)) attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None hidden_states = attn.attn(query, key, value, attn_metadata) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) # LTX-2.3: per-head gated attention if attn.to_gate_logits is not None: hidden_states = hidden_states.unflatten(2, (attn.heads, attn.head_dim)) # [B, T, H, D] # 2.0 * sigmoid so zero-init gates produce 1.0 (identity) gates = 2.0 * torch.sigmoid(gate_logits) # [B, T, H] hidden_states = hidden_states * gates.unsqueeze(-1) hidden_states = hidden_states.flatten(2, 3) hidden_states = attn.to_out[0](hidden_states) if isinstance(hidden_states, tuple): hidden_states = hidden_states[0] hidden_states = attn.to_out[1](hidden_states) return hidden_states class LTX2Attention(torch.nn.Module): r""" Attention class for all LTX-2.0 attention layers. Compared to LTX-1.0, this supports specifying the query and key RoPE embeddings separately for audio-to-video (a2v) and video-to-audio (v2a) cross-attention. """ _default_processor_cls = LTX2AudioVideoAttnProcessor _available_processors = [LTX2AudioVideoAttnProcessor] def __init__( self, query_dim: int, heads: int = 8, kv_heads: int = 8, dim_head: int = 64, dropout: float = 0.0, bias: bool = True, cross_attention_dim: int | None = None, out_bias: bool = True, qk_norm: str = "rms_norm_across_heads", norm_eps: float = 1e-6, norm_elementwise_affine: bool = True, rope_type: str = "interleaved", apply_gated_attention: bool = False, processor=None, quant_config: "QuantizationConfig | None" = None, prefix: str = "", disable_kv_quant: bool = False, ): super().__init__() # LTX-2 uses "rms_norm_across_heads", LTX-2.3 uses "rms_norm" -- both # map to the same RMSNorm implementation applied across Q/K heads. if qk_norm not in ("rms_norm_across_heads", "rms_norm"): raise NotImplementedError( f"Only 'rms_norm_across_heads' and 'rms_norm' are supported for `qk_norm`, got {qk_norm!r}." ) kv_heads = heads if kv_heads is None else kv_heads self.head_dim = dim_head self.inner_dim = dim_head * heads self.inner_kv_dim = dim_head * kv_heads self.query_dim = query_dim self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim self.use_bias = bias self.dropout = dropout self.out_dim = query_dim self.total_num_heads = heads self.total_num_kv_heads = kv_heads self.rope_type = rope_type self.to_qkv = None self.to_q = None self.to_k = None self.to_v = None if cross_attention_dim is None: self.to_qkv = QKVParallelLinear( hidden_size=query_dim, head_size=self.head_dim, total_num_heads=heads, bias=bias, quant_config=quant_config, prefix=f"{prefix}.to_qkv" if prefix else "to_qkv", ) self.query_num_heads = self.to_qkv.num_heads self.kv_num_heads = self.to_qkv.num_kv_heads else: tp_size = get_tensor_model_parallel_world_size() self.query_num_heads = heads // tp_size self.kv_num_heads = kv_heads // tp_size self.to_q = ColumnParallelLinear( query_dim, self.inner_dim, bias=bias, gather_output=False, return_bias=False, quant_config=quant_config, prefix=f"{prefix}.to_q" if prefix else "to_q", ) self.to_k = ColumnParallelLinear( self.cross_attention_dim, self.inner_kv_dim, bias=bias, gather_output=False, return_bias=False, quant_config=quant_config, prefix=f"{prefix}.to_k" if prefix else "to_k", ) self.to_v = ColumnParallelLinear( self.cross_attention_dim, self.inner_kv_dim, bias=bias, gather_output=False, return_bias=False, quant_config=quant_config, prefix=f"{prefix}.to_v" if prefix else "to_v", ) self.heads = self.query_num_heads tp_size = get_tensor_model_parallel_world_size() # At TP > 1 with rms_norm_across_heads, use TensorParallelRMSNorm # which all-reduces squared sums to match global RMS statistics. # At TP=1, use torch.nn.RMSNorm which is numerically identical to # the diffusers reference (verified via hook-based comparison). if tp_size > 1 and qk_norm == "rms_norm_across_heads": self.norm_q = TensorParallelRMSNorm( dim_head * self.query_num_heads, eps=norm_eps, elementwise_affine=norm_elementwise_affine, tp_size=tp_size, ) self.norm_k = TensorParallelRMSNorm( dim_head * self.kv_num_heads, eps=norm_eps, elementwise_affine=norm_elementwise_affine, tp_size=tp_size, ) else: self.norm_q = torch.nn.RMSNorm( dim_head * self.query_num_heads, eps=norm_eps, elementwise_affine=norm_elementwise_affine, ) self.norm_k = torch.nn.RMSNorm( dim_head * self.kv_num_heads, eps=norm_eps, elementwise_affine=norm_elementwise_affine, ) self.to_out = torch.nn.ModuleList( [ RowParallelLinear( self.inner_dim, self.out_dim, bias=out_bias, input_is_parallel=True, return_bias=False, quant_config=quant_config, prefix=f"{prefix}.to_out.0" if prefix else "to_out.0", ), torch.nn.Dropout(dropout) if dropout > 0 else torch.nn.Identity(), ] ) self.attn = Attention( num_heads=self.query_num_heads, head_size=dim_head, num_kv_heads=self.kv_num_heads, softmax_scale=1.0 / (dim_head**0.5), causal=False, prefix=prefix, disable_kv_quant=disable_kv_quant, ) # LTX-2.3: per-head gated attention # leave unquantized for this linear if apply_gated_attention: self.to_gate_logits = nn.Linear(query_dim, heads, bias=True) else: self.to_gate_logits = None if processor is None: processor = self._default_processor_cls() self.set_processor(processor) def set_processor(self, processor: Any) -> None: if processor is None: raise ValueError("processor must not be None.") self.processor = processor def get_processor(self) -> Any: return self.processor def prepare_attention_mask( self, attention_mask: torch.Tensor | None, target_length: int, batch_size: int, out_dim: int = 3, ) -> torch.Tensor | None: if attention_mask is None: return None current_length = attention_mask.shape[-1] if current_length != target_length: pad_length = target_length - current_length if pad_length > 0: attention_mask = F.pad(attention_mask, (0, pad_length), value=0.0) else: attention_mask = attention_mask[..., :target_length] if out_dim == 3: expected_batch = batch_size * self.heads if attention_mask.shape[0] != expected_batch: repeat_factor = expected_batch // attention_mask.shape[0] if repeat_factor * attention_mask.shape[0] != expected_batch: raise ValueError( "attention_mask batch dimension is incompatible with the requested batch/head expansion: " f"got {attention_mask.shape[0]}, expected a divisor of {expected_batch}." ) attention_mask = attention_mask.repeat_interleave(repeat_factor, dim=0) elif out_dim == 4: attention_mask = attention_mask.unsqueeze(1) attention_mask = attention_mask.repeat_interleave(self.heads, dim=1) else: raise ValueError(f"Unsupported out_dim={out_dim}; expected 3 or 4.") return attention_mask def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, **kwargs, ) -> torch.Tensor: attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters] if len(unused_kwargs) > 0: logger.warning( "attention_kwargs %s are not expected by %s and will be ignored.", unused_kwargs, self.processor.__class__.__name__, ) kwargs = {k: w for k, w in kwargs.items() if k in attn_parameters} hidden_states = self.processor( self, hidden_states, encoder_hidden_states, attention_mask, query_rotary_emb, key_rotary_emb, **kwargs ) return hidden_states class LTX2VideoTransformerBlock(nn.Module): r""" Transformer block used in [LTX-2.0](https://huggingface.co/Lightricks/LTX-Video). Args: dim (`int`): The number of channels in the input and output. num_attention_heads (`int`): The number of heads to use for multi-head attention. attention_head_dim (`int`): The number of channels in each head. qk_norm (`str`, defaults to `"rms_norm"`): The normalization layer to use. activation_fn (`str`, defaults to `"gelu-approximate"`): Activation function to use in feed-forward. eps (`float`, defaults to `1e-6`): Epsilon value for normalization layers. """ def __init__( self, dim: int, num_attention_heads: int, attention_head_dim: int, cross_attention_dim: int, audio_dim: int, audio_num_attention_heads: int, audio_attention_head_dim, audio_cross_attention_dim: int, video_gated_attn: bool = False, video_cross_attn_adaln: bool = False, audio_gated_attn: bool = False, audio_cross_attn_adaln: bool = False, qk_norm: str = "rms_norm_across_heads", activation_fn: str = "gelu-approximate", attention_bias: bool = True, attention_out_bias: bool = True, eps: float = 1e-6, elementwise_affine: bool = False, rope_type: str = "interleaved", perturbed_attn: bool = False, quant_config: "QuantizationConfig | None" = None, prefix: str = "", ): super().__init__() self.video_cross_attn_adaln = video_cross_attn_adaln self.audio_cross_attn_adaln = audio_cross_attn_adaln self.perturbed_attn = perturbed_attn # 1. Self-Attention (video and audio) self.norm1 = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine) self.attn1 = LTX2Attention( query_dim=dim, heads=num_attention_heads, kv_heads=num_attention_heads, dim_head=attention_head_dim, bias=attention_bias, cross_attention_dim=None, out_bias=attention_out_bias, qk_norm=qk_norm, rope_type=rope_type, apply_gated_attention=video_gated_attn, quant_config=quant_config, prefix=f"{prefix}.attn1" if prefix else "attn1", ) self.audio_norm1 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) self.audio_attn1 = LTX2Attention( query_dim=audio_dim, heads=audio_num_attention_heads, kv_heads=audio_num_attention_heads, dim_head=audio_attention_head_dim, bias=attention_bias, cross_attention_dim=None, out_bias=attention_out_bias, qk_norm=qk_norm, rope_type=rope_type, apply_gated_attention=audio_gated_attn, quant_config=quant_config, prefix=f"{prefix}.audio_attn1" if prefix else "audio_attn1", ) # 2. Prompt Cross-Attention self.norm2 = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine) self.attn2 = LTX2Attention( query_dim=dim, cross_attention_dim=cross_attention_dim, heads=num_attention_heads, kv_heads=num_attention_heads, dim_head=attention_head_dim, bias=attention_bias, out_bias=attention_out_bias, qk_norm=qk_norm, rope_type=rope_type, apply_gated_attention=video_gated_attn, quant_config=quant_config, prefix=f"{prefix}.attn2" if prefix else "attn2", disable_kv_quant=True, ) self.audio_norm2 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) self.audio_attn2 = LTX2Attention( query_dim=audio_dim, cross_attention_dim=audio_cross_attention_dim, heads=audio_num_attention_heads, kv_heads=audio_num_attention_heads, dim_head=audio_attention_head_dim, bias=attention_bias, out_bias=attention_out_bias, qk_norm=qk_norm, rope_type=rope_type, apply_gated_attention=audio_gated_attn, quant_config=quant_config, prefix=f"{prefix}.audio_attn2" if prefix else "audio_attn2", disable_kv_quant=True, ) # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention # K/V here come from the other modality's latents (long sequences), # so KV-cache quant remains enabled. self.audio_to_video_norm = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine) self.audio_to_video_attn = LTX2Attention( query_dim=dim, cross_attention_dim=audio_dim, heads=audio_num_attention_heads, kv_heads=audio_num_attention_heads, dim_head=audio_attention_head_dim, bias=attention_bias, out_bias=attention_out_bias, qk_norm=qk_norm, rope_type=rope_type, apply_gated_attention=video_gated_attn, quant_config=quant_config, prefix=f"{prefix}.audio_to_video_attn" if prefix else "audio_to_video_attn", ) self.video_to_audio_norm = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) self.video_to_audio_attn = LTX2Attention( query_dim=audio_dim, cross_attention_dim=dim, heads=audio_num_attention_heads, kv_heads=audio_num_attention_heads, dim_head=audio_attention_head_dim, bias=attention_bias, out_bias=attention_out_bias, qk_norm=qk_norm, rope_type=rope_type, apply_gated_attention=audio_gated_attn, quant_config=quant_config, prefix=f"{prefix}.video_to_audio_attn" if prefix else "video_to_audio_attn", ) # 4. Feedforward layers self.norm3 = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine) self.ff = LTX2FeedForward( dim, activation_fn=activation_fn, quant_config=quant_config, prefix=f"{prefix}.ff" if prefix else "ff", ) self.audio_norm3 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) self.audio_ff = LTX2FeedForward( audio_dim, activation_fn=activation_fn, quant_config=quant_config, prefix=f"{prefix}.audio_ff" if prefix else "audio_ff", ) # 5. Per-Layer Modulation Parameters # LTX-2.3 with cross_attn_adaln uses 9 params (extra 3 for cross-attn modulation); # LTX-2 uses 6. video_mod_param_num = 9 if self.video_cross_attn_adaln else 6 audio_mod_param_num = 9 if self.audio_cross_attn_adaln else 6 self.scale_shift_table = nn.Parameter(torch.randn(video_mod_param_num, dim) / dim**0.5) self.audio_scale_shift_table = nn.Parameter(torch.randn(audio_mod_param_num, audio_dim) / audio_dim**0.5) # Prompt cross-attn additional modulation params (LTX-2.3) self.cross_attn_adaln = video_cross_attn_adaln or audio_cross_attn_adaln if self.cross_attn_adaln: self.prompt_scale_shift_table = nn.Parameter(torch.randn(2, dim)) self.audio_prompt_scale_shift_table = nn.Parameter(torch.randn(2, audio_dim)) # Per-layer a2v, v2a Cross-Attention mod params self.video_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, dim)) self.audio_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, audio_dim)) @staticmethod def get_mod_params( scale_shift_table: torch.Tensor, temb: torch.Tensor, batch_size: int ) -> tuple[torch.Tensor, ...]: num_ada_params = scale_shift_table.shape[0] ada_values = scale_shift_table[None, None].to(temb.device) + temb.reshape( batch_size, temb.shape[1], num_ada_params, -1 ) return ada_values.unbind(dim=2) def forward( self, hidden_states: torch.Tensor, audio_hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, audio_encoder_hidden_states: torch.Tensor, temb: torch.Tensor, temb_audio: torch.Tensor, temb_ca_scale_shift: torch.Tensor, temb_ca_audio_scale_shift: torch.Tensor, temb_ca_gate: torch.Tensor, temb_ca_audio_gate: torch.Tensor, temb_prompt: torch.Tensor | None = None, temb_prompt_audio: torch.Tensor | None = None, video_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, audio_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, ca_video_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, ca_audio_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, encoder_attention_mask: torch.Tensor | None = None, audio_encoder_attention_mask: torch.Tensor | None = None, self_attention_mask: torch.Tensor | None = None, audio_self_attention_mask: torch.Tensor | None = None, a2v_cross_attention_mask: torch.Tensor | None = None, v2a_cross_attention_mask: torch.Tensor | None = None, use_a2v_cross_attention: bool = True, use_v2a_cross_attention: bool = True, perturbation_mask: torch.Tensor | None = None, all_perturbed: bool | None = None, ) -> torch.Tensor: batch_size = hidden_states.size(0) # 1. Video and Audio Self-Attention # 1.1. Video Self-Attention video_ada_params = self.get_mod_params(self.scale_shift_table, temb, batch_size) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = video_ada_params[:6] if self.video_cross_attn_adaln: shift_text_q, scale_text_q, gate_text_q = video_ada_params[6:9] norm_hidden_states = self.norm1(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa attn_hidden_states = self.attn1( hidden_states=norm_hidden_states, encoder_hidden_states=None, query_rotary_emb=video_rotary_emb, attention_mask=self_attention_mask, ) hidden_states = hidden_states + attn_hidden_states * gate_msa # 1.2. Audio Self-Attention audio_ada_params = self.get_mod_params(self.audio_scale_shift_table, temb_audio, batch_size) audio_shift_msa, audio_scale_msa, audio_gate_msa, audio_shift_mlp, audio_scale_mlp, audio_gate_mlp = ( audio_ada_params[:6] ) if self.audio_cross_attn_adaln: audio_shift_text_q, audio_scale_text_q, audio_gate_text_q = audio_ada_params[6:9] norm_audio_hidden_states = self.audio_norm1(audio_hidden_states) norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_msa) + audio_shift_msa attn_audio_hidden_states = self.audio_attn1( hidden_states=norm_audio_hidden_states, encoder_hidden_states=None, query_rotary_emb=audio_rotary_emb, attention_mask=audio_self_attention_mask, ) audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * audio_gate_msa # 2. Video and Audio Cross-Attention with text embeddings (Q: Video/Audio; K,V: Text) # LTX-2.3: compute prompt modulation params for K/V if self.cross_attn_adaln and temb_prompt is not None: video_prompt_ada_params = self.get_mod_params(self.prompt_scale_shift_table, temb_prompt, batch_size) shift_text_kv, scale_text_kv = video_prompt_ada_params audio_prompt_ada_params = self.get_mod_params( self.audio_prompt_scale_shift_table, temb_prompt_audio, batch_size ) audio_shift_text_kv, audio_scale_text_kv = audio_prompt_ada_params # 2.1. Video-Text Cross-Attention norm_hidden_states = self.norm2(hidden_states) if self.video_cross_attn_adaln: norm_hidden_states = norm_hidden_states * (1 + scale_text_q) + shift_text_q if self.cross_attn_adaln and temb_prompt is not None: encoder_hidden_states = encoder_hidden_states * (1 + scale_text_kv) + shift_text_kv attn_hidden_states = self.attn2( norm_hidden_states, encoder_hidden_states=encoder_hidden_states, query_rotary_emb=None, attention_mask=encoder_attention_mask, ) if self.video_cross_attn_adaln: attn_hidden_states = attn_hidden_states * gate_text_q hidden_states = hidden_states + attn_hidden_states # 2.2. Audio-Text Cross-Attention norm_audio_hidden_states = self.audio_norm2(audio_hidden_states) if self.audio_cross_attn_adaln: norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_text_q) + audio_shift_text_q if self.cross_attn_adaln and temb_prompt is not None: audio_encoder_hidden_states = audio_encoder_hidden_states * (1 + audio_scale_text_kv) + audio_shift_text_kv attn_audio_hidden_states = self.audio_attn2( norm_audio_hidden_states, encoder_hidden_states=audio_encoder_hidden_states, query_rotary_emb=None, attention_mask=audio_encoder_attention_mask, ) if self.audio_cross_attn_adaln: attn_audio_hidden_states = attn_audio_hidden_states * audio_gate_text_q audio_hidden_states = audio_hidden_states + attn_audio_hidden_states # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention if use_a2v_cross_attention or use_v2a_cross_attention: norm_hidden_states = self.audio_to_video_norm(hidden_states) norm_audio_hidden_states = self.video_to_audio_norm(audio_hidden_states) # Combine global and per-layer cross attention modulation parameters # Video video_per_layer_ca_scale_shift = self.video_a2v_cross_attn_scale_shift_table[:4, :] video_per_layer_ca_gate = self.video_a2v_cross_attn_scale_shift_table[4:, :] video_ca_ada_params = self.get_mod_params(video_per_layer_ca_scale_shift, temb_ca_scale_shift, batch_size) video_ca_gate_param = self.get_mod_params(video_per_layer_ca_gate, temb_ca_gate, batch_size) video_a2v_ca_scale, video_a2v_ca_shift, video_v2a_ca_scale, video_v2a_ca_shift = video_ca_ada_params a2v_gate = video_ca_gate_param[0].squeeze(2) # Audio audio_per_layer_ca_scale_shift = self.audio_a2v_cross_attn_scale_shift_table[:4, :] audio_per_layer_ca_gate = self.audio_a2v_cross_attn_scale_shift_table[4:, :] audio_ca_ada_params = self.get_mod_params( audio_per_layer_ca_scale_shift, temb_ca_audio_scale_shift, batch_size ) audio_ca_gate_param = self.get_mod_params(audio_per_layer_ca_gate, temb_ca_audio_gate, batch_size) audio_a2v_ca_scale, audio_a2v_ca_shift, audio_v2a_ca_scale, audio_v2a_ca_shift = audio_ca_ada_params v2a_gate = audio_ca_gate_param[0].squeeze(2) # 3.2. Audio-to-Video Cross Attention: Q: Video; K,V: Audio if use_a2v_cross_attention: mod_norm_hidden_states = norm_hidden_states * ( 1 + video_a2v_ca_scale.squeeze(2) ) + video_a2v_ca_shift.squeeze(2) mod_norm_audio_hidden_states = norm_audio_hidden_states * ( 1 + audio_a2v_ca_scale.squeeze(2) ) + audio_a2v_ca_shift.squeeze(2) a2v_attn_hidden_states = self.audio_to_video_attn( mod_norm_hidden_states, encoder_hidden_states=mod_norm_audio_hidden_states, query_rotary_emb=ca_video_rotary_emb, key_rotary_emb=ca_audio_rotary_emb, attention_mask=a2v_cross_attention_mask, ) hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states # 3.3. Video-to-Audio Cross Attention: Q: Audio; K,V: Video if use_v2a_cross_attention: mod_norm_hidden_states = norm_hidden_states * ( 1 + video_v2a_ca_scale.squeeze(2) ) + video_v2a_ca_shift.squeeze(2) mod_norm_audio_hidden_states = norm_audio_hidden_states * ( 1 + audio_v2a_ca_scale.squeeze(2) ) + audio_v2a_ca_shift.squeeze(2) v2a_attn_hidden_states = self.video_to_audio_attn( mod_norm_audio_hidden_states, encoder_hidden_states=mod_norm_hidden_states, query_rotary_emb=ca_audio_rotary_emb, key_rotary_emb=ca_video_rotary_emb, attention_mask=v2a_cross_attention_mask, ) audio_hidden_states = audio_hidden_states + v2a_gate * v2a_attn_hidden_states # 4. Feedforward norm_hidden_states = self.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp ff_output = self.ff(norm_hidden_states) hidden_states = hidden_states + ff_output * gate_mlp norm_audio_hidden_states = self.audio_norm3(audio_hidden_states) * (1 + audio_scale_mlp) + audio_shift_mlp audio_ff_output = self.audio_ff(norm_audio_hidden_states) audio_hidden_states = audio_hidden_states + audio_ff_output * audio_gate_mlp return hidden_states, audio_hidden_states class LTX2AudioVideoRotaryPosEmbed(nn.Module): """ Video and audio rotary positional embeddings (RoPE) for the LTX-2.0 model. Args: causal_offset (`int`, *optional*, defaults to `1`): Offset in the temporal axis for causal VAE modeling. This is typically 1 (for causal modeling where the VAE treats the very first frame differently), but could also be 0 (for non-causal modeling). """ def __init__( self, dim: int, patch_size: int = 1, patch_size_t: int = 1, base_num_frames: int = 20, base_height: int = 2048, base_width: int = 2048, sampling_rate: int = 16000, hop_length: int = 160, scale_factors: tuple[int, ...] = (8, 32, 32), theta: float = 10000.0, causal_offset: int = 1, modality: str = "video", double_precision: bool = True, rope_type: str = "interleaved", num_attention_heads: int = 32, ) -> None: super().__init__() self.dim = dim self.patch_size = patch_size self.patch_size_t = patch_size_t if rope_type not in ["interleaved", "split"]: raise ValueError(f"{rope_type=} not supported. Choose between 'interleaved' and 'split'.") self.rope_type = rope_type self.base_num_frames = base_num_frames self.num_attention_heads = num_attention_heads # Video-specific self.base_height = base_height self.base_width = base_width # Audio-specific self.sampling_rate = sampling_rate self.hop_length = hop_length self.audio_latents_per_second = float(sampling_rate) / float(hop_length) / float(scale_factors[0]) self.scale_factors = scale_factors self.theta = theta self.causal_offset = causal_offset self.modality = modality if self.modality not in ["video", "audio"]: raise ValueError(f"Modality {modality} is not supported. Supported modalities are `video` and `audio`.") self.double_precision = double_precision def prepare_video_coords( self, batch_size: int, num_frames: int, height: int, width: int, device: torch.device, fps: float = 24.0, ) -> torch.Tensor: """ Create per-dimension bounds [inclusive start, exclusive end) for each patch with respect to the original pixel space video grid (num_frames, height, width). This will ultimately have shape (batch_size, 3, num_patches, 2) where - axis 1 (size 3) enumerates (frame, height, width) dimensions (e.g. idx 0 corresponds to frames) - axis 3 (size 2) stores `[start, end)` indices within each dimension Args: batch_size (`int`): Batch size of the video latents. num_frames (`int`): Number of latent frames in the video latents. height (`int`): Latent height of the video latents. width (`int`): Latent width of the video latents. device (`torch.device`): Device on which to create the video grid. Returns: `torch.Tensor`: Per-dimension patch boundaries tensor of shape [batch_size, 3, num_patches, 2]. """ # 1. Generate grid coordinates for each spatiotemporal dimension (frames, height, width) # Always compute rope in fp32 grid_f = torch.arange(start=0, end=num_frames, step=self.patch_size_t, dtype=torch.float32, device=device) grid_h = torch.arange(start=0, end=height, step=self.patch_size, dtype=torch.float32, device=device) grid_w = torch.arange(start=0, end=width, step=self.patch_size, dtype=torch.float32, device=device) # indexing='ij' ensures that the dimensions are kept in order as (frames, height, width) grid = torch.meshgrid(grid_f, grid_h, grid_w, indexing="ij") grid = torch.stack(grid, dim=0) # [3, N_F, N_H, N_W], where e.g. N_F is the number of temporal patches # 2. Get the patch boundaries with respect to the latent video grid patch_size = (self.patch_size_t, self.patch_size, self.patch_size) patch_size_delta = torch.tensor(patch_size, dtype=grid.dtype, device=grid.device) patch_ends = grid + patch_size_delta.view(3, 1, 1, 1) # Combine the start (grid) and end (patch_ends) coordinates along new trailing dimension latent_coords = torch.stack([grid, patch_ends], dim=-1) # [3, N_F, N_H, N_W, 2] # Reshape to (batch_size, 3, num_patches, 2) latent_coords = latent_coords.flatten(1, 3) latent_coords = latent_coords.unsqueeze(0).repeat(batch_size, 1, 1, 1) # 3. Calculate the pixel space patch boundaries from the latent boundaries. scale_tensor = torch.tensor(self.scale_factors, device=latent_coords.device) # Broadcast the VAE scale factors such that they are compatible with latent_coords's shape broadcast_shape = [1] * latent_coords.ndim broadcast_shape[1] = -1 # This is the (frame, height, width) dim # Apply per-axis scaling to convert latent coordinates to pixel space coordinates pixel_coords = latent_coords * scale_tensor.view(*broadcast_shape) # As the VAE temporal stride for the first frame is 1 instead of self.vae_scale_factors[0], we need to shift # and clamp to keep the first-frame timestamps causal and non-negative. pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + self.causal_offset - self.scale_factors[0]).clamp(min=0) # Scale the temporal coordinates by the video FPS pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps return pixel_coords def prepare_audio_coords( self, batch_size: int, num_frames: int, device: torch.device, shift: int = 0, ) -> torch.Tensor: """ Create per-dimension bounds [inclusive start, exclusive end) of start and end timestamps for each latent frame. This will ultimately have shape (batch_size, 3, num_patches, 2) where - axis 1 (size 1) represents the temporal dimension - axis 3 (size 2) stores `[start, end)` indices within each dimension Args: batch_size (`int`): Batch size of the audio latents. num_frames (`int`): Number of latent frames in the audio latents. device (`torch.device`): Device on which to create the audio grid. shift (`int`, *optional*, defaults to `0`): Offset on the latent indices. Different shift values correspond to different overlapping windows with respect to the same underlying latent grid. Returns: `torch.Tensor`: Per-dimension patch boundaries tensor of shape [batch_size, 1, num_patches, 2]. """ # 1. Generate coordinates in the frame (time) dimension. # Always compute rope in fp32 grid_f = torch.arange( start=shift, end=num_frames + shift, step=self.patch_size_t, dtype=torch.float32, device=device ) # 2. Calculate start timestamps in seconds with respect to the original spectrogram grid audio_scale_factor = self.scale_factors[0] # Scale back to mel spectrogram space grid_start_mel = grid_f * audio_scale_factor # Handle first frame causal offset, ensuring non-negative timestamps grid_start_mel = (grid_start_mel + self.causal_offset - audio_scale_factor).clip(min=0) # Convert mel bins back into seconds grid_start_s = grid_start_mel * self.hop_length / self.sampling_rate # 3. Calculate start timestamps in seconds with respect to the original spectrogram grid grid_end_mel = (grid_f + self.patch_size_t) * audio_scale_factor grid_end_mel = (grid_end_mel + self.causal_offset - audio_scale_factor).clip(min=0) grid_end_s = grid_end_mel * self.hop_length / self.sampling_rate audio_coords = torch.stack([grid_start_s, grid_end_s], dim=-1) # [num_patches, 2] audio_coords = audio_coords.unsqueeze(0).expand(batch_size, -1, -1) # [batch_size, num_patches, 2] audio_coords = audio_coords.unsqueeze(1) # [batch_size, 1, num_patches, 2] return audio_coords def prepare_coords(self, *args, **kwargs): if self.modality == "video": return self.prepare_video_coords(*args, **kwargs) elif self.modality == "audio": return self.prepare_audio_coords(*args, **kwargs) def forward( self, coords: torch.Tensor, device: str | torch.device | None = None ) -> tuple[torch.Tensor, torch.Tensor]: device = device or coords.device # Number of spatiotemporal dimensions (3 for video, 1 (temporal) for audio and cross attn) num_pos_dims = coords.shape[1] # 1. If the coords are patch boundaries [start, end), use the midpoint of these boundaries as the patch # position index if coords.ndim == 4: coords_start, coords_end = coords.chunk(2, dim=-1) coords = (coords_start + coords_end) / 2.0 coords = coords.squeeze(-1) # [B, num_pos_dims, num_patches] # 2. Get coordinates as a fraction of the base data shape if self.modality == "video": max_positions = (self.base_num_frames, self.base_height, self.base_width) elif self.modality == "audio": max_positions = (self.base_num_frames,) # [B, num_pos_dims, num_patches] --> [B, num_patches, num_pos_dims] grid = torch.stack([coords[:, i] / max_positions[i] for i in range(num_pos_dims)], dim=-1).to(device) # Number of spatiotemporal dimensions (3 for video, 1 for audio and cross attn) times 2 for cos, sin num_rope_elems = num_pos_dims * 2 # 3. Create a 1D grid of frequencies for RoPE freqs_dtype = torch.float64 if self.double_precision else torch.float32 pow_indices = torch.pow( self.theta, torch.linspace(start=0.0, end=1.0, steps=self.dim // num_rope_elems, dtype=freqs_dtype, device=device), ) freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32) # 4. Tensor-vector outer product between pos ids tensor of shape (B, 3, num_patches) and freqs vector of shape # (self.dim // num_elems,) freqs = (grid.unsqueeze(-1) * 2 - 1) * freqs # [B, num_patches, num_pos_dims, self.dim // num_elems] freqs = freqs.transpose(-1, -2).flatten(2) # [B, num_patches, self.dim // 2] # 5. Get real, interleaved (cos, sin) frequencies, padded to self.dim. # Kept inline here to stay aligned with the upstream diffusers connector logic. # src/diffusers/pipelines/ltx2/connectors.py if self.rope_type == "interleaved": cos_freqs = freqs.cos().repeat_interleave(2, dim=-1) sin_freqs = freqs.sin().repeat_interleave(2, dim=-1) if self.dim % num_rope_elems != 0: cos_padding = torch.ones_like(cos_freqs[:, :, : self.dim % num_rope_elems]) sin_padding = torch.zeros_like(cos_freqs[:, :, : self.dim % num_rope_elems]) cos_freqs = torch.cat([cos_padding, cos_freqs], dim=-1) sin_freqs = torch.cat([sin_padding, sin_freqs], dim=-1) elif self.rope_type == "split": expected_freqs = self.dim // 2 current_freqs = freqs.shape[-1] pad_size = expected_freqs - current_freqs cos_freq = freqs.cos() sin_freq = freqs.sin() if pad_size != 0: cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size]) cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1) sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1) # Reshape freqs to be compatible with multi-head attention b = cos_freq.shape[0] t = cos_freq.shape[1] # Concrete per-head dim instead of -1 — see apply_split_rotary_emb r = self.dim // self.num_attention_heads // 2 cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, r) sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, r) cos_freqs = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2) sin_freqs = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2) return cos_freqs, sin_freqs class LTX2VideoTransformer3DModel(nn.Module): r""" A Transformer model for video-like data used in [LTX](https://huggingface.co/Lightricks/LTX-Video). Args: in_channels (`int`, defaults to `128`): The number of channels in the input. out_channels (`int`, defaults to `128`): The number of channels in the output. patch_size (`int`, defaults to `1`): The size of the spatial patches to use in the patch embedding layer. patch_size_t (`int`, defaults to `1`): The size of the tmeporal patches to use in the patch embedding layer. num_attention_heads (`int`, defaults to `32`): The number of heads to use for multi-head attention. attention_head_dim (`int`, defaults to `64`): The number of channels in each head. cross_attention_dim (`int`, defaults to `2048 `): The number of channels for cross attention heads. num_layers (`int`, defaults to `28`): The number of layers of Transformer blocks to use. activation_fn (`str`, defaults to `"gelu-approximate"`): Activation function to use in feed-forward. qk_norm (`str`, defaults to `"rms_norm_across_heads"`): The normalization layer to use. """ _supports_gradient_checkpointing = True _skip_layerwise_casting_patterns = ["norm"] _repeated_blocks = ["LTX2VideoTransformerBlock"] _layerwise_offload_blocks_attrs = ["transformer_blocks"] _hsdp_shard_conditions = [is_transformer_block_module] _sp_plan: dict[str, Any] | None = None packed_modules_mapping = { "to_qkv": ["to_q", "to_k", "to_v"], } @staticmethod def _build_sp_plan(rope_type: str) -> dict[str, Any]: if rope_type == "split": # split RoPE returns (B, H, T, D/2) -> shard along T dim rope_expected_dims = 4 rope_split_dim = 2 else: # interleaved RoPE returns (B, T, D) -> shard along T dim rope_expected_dims = 3 rope_split_dim = 1 return { "": { # Shard video/audio latents across sequence "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), "audio_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), # Shard prompt embeds across sequence "encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), "audio_encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), # Shard video timestep when provided as (B, seq_len) "timestep": SequenceParallelInput(split_dim=1, expected_dims=2, split_output=False), }, "rope": { 0: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), 1: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), }, "audio_rope": { 0: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), 1: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), }, "cross_attn_rope": { 0: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), 1: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), }, "cross_attn_audio_rope": { 0: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), 1: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True), }, # Gather outputs before returning "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), "audio_proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), } def __init__( self, in_channels: int = 128, # Video Arguments out_channels: int | None = 128, patch_size: int = 1, patch_size_t: int = 1, num_attention_heads: int = 32, attention_head_dim: int = 128, cross_attention_dim: int = 4096, vae_scale_factors: tuple[int, int, int] = (8, 32, 32), pos_embed_max_pos: int = 20, base_height: int = 2048, base_width: int = 2048, audio_in_channels: int = 128, # Audio Arguments audio_out_channels: int | None = 128, audio_patch_size: int = 1, audio_patch_size_t: int = 1, audio_num_attention_heads: int = 32, audio_attention_head_dim: int = 64, audio_cross_attention_dim: int = 2048, audio_scale_factor: int = 4, audio_pos_embed_max_pos: int = 20, audio_sampling_rate: int = 16000, audio_hop_length: int = 160, num_layers: int = 48, # Shared arguments activation_fn: str = "gelu-approximate", qk_norm: str = "rms_norm_across_heads", norm_elementwise_affine: bool = False, norm_eps: float = 1e-6, caption_channels: int = 3840, attention_bias: bool = True, attention_out_bias: bool = True, rope_theta: float = 10000.0, rope_double_precision: bool = True, causal_offset: int = 1, timestep_scale_multiplier: int = 1000, cross_attn_timestep_scale_multiplier: int = 1000, rope_type: str = "interleaved", use_prompt_embeddings: bool = True, perturbed_attn: bool = False, gated_attn: bool = False, cross_attn_mod: bool = False, audio_gated_attn: bool = False, audio_cross_attn_mod: bool = False, quant_config: "QuantizationConfig | None" = None, ) -> None: super().__init__() self.perturbed_attn = perturbed_attn out_channels = out_channels or in_channels audio_out_channels = audio_out_channels or audio_in_channels inner_dim = num_attention_heads * attention_head_dim audio_inner_dim = audio_num_attention_heads * audio_attention_head_dim self.config = SimpleNamespace( in_channels=in_channels, out_channels=out_channels, patch_size=patch_size, patch_size_t=patch_size_t, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, cross_attention_dim=cross_attention_dim, vae_scale_factors=vae_scale_factors, pos_embed_max_pos=pos_embed_max_pos, base_height=base_height, base_width=base_width, audio_in_channels=audio_in_channels, audio_out_channels=audio_out_channels, audio_patch_size=audio_patch_size, audio_patch_size_t=audio_patch_size_t, audio_num_attention_heads=audio_num_attention_heads, audio_attention_head_dim=audio_attention_head_dim, audio_cross_attention_dim=audio_cross_attention_dim, audio_scale_factor=audio_scale_factor, audio_pos_embed_max_pos=audio_pos_embed_max_pos, audio_sampling_rate=audio_sampling_rate, audio_hop_length=audio_hop_length, num_layers=num_layers, activation_fn=activation_fn, qk_norm=qk_norm, norm_elementwise_affine=norm_elementwise_affine, norm_eps=norm_eps, caption_channels=caption_channels, attention_bias=attention_bias, attention_out_bias=attention_out_bias, rope_theta=rope_theta, rope_double_precision=rope_double_precision, causal_offset=causal_offset, timestep_scale_multiplier=timestep_scale_multiplier, cross_attn_timestep_scale_multiplier=cross_attn_timestep_scale_multiplier, rope_type=rope_type, ) # 1. Patchification input projections self.proj_in = nn.Linear(in_channels, inner_dim) self.audio_proj_in = nn.Linear(audio_in_channels, audio_inner_dim) # 2. Prompt embeddings # LTX-2 (use_prompt_embeddings=True): caption projection in the transformer # LTX-2.3 (use_prompt_embeddings=False): caption projection in the connectors if use_prompt_embeddings: self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) self.audio_caption_projection = PixArtAlphaTextProjection( in_features=caption_channels, hidden_size=audio_inner_dim ) # 3. Timestep Modulation Params and Embedding # 3.1. Global Timestep Modulation Parameters # LTX-2.3 with cross_attn_mod uses 9 mod params (extra 3 for cross-attn); LTX-2 uses 6. video_num_mod_params = 9 if cross_attn_mod else 6 audio_num_mod_params = 9 if audio_cross_attn_mod else 6 self.time_embed = LTX2AdaLayerNormSingle( inner_dim, num_mod_params=video_num_mod_params, use_additional_conditions=False ) self.audio_time_embed = LTX2AdaLayerNormSingle( audio_inner_dim, num_mod_params=audio_num_mod_params, use_additional_conditions=False ) # 3.3. LTX-2.3: Prompt modulation from sigma self.prompt_modulation = cross_attn_mod or audio_cross_attn_mod if self.prompt_modulation: self.prompt_adaln = LTX2AdaLayerNormSingle(inner_dim, num_mod_params=2, use_additional_conditions=False) self.audio_prompt_adaln = LTX2AdaLayerNormSingle( audio_inner_dim, num_mod_params=2, use_additional_conditions=False ) # 3.2. Global Cross Attention Modulation Parameters # Used in the audio-to-video and video-to-audio cross attention layers as a global set of modulation params, # which are then further modified by per-block modulaton params in each transformer block. # There are 2 sets of scale/shift parameters for each modality, 1 each for audio-to-video (a2v) and # video-to-audio (v2a) cross attention self.av_cross_attn_video_scale_shift = LTX2AdaLayerNormSingle( inner_dim, num_mod_params=4, use_additional_conditions=False ) self.av_cross_attn_audio_scale_shift = LTX2AdaLayerNormSingle( audio_inner_dim, num_mod_params=4, use_additional_conditions=False ) # Gate param for audio-to-video (a2v) cross attn (where the video is the queries (Q) and the audio is the keys # and values (KV)) self.av_cross_attn_video_a2v_gate = LTX2AdaLayerNormSingle( inner_dim, num_mod_params=1, use_additional_conditions=False ) # Gate param for video-to-audio (v2a) cross attn (where the audio is the queries (Q) and the video is the keys # and values (KV)) self.av_cross_attn_audio_v2a_gate = LTX2AdaLayerNormSingle( audio_inner_dim, num_mod_params=1, use_additional_conditions=False ) # 3.3. Output Layer Scale/Shift Modulation parameters self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5) self.audio_scale_shift_table = nn.Parameter(torch.randn(2, audio_inner_dim) / audio_inner_dim**0.5) # 4. Rotary Positional Embeddings (RoPE) # Self-Attention self.rope = LTX2AudioVideoRotaryPosEmbed( dim=inner_dim, patch_size=patch_size, patch_size_t=patch_size_t, base_num_frames=pos_embed_max_pos, base_height=base_height, base_width=base_width, scale_factors=vae_scale_factors, theta=rope_theta, causal_offset=causal_offset, modality="video", double_precision=rope_double_precision, rope_type=rope_type, num_attention_heads=num_attention_heads, ) self.audio_rope = LTX2AudioVideoRotaryPosEmbed( dim=audio_inner_dim, patch_size=audio_patch_size, patch_size_t=audio_patch_size_t, base_num_frames=audio_pos_embed_max_pos, sampling_rate=audio_sampling_rate, hop_length=audio_hop_length, scale_factors=[audio_scale_factor], theta=rope_theta, causal_offset=causal_offset, modality="audio", double_precision=rope_double_precision, rope_type=rope_type, num_attention_heads=audio_num_attention_heads, ) # Audio-to-Video, Video-to-Audio Cross-Attention cross_attn_pos_embed_max_pos = max(pos_embed_max_pos, audio_pos_embed_max_pos) self.cross_attn_rope = LTX2AudioVideoRotaryPosEmbed( dim=audio_cross_attention_dim, patch_size=patch_size, patch_size_t=patch_size_t, base_num_frames=cross_attn_pos_embed_max_pos, base_height=base_height, base_width=base_width, theta=rope_theta, causal_offset=causal_offset, modality="video", double_precision=rope_double_precision, rope_type=rope_type, num_attention_heads=num_attention_heads, ) self.cross_attn_audio_rope = LTX2AudioVideoRotaryPosEmbed( dim=audio_cross_attention_dim, patch_size=audio_patch_size, patch_size_t=audio_patch_size_t, base_num_frames=cross_attn_pos_embed_max_pos, sampling_rate=audio_sampling_rate, hop_length=audio_hop_length, theta=rope_theta, causal_offset=causal_offset, modality="audio", double_precision=rope_double_precision, rope_type=rope_type, num_attention_heads=audio_num_attention_heads, ) # 5. Transformer Blocks self.transformer_blocks = nn.ModuleList( [ LTX2VideoTransformerBlock( dim=inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, cross_attention_dim=cross_attention_dim, audio_dim=audio_inner_dim, audio_num_attention_heads=audio_num_attention_heads, audio_attention_head_dim=audio_attention_head_dim, audio_cross_attention_dim=audio_cross_attention_dim, video_gated_attn=gated_attn, video_cross_attn_adaln=cross_attn_mod, audio_gated_attn=audio_gated_attn, audio_cross_attn_adaln=audio_cross_attn_mod, qk_norm=qk_norm, activation_fn=activation_fn, attention_bias=attention_bias, attention_out_bias=attention_out_bias, eps=norm_eps, elementwise_affine=norm_elementwise_affine, rope_type=rope_type, perturbed_attn=perturbed_attn, quant_config=quant_config, prefix=f"transformer_blocks.{layer_idx}", ) for layer_idx in range(num_layers) ] ) # 6. Output layers self.norm_out = nn.LayerNorm(inner_dim, eps=1e-6, elementwise_affine=False) self.proj_out = nn.Linear(inner_dim, out_channels) self.audio_norm_out = nn.LayerNorm(audio_inner_dim, eps=1e-6, elementwise_affine=False) self.audio_proj_out = nn.Linear(audio_inner_dim, audio_out_channels) self.gradient_checkpointing = False self._sp_plan = self._build_sp_plan(rope_type) def _gradient_checkpointing_func(self, module: nn.Module, *inputs: Any): def custom_forward(*checkpoint_inputs: Any): return module(*checkpoint_inputs) return checkpoint(custom_forward, *inputs, use_reentrant=False) def enable_gradient_checkpointing(self) -> None: self.gradient_checkpointing = True def disable_gradient_checkpointing(self) -> None: self.gradient_checkpointing = False def forward( self, hidden_states: torch.Tensor, audio_hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, audio_encoder_hidden_states: torch.Tensor, timestep: torch.LongTensor, audio_timestep: torch.LongTensor | None = None, sigma: torch.Tensor | None = None, audio_sigma: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, audio_encoder_attention_mask: torch.Tensor | None = None, num_frames: int | None = None, height: int | None = None, width: int | None = None, fps: float = 24.0, audio_num_frames: int | None = None, video_coords: torch.Tensor | None = None, audio_coords: torch.Tensor | None = None, attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, **kwargs, # Accept extra diffusers kwargs (isolate_modalities, perturbation_mask, etc.) ) -> torch.Tensor: """ Forward pass for LTX-2.0 audiovisual video transformer. Args: hidden_states (`torch.Tensor`): Input patchified video latents of shape `(batch_size, num_video_tokens, in_channels)`. audio_hidden_states (`torch.Tensor`): Input patchified audio latents of shape `(batch_size, num_audio_tokens, audio_in_channels)`. encoder_hidden_states (`torch.Tensor`): Input video text embeddings of shape `(batch_size, text_seq_len, self.config.caption_channels)`. audio_encoder_hidden_states (`torch.Tensor`): Input audio text embeddings of shape `(batch_size, text_seq_len, self.config.caption_channels)`. timestep (`torch.Tensor`): Input timestep of shape `(batch_size, num_video_tokens)`. These should already be scaled by `self.config.timestep_scale_multiplier`. audio_timestep (`torch.Tensor`, *optional*): Input timestep of shape `(batch_size,)` or `(batch_size, num_audio_tokens)` for audio modulation params. This is only used by certain pipelines such as the I2V pipeline. encoder_attention_mask (`torch.Tensor`, *optional*): Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)`. audio_encoder_attention_mask (`torch.Tensor`, *optional*): Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)` for audio modeling. num_frames (`int`, *optional*): The number of latent video frames. Used if calculating the video coordinates for RoPE. height (`int`, *optional*): The latent video height. Used if calculating the video coordinates for RoPE. width (`int`, *optional*): The latent video width. Used if calculating the video coordinates for RoPE. fps: (`float`, *optional*, defaults to `24.0`): The desired frames per second of the generated video. Used if calculating the video coordinates for RoPE. audio_num_frames: (`int`, *optional*): The number of latent audio frames. Used if calculating the audio coordinates for RoPE. video_coords (`torch.Tensor`, *optional*): The video coordinates to be used when calculating the rotary positional embeddings (RoPE) of shape `(batch_size, 3, num_video_tokens, 2)`. If not supplied, this will be calculated inside `forward`. audio_coords (`torch.Tensor`, *optional*): The audio coordinates to be used when calculating the rotary positional embeddings (RoPE) of shape `(batch_size, 1, num_audio_tokens, 2)`. If not supplied, this will be calculated inside `forward`. attention_kwargs (`Dict[str, Any]`, *optional*): Optional dict of keyword args to be passed to the attention processor. return_dict (`bool`, *optional*, defaults to `True`): Whether to return a dict-like structured output of type `AudioVisualModelOutput` or a tuple. Returns: `AudioVisualModelOutput` or `tuple`: If `return_dict` is `True`, returns a structured output of type `AudioVisualModelOutput`, otherwise a `tuple` is returned where the first element is the denoised video latent patch sequence and the second element is the denoised audio latent patch sequence. """ # Determine timestep for audio. audio_timestep = audio_timestep if audio_timestep is not None else timestep # convert encoder_attention_mask to a bias the same way we do for attention_mask if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2: encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0 encoder_attention_mask = encoder_attention_mask.unsqueeze(1) if audio_encoder_attention_mask is not None and audio_encoder_attention_mask.ndim == 2: audio_encoder_attention_mask = (1 - audio_encoder_attention_mask.to(audio_hidden_states.dtype)) * -10000.0 audio_encoder_attention_mask = audio_encoder_attention_mask.unsqueeze(1) batch_size = hidden_states.size(0) # 1. Prepare RoPE positional embeddings if video_coords is None: video_coords = self.rope.prepare_video_coords( batch_size, num_frames, height, width, hidden_states.device, fps=fps ) if audio_coords is None: audio_coords = self.audio_rope.prepare_audio_coords( batch_size, audio_num_frames, audio_hidden_states.device ) video_rotary_emb = self.rope(video_coords, device=hidden_states.device) audio_rotary_emb = self.audio_rope(audio_coords, device=audio_hidden_states.device) video_cross_attn_rotary_emb = self.cross_attn_rope(video_coords[:, 0:1, :], device=hidden_states.device) audio_cross_attn_rotary_emb = self.cross_attn_audio_rope( audio_coords[:, 0:1, :], device=audio_hidden_states.device ) # 2. Patchify input projections hidden_states = self.proj_in(hidden_states) audio_hidden_states = self.audio_proj_in(audio_hidden_states) # 3. Prepare timestep embeddings and modulation parameters timestep_cross_attn_gate_scale_factor = ( self.config.cross_attn_timestep_scale_multiplier / self.config.timestep_scale_multiplier ) # 3.1. Prepare global modality (video and audio) timestep embedding and modulation parameters # temb is used in the transformer blocks (as expected), while embedded_timestep is used for the output layer # modulation with scale_shift_table (and similarly for audio) temb, embedded_timestep = self.time_embed( timestep.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype, ) temb = temb.view(batch_size, -1, temb.size(-1)) embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.size(-1)) temb_audio, audio_embedded_timestep = self.audio_time_embed( audio_timestep.flatten(), batch_size=batch_size, hidden_dtype=audio_hidden_states.dtype, ) temb_audio = temb_audio.view(batch_size, -1, temb_audio.size(-1)) audio_embedded_timestep = audio_embedded_timestep.view(batch_size, -1, audio_embedded_timestep.size(-1)) # 3.2. Prepare global modality cross attention modulation parameters video_cross_attn_scale_shift, _ = self.av_cross_attn_video_scale_shift( timestep.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype, ) video_cross_attn_a2v_gate, _ = self.av_cross_attn_video_a2v_gate( timestep.flatten() * timestep_cross_attn_gate_scale_factor, batch_size=batch_size, hidden_dtype=hidden_states.dtype, ) video_cross_attn_scale_shift = video_cross_attn_scale_shift.view( batch_size, -1, video_cross_attn_scale_shift.shape[-1] ) video_cross_attn_a2v_gate = video_cross_attn_a2v_gate.view(batch_size, -1, video_cross_attn_a2v_gate.shape[-1]) audio_cross_attn_scale_shift, _ = self.av_cross_attn_audio_scale_shift( audio_timestep.flatten(), batch_size=batch_size, hidden_dtype=audio_hidden_states.dtype, ) audio_cross_attn_v2a_gate, _ = self.av_cross_attn_audio_v2a_gate( audio_timestep.flatten() * timestep_cross_attn_gate_scale_factor, batch_size=batch_size, hidden_dtype=audio_hidden_states.dtype, ) audio_cross_attn_scale_shift = audio_cross_attn_scale_shift.view( batch_size, -1, audio_cross_attn_scale_shift.shape[-1] ) audio_cross_attn_v2a_gate = audio_cross_attn_v2a_gate.view(batch_size, -1, audio_cross_attn_v2a_gate.shape[-1]) # 3.3. LTX-2.3: Compute prompt modulation from sigma audio_sigma = audio_sigma if audio_sigma is not None else sigma if self.prompt_modulation and sigma is not None: temb_prompt, _ = self.prompt_adaln(sigma.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype) temb_prompt_audio, _ = self.audio_prompt_adaln( audio_sigma.flatten(), batch_size=batch_size, hidden_dtype=audio_hidden_states.dtype ) temb_prompt = temb_prompt.view(batch_size, -1, temb_prompt.size(-1)) temb_prompt_audio = temb_prompt_audio.view(batch_size, -1, temb_prompt_audio.size(-1)) else: temb_prompt = temb_prompt_audio = None # 4. Prepare prompt embeddings # LTX-2: caption projection is in the transformer # LTX-2.3: caption projection is in the connectors (already applied) if hasattr(self, "caption_projection"): encoder_hidden_states = self.caption_projection(encoder_hidden_states) encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1)) if hasattr(self, "audio_caption_projection"): audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states) audio_encoder_hidden_states = audio_encoder_hidden_states.view(batch_size, -1, audio_hidden_states.size(-1)) # 5. Run transformer blocks for block in self.transformer_blocks: block_kwargs = { "temb": temb, "temb_audio": temb_audio, "temb_ca_scale_shift": video_cross_attn_scale_shift, "temb_ca_audio_scale_shift": audio_cross_attn_scale_shift, "temb_ca_gate": video_cross_attn_a2v_gate, "temb_ca_audio_gate": audio_cross_attn_v2a_gate, "temb_prompt": temb_prompt, "temb_prompt_audio": temb_prompt_audio, "video_rotary_emb": video_rotary_emb, "audio_rotary_emb": audio_rotary_emb, "ca_video_rotary_emb": video_cross_attn_rotary_emb, "ca_audio_rotary_emb": audio_cross_attn_rotary_emb, "encoder_attention_mask": encoder_attention_mask, "audio_encoder_attention_mask": audio_encoder_attention_mask, } # NOTE: We intentionally pass the first 4 args here positionally, # because CacheDiT expects the first 2 args to be hidden_states # and encoder_hidden_states, so passing them as kwargs will cause # positional / keywords arg collisions. if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states, audio_hidden_states = self._gradient_checkpointing_func( block, hidden_states, audio_hidden_states, encoder_hidden_states, audio_encoder_hidden_states, **block_kwargs, ) else: hidden_states, audio_hidden_states = block( hidden_states, audio_hidden_states, encoder_hidden_states, audio_encoder_hidden_states, **block_kwargs, ) # 6. Output layers (including unpatchification) scale_shift_values = self.scale_shift_table[None, None] + embedded_timestep[:, :, None] shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] hidden_states = self.norm_out(hidden_states) hidden_states = hidden_states * (1 + scale) + shift output = self.proj_out(hidden_states) audio_scale_shift_values = self.audio_scale_shift_table[None, None] + audio_embedded_timestep[:, :, None] audio_shift, audio_scale = audio_scale_shift_values[:, :, 0], audio_scale_shift_values[:, :, 1] audio_hidden_states = self.audio_norm_out(audio_hidden_states) audio_hidden_states = audio_hidden_states * (1 + audio_scale) + audio_shift audio_output = self.audio_proj_out(audio_hidden_states) if not return_dict: return (output, audio_output) return AudioVisualModelOutput(sample=output, audio_sample=audio_output) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """ Load weights from a pretrained model, mapping separate Q/K/V projections into fused QKV projections for self-attention blocks. Returns: Set of parameter names that were successfully loaded. """ stacked_params_mapping = [ (".attn1.to_qkv", ".attn1.to_q", "q"), (".attn1.to_qkv", ".attn1.to_k", "k"), (".attn1.to_qkv", ".attn1.to_v", "v"), (".audio_attn1.to_qkv", ".audio_attn1.to_q", "q"), (".audio_attn1.to_qkv", ".audio_attn1.to_k", "k"), (".audio_attn1.to_qkv", ".audio_attn1.to_v", "v"), ] params_dict = dict(self.named_parameters()) tp_size = get_tensor_model_parallel_world_size() tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0 loaded_params: set[str] = set() def _maybe_shard_weight(weight: torch.Tensor, param: torch.Tensor) -> torch.Tensor: if tp_size <= 1 or weight.shape == param.shape: return weight if weight.ndim == 1 and weight.numel() == param.numel() * tp_size: return weight.chunk(tp_size, dim=0)[tp_rank] if weight.ndim == 2: if weight.shape[0] == param.shape[0] * tp_size: return weight.chunk(tp_size, dim=0)[tp_rank] if weight.shape[1] == param.shape[1] * tp_size: return weight.chunk(tp_size, dim=1)[tp_rank] return weight for name, loaded_weight in weights: for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) break else: if name not in params_dict: logger.warning( "Skipping transformer weight %s -- not found in model " "parameters. This may indicate an incomplete " "implementation or checkpoint mismatch.", name, ) continue param = params_dict[name] weight_loader = getattr(param, "weight_loader", None) if weight_loader is not None: weight_loader(param, loaded_weight) else: loaded_weight = _maybe_shard_weight(loaded_weight, param) default_weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params