# Copied and adapted from LTX-2 and WanVideo implementations. # # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations from typing import Any, Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F from sglang.jit_kernel.diffusion.ltx2_qknorm_split_rope import ( can_use_ltx2_qknorm_split_rope_cuda, ltx2_qknorm_split_rope_cuda, ) from sglang.jit_kernel.diffusion.residual_gate_add import ( can_use_residual_gate_add_cuda, residual_gate_add_cuda, ) from sglang.multimodal_gen.configs.models.dits.ltx_2 import LTX2ArchConfig, LTX2Config from sglang.multimodal_gen.runtime.distributed import ( get_sp_parallel_rank, get_sp_world_size, get_tp_rank, get_tp_world_size, model_parallel_is_initialized, ) from sglang.multimodal_gen.runtime.distributed.communication_op import ( sequence_model_parallel_all_gather, tensor_model_parallel_all_reduce, ) from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAttention from sglang.multimodal_gen.runtime.layers.layernorm import RMSNormNoWeight from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.visual_embedding import timestep_embedding from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger _is_npu = current_platform.is_npu() logger = init_logger(__name__) ADALN_NUM_BASE_PARAMS = 6 ADALN_NUM_CROSS_ATTN_PARAMS = 3 _LTX2_RESIDUAL_GATE_CUDA_DISABLED = False _LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED = False def _ltx2_residual_gate_add( residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor, ) -> torch.Tensor: global _LTX2_RESIDUAL_GATE_CUDA_DISABLED if not _LTX2_RESIDUAL_GATE_CUDA_DISABLED and can_use_residual_gate_add_cuda( residual, update, gate ): try: return residual_gate_add_cuda(residual, update, gate) except Exception as exc: if torch.compiler.is_compiling(): raise logger.warning_once(f"Disabling LTX2 residual-gate CUDA fast path: {exc}") _LTX2_RESIDUAL_GATE_CUDA_DISABLED = True return residual + update * gate def _ltx2_try_fused_qknorm_split_rope( q: torch.Tensor, k: torch.Tensor, q_norm: nn.Module, k_norm: nn.Module, q_cos: torch.Tensor, q_sin: torch.Tensor, k_cos: torch.Tensor, k_sin: torch.Tensor, *, eps: float, num_heads: int, head_dim: int, ) -> tuple[torch.Tensor, torch.Tensor] | None: global _LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED if ( _LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED or get_tp_world_size() != 1 or not isinstance(q_norm, nn.RMSNorm) or not isinstance(k_norm, nn.RMSNorm) or float(q_norm.eps) != float(eps) or float(k_norm.eps) != float(eps) or not can_use_ltx2_qknorm_split_rope_cuda( q, q_cos, q_sin, q_norm.weight, k, k_cos, k_sin, k_norm.weight, num_heads=num_heads, head_dim=head_dim, ) ): return None try: return ltx2_qknorm_split_rope_cuda( q, q_cos, q_sin, q_norm.weight, k, k_cos, k_sin, k_norm.weight, eps=eps, num_heads=num_heads, head_dim=head_dim, ) except Exception as exc: if torch.compiler.is_compiling(): raise logger.warning_once(f"Disabling LTX2 QKNorm split-RoPE CUDA fast path: {exc}") _LTX2_QKNORM_SPLIT_ROPE_CUDA_DISABLED = True return None _LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED = False 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 ) def _ltx2_disable_fused_ada_values(exc: Exception) -> None: global _LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED _LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED = True logger.warning_once(f"Disabling LTX2 fused Ada values fast path: {exc}") def _ltx2_try_fused_ada_values9( scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, ) -> tuple[torch.Tensor, ...] | None: if ( _LTX2_FUSED_ADA_VALUES_RUNTIME_DISABLED or get_tp_world_size() != 1 or not timestep.is_cuda or timestep.dtype != torch.bfloat16 or timestep.ndim != 3 or int(timestep.shape[0]) != int(batch_size) or not timestep.is_contiguous() or not scale_shift_table.is_cuda or scale_shift_table.dtype not in (torch.bfloat16, torch.float32) or scale_shift_table.ndim != 2 or int(scale_shift_table.shape[0]) != 9 or scale_shift_table.stride(-1) != 1 ): return None hidden = int(scale_shift_table.shape[1]) if hidden % 256 != 0 or hidden > 8192 or timestep.shape[-1] != 9 * hidden: return None try: from sglang.jit_kernel.diffusion.triton.ltx2_ada_values import ( ltx2_ada_values9, ) return ltx2_ada_values9(scale_shift_table, timestep) except Exception as exc: _ltx2_disable_fused_ada_values(exc) return None def _ltx2_is_perturbed( perturbation_config: dict[str, object], key: str, block_idx: int, ) -> bool: value = perturbation_config.get(key) if value is None: return False if key.endswith("_blocks"): return block_idx in value return bool(value) def _ltx2_build_batched_perturbation_states( perturbation_configs: tuple[dict[str, object], ...], key: str, block_indices: tuple[int, ...], values: torch.Tensor, ) -> dict[int, tuple[torch.Tensor | None, bool]]: mask_cache: dict[tuple[int, ...], torch.Tensor] = {} states: dict[int, tuple[torch.Tensor | None, bool]] = {} for block_idx in block_indices: keep_values = [] any_perturbed = False all_perturbed = True for config in perturbation_configs: perturbed = _ltx2_is_perturbed(config, key, block_idx) any_perturbed = any_perturbed or perturbed all_perturbed = all_perturbed and perturbed keep_values.append(0 if perturbed else 1) if not any_perturbed: states[block_idx] = (None, False) elif all_perturbed: states[block_idx] = (None, True) else: cache_key = tuple(keep_values) mask = mask_cache.get(cache_key) if mask is None: mask = torch.tensor( keep_values, device=values.device, dtype=values.dtype ).view(len(keep_values), *([1] * (values.ndim - 1))) mask_cache[cache_key] = mask states[block_idx] = (mask, False) return states def apply_interleaved_rotary_emb( x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor] ) -> torch.Tensor: cos, sin = freqs x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1) x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2) return x * cos + x_rotated * sin def apply_split_rotary_emb( x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor] ) -> torch.Tensor: cos, sin = freqs if ( x.ndim == 3 and cos.ndim == 4 and sin.ndim == 4 and x.dtype == torch.bfloat16 and cos.dtype == torch.bfloat16 and sin.dtype == torch.bfloat16 and x.is_cuda and x.is_contiguous() and cos.is_cuda and sin.is_cuda ): from sglang.jit_kernel.diffusion.triton.ltx2_rotary import ( apply_ltx2_split_rotary_emb, ) return apply_ltx2_split_rotary_emb(x, cos, sin) x_dtype = x.dtype needs_reshape = False if x.ndim != 4 and cos.ndim == 4: b = x.shape[0] _, h, t, _ = cos.shape x = x.reshape(b, t, h, -1).swapaxes(1, 2) needs_reshape = True 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 split_x = x.reshape(*x.shape[:-1], 2, r) first_x = split_x[..., :1, :] second_x = split_x[..., 1:, :] cos_u = cos.unsqueeze(-2) 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, -1) return out.to(dtype=x_dtype) # ============================================================================== # Layers and Embeddings # ============================================================================== class LTX2AudioVideoRotaryPosEmbed(nn.Module): 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 = int(dim) self.patch_size = int(patch_size) self.patch_size_t = int(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 = int(base_num_frames) self.num_attention_heads = int(num_attention_heads) self.base_height = int(base_height) self.base_width = int(base_width) self.sampling_rate = int(sampling_rate) self.hop_length = int(hop_length) self.audio_latents_per_second = ( float(self.sampling_rate) / float(self.hop_length) / float(scale_factors[0]) ) self.scale_factors = tuple(int(x) for x in scale_factors) self.theta = float(theta) self.causal_offset = int(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 = bool(double_precision) def prepare_video_coords( self, batch_size: int, num_frames: int, height: int, width: int, device: torch.device, fps: float = 24.0, *, start_frame: int = 0, ) -> torch.Tensor: grid_f = torch.arange( start=int(start_frame), end=int(num_frames) + int(start_frame), 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, ) grid = torch.meshgrid(grid_f, grid_h, grid_w, indexing="ij") grid = torch.stack(grid, dim=0) 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) latent_coords = torch.stack([grid, patch_ends], dim=-1) latent_coords = latent_coords.flatten(1, 3) latent_coords = latent_coords.unsqueeze(0).repeat(batch_size, 1, 1, 1) scale_tensor = torch.tensor(self.scale_factors, device=latent_coords.device) broadcast_shape = [1] * latent_coords.ndim broadcast_shape[1] = -1 pixel_coords = latent_coords * scale_tensor.view(*broadcast_shape) pixel_coords[:, 0, ...] = ( pixel_coords[:, 0, ...] + self.causal_offset - self.scale_factors[0] ).clamp(min=0) pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps return pixel_coords def prepare_audio_coords( self, batch_size: int, num_frames: int, device: torch.device, *, start_frame: int = 0, ) -> torch.Tensor: grid_f = torch.arange( start=int(start_frame), end=int(num_frames) + int(start_frame), step=self.patch_size_t, dtype=torch.float32, device=device, ) audio_scale_factor = self.scale_factors[0] grid_start_mel = grid_f * audio_scale_factor grid_start_mel = ( grid_start_mel + self.causal_offset - audio_scale_factor ).clip(min=0) grid_start_s = grid_start_mel * self.hop_length / self.sampling_rate 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) audio_coords = audio_coords.unsqueeze(0).expand(batch_size, -1, -1) audio_coords = audio_coords.unsqueeze(1) return audio_coords def prepare_coords(self, *args, **kwargs): if self.modality == "video": return self.prepare_video_coords(*args, **kwargs) return self.prepare_audio_coords(*args, **kwargs) def forward( self, coords: torch.Tensor, device: Optional[Union[str, torch.device]] = None, out_dtype: Optional[torch.dtype] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: device = device or coords.device out_dtype = out_dtype or coords.dtype num_pos_dims = coords.shape[1] if coords.ndim == 4: coords_start, coords_end = coords.chunk(2, dim=-1) coords = (coords_start + coords_end) / 2.0 coords = coords.squeeze(-1) if self.modality == "video": max_positions = (self.base_num_frames, self.base_height, self.base_width) else: max_positions = (self.base_num_frames,) grid = torch.stack( [coords[:, i] / max_positions[i] for i in range(num_pos_dims)], dim=-1 ).to(device) num_rope_elems = num_pos_dims * 2 # LTX-2.3 HQ is sensitive to RoPE rounding; keep frequency generation on # the target device instead of caching a CPU/NumPy tensor. 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) freqs = (grid.unsqueeze(-1) * 2 - 1) * freqs freqs = freqs.transpose(-1, -2).flatten(2) 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) else: 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.cat([cos_padding, cos_freq], dim=-1) sin_freq = torch.cat([sin_padding, sin_freq], dim=-1) b = cos_freq.shape[0] t = cos_freq.shape[1] cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, -1) sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, -1) cos_freqs = torch.swapaxes(cos_freq, 1, 2) sin_freqs = torch.swapaxes(sin_freq, 1, 2) return cos_freqs.to(dtype=out_dtype), sin_freqs.to(dtype=out_dtype) class LTX2TextProjection(nn.Module): def __init__( self, in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = "gelu_tanh", ) -> None: super().__init__() if out_features is None: out_features = hidden_size self.linear_1 = ColumnParallelLinear( in_features, hidden_size, bias=True, gather_output=True ) if act_fn == "gelu_tanh": self.act_1 = nn.GELU(approximate="tanh") elif act_fn == "silu": self.act_1 = nn.SiLU() else: raise ValueError(f"Unknown activation function: {act_fn}") self.linear_2 = ColumnParallelLinear( hidden_size, out_features, bias=True, gather_output=True ) def forward(self, caption: torch.Tensor) -> torch.Tensor: hidden_states, _ = self.linear_1(caption) hidden_states = self.act_1(hidden_states) hidden_states, _ = self.linear_2(hidden_states) return hidden_states class LTX2TimestepEmbedder(nn.Module): def __init__(self, embedding_dim: int, in_channels: int = 256) -> None: super().__init__() self.linear_1 = ColumnParallelLinear( in_channels, embedding_dim, bias=True, gather_output=True ) self.linear_2 = ColumnParallelLinear( embedding_dim, embedding_dim, bias=True, gather_output=True ) def forward(self, t_emb: torch.Tensor) -> torch.Tensor: x, _ = self.linear_1(t_emb) x = F.silu(x) x, _ = self.linear_2(x) return x class LTX2PixArtAlphaCombinedTimestepSizeEmbeddings(nn.Module): def __init__(self, embedding_dim: int) -> None: super().__init__() self.timestep_embedder = LTX2TimestepEmbedder(embedding_dim, in_channels=256) def forward( self, timestep: torch.Tensor, hidden_dtype: torch.dtype | None = None ) -> torch.Tensor: t = timestep.reshape(-1).to(dtype=torch.float32) t_emb = timestep_embedding(t, dim=256, max_period=10000, dtype=torch.float32) if hidden_dtype is not None: t_emb = t_emb.to(dtype=hidden_dtype) return self.timestep_embedder(t_emb) class LTX2AdaLayerNormSingle(nn.Module): def __init__(self, embedding_dim: int, embedding_coefficient: int = 6) -> None: super().__init__() self.emb = LTX2PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim) self.silu = nn.SiLU() self.linear = ColumnParallelLinear( embedding_dim, embedding_coefficient * embedding_dim, bias=True, gather_output=True, ) def forward( self, timestep: torch.Tensor, hidden_dtype: torch.dtype | None = None ) -> tuple[torch.Tensor, torch.Tensor]: embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype).to( dtype=self.linear.weight.dtype ) out, _ = self.linear(self.silu(embedded_timestep)) return out, embedded_timestep class LTX2TPRMSNormAcrossHeads(nn.Module): def __init__( self, full_hidden_size: int, local_hidden_size: int, eps: float ) -> None: super().__init__() self.full_hidden_size = full_hidden_size self.local_hidden_size = local_hidden_size self.eps = eps self.weight = nn.Parameter(torch.ones(local_hidden_size)) tp_rank = get_tp_rank() def _weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: shard = loaded_weight.narrow( 0, tp_rank * local_hidden_size, local_hidden_size ) param.data.copy_(shard.to(dtype=param.dtype, device=param.device)) setattr(self.weight, "weight_loader", _weight_loader) def forward(self, x: torch.Tensor) -> torch.Tensor: # Keep track of the original dtype. We do the statistics in fp32 for # numerical stability, but cast the output back to the input dtype to orig_dtype = x.dtype if get_tp_world_size() == 1: var = x.float().pow(2).mean(dim=-1, keepdim=True) else: local_sumsq = x.float().pow(2).sum(dim=-1, keepdim=True) global_sumsq = tensor_model_parallel_all_reduce(local_sumsq) var = global_sumsq / float(self.full_hidden_size) inv_rms_fp32 = torch.rsqrt(var + self.eps) y = (x.float() * inv_rms_fp32).to(dtype=orig_dtype) return y * self.weight.to(dtype=orig_dtype) class LTX2Attention(nn.Module): def __init__( self, query_dim: int, context_dim: int | None = None, heads: int = 8, dim_head: int = 64, norm_eps: float = 1e-6, qk_norm: bool = True, use_local_attention: bool = False, apply_gated_attention: bool = False, enable_packed_qkv_input_a2a: bool = False, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", quant_config: QuantizationConfig | None = None, ) -> None: super().__init__() self.query_dim = int(query_dim) self.context_dim = int(query_dim if context_dim is None else context_dim) self.heads = int(heads) self.dim_head = int(dim_head) self.inner_dim = self.heads * self.dim_head self.norm_eps = float(norm_eps) self.qk_norm = bool(qk_norm) self.use_local_attention = bool(use_local_attention) self.apply_gated_attention = bool(apply_gated_attention) self.enable_packed_qkv_input_a2a = bool(enable_packed_qkv_input_a2a) self.prefix = prefix tp_size = get_tp_world_size() if tp_size <= 0: raise ValueError(f"Invalid {tp_size=}. Expected tp_size >= 1.") if self.heads % tp_size != 0: raise ValueError( f"LTX2Attention requires heads divisible by tp_size, got " f"{self.heads=} {tp_size=}." ) if self.inner_dim % tp_size != 0: # This should follow from heads % tp_size, but keep explicit for clarity. raise ValueError( f"LTX2Attention requires inner_dim divisible by tp_size, got " f"{self.inner_dim=} {tp_size=}." ) self.local_heads = self.heads // tp_size self.to_q = ColumnParallelLinear( self.query_dim, self.inner_dim, bias=True, gather_output=False, quant_config=quant_config, ) self.to_k = ColumnParallelLinear( self.context_dim, self.inner_dim, bias=True, gather_output=False, quant_config=quant_config, ) self.to_v = ColumnParallelLinear( self.context_dim, self.inner_dim, bias=True, gather_output=False, quant_config=quant_config, ) self.to_gate_logits: ColumnParallelLinear | None = None if self.apply_gated_attention: self.to_gate_logits = ColumnParallelLinear( self.query_dim, self.heads, bias=True, gather_output=False, quant_config=quant_config, ) self.q_norm: nn.Module | None = None self.k_norm: nn.Module | None = None if self.qk_norm: if tp_size == 1: self.q_norm = torch.nn.RMSNorm(self.inner_dim, eps=self.norm_eps) self.k_norm = torch.nn.RMSNorm(self.inner_dim, eps=self.norm_eps) else: self.q_norm = LTX2TPRMSNormAcrossHeads( full_hidden_size=self.inner_dim, local_hidden_size=self.inner_dim // tp_size, eps=self.norm_eps, ) self.k_norm = LTX2TPRMSNormAcrossHeads( full_hidden_size=self.inner_dim, local_hidden_size=self.inner_dim // tp_size, eps=self.norm_eps, ) self.to_out = nn.Sequential( RowParallelLinear( self.inner_dim, self.query_dim, bias=True, input_is_parallel=True, quant_config=quant_config, ), nn.Identity(), ) if self.use_local_attention: self.attn = LocalAttention( num_heads=self.local_heads, head_size=self.dim_head, num_kv_heads=self.local_heads, softmax_scale=None, causal=False, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn", enable_packed_qkv_input_a2a=self.enable_packed_qkv_input_a2a, # official LTX2 torch_sdpa uses cuDNN; cuda setup disables it allow_cudnn_sdp=True, ) else: self.attn = USPAttention( num_heads=self.local_heads, head_size=self.dim_head, num_kv_heads=self.local_heads, dropout_rate=0, softmax_scale=None, causal=False, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn", # official LTX2 torch_sdpa uses cuDNN; cuda setup disables it allow_cudnn_sdp=True, ) 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, perturbation_mask: torch.Tensor | None = None, all_perturbed: bool = False, skip_sequence_parallel_override: bool = False, gather_context_kv_for_sp: bool = False, context_replicated_prefix_len: int = 0, ) -> torch.Tensor: gate_input = x context_ = x if context is None else context v, _ = self.to_v(context_) use_attention = not all_perturbed if use_attention: q, _ = self.to_q(x) k, _ = self.to_k(context_) fused_qk = None if pe is not None: cos, sin = pe k_cos, k_sin = pe if k_pe is None else k_pe tp_size = get_tp_world_size() if tp_size > 1: tp_rank = get_tp_rank() cos, sin = self._slice_rope_for_tp( cos, sin, tp_rank=tp_rank, tp_size=tp_size ) k_cos, k_sin = self._slice_rope_for_tp( k_cos, k_sin, tp_rank=tp_rank, tp_size=tp_size ) if self.qk_norm and cos.dim() != 3: assert self.q_norm is not None and self.k_norm is not None fused_qk = _ltx2_try_fused_qknorm_split_rope( q, k, self.q_norm, self.k_norm, cos, sin, k_cos, k_sin, eps=self.norm_eps, num_heads=self.local_heads, head_dim=self.dim_head, ) if fused_qk is not None: q, k = fused_qk else: if self.qk_norm: assert self.q_norm is not None and self.k_norm is not None q = self.q_norm(q) k = self.k_norm(k) if pe is not None and cos.dim() == 3: q = apply_interleaved_rotary_emb(q, (cos, sin)) k = apply_interleaved_rotary_emb(k, (k_cos, k_sin)) elif pe is not None: q = apply_split_rotary_emb(q, (cos, sin)) k = apply_split_rotary_emb(k, (k_cos, k_sin)) v = v.view(*v.shape[:-1], self.local_heads, self.dim_head) if use_attention: q = q.view(*q.shape[:-1], self.local_heads, self.dim_head) k = k.view(*k.shape[:-1], self.local_heads, self.dim_head) if gather_context_kv_for_sp: # Replicated prefix (e.g. JoyEcho memory) is identical on every rank; only gather the sharded suffix. if context_replicated_prefix_len > 0: prefix = int(context_replicated_prefix_len) k_prefix, k_suffix = k[:, :prefix], k[:, prefix:] v_prefix, v_suffix = v[:, :prefix], v[:, prefix:] k_full = torch.cat( [ k_prefix, sequence_model_parallel_all_gather( k_suffix.contiguous(), dim=1 ), ], dim=1, ) v_full = torch.cat( [ v_prefix, sequence_model_parallel_all_gather( v_suffix.contiguous(), dim=1 ), ], dim=1, ) gathered_mask = mask else: k_full = sequence_model_parallel_all_gather(k.contiguous(), dim=1) v_full = sequence_model_parallel_all_gather(v.contiguous(), dim=1) gathered_mask = None if mask is not None: gathered_mask = sequence_model_parallel_all_gather( mask.contiguous(), dim=1 ) if self.use_local_attention: out = self.attn(q, k_full, v_full, attn_mask=gathered_mask) else: out = self.attn( q, k_full, v_full, attn_mask=gathered_mask, skip_sequence_parallel_override=True, ) elif self.use_local_attention: out = self.attn(q, k, v, attn_mask=mask) else: out = self.attn( q, k, v, attn_mask=mask, skip_sequence_parallel_override=skip_sequence_parallel_override, ) if perturbation_mask is not None: if perturbation_mask.ndim == out.ndim - 1: perturbation_mask = perturbation_mask.unsqueeze(-1) out = out * perturbation_mask + v * (1 - perturbation_mask) if not use_attention: out = v if self.to_gate_logits is not None: gate_logits, _ = self.to_gate_logits(gate_input) b, t = out.shape[:2] out = out.view(b, t, self.local_heads, self.dim_head) out = out * (2.0 * torch.sigmoid(gate_logits).unsqueeze(-1)) out = out.view(b, t, self.local_heads * self.dim_head) out_flat = out.flatten(2) out_proj, _ = self.to_out[0](out_flat) return out_proj def _slice_rope_for_tp( self, cos: torch.Tensor, sin: torch.Tensor, *, tp_rank: int, tp_size: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Slice RoPE tensors to the local TP shard. - split-rope: cos/sin are shaped [B, H, T, R] (head-major), slice by heads. - interleaved-rope: cos/sin are shaped [B, T, D], where D matches the projected feature dimension and is sharded by TP. """ if cos.ndim == 4: # [B, H, T, R] start = tp_rank * self.local_heads end = start + self.local_heads return cos[:, start:end, :, :], sin[:, start:end, :, :] elif cos.ndim == 3: # [B, T, D] d = cos.shape[-1] if d % tp_size != 0: raise ValueError( f"RoPE dim must be divisible by tp_size, got {d=} {tp_size=}." ) local_d = d // tp_size start = tp_rank * local_d end = start + local_d return cos[:, :, start:end], sin[:, :, start:end] raise ValueError(f"Unexpected RoPE tensor rank: {cos.ndim}. Expected 3 or 4.") class LTX2FeedForward(nn.Module): def __init__( self, dim: int, dim_out: int | None = None, mult: int = 4, quant_config: QuantizationConfig | None = None, ) -> None: super().__init__() if dim_out is None: dim_out = dim inner_dim = int(dim * mult) self.proj_in = ColumnParallelLinear( dim, inner_dim, bias=True, gather_output=False, quant_config=quant_config ) self.act = nn.GELU(approximate="tanh") self.proj_out = RowParallelLinear( inner_dim, dim_out, bias=True, input_is_parallel=True, quant_config=quant_config, ) def forward(self, x: torch.Tensor) -> torch.Tensor: x, _ = self.proj_in(x) x = self.act(x) x, _ = self.proj_out(x) return x class LTX2TransformerBlock(nn.Module): def __init__( self, idx: int, 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: int, audio_cross_attention_dim: int, qk_norm: bool = True, norm_eps: float = 1e-6, apply_gated_attention: bool = False, cross_attention_adaln: bool = False, use_local_av_cross_attention: bool = False, force_sdpa_v2a_cross_attention: bool = False, enable_packed_qkv_input_a2a: bool = False, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", quant_config: QuantizationConfig | None = None, ): super().__init__() self.idx = idx self.norm_eps = norm_eps self.rms_norm = RMSNormNoWeight() # LTX2.3 self.cross_attention_adaln = cross_attention_adaln self.use_local_av_cross_attention = use_local_av_cross_attention # 1. Self-Attention (video and audio) self.attn1 = LTX2Attention( query_dim=dim, heads=num_attention_heads, dim_head=attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, apply_gated_attention=apply_gated_attention, enable_packed_qkv_input_a2a=enable_packed_qkv_input_a2a, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn1", quant_config=quant_config, ) self.audio_attn1 = LTX2Attention( query_dim=audio_dim, heads=audio_num_attention_heads, dim_head=audio_attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, apply_gated_attention=apply_gated_attention, enable_packed_qkv_input_a2a=enable_packed_qkv_input_a2a, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.audio_attn1", quant_config=quant_config, ) # 2. Prompt Cross-Attention # Prompt KV is replicated across SP ranks, so prompt cross-attn should # stay local and preserve the explicit KV mask semantics from official. self.attn2 = LTX2Attention( query_dim=dim, context_dim=cross_attention_dim, heads=num_attention_heads, dim_head=attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, use_local_attention=True, apply_gated_attention=apply_gated_attention, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn2", quant_config=quant_config, ) self.audio_attn2 = LTX2Attention( query_dim=audio_dim, context_dim=audio_cross_attention_dim, heads=audio_num_attention_heads, dim_head=audio_attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, use_local_attention=True, apply_gated_attention=apply_gated_attention, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.audio_attn2", quant_config=quant_config, ) # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention self.audio_to_video_attn = LTX2Attention( query_dim=dim, context_dim=audio_dim, heads=audio_num_attention_heads, dim_head=audio_attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, use_local_attention=use_local_av_cross_attention, apply_gated_attention=apply_gated_attention, enable_packed_qkv_input_a2a=enable_packed_qkv_input_a2a, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.audio_to_video_attn", quant_config=quant_config, ) self.video_to_audio_attn = LTX2Attention( query_dim=audio_dim, context_dim=dim, heads=audio_num_attention_heads, dim_head=audio_attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, use_local_attention=use_local_av_cross_attention, apply_gated_attention=apply_gated_attention, enable_packed_qkv_input_a2a=enable_packed_qkv_input_a2a, supported_attention_backends=( {AttentionBackendEnum.TORCH_SDPA} if force_sdpa_v2a_cross_attention else supported_attention_backends ), prefix=f"{prefix}.video_to_audio_attn", quant_config=quant_config, ) # 4. Feedforward layers self.ff = LTX2FeedForward(dim, dim_out=dim, quant_config=quant_config) self.audio_ff = LTX2FeedForward( audio_dim, dim_out=audio_dim, quant_config=quant_config ) # 5. Modulation Parameters num_ada_params = adaln_embedding_coefficient(cross_attention_adaln) self.scale_shift_table = nn.Parameter( torch.randn(num_ada_params, dim) / dim**0.5 ) self.audio_scale_shift_table = nn.Parameter( torch.randn(num_ada_params, audio_dim) / audio_dim**0.5 ) 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) ) if self.cross_attention_adaln: # LTX2.3 self.prompt_scale_shift_table = nn.Parameter(torch.randn(2, dim)) self.audio_prompt_scale_shift_table = nn.Parameter( torch.randn(2, audio_dim) ) def get_ada_values( self, scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, indices: slice, ) -> tuple[torch.Tensor, ...]: num_ada_params = int(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 [t.squeeze(2) for t in ada_values] 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_prompt: torch.Tensor | None, temb_audio_prompt: torch.Tensor | None, temb_ca_scale_shift: torch.Tensor, temb_ca_audio_scale_shift: torch.Tensor, temb_ca_gate: torch.Tensor, temb_ca_audio_gate: torch.Tensor, video_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, audio_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ca_video_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ca_audio_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, encoder_attention_mask: Optional[torch.Tensor] = None, audio_encoder_attention_mask: Optional[torch.Tensor] = None, video_self_attention_mask: Optional[torch.Tensor] = None, audio_self_attention_mask: Optional[torch.Tensor] = None, a2v_cross_attention_mask: Optional[torch.Tensor] = None, v2a_cross_attention_mask: Optional[torch.Tensor] = None, skip_video_self_attn: bool = False, skip_audio_self_attn: bool = False, skip_a2v_cross_attn: bool = False, skip_v2a_cross_attn: bool = False, video_self_attn_perturbation_mask: Optional[torch.Tensor] = None, audio_self_attn_perturbation_mask: Optional[torch.Tensor] = None, a2v_cross_attn_perturbation_mask: Optional[torch.Tensor] = None, v2a_cross_attn_perturbation_mask: Optional[torch.Tensor] = None, audio_replicated_for_sp: bool = False, video_memory_prefix_len: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: batch_size = hidden_states.size(0) video_ada_values = _ltx2_try_fused_ada_values9( self.scale_shift_table, batch_size, temb ) audio_ada_values = _ltx2_try_fused_ada_values9( self.audio_scale_shift_table, batch_size, temb_audio ) # 1. Video and Audio Self-Attention if video_ada_values is None: vshift_msa, vscale_msa, vgate_msa = self.get_ada_values( self.scale_shift_table, batch_size, temb, slice(0, 3) ) else: vshift_msa, vscale_msa, vgate_msa = video_ada_values[0:3] norm_hidden_states = ( self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_msa) + vshift_msa ) attn_hidden_states = self.attn1( norm_hidden_states, mask=video_self_attention_mask, pe=video_rotary_emb, perturbation_mask=video_self_attn_perturbation_mask, all_perturbed=skip_video_self_attn, gather_context_kv_for_sp=audio_replicated_for_sp, context_replicated_prefix_len=video_memory_prefix_len, ) hidden_states = _ltx2_residual_gate_add( hidden_states, attn_hidden_states, vgate_msa ) if audio_ada_values is None: ashift_msa, ascale_msa, agate_msa = self.get_ada_values( self.audio_scale_shift_table, batch_size, temb_audio, slice(0, 3) ) else: ashift_msa, ascale_msa, agate_msa = audio_ada_values[0:3] norm_audio_hidden_states = ( self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_msa) + ashift_msa ) attn_audio_hidden_states = self.audio_attn1( norm_audio_hidden_states, mask=audio_self_attention_mask, pe=audio_rotary_emb, perturbation_mask=audio_self_attn_perturbation_mask, all_perturbed=skip_audio_self_attn, skip_sequence_parallel_override=audio_replicated_for_sp, ) audio_hidden_states = _ltx2_residual_gate_add( audio_hidden_states, attn_audio_hidden_states, agate_msa ) # 2. Prompt Cross-Attention if self.cross_attention_adaln: # LTX2.3 if temb_prompt is None or temb_audio_prompt is None: raise ValueError( "cross_attention_adaln requires prompt modulation tensors." ) if video_ada_values is None: vshift_q, vscale_q, vgate_q = self.get_ada_values( self.scale_shift_table, batch_size, temb, slice(6, 9) ) else: vshift_q, vscale_q, vgate_q = video_ada_values[6:9] v_prompt_shift, v_prompt_scale = self.get_ada_values( self.prompt_scale_shift_table, batch_size, temb_prompt, slice(None) ) norm_hidden_states = ( self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_q) + vshift_q ) mod_encoder_hidden_states = ( encoder_hidden_states * (1 + v_prompt_scale) + v_prompt_shift ) attn_hidden_states = self.attn2( norm_hidden_states, context=mod_encoder_hidden_states, mask=encoder_attention_mask, ) hidden_states = _ltx2_residual_gate_add( hidden_states, attn_hidden_states, vgate_q ) if audio_ada_values is None: ashift_q, ascale_q, agate_q = self.get_ada_values( self.audio_scale_shift_table, batch_size, temb_audio, slice(6, 9) ) else: ashift_q, ascale_q, agate_q = audio_ada_values[6:9] a_prompt_shift, a_prompt_scale = self.get_ada_values( self.audio_prompt_scale_shift_table, batch_size, temb_audio_prompt, slice(None), ) norm_audio_hidden_states = ( self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_q) + ashift_q ) mod_audio_encoder_hidden_states = ( audio_encoder_hidden_states * (1 + a_prompt_scale) + a_prompt_shift ) attn_audio_hidden_states = self.audio_attn2( norm_audio_hidden_states, context=mod_audio_encoder_hidden_states, mask=audio_encoder_attention_mask, ) audio_hidden_states = _ltx2_residual_gate_add( audio_hidden_states, attn_audio_hidden_states, agate_q ) else: norm_hidden_states = self.rms_norm(hidden_states, self.norm_eps) attn_hidden_states = self.attn2( norm_hidden_states, context=encoder_hidden_states, mask=encoder_attention_mask, ) hidden_states = hidden_states + attn_hidden_states norm_audio_hidden_states = self.rms_norm(audio_hidden_states, self.norm_eps) attn_audio_hidden_states = self.audio_attn2( norm_audio_hidden_states, context=audio_encoder_hidden_states, mask=audio_encoder_attention_mask, ) audio_hidden_states = audio_hidden_states + attn_audio_hidden_states # 3. Audio-to-Video and Video-to-Audio Cross-Attention norm_hidden_states = self.rms_norm(hidden_states, self.norm_eps) norm_audio_hidden_states = self.rms_norm(audio_hidden_states, self.norm_eps) # Compute combined ada params 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_scale_shift_table = ( video_per_layer_ca_scale_shift[None, None, :, :].to( dtype=temb_ca_scale_shift.dtype, device=temb_ca_scale_shift.device ) + temb_ca_scale_shift.reshape( batch_size, temb_ca_scale_shift.shape[1], 4, -1 ) ).unbind(dim=2) video_ca_gate = ( video_per_layer_ca_gate[None, None, :, :].to( dtype=temb_ca_gate.dtype, device=temb_ca_gate.device ) + temb_ca_gate.reshape(batch_size, temb_ca_gate.shape[1], 1, -1) ).unbind(dim=2) ( video_a2v_ca_scale, video_a2v_ca_shift, video_v2a_ca_scale, video_v2a_ca_shift, ) = [t.squeeze(2) for t in video_ca_scale_shift_table] a2v_gate = video_ca_gate[0].squeeze(2) 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_scale_shift_table = ( audio_per_layer_ca_scale_shift[None, None, :, :].to( dtype=temb_ca_audio_scale_shift.dtype, device=temb_ca_audio_scale_shift.device, ) + temb_ca_audio_scale_shift.reshape( batch_size, temb_ca_audio_scale_shift.shape[1], 4, -1 ) ).unbind(dim=2) audio_ca_gate = ( audio_per_layer_ca_gate[None, None, :, :].to( dtype=temb_ca_audio_gate.dtype, device=temb_ca_audio_gate.device ) + temb_ca_audio_gate.reshape(batch_size, temb_ca_audio_gate.shape[1], 1, -1) ).unbind(dim=2) ( audio_a2v_ca_scale, audio_a2v_ca_shift, audio_v2a_ca_scale, audio_v2a_ca_shift, ) = [t.squeeze(2) for t in audio_ca_scale_shift_table] v2a_gate = audio_ca_gate[0].squeeze(2) # A2V mod_norm_hidden_states = ( norm_hidden_states * (1 + video_a2v_ca_scale) + video_a2v_ca_shift ) mod_norm_audio_hidden_states = ( norm_audio_hidden_states * (1 + audio_a2v_ca_scale) + audio_a2v_ca_shift ) if not skip_a2v_cross_attn: a2v_attn_hidden_states = self.audio_to_video_attn( mod_norm_hidden_states, context=mod_norm_audio_hidden_states, pe=ca_video_rotary_emb, k_pe=ca_audio_rotary_emb, mask=a2v_cross_attention_mask, skip_sequence_parallel_override=audio_replicated_for_sp, ) if a2v_cross_attn_perturbation_mask is not None: a2v_attn_hidden_states = ( a2v_attn_hidden_states * a2v_cross_attn_perturbation_mask ) hidden_states = _ltx2_residual_gate_add( hidden_states, a2v_attn_hidden_states, a2v_gate ) # V2A mod_norm_hidden_states = ( norm_hidden_states * (1 + video_v2a_ca_scale) + video_v2a_ca_shift ) mod_norm_audio_hidden_states = ( norm_audio_hidden_states * (1 + audio_v2a_ca_scale) + audio_v2a_ca_shift ) if not skip_v2a_cross_attn: v2a_attn_hidden_states = self.video_to_audio_attn( mod_norm_audio_hidden_states, context=mod_norm_hidden_states, pe=ca_audio_rotary_emb, k_pe=ca_video_rotary_emb, mask=v2a_cross_attention_mask, gather_context_kv_for_sp=audio_replicated_for_sp, context_replicated_prefix_len=video_memory_prefix_len, ) if v2a_cross_attn_perturbation_mask is not None: v2a_attn_hidden_states = ( v2a_attn_hidden_states * v2a_cross_attn_perturbation_mask ) audio_hidden_states = _ltx2_residual_gate_add( audio_hidden_states, v2a_attn_hidden_states, v2a_gate ) # 4. Feedforward if video_ada_values is None: vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values( self.scale_shift_table, batch_size, temb, slice(3, 6) ) else: vshift_mlp, vscale_mlp, vgate_mlp = video_ada_values[3:6] norm_hidden_states = ( self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_mlp) + vshift_mlp ) ff_output = self.ff(norm_hidden_states) hidden_states = _ltx2_residual_gate_add(hidden_states, ff_output, vgate_mlp) if audio_ada_values is None: ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values( self.audio_scale_shift_table, batch_size, temb_audio, slice(3, 6) ) else: ashift_mlp, ascale_mlp, agate_mlp = audio_ada_values[3:6] norm_audio_hidden_states = ( self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_mlp) + ashift_mlp ) audio_ff_output = self.audio_ff(norm_audio_hidden_states) audio_hidden_states = _ltx2_residual_gate_add( audio_hidden_states, audio_ff_output, agate_mlp ) return hidden_states, audio_hidden_states class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin): _fsdp_shard_conditions = LTX2ArchConfig()._fsdp_shard_conditions _compile_conditions = LTX2ArchConfig()._compile_conditions _supported_attention_backends = LTX2ArchConfig()._supported_attention_backends param_names_mapping = LTX2ArchConfig().param_names_mapping reverse_param_names_mapping = LTX2ArchConfig().reverse_param_names_mapping lora_param_names_mapping = LTX2ArchConfig().lora_param_names_mapping @staticmethod def _collapse_prompt_timestep(timestep: torch.Tensor) -> torch.Tensor: if timestep.ndim <= 1: return timestep return timestep.amax(dim=tuple(range(1, timestep.ndim))) def _scale_timestep_for_adaln(self, timestep: torch.Tensor) -> torch.Tensor: ltx_variant = str(getattr(self.config.arch_config, "ltx_variant", "ltx_2")) if ltx_variant == "ltx_2_3" and bool( getattr(self, "_sglang_use_ltx23_hq_timestep_semantics", False) ): return timestep * float(self.timestep_scale_multiplier) return timestep def _validate_tp_config(self, *, arch: LTX2ArchConfig, tp_size: int) -> None: """Validate TP-related dimension constraints (fail-fast).""" if tp_size < 1: raise ValueError(f"Invalid tp_size={tp_size}. Expected tp_size >= 1.") if self.hidden_size % self.num_attention_heads != 0: raise ValueError( "video hidden_size must be divisible by num_attention_heads, got " f"{self.hidden_size=} {self.num_attention_heads=}." ) if self.audio_hidden_size % self.audio_num_attention_heads != 0: raise ValueError( "audio_hidden_size must be divisible by audio_num_attention_heads, got " f"{self.audio_hidden_size=} {self.audio_num_attention_heads=}." ) if tp_size == 1: return if self.num_attention_heads % tp_size != 0: raise ValueError( "num_attention_heads must be divisible by tp_size, got " f"{self.num_attention_heads=} {tp_size=}." ) if self.audio_num_attention_heads % tp_size != 0: raise ValueError( "audio_num_attention_heads must be divisible by tp_size, got " f"{self.audio_num_attention_heads=} {tp_size=}." ) if self.hidden_size % tp_size != 0: raise ValueError( "hidden_size must be divisible by tp_size for TP-sharded projections, got " f"{self.hidden_size=} {tp_size=}." ) if self.audio_hidden_size % tp_size != 0: raise ValueError( "audio_hidden_size must be divisible by tp_size for TP-sharded projections, got " f"{self.audio_hidden_size=} {tp_size=}." ) if int(arch.out_channels) % tp_size != 0: raise ValueError( "out_channels must be divisible by tp_size for TP-sharded output projection, got " f"{arch.out_channels=} {tp_size=}." ) if int(arch.audio_out_channels) % tp_size != 0: raise ValueError( "audio_out_channels must be divisible by tp_size for TP-sharded output projection, got " f"{arch.audio_out_channels=} {tp_size=}." ) def __init__( self, config: LTX2Config, hf_config: dict[str, Any], quant_config: QuantizationConfig | None = None, ) -> None: super().__init__(config=config, hf_config=hf_config) arch = config.arch_config self.hidden_size = arch.hidden_size self.num_attention_heads = arch.num_attention_heads self.audio_hidden_size = arch.audio_hidden_size self.audio_num_attention_heads = arch.audio_num_attention_heads self.norm_eps = arch.norm_eps tp_size = get_tp_world_size() self._validate_tp_config(arch=arch, tp_size=tp_size) # 1. Patchification input projections # Matches LTX2Config().param_names_mapping self.patchify_proj = ColumnParallelLinear( arch.in_channels, self.hidden_size, bias=True, gather_output=True, quant_config=quant_config, ) self.audio_patchify_proj = ColumnParallelLinear( arch.audio_in_channels, self.audio_hidden_size, bias=True, gather_output=True, quant_config=quant_config, ) # 2. Prompt embeddings self.caption_projection: LTX2TextProjection | None = None self.audio_caption_projection: LTX2TextProjection | None = None if not arch.caption_proj_before_connector: self.caption_projection = LTX2TextProjection( in_features=arch.caption_channels, hidden_size=self.hidden_size ) self.audio_caption_projection = LTX2TextProjection( in_features=arch.caption_channels, hidden_size=self.audio_hidden_size ) # 3. Timestep Modulation Params and Embedding self.adaln_single = LTX2AdaLayerNormSingle( self.hidden_size, embedding_coefficient=adaln_embedding_coefficient( arch.cross_attention_adaln ), ) self.audio_adaln_single = LTX2AdaLayerNormSingle( self.audio_hidden_size, embedding_coefficient=adaln_embedding_coefficient( arch.cross_attention_adaln ), ) self.prompt_adaln_single: LTX2AdaLayerNormSingle | None = None self.audio_prompt_adaln_single: LTX2AdaLayerNormSingle | None = None if arch.cross_attention_adaln: self.prompt_adaln_single = LTX2AdaLayerNormSingle( self.hidden_size, embedding_coefficient=2 ) self.audio_prompt_adaln_single = LTX2AdaLayerNormSingle( self.audio_hidden_size, embedding_coefficient=2 ) # Global Cross Attention Modulation Parameters self.av_ca_video_scale_shift_adaln_single = LTX2AdaLayerNormSingle( self.hidden_size, embedding_coefficient=4 ) self.av_ca_a2v_gate_adaln_single = LTX2AdaLayerNormSingle( self.hidden_size, embedding_coefficient=1 ) self.av_ca_audio_scale_shift_adaln_single = LTX2AdaLayerNormSingle( self.audio_hidden_size, embedding_coefficient=4 ) self.av_ca_v2a_gate_adaln_single = LTX2AdaLayerNormSingle( self.audio_hidden_size, embedding_coefficient=1 ) # Output Layer Scale/Shift Modulation parameters self.scale_shift_table = nn.Parameter( torch.randn(2, self.hidden_size) / self.hidden_size**0.5 ) self.audio_scale_shift_table = nn.Parameter( torch.randn(2, self.audio_hidden_size) / self.audio_hidden_size**0.5 ) hf_patch_size = int(hf_config.get("patch_size", 1)) hf_patch_size_t = int(hf_config.get("patch_size_t", 1)) self.patch_size = (hf_patch_size_t, hf_patch_size, hf_patch_size) hf_audio_patch_size = int(hf_config.get("audio_patch_size", 1)) hf_audio_patch_size_t = int(hf_config.get("audio_patch_size_t", 1)) rope_type = ( arch.rope_type.value if hasattr(arch.rope_type, "value") else str(arch.rope_type) ) frequencies_precision = hf_config.get("frequencies_precision") if frequencies_precision is None: frequencies_precision = getattr(arch, "frequencies_precision", None) # diffusers/LTX configs use `frequencies_precision` for this RoPE switch rope_double_precision = ( str(frequencies_precision) == "float64" if frequencies_precision is not None else bool( hf_config.get("rope_double_precision", arch.double_precision_rope) ) ) self.quantize_video_rope_coords_to_hidden_dtype = bool( hf_config.get("quantize_video_rope_coords_to_hidden_dtype", False) ) causal_offset = int(hf_config.get("causal_offset", 1)) pos_embed_max_pos = int(arch.positional_embedding_max_pos[0]) base_height = int(arch.positional_embedding_max_pos[1]) base_width = int(arch.positional_embedding_max_pos[2]) audio_pos_embed_max_pos = int(arch.audio_positional_embedding_max_pos[0]) self.video_scale_factors = (8, 32, 32) self.audio_scale_factors = (4,) self.rope = LTX2AudioVideoRotaryPosEmbed( dim=self.hidden_size, patch_size=hf_patch_size, patch_size_t=hf_patch_size_t, base_num_frames=pos_embed_max_pos, base_height=base_height, base_width=base_width, scale_factors=self.video_scale_factors, theta=float(arch.positional_embedding_theta), causal_offset=causal_offset, modality="video", double_precision=rope_double_precision, rope_type=rope_type, num_attention_heads=self.num_attention_heads, ) self.audio_rope = LTX2AudioVideoRotaryPosEmbed( dim=self.audio_hidden_size, patch_size=hf_audio_patch_size, patch_size_t=hf_audio_patch_size_t, base_num_frames=audio_pos_embed_max_pos, sampling_rate=16000, hop_length=160, scale_factors=self.audio_scale_factors, theta=float(arch.positional_embedding_theta), causal_offset=causal_offset, modality="audio", double_precision=rope_double_precision, rope_type=rope_type, num_attention_heads=self.audio_num_attention_heads, ) cross_attn_pos_embed_max_pos = max(pos_embed_max_pos, audio_pos_embed_max_pos) self.cross_attn_rope = LTX2AudioVideoRotaryPosEmbed( dim=int(arch.audio_cross_attention_dim), patch_size=hf_patch_size, patch_size_t=hf_patch_size_t, base_num_frames=cross_attn_pos_embed_max_pos, base_height=base_height, base_width=base_width, theta=float(arch.positional_embedding_theta), causal_offset=causal_offset, modality="video", double_precision=rope_double_precision, rope_type=rope_type, num_attention_heads=self.num_attention_heads, ) self.cross_attn_audio_rope = LTX2AudioVideoRotaryPosEmbed( dim=int(arch.audio_cross_attention_dim), patch_size=hf_audio_patch_size, patch_size_t=hf_audio_patch_size_t, base_num_frames=cross_attn_pos_embed_max_pos, sampling_rate=16000, hop_length=160, scale_factors=self.audio_scale_factors, theta=float(arch.positional_embedding_theta), causal_offset=causal_offset, modality="audio", double_precision=rope_double_precision, rope_type=rope_type, num_attention_heads=self.audio_num_attention_heads, ) self.cross_pe_max_pos = cross_attn_pos_embed_max_pos # 5. Transformer Blocks self.transformer_blocks = nn.ModuleList( [ LTX2TransformerBlock( idx=idx, dim=self.hidden_size, num_attention_heads=self.num_attention_heads, attention_head_dim=self.hidden_size // self.num_attention_heads, cross_attention_dim=arch.cross_attention_dim, audio_dim=self.audio_hidden_size, audio_num_attention_heads=self.audio_num_attention_heads, audio_attention_head_dim=self.audio_hidden_size // self.audio_num_attention_heads, audio_cross_attention_dim=arch.audio_cross_attention_dim, norm_eps=self.norm_eps, qk_norm=True, # Always True in LTX2 apply_gated_attention=arch.apply_gated_attention, cross_attention_adaln=arch.cross_attention_adaln, use_local_av_cross_attention=bool( getattr(arch, "use_local_av_cross_attention", False) ), force_sdpa_v2a_cross_attention=bool( getattr(arch, "force_sdpa_v2a_cross_attention", False) ), enable_packed_qkv_input_a2a=arch.enable_packed_qkv_input_a2a, supported_attention_backends=self._supported_attention_backends, prefix=config.prefix, quant_config=quant_config, ) for idx in range(arch.num_layers) ] ) # 6. Output layers self.norm_out = nn.LayerNorm( self.hidden_size, eps=self.norm_eps, elementwise_affine=False ) self.proj_out = ColumnParallelLinear( self.hidden_size, arch.out_channels, bias=True, gather_output=True, quant_config=quant_config, ) self.audio_norm_out = nn.LayerNorm( self.audio_hidden_size, eps=self.norm_eps, elementwise_affine=False ) self.audio_proj_out = ColumnParallelLinear( self.audio_hidden_size, arch.audio_out_channels, bias=True, gather_output=True, quant_config=quant_config, ) self.out_channels_raw = arch.out_channels // ( self.patch_size[0] * self.patch_size[1] * self.patch_size[2] ) self.audio_out_channels = arch.audio_out_channels self.timestep_scale_multiplier = arch.timestep_scale_multiplier self.av_ca_timestep_scale_multiplier = arch.av_ca_timestep_scale_multiplier self.layer_names = ["transformer_blocks"] def _maybe_quantize_video_rope_coords( self, video_coords: torch.Tensor, hidden_device: torch.device, hidden_dtype: torch.dtype, ) -> torch.Tensor: if self.quantize_video_rope_coords_to_hidden_dtype: return video_coords.to(device=hidden_device, dtype=hidden_dtype) return video_coords.to(device=hidden_device) def _get_av_ca_gate_timestep_factor(self) -> float: ltx_variant = str(getattr(self.config.arch_config, "ltx_variant", "ltx_2")) if ltx_variant == "ltx_2_3": return self.av_ca_timestep_scale_multiplier / self.timestep_scale_multiplier return float(self.av_ca_timestep_scale_multiplier) def _get_av_ca_timesteps( self, timestep: torch.Tensor, audio_timestep: torch.Tensor, prompt_timestep: torch.Tensor | None, audio_prompt_timestep: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: ltx_variant = str(getattr(self.config.arch_config, "ltx_variant", "ltx_2")) if ltx_variant != "ltx_2_3": return timestep, audio_timestep video_timestep = ( self._collapse_prompt_timestep(timestep) if prompt_timestep is None else prompt_timestep ) audio_timestep_for_ca = ( self._collapse_prompt_timestep(audio_timestep) if audio_prompt_timestep is None else audio_prompt_timestep ) return video_timestep, audio_timestep_for_ca 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: Optional[torch.LongTensor] = None, prompt_timestep: Optional[torch.Tensor] = None, audio_prompt_timestep: Optional[torch.Tensor] = None, encoder_attention_mask: Optional[torch.Tensor] = None, audio_encoder_attention_mask: Optional[torch.Tensor] = None, num_frames: Optional[int] = None, height: Optional[int] = None, width: Optional[int] = None, fps: float = 24.0, audio_num_frames: Optional[int] = None, video_coords: Optional[torch.Tensor] = None, audio_coords: Optional[torch.Tensor] = None, video_self_attention_mask: Optional[torch.Tensor] = None, audio_self_attention_mask: Optional[torch.Tensor] = None, a2v_cross_attention_mask: Optional[torch.Tensor] = None, v2a_cross_attention_mask: Optional[torch.Tensor] = None, skip_video_self_attn_blocks: Optional[tuple[int, ...]] = None, skip_audio_self_attn_blocks: Optional[tuple[int, ...]] = None, disable_a2v_cross_attn: bool = False, disable_v2a_cross_attn: bool = False, audio_replicated_for_sp: bool = False, video_memory_prefix_len: int = 0, late_layer_ratio: float = 1.0, late_audio_self_attention_mask: Optional[torch.Tensor] = None, **kwargs, ) -> tuple[torch.Tensor | None, torch.Tensor | None]: batch_size = hidden_states.size(0) audio_timestep = audio_timestep if audio_timestep is not None else timestep if num_frames is None or height is None or width is None: raise ValueError( "num_frames/height/width must be provided for RoPE coordinate generation." ) if audio_num_frames is None: raise ValueError( "audio_num_frames must be provided for RoPE coordinate generation." ) perturbation_configs = kwargs.get("perturbation_configs") if perturbation_configs is not None and len(perturbation_configs) != batch_size: raise ValueError( "perturbation_configs length must match batch size, got " f"{len(perturbation_configs)=} {batch_size=}." ) if video_coords is None: # Wan-style SP-RoPE: when SP is enabled, each rank runs on its local # time shard but RoPE positions must be offset to global time. # # We assume equal time sharding across SP ranks. if model_parallel_is_initialized(): sp_world_size = get_sp_world_size() sp_rank = get_sp_parallel_rank() else: sp_world_size = 1 sp_rank = 0 video_shift = int(sp_rank) * int(num_frames) if sp_world_size > 1 else 0 video_coords = self.rope.prepare_video_coords( batch_size=batch_size, num_frames=num_frames, height=height, width=width, device=hidden_states.device, fps=fps, start_frame=video_shift, ) if audio_coords is None: audio_coords = self.audio_rope.prepare_audio_coords( batch_size=batch_size, num_frames=audio_num_frames, device=audio_hidden_states.device, ) video_coords = self._maybe_quantize_video_rope_coords( video_coords, hidden_states.device, hidden_states.dtype ) audio_coords = audio_coords.to(device=audio_hidden_states.device) video_rotary_emb = self.rope( video_coords, device=hidden_states.device, out_dtype=hidden_states.dtype, ) audio_rotary_emb = self.audio_rope( audio_coords, device=audio_hidden_states.device, out_dtype=audio_hidden_states.dtype, ) ca_video_rotary_emb = self.cross_attn_rope( video_coords[:, 0:1, :], device=hidden_states.device, out_dtype=hidden_states.dtype, ) ca_audio_rotary_emb = self.cross_attn_audio_rope( audio_coords[:, 0:1, :], device=audio_hidden_states.device, out_dtype=audio_hidden_states.dtype, ) # 2. Patchify input projections hidden_states, _ = self.patchify_proj(hidden_states) audio_hidden_states, _ = self.audio_patchify_proj(audio_hidden_states) # 3. Prepare timestep embeddings # 3.1. Prepare global modality (video and audio) timestep embedding and modulation parameters timestep_for_adaln = self._scale_timestep_for_adaln(timestep) audio_timestep_for_adaln = self._scale_timestep_for_adaln(audio_timestep) temb, embedded_timestep = self.adaln_single( timestep_for_adaln.flatten(), 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_adaln_single( audio_timestep_for_adaln.flatten(), 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) ) temb_prompt = None temb_audio_prompt = None if self.prompt_adaln_single is not None: prompt_timestep = ( self._collapse_prompt_timestep(timestep) if prompt_timestep is None else prompt_timestep ) prompt_timestep_for_adaln = self._scale_timestep_for_adaln(prompt_timestep) temb_prompt, _ = self.prompt_adaln_single( prompt_timestep_for_adaln.flatten(), hidden_dtype=hidden_states.dtype ) temb_prompt = temb_prompt.view(batch_size, -1, temb_prompt.size(-1)) if self.audio_prompt_adaln_single is not None: audio_prompt_timestep = ( self._collapse_prompt_timestep(audio_timestep) if audio_prompt_timestep is None else audio_prompt_timestep ) audio_prompt_timestep_for_adaln = self._scale_timestep_for_adaln( audio_prompt_timestep ) temb_audio_prompt, _ = self.audio_prompt_adaln_single( audio_prompt_timestep_for_adaln.flatten(), hidden_dtype=audio_hidden_states.dtype, ) temb_audio_prompt = temb_audio_prompt.view( batch_size, -1, temb_audio_prompt.size(-1) ) # 3.2. Prepare global modality cross attention modulation parameters hidden_dtype = hidden_states.dtype av_ca_video_timestep, av_ca_audio_timestep = self._get_av_ca_timesteps( timestep, audio_timestep, prompt_timestep, audio_prompt_timestep, ) av_ca_video_timestep_for_adaln = self._scale_timestep_for_adaln( av_ca_video_timestep ) av_ca_audio_timestep_for_adaln = self._scale_timestep_for_adaln( av_ca_audio_timestep ) temb_ca_scale_shift, _ = self.av_ca_video_scale_shift_adaln_single( av_ca_video_timestep_for_adaln.flatten(), hidden_dtype=hidden_dtype ) temb_ca_scale_shift = temb_ca_scale_shift.view( batch_size, -1, temb_ca_scale_shift.shape[-1] ) av_ca_gate_factor = self._get_av_ca_gate_timestep_factor() temb_ca_gate, _ = self.av_ca_a2v_gate_adaln_single( av_ca_video_timestep_for_adaln.flatten() * av_ca_gate_factor, hidden_dtype=hidden_dtype, ) temb_ca_gate = temb_ca_gate.view(batch_size, -1, temb_ca_gate.shape[-1]) temb_ca_audio_scale_shift, _ = self.av_ca_audio_scale_shift_adaln_single( av_ca_audio_timestep_for_adaln.flatten(), hidden_dtype=audio_hidden_states.dtype, ) temb_ca_audio_scale_shift = temb_ca_audio_scale_shift.view( batch_size, -1, temb_ca_audio_scale_shift.shape[-1] ) temb_ca_audio_gate, _ = self.av_ca_v2a_gate_adaln_single( av_ca_audio_timestep_for_adaln.flatten() * av_ca_gate_factor, hidden_dtype=audio_hidden_states.dtype, ) temb_ca_audio_gate = temb_ca_audio_gate.view( batch_size, -1, temb_ca_audio_gate.shape[-1] ) # 4. Prepare prompt embeddings if self.caption_projection is not None: encoder_hidden_states = self.caption_projection(encoder_hidden_states) if self.audio_caption_projection is not None: audio_encoder_hidden_states = self.audio_caption_projection( audio_encoder_hidden_states ) if _is_npu: # If the 'encoder_attention_mask' is provided and it is all ones, # it can be set to 'None' to avoid the degradation of performance on the NPU side, # where the mask, even though it has no affect, # can lead to the introduction of multiple small operators. if encoder_attention_mask is not None and torch.all( encoder_attention_mask == 1 ): encoder_attention_mask = None # 5. Run blocks skip_video_self_attn_blocks = set(skip_video_self_attn_blocks or ()) skip_audio_self_attn_blocks = set(skip_audio_self_attn_blocks or ()) video_self_attn_perturbation_states = None audio_self_attn_perturbation_states = None a2v_cross_attn_perturbation_states = None v2a_cross_attn_perturbation_states = None if perturbation_configs is not None: block_indices = tuple( getattr(block, "idx", -1) for block in self.transformer_blocks ) video_self_attn_perturbation_states = ( _ltx2_build_batched_perturbation_states( perturbation_configs, "skip_video_self_attn_blocks", block_indices, hidden_states, ) ) audio_self_attn_perturbation_states = ( _ltx2_build_batched_perturbation_states( perturbation_configs, "skip_audio_self_attn_blocks", block_indices, audio_hidden_states, ) ) a2v_cross_attn_perturbation_states = ( _ltx2_build_batched_perturbation_states( perturbation_configs, "skip_a2v_cross_attn", block_indices, hidden_states, ) ) v2a_cross_attn_perturbation_states = ( _ltx2_build_batched_perturbation_states( perturbation_configs, "skip_v2a_cross_attn", block_indices, audio_hidden_states, ) ) late_layer_start = int(len(self.transformer_blocks) * float(late_layer_ratio)) for block in self.transformer_blocks: block_idx = getattr(block, "idx", -1) video_self_attn_perturbation_mask = None audio_self_attn_perturbation_mask = None a2v_cross_attn_perturbation_mask = None v2a_cross_attn_perturbation_mask = None skip_video_self_attn = block_idx in skip_video_self_attn_blocks skip_audio_self_attn = block_idx in skip_audio_self_attn_blocks skip_a2v_cross_attn = disable_a2v_cross_attn skip_v2a_cross_attn = disable_v2a_cross_attn block_audio_self_attention_mask = audio_self_attention_mask if ( block_idx >= late_layer_start and late_audio_self_attention_mask is not None ): block_audio_self_attention_mask = late_audio_self_attention_mask elif block_idx >= late_layer_start and late_layer_ratio < 1.0: block_audio_self_attention_mask = None if perturbation_configs is not None: if not skip_video_self_attn: assert video_self_attn_perturbation_states is not None state = video_self_attn_perturbation_states[block_idx] video_self_attn_perturbation_mask, skip_video_self_attn = state if not skip_audio_self_attn: assert audio_self_attn_perturbation_states is not None state = audio_self_attn_perturbation_states[block_idx] audio_self_attn_perturbation_mask, skip_audio_self_attn = state if not skip_a2v_cross_attn: assert a2v_cross_attn_perturbation_states is not None state = a2v_cross_attn_perturbation_states[block_idx] a2v_cross_attn_perturbation_mask, skip_a2v_cross_attn = state if not skip_v2a_cross_attn: assert v2a_cross_attn_perturbation_states is not None state = v2a_cross_attn_perturbation_states[block_idx] v2a_cross_attn_perturbation_mask, skip_v2a_cross_attn = state hidden_states, audio_hidden_states = block( hidden_states, audio_hidden_states, encoder_hidden_states, audio_encoder_hidden_states, # Keep the first 4 args positional to stay compatible with cache-dit's # LTX2 adapter, which treats `audio_hidden_states` as `encoder_hidden_states` # under ForwardPattern.Pattern_0. temb=temb, temb_audio=temb_audio, temb_prompt=temb_prompt, temb_audio_prompt=temb_audio_prompt, temb_ca_scale_shift=temb_ca_scale_shift, temb_ca_audio_scale_shift=temb_ca_audio_scale_shift, temb_ca_gate=temb_ca_gate, temb_ca_audio_gate=temb_ca_audio_gate, video_rotary_emb=video_rotary_emb, audio_rotary_emb=audio_rotary_emb, ca_video_rotary_emb=ca_video_rotary_emb, ca_audio_rotary_emb=ca_audio_rotary_emb, encoder_attention_mask=encoder_attention_mask, audio_encoder_attention_mask=audio_encoder_attention_mask, video_self_attention_mask=video_self_attention_mask, audio_self_attention_mask=block_audio_self_attention_mask, a2v_cross_attention_mask=a2v_cross_attention_mask, v2a_cross_attention_mask=v2a_cross_attention_mask, skip_video_self_attn=skip_video_self_attn, skip_audio_self_attn=skip_audio_self_attn, skip_a2v_cross_attn=skip_a2v_cross_attn, skip_v2a_cross_attn=skip_v2a_cross_attn, video_self_attn_perturbation_mask=video_self_attn_perturbation_mask, audio_self_attn_perturbation_mask=audio_self_attn_perturbation_mask, a2v_cross_attn_perturbation_mask=a2v_cross_attn_perturbation_mask, v2a_cross_attn_perturbation_mask=v2a_cross_attn_perturbation_mask, audio_replicated_for_sp=audio_replicated_for_sp, video_memory_prefix_len=video_memory_prefix_len, ) # 6. Output layers # Video scale_shift_values = self.scale_shift_table[None, None].to( device=hidden_states.device, dtype=hidden_states.dtype ) + embedded_timestep[:, :, None].to(dtype=hidden_states.dtype) shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] with torch.autocast(device_type=hidden_states.device.type, enabled=False): hidden_states = self.norm_out(hidden_states) hidden_states = hidden_states * (1 + scale) + shift hidden_states, _ = self.proj_out(hidden_states) # Audio audio_scale_shift_values = self.audio_scale_shift_table[None, None].to( device=audio_hidden_states.device, dtype=audio_hidden_states.dtype ) + audio_embedded_timestep[:, :, None].to(dtype=audio_hidden_states.dtype) audio_shift, audio_scale = ( audio_scale_shift_values[:, :, 0], audio_scale_shift_values[:, :, 1], ) with torch.autocast(device_type=audio_hidden_states.device.type, enabled=False): audio_hidden_states = self.audio_norm_out(audio_hidden_states) audio_hidden_states = audio_hidden_states * (1 + audio_scale) + audio_shift audio_hidden_states, _ = self.audio_proj_out(audio_hidden_states) # Unpatchify if requested (default True for pipeline compatibility) return_latents = kwargs.get("return_latents", True) if return_latents: # Unpatchify Video # [B, N, C_out_raw*patch_vol] -> [B, C_out_raw, T, H, W] # Requires num_frames, height, width to be known if num_frames is not None and height is not None and width is not None: p_t, p_h, p_w = self.patch_size post_t, post_h, post_w = num_frames // p_t, height // p_h, width // p_w b = batch_size hidden_states = hidden_states.reshape( b, post_t, post_h, post_w, self.out_channels_raw, p_t, p_h, p_w ) hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape( b, self.out_channels_raw, num_frames, height, width ) # Unpatchify Audio # [B, N, C_out] -> [B, C_out, T] (or 4D/5D) if audio_num_frames is not None: b = batch_size # simple reshape for 1D patch audio_hidden_states = audio_hidden_states.permute(0, 2, 1) # [B, C, T] return hidden_states, audio_hidden_states # Backward-compatible alias (older internal name). LTXModel = LTX2VideoTransformer3DModel EntryClass = LTX2VideoTransformer3DModel