Spaces:
Sleeping
Sleeping
| # Copyright 2024-2025 LTX-2 Refactored (WAN-style) | |
| """ | |
| LTX-2 Diffusion Model following WAN style with Audio-Video support. | |
| All core components in one file as per WAN convention. | |
| """ | |
| import math | |
| import os | |
| from dataclasses import dataclass | |
| from enum import Enum | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from diffusers.configuration_utils import ConfigMixin, register_to_config | |
| from diffusers.models.modeling_utils import ModelMixin | |
| from typing import Optional, Tuple | |
| from ltx2.modules.attention import AttentionFunction, AttentionCallable | |
| from ltx2.modules.perturbations import BatchedPerturbationConfig, PerturbationType | |
| from ltx2.modules.timestep_embedding import get_timestep_embedding, TimestepEmbedding | |
| from ltx2.modules.rope import apply_rotary_emb, LTXRopeType, precompute_freqs_cis,generate_freq_grid_np,generate_freq_grid_pytorch | |
| from fastvideo.utils.context_parallel import ( | |
| is_cp_enabled, get_cp_world_size, | |
| scatter_sequence, gather_sequence, gather_for_loss, | |
| apply_ulysses_attention, | |
| pad_to_cp_divisible, unpad_from_cp, | |
| ) | |
| __all__ = [ | |
| 'LTX2Model', 'LTX2ModelType', 'VideoLatentShape', 'X0Model', 'to_denoised', 'Modality', | |
| # Re-exports for backward compatibility (used by text_encoder.py) | |
| 'Attention', 'FeedForward', 'LTXRopeType', 'precompute_freqs_cis', 'rms_norm', | |
| ] | |
| def rms_norm(x: torch.Tensor, weight: torch.Tensor | None = None, eps: float = 1e-6) -> torch.Tensor: | |
| """Root-mean-square (RMS) normalize `x` over its last dimension. | |
| Thin wrapper around `torch.nn.functional.rms_norm` that infers the normalized | |
| shape and forwards `weight` and `eps`. | |
| """ | |
| return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight=weight, eps=eps) | |
| class Modality: | |
| latent: ( | |
| torch.Tensor | |
| ) # Shape: (B, T, D) where B is the batch size, T is the number of tokens, and D is input dimension | |
| timesteps: torch.Tensor # Shape: (B, T) where T is the number of timesteps | |
| positions: ( | |
| torch.Tensor | |
| ) # Shape: (B, 3, T) for video, where 3 is the number of dimensions and T is the number of tokens | |
| context: torch.Tensor | |
| enabled: bool = True | |
| context_mask: torch.Tensor | None = None | |
| class TransformerConfig: | |
| dim: int | |
| heads: int | |
| d_head: int | |
| context_dim: int | |
| class PixArtAlphaTextProjection(torch.nn.Module): | |
| """ | |
| Projects caption embeddings. Also handles dropout for classifier-free guidance. | |
| Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py | |
| """ | |
| def __init__(self, in_features: int, hidden_size: int, out_features: int | None = None, act_fn: str = "gelu_tanh"): | |
| super().__init__() | |
| if out_features is None: | |
| out_features = hidden_size | |
| self.linear_1 = torch.nn.Linear(in_features=in_features, out_features=hidden_size, bias=True) | |
| if act_fn == "gelu_tanh": | |
| self.act_1 = torch.nn.GELU(approximate="tanh") | |
| elif act_fn == "silu": | |
| self.act_1 = torch.nn.SiLU() | |
| else: | |
| raise ValueError(f"Unknown activation function: {act_fn}") | |
| self.linear_2 = torch.nn.Linear(in_features=hidden_size, out_features=out_features, bias=True) | |
| def forward(self, caption: torch.Tensor) -> torch.Tensor: | |
| hidden_states = self.linear_1(caption) | |
| hidden_states = self.act_1(hidden_states) | |
| hidden_states = self.linear_2(hidden_states) | |
| return hidden_states | |
| class Timesteps(torch.nn.Module): | |
| def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1): | |
| super().__init__() | |
| self.num_channels = num_channels | |
| self.flip_sin_to_cos = flip_sin_to_cos | |
| self.downscale_freq_shift = downscale_freq_shift | |
| self.scale = scale | |
| def forward(self, timesteps: torch.Tensor) -> torch.Tensor: | |
| t_emb = get_timestep_embedding( | |
| timesteps, | |
| self.num_channels, | |
| flip_sin_to_cos=self.flip_sin_to_cos, | |
| downscale_freq_shift=self.downscale_freq_shift, | |
| scale=self.scale, | |
| ) | |
| return t_emb | |
| class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module): | |
| """ | |
| For PixArt-Alpha. | |
| Reference: | |
| https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L164C9-L168C29 | |
| """ | |
| def __init__( | |
| self, | |
| embedding_dim: int, | |
| size_emb_dim: int, | |
| ): | |
| super().__init__() | |
| self.outdim = size_emb_dim | |
| self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) | |
| self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) | |
| def forward( | |
| self, | |
| timestep: torch.Tensor, | |
| hidden_dtype: torch.dtype, | |
| ) -> torch.Tensor: | |
| timesteps_proj = self.time_proj(timestep) | |
| timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) | |
| return timesteps_emb | |
| class AdaLayerNormSingle(torch.nn.Module): | |
| r""" | |
| Norm layer adaptive layer norm single (adaLN-single). | |
| As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3). | |
| Parameters: | |
| embedding_dim (`int`): The size of each embedding vector. | |
| use_additional_conditions (`bool`): To use additional conditions for normalization or not. | |
| """ | |
| def __init__(self, embedding_dim: int, embedding_coefficient: int = 6): | |
| super().__init__() | |
| self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings( | |
| embedding_dim, | |
| size_emb_dim=embedding_dim // 3, | |
| ) | |
| self.silu = torch.nn.SiLU() | |
| self.linear = torch.nn.Linear(embedding_dim, embedding_coefficient * embedding_dim, bias=True) | |
| def forward( | |
| self, | |
| timestep: torch.Tensor, | |
| hidden_dtype: Optional[torch.dtype] = None, | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype) | |
| return self.linear(self.silu(embedded_timestep)), embedded_timestep | |
| class VideoLatentShape: | |
| """Helper class for video latent shape calculations.""" | |
| def __init__(self, batch_size, num_frames, num_channels, height, width): | |
| self.batch_size = batch_size | |
| self.num_frames = num_frames | |
| self.num_channels = num_channels | |
| self.height = height | |
| self.width = width | |
| def to_tuple(self): | |
| return (self.batch_size, self.num_frames, self.num_channels, self.height, self.width) | |
| class LTX2ModelType(Enum): | |
| """Model type for LTX-2.""" | |
| AudioVideo = "ltx av model" | |
| VideoOnly = "ltx video only model" | |
| AudioOnly = "ltx audio only model" | |
| def is_video_enabled(self) -> bool: | |
| return self in (LTX2ModelType.AudioVideo, LTX2ModelType.VideoOnly) | |
| def is_audio_enabled(self) -> bool: | |
| return self in (LTX2ModelType.AudioVideo, LTX2ModelType.AudioOnly) | |
| def sinusoidal_embedding_1d(dim, position): | |
| """ | |
| Sinusoidal positional embedding for timesteps (WAN-style). | |
| Args: | |
| dim: Embedding dimension (must be even) | |
| position: Position tensor [B] | |
| Returns: | |
| Embeddings [B, dim] | |
| """ | |
| assert dim % 2 == 0 | |
| half = dim // 2 | |
| position = position.type(torch.float64) | |
| sinusoid = torch.outer( | |
| position, torch.pow(10000, -torch.arange(half).to(position).div(half))) | |
| x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) | |
| return x | |
| class LTX2RMSNorm(nn.Module): | |
| """ | |
| RMS Normalization (WAN-style). | |
| """ | |
| def __init__(self, dim, eps=1e-5): | |
| super().__init__() | |
| self.dim = dim | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| def forward(self, x): | |
| """ | |
| Args: | |
| x(Tensor): Shape [B, L, C] | |
| """ | |
| return self._norm(x.float()).type_as(x) * self.weight | |
| def _norm(self, x): | |
| return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) | |
| class LTX2LayerNorm(nn.LayerNorm): | |
| """ | |
| Layer Normalization (WAN-style). | |
| """ | |
| def __init__(self, dim, eps=1e-6, elementwise_affine=False): | |
| super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps) | |
| def forward(self, x): | |
| """ | |
| Args: | |
| x(Tensor): Shape [B, L, C] | |
| """ | |
| return super().forward(x).type_as(x) | |
| class GELUApprox(torch.nn.Module): | |
| def __init__(self, dim_in: int, dim_out: int) -> None: | |
| super().__init__() | |
| self.proj = torch.nn.Linear(dim_in, dim_out) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return torch.nn.functional.gelu(self.proj(x), approximate="tanh") | |
| class FeedForward(torch.nn.Module): | |
| def __init__(self, dim: int, dim_out: int, mult: int = 4) -> None: | |
| super().__init__() | |
| inner_dim = int(dim * mult) | |
| project_in = GELUApprox(dim, inner_dim) | |
| self.net = torch.nn.Sequential(project_in, torch.nn.Identity(), torch.nn.Linear(inner_dim, dim_out)) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return self.net(x) | |
| class Attention(torch.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, | |
| rope_type: LTXRopeType = LTXRopeType.SPLIT, | |
| attention_function: AttentionCallable | AttentionFunction = AttentionFunction.DEFAULT, | |
| ) -> None: | |
| super().__init__() | |
| self.rope_type = rope_type | |
| self.attention_function = attention_function | |
| inner_dim = dim_head * heads | |
| context_dim = query_dim if context_dim is None else context_dim | |
| self.heads = heads | |
| self.dim_head = dim_head | |
| self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) | |
| self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) | |
| self.to_q = torch.nn.Linear(query_dim, inner_dim, bias=True) | |
| self.to_k = torch.nn.Linear(context_dim, inner_dim, bias=True) | |
| self.to_v = torch.nn.Linear(context_dim, inner_dim, bias=True) | |
| self.to_out = torch.nn.Sequential(torch.nn.Linear(inner_dim, query_dim, bias=True), torch.nn.Identity()) | |
| def forward( | |
| self, | |
| x: torch.Tensor, | |
| context: torch.Tensor | None = None, | |
| mask: torch.Tensor | None = None, | |
| pe: torch.Tensor | None = None, | |
| k_pe: torch.Tensor | None = None, | |
| ) -> torch.Tensor: | |
| q = self.to_q(x) | |
| context = x if context is None else context | |
| k = self.to_k(context) | |
| v = self.to_v(context) | |
| q = self.q_norm(q) | |
| k = self.k_norm(k) | |
| if pe is not None: | |
| q = apply_rotary_emb(q, pe, self.rope_type) | |
| k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type) | |
| # attention_function can be an enum *or* a custom callable | |
| out = self.attention_function(q, k, v, self.heads, mask) | |
| return self.to_out(out) | |
| class LTX2SelfAttention(torch.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, | |
| rope_type: LTXRopeType = LTXRopeType.SPLIT, | |
| attention_function: AttentionCallable | AttentionFunction = AttentionFunction.DEFAULT, | |
| # Sparse attention parameters | |
| enable_sparse_attention: bool = False, | |
| sparse_block_size: tuple = (4, 4, 4), | |
| sparse_ratio: float = 0.125, | |
| ) -> None: | |
| super().__init__() | |
| self.rope_type = rope_type | |
| self.attention_function = attention_function | |
| inner_dim = dim_head * heads | |
| context_dim = query_dim if context_dim is None else context_dim | |
| self.heads = heads | |
| self.dim_head = dim_head | |
| self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) | |
| self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps) | |
| self.to_q = torch.nn.Linear(query_dim, inner_dim, bias=True) | |
| self.to_k = torch.nn.Linear(context_dim, inner_dim, bias=True) | |
| self.to_v = torch.nn.Linear(context_dim, inner_dim, bias=True) | |
| self.to_out = torch.nn.Sequential(torch.nn.Linear(inner_dim, query_dim, bias=True), torch.nn.Identity()) | |
| # Sparse attention support | |
| self.enable_sparse_attention = enable_sparse_attention | |
| if enable_sparse_attention: | |
| from ltx2.modules.sparse_attention import BlockSparseAttention | |
| self.sparse_attn = BlockSparseAttention( | |
| block_size=sparse_block_size, | |
| sparsity_ratio=sparse_ratio, | |
| num_heads=heads, | |
| head_dim=dim_head, | |
| ) | |
| def forward( | |
| self, | |
| x: torch.Tensor, | |
| context: torch.Tensor | None = None, | |
| mask: torch.Tensor | None = None, | |
| pe: torch.Tensor | None = None, | |
| k_pe: torch.Tensor | None = None, | |
| video_shape: tuple = None, # (T, H, W) for sparse attention | |
| ) -> torch.Tensor: | |
| is_self_attn = context is None | |
| q = self.to_q(x) | |
| context = x if context is None else context | |
| k = self.to_k(context) | |
| v = self.to_v(context) | |
| q = self.q_norm(q) | |
| k = self.k_norm(k) | |
| if pe is not None: | |
| q = apply_rotary_emb(q, pe, self.rope_type) | |
| k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type) | |
| # Choose attention implementation | |
| if self.enable_sparse_attention and video_shape is not None: | |
| # Use sparse attention | |
| # Reshape to [B, seq_len, num_heads, head_dim] | |
| q = q.view(q.shape[0], q.shape[1], self.heads, self.dim_head) | |
| k = k.view(k.shape[0], k.shape[1], self.heads, self.dim_head) | |
| v = v.view(v.shape[0], v.shape[1], self.heads, self.dim_head) | |
| out = self.sparse_attn(q, k, v, video_shape=video_shape, mask=mask) | |
| # Flatten back: [B, seq_len, num_heads * head_dim] | |
| out = out.reshape(q.shape[0], q.shape[1], -1) | |
| elif is_self_attn and is_cp_enabled(): | |
| # Ulysses-style context parallel for self-attention | |
| out = apply_ulysses_attention(q, k, v, self.heads, self.attention_function, mask) | |
| else: | |
| # Use standard attention (cross-attention or CP disabled) | |
| out = self.attention_function(q, k, v, self.heads, mask) | |
| return self.to_out(out) | |
| class LTX2CrossAttention(LTX2SelfAttention): | |
| """ | |
| Cross attention layer - inherits from LTX2SelfAttention. | |
| Uses the parent class's forward method by simply passing context. | |
| """ | |
| pass # All functionality is inherited from parent class | |
| # Alias for backward compatibility (used by text_encoder.py) | |
| # Attention = LTX2SelfAttention | |
| class LTX2AttentionBlock(torch.nn.Module): | |
| def __init__( | |
| self, | |
| idx: int, | |
| video: TransformerConfig | None = None, | |
| rope_type: LTXRopeType = LTXRopeType.SPLIT, | |
| norm_eps: float = 1e-6, | |
| attention_function: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT, | |
| enable_camera_injection: bool = False, | |
| # Sparse attention parameters | |
| enable_sparse_attention: bool = False, | |
| sparse_block_size: tuple = (4, 4, 4), | |
| sparse_ratio: float = 0.125, | |
| ): | |
| super().__init__() | |
| self.idx = idx | |
| self.enable_camera_injection = enable_camera_injection | |
| self.attn1 = LTX2SelfAttention( | |
| query_dim=video.dim, | |
| heads=video.heads, | |
| dim_head=video.d_head, | |
| context_dim=None, | |
| rope_type=rope_type, | |
| norm_eps=norm_eps, | |
| attention_function=attention_function, | |
| # Pass sparse parameters to self-attention | |
| enable_sparse_attention=enable_sparse_attention, | |
| sparse_block_size=sparse_block_size, | |
| sparse_ratio=sparse_ratio, | |
| ) | |
| # Cross-attention does not use sparse attention | |
| self.attn2 = LTX2CrossAttention( | |
| query_dim=video.dim, | |
| context_dim=video.context_dim, | |
| heads=video.heads, | |
| dim_head=video.d_head, | |
| rope_type=rope_type, | |
| norm_eps=norm_eps, | |
| attention_function=attention_function, | |
| ) | |
| self.ff = FeedForward(video.dim, dim_out=video.dim) | |
| self.scale_shift_table = torch.nn.Parameter(torch.empty(6, video.dim)) | |
| self.norm_eps = norm_eps | |
| # Camera injection layers (scale/shift modulation, zero-initialized) | |
| if enable_camera_injection: | |
| self.cam_injector_layer1 = nn.Linear(video.dim, video.dim) | |
| self.cam_injector_layer2 = nn.Linear(video.dim, video.dim) | |
| self.cam_scale_layer = nn.Linear(video.dim, video.dim) | |
| self.cam_shift_layer = nn.Linear(video.dim, video.dim) | |
| self._init_camera_weights() | |
| def _init_camera_weights(self): | |
| """Initialize camera injection layers with small values for stable training.""" | |
| nn.init.xavier_uniform_(self.cam_injector_layer1.weight) | |
| nn.init.zeros_(self.cam_injector_layer1.bias) | |
| nn.init.xavier_uniform_(self.cam_injector_layer2.weight) | |
| nn.init.zeros_(self.cam_injector_layer2.bias) | |
| nn.init.constant_(self.cam_scale_layer.weight, 1e-6) | |
| nn.init.constant_(self.cam_scale_layer.bias, 0.0) | |
| nn.init.constant_(self.cam_shift_layer.weight, 1e-6) | |
| nn.init.constant_(self.cam_shift_layer.bias, 0.0) | |
| def get_ada_values( | |
| self, scale_shift_table: torch.Tensor, batch_size: int, timestep: torch.Tensor, indices: slice | |
| ) -> tuple[torch.Tensor, ...]: | |
| num_ada_params = scale_shift_table.shape[0] | |
| ada_values = ( | |
| scale_shift_table[indices].unsqueeze(0).unsqueeze(0).to(device=timestep.device, dtype=timestep.dtype) | |
| + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[:, :, indices, :] | |
| ).unbind(dim=2) | |
| return ada_values | |
| def forward( # noqa: PLR0915 | |
| self, | |
| x: torch.Tensor | None, | |
| timesteps: torch.Tensor | None = None, | |
| freqs: tuple[torch.Tensor, torch.Tensor] | None = None, | |
| context: torch.Tensor | None = None, | |
| context_mask: torch.Tensor | None = None, | |
| perturbations: BatchedPerturbationConfig | None = None, | |
| plucker_emb: Optional[torch.Tensor] = None, | |
| video_shape: tuple = None, # (T, H, W) for sparse attention | |
| ) -> torch.Tensor: | |
| batch_size = x.shape[0] if x is not None else 1 | |
| if perturbations is None: | |
| perturbations = BatchedPerturbationConfig.empty(batch_size) | |
| vshift_msa, vscale_msa, vgate_msa = self.get_ada_values( | |
| self.scale_shift_table, x.shape[0], timesteps, slice(0, 3) | |
| ) | |
| if not perturbations.all_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx): | |
| norm_vx = rms_norm(x, eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa | |
| v_mask = perturbations.mask_like(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx, x) | |
| x = x + self.attn1(norm_vx, pe=freqs, video_shape=video_shape) * vgate_msa * v_mask | |
| # Camera signal injection: scale/shift after self-attention, before cross-attention | |
| if self.enable_camera_injection: | |
| if plucker_emb is not None: | |
| cam_hidden = self.cam_injector_layer1(plucker_emb) | |
| cam_hidden = torch.nn.functional.silu(cam_hidden) | |
| cam_hidden = self.cam_injector_layer2(cam_hidden) | |
| cam_hidden = cam_hidden + plucker_emb # residual | |
| cam_scale = self.cam_scale_layer(cam_hidden) | |
| cam_shift = self.cam_shift_layer(cam_hidden) | |
| del cam_hidden # free immediately | |
| x = (1.0 + cam_scale) * x + cam_shift | |
| del cam_scale, cam_shift # free immediately | |
| else: | |
| _dummy = torch.zeros(x.shape[0], 1, x.shape[-1], device=x.device, dtype=x.dtype) | |
| _ = self.cam_injector_layer1(_dummy) | |
| _ = self.cam_injector_layer2(_dummy) | |
| _ = self.cam_scale_layer(_dummy) | |
| _ = self.cam_shift_layer(_dummy) | |
| x = x + self.attn2(rms_norm(x, eps=self.norm_eps), context=context, mask=context_mask) | |
| del vshift_msa, vscale_msa, vgate_msa | |
| vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values( | |
| self.scale_shift_table, x.shape[0], timesteps, slice(3, None) | |
| ) | |
| x_scaled = rms_norm(x, eps=self.norm_eps) * (1 + vscale_mlp) + vshift_mlp | |
| x = x + self.ff(x_scaled) * vgate_mlp | |
| del vshift_mlp, vscale_mlp, vgate_mlp | |
| return x | |
| class LTX2Model(ModelMixin, ConfigMixin): | |
| ignore_for_config = [ | |
| 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size' | |
| ] | |
| _no_split_modules = ['LTX2AttentionBlock'] | |
| _supports_gradient_checkpointing = True | |
| def __init__(self, | |
| # Model type | |
| model_type='video_only', | |
| # Video params from LTX2 | |
| num_attention_heads: int = 32, | |
| attention_head_dim: int = 128, | |
| in_channels: int = 128, | |
| out_channels: int = 128, | |
| num_layers: int = 48, | |
| cross_attention_dim: int = 4096, | |
| norm_eps: float = 1e-06, | |
| attention_type: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT, | |
| caption_channels: int = 3840, | |
| positional_embedding_theta: float = 10000.0, | |
| positional_embedding_max_pos: list[int] | None = None, | |
| timestep_scale_multiplier: int = 1000, | |
| use_middle_indices_grid: bool = True, | |
| rope_type: LTXRopeType = LTXRopeType.SPLIT, | |
| double_precision_rope: bool = False, | |
| normalize_rope_positions: bool = True, # Whether to normalize positions to [0, 1] in RoPE | |
| normalize_time_by_fps: bool = True, # Whether to normalize time positions by FPS (divide by 25) | |
| # Legacy WAN params (kept for compatibility) | |
| patch_size=(1, 1, 1), | |
| text_len=1024, | |
| in_dim=128, | |
| dim=4096, | |
| ffn_dim=16384, | |
| freq_dim=256, | |
| text_dim=3840, | |
| out_dim=128, | |
| num_heads=32, | |
| window_size=(-1, -1), | |
| qk_norm=True, | |
| cross_attn_norm=True, | |
| eps=1e-6, | |
| # Camera control params (optional) | |
| enable_camera_control: bool = False, | |
| camera_injection_mode: str = "scale_shift", # "scale_shift", "additive", or "wan_inject" | |
| camera_split_plucker: bool = False, # split moment/direction branches | |
| num_actions: int = 81, | |
| plucker_hidden_dim: int = 64, | |
| continuous_camera_dropout_prob: float = 0.1, | |
| discrete_camera_dropout_prob: float = 0.1, | |
| enable_continuous_camera: bool = True, | |
| enable_discrete_camera: bool = True, | |
| # Sparse Attention params (optional, default off) | |
| enable_sparse_attention: bool = False, | |
| sparse_block_size: tuple = (4, 4, 4), | |
| sparse_ratio_train: float = 0.125, # 1/8 for training | |
| sparse_ratio_inference: float = 0.0625, # 1/16 for inference | |
| ): | |
| super().__init__() | |
| # Parse model type | |
| attention_type = AttentionFunction.FLASH_ATTENTION_3 | |
| print("forcing flash-attn3") | |
| self.model_type = LTX2ModelType.VideoOnly | |
| self.patch_size = patch_size | |
| self.text_len = text_len | |
| self.in_dim = in_dim | |
| self.dim = dim | |
| self.ffn_dim = ffn_dim | |
| self.freq_dim = freq_dim | |
| self.text_dim = text_dim | |
| self.out_dim = out_dim | |
| self.num_heads = num_heads | |
| self.num_layers = num_layers | |
| self.window_size = window_size | |
| self.qk_norm = qk_norm | |
| self.cross_attn_norm = cross_attn_norm | |
| self.eps = eps | |
| # LTX2 specific params | |
| self.timestep_scale_multiplier = timestep_scale_multiplier | |
| self.use_middle_indices_grid = use_middle_indices_grid | |
| self.positional_embedding_theta = positional_embedding_theta | |
| self.rope_type = rope_type | |
| self.double_precision_rope = double_precision_rope | |
| self.normalize_rope_positions = normalize_rope_positions # Whether to normalize positions to [0, 1] in RoPE | |
| self.normalize_time_by_fps = normalize_time_by_fps # Whether to normalize time positions by FPS | |
| self.yarn_rope = False # YaRN RoPE extrapolation for long videos (off by default) | |
| self.yarn_max_train_seconds = 20.0 # maximum duration the pretrained model supports | |
| if positional_embedding_max_pos is None: | |
| positional_embedding_max_pos = [2048, 2048, 2048] | |
| self.positional_embedding_max_pos = positional_embedding_max_pos | |
| self.num_attention_heads = num_attention_heads | |
| self.inner_dim = num_attention_heads * attention_head_dim | |
| ## VIDEO INPUT COMPONENTS | |
| self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True) | |
| self.adaln_single = AdaLayerNormSingle(self.inner_dim) | |
| self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels,hidden_size=self.inner_dim) | |
| self.scale_shift_table = torch.nn.Parameter(torch.empty(2, self.inner_dim)) | |
| self.norm_out = torch.nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=norm_eps) | |
| self.proj_out = torch.nn.Linear(self.inner_dim, out_channels) | |
| # ===== Sparse Attention configuration ===== | |
| self.enable_sparse_attention = enable_sparse_attention | |
| self.sparse_block_size = sparse_block_size | |
| self.sparse_ratio = sparse_ratio_train # Default to training ratio | |
| self.sparse_ratio_train = sparse_ratio_train | |
| self.sparse_ratio_inference = sparse_ratio_inference | |
| # ===== Camera control modules (optional) ===== | |
| self.enable_camera_control = enable_camera_control | |
| self.camera_injection_mode = camera_injection_mode | |
| self.camera_split_plucker = camera_split_plucker | |
| self.enable_continuous_camera = enable_continuous_camera | |
| self.enable_discrete_camera = enable_discrete_camera | |
| self._camera_use_cross_norm = os.environ.get('LTX_CAMERA_USE_CROSS_NORM', '0') == '1' | |
| self._cross_norm_scale = float(os.environ.get('LTX_CROSS_NORM_SCALE', '0.2')) | |
| if enable_camera_control: | |
| from ltx2.modules.camera_control import ( | |
| WanCameraAdapter, | |
| DiscreteActionEmbedder, CameraControlDropout, DiscreteActionDropout, | |
| ) | |
| print(f"[LTX2-CameraControl] Enabled: mode={camera_injection_mode}, " | |
| f"continuous={enable_continuous_camera}, discrete={enable_discrete_camera}") | |
| if enable_discrete_camera: | |
| self.action_embedder = DiscreteActionEmbedder( | |
| dim=self.inner_dim, freq_dim=freq_dim, num_actions=num_actions) | |
| self.discrete_action_dropout = DiscreteActionDropout(discrete_camera_dropout_prob) | |
| if enable_continuous_camera: | |
| self.camera_dropout = CameraControlDropout(continuous_camera_dropout_prob) | |
| # All modes now use WanCameraAdapter with pixel-resolution Plucker | |
| # (PixelUnshuffle(8) + Conv2d(k=4,s=4) = 32x lossless spatial downsample) | |
| _adapter_kwargs = dict(in_channels=6, out_dim=self.inner_dim, | |
| vae_temporal_stride=8, split_plucker=camera_split_plucker) | |
| if camera_injection_mode == "scale_shift": | |
| self.plucker_processor = WanCameraAdapter(**_adapter_kwargs) | |
| elif camera_injection_mode == "additive": | |
| self.additive_camera_adapter = WanCameraAdapter(**_adapter_kwargs) | |
| elif camera_injection_mode == "wan_inject": | |
| self.wan_camera_adapter = WanCameraAdapter(**_adapter_kwargs) | |
| # ===== Transformer blocks ===== | |
| # Note: No need for preprocessors in WAN-style - all processing happens directly in forward() | |
| video_config = ( | |
| TransformerConfig( | |
| dim=self.inner_dim, | |
| heads=self.num_attention_heads, | |
| d_head=attention_head_dim, | |
| context_dim=cross_attention_dim, | |
| ) | |
| if self.model_type.is_video_enabled() | |
| else None | |
| ) | |
| _enable_block_camera = (enable_camera_control and enable_continuous_camera | |
| and camera_injection_mode == "scale_shift") | |
| self.blocks = torch.nn.ModuleList( | |
| [ | |
| LTX2AttentionBlock( | |
| idx=idx, | |
| video=video_config, | |
| rope_type=self.rope_type, | |
| norm_eps=norm_eps, | |
| attention_function=attention_type, | |
| enable_camera_injection=_enable_block_camera, | |
| # Sparse attention configuration | |
| enable_sparse_attention=enable_sparse_attention, | |
| sparse_block_size=sparse_block_size, | |
| sparse_ratio=self.sparse_ratio, | |
| ) | |
| for idx in range(num_layers) | |
| ] | |
| ) | |
| # ===== RoPE frequencies ===== | |
| # Note: RoPE frequencies are computed dynamically in forward() using LTX2's precompute_freqs_cis | |
| # This allows for flexible video sizes and maintains LTX2's original behavior | |
| # initialize weights (skip if loading from checkpoint to save time) | |
| # self.init_weights() # Will be called manually if needed | |
| self.gradient_checkpointing = False | |
| def _set_gradient_checkpointing(self, module, value=False): | |
| self.gradient_checkpointing = value | |
| def enable_gradient_checkpointing(self) -> None: | |
| """Enable gradient checkpointing for transformer blocks (WAN-style).""" | |
| self.gradient_checkpointing = True | |
| def set_gradient_checkpointing(self, enable: bool) -> None: | |
| """Enable or disable gradient checkpointing for transformer blocks.""" | |
| self.gradient_checkpointing = enable | |
| def set_sparse_ratio(self, ratio: float): | |
| """ | |
| Dynamically adjust sparsity ratio (for switching between training/inference). | |
| Args: | |
| ratio: float - new sparsity ratio (e.g., 0.125 for training, 0.0625 for inference) | |
| """ | |
| self.sparse_ratio = ratio | |
| for block in self.blocks: | |
| if hasattr(block.attn1, 'sparse_attn'): | |
| block.attn1.sparse_attn.sparsity_ratio = ratio | |
| def _process_output( | |
| self, | |
| scale_shift_table: torch.Tensor, | |
| norm_out: torch.nn.LayerNorm, | |
| proj_out: torch.nn.Linear, | |
| x: torch.Tensor, | |
| embedded_timestep: torch.Tensor, | |
| ) -> torch.Tensor: | |
| """Process output for LTX2.""" | |
| # Apply scale-shift modulation | |
| scale_shift_values = ( | |
| scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + embedded_timestep[:, :, None] | |
| ) | |
| shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] | |
| x = norm_out(x) | |
| x = x * (1 + scale) + shift | |
| x = proj_out(x) | |
| return x | |
| def forward( | |
| self, | |
| x, | |
| t, | |
| context, | |
| seq_len, | |
| action_labels: Optional[torch.Tensor] = None, | |
| plucker_coords: Optional[torch.Tensor] = None, | |
| **kwargs, | |
| ): | |
| """ | |
| Forward pass through the diffusion model (WAN-style interface). | |
| Args: | |
| x (List[Tensor]): List of input video tensors, each with shape [C_in, F, H, W] | |
| t (Tensor): Diffusion timesteps tensor of shape [B] | |
| context (List[Tensor]): List of text embeddings each with shape [L, C] | |
| seq_len (int): Maximum sequence length for positional encoding | |
| action_labels (Tensor, optional): Discrete action labels [B] or [B, T] (0~80) | |
| plucker_coords (Tensor, optional): Plucker coordinates [B, 6, F, H, W] | |
| Returns: | |
| Tensor: Denoised video tensors with shape [B, C_out, F, H, W] | |
| """ | |
| return self._forward(x, t, context, seq_len, | |
| action_labels=action_labels, | |
| plucker_coords=plucker_coords, | |
| cond_latent_frames=kwargs.get('cond_latent_frames', 0), | |
| **{k: v for k, v in kwargs.items() if k != 'cond_latent_frames'}) | |
| def _forward( | |
| self, | |
| x, | |
| t, | |
| context, | |
| seq_len, | |
| fps: float = 25.0, | |
| action_labels: Optional[torch.Tensor] = None, | |
| plucker_coords: Optional[torch.Tensor] = None, | |
| cond_latent_frames: int = 0, | |
| **kwargs, | |
| ): | |
| """ | |
| Internal forward pass (WAN-style). | |
| Args: | |
| x: List[[C, F, H, W]] or [B, C, F, H, W] input latents | |
| t: [B] timesteps (0-1000 range) | |
| context: List[[L, C]] or [B, L, C] text embeddings | |
| seq_len: max sequence length | |
| fps: frames per second for position normalization (default 25.0) | |
| action_labels: Optional discrete action labels [B] or [B, T] | |
| plucker_coords: Optional Plucker coordinates [B, 6, F, H, W] | |
| cond_latent_frames: number of i2v/v2v condition latents; their timestep is forced to 0 | |
| """ | |
| from ltx2.modules.patchifier import VideoLatentPatchifier, VideoLatentShape, SpatioTemporalScaleFactors, get_pixel_coords | |
| device = self.patchify_proj.weight.device | |
| dtype = self.patchify_proj.weight.dtype | |
| # x is List[[C, F, H, W]], convert to [B, C, F, H, W] | |
| if isinstance(x, list): | |
| x_stacked = torch.stack([u for u in x]) # [B, C, F, H, W] | |
| else: | |
| x_stacked = x | |
| batch_size, channels, num_frames, height, width = x_stacked.shape | |
| # Video shape for sparse attention (T, H, W) | |
| video_shape = (num_frames, height, width) if self.enable_sparse_attention else None | |
| # Patchify: [B, C, F, H, W] -> [B, num_tokens, C] | |
| patchifier = VideoLatentPatchifier(patch_size=1) | |
| target_shape = VideoLatentShape( | |
| batch=batch_size, channels=channels, frames=num_frames, | |
| height=height, width=width | |
| ) | |
| latent = patchifier.patchify(x_stacked).to(dtype) | |
| num_tokens = latent.shape[1] | |
| # Time embedding: t can be either timestep (0-1000) or sigma (0-1) | |
| # Detect based on range: if max > 1, it's timesteps, else it's sigma | |
| t_float = t.float() | |
| if t_float.max() > 1.0: | |
| sigma_val = t_float / 1000.0 # Convert timestep to sigma | |
| else: | |
| sigma_val = t_float # Already sigma | |
| timesteps = sigma_val.unsqueeze(-1).unsqueeze(-1).expand(-1, num_tokens, 1).clone().to(dtype) | |
| if cond_latent_frames > 0: | |
| tokens_per_frame = height * width | |
| cond_tokens = cond_latent_frames * tokens_per_frame | |
| cond_tokens = min(cond_tokens, num_tokens) # safety clamp | |
| timesteps[:, :cond_tokens, :] = 0.0 # condition frames get sigma=0 | |
| if not hasattr(self, '_ts_print_count'): | |
| self._ts_print_count = 0 | |
| if self._ts_print_count < 2: | |
| sigma_gen = timesteps[0, -1, 0].item() if num_tokens > 0 else 0 | |
| sigma_cond = timesteps[0, 0, 0].item() if num_tokens > 0 else 0 | |
| print(f"[Model.forward] sigma_cond={sigma_cond:.4f}, sigma_gen={sigma_gen:.4f}, " | |
| f"cond_frames={cond_latent_frames}, x={x_stacked.shape}, cam={getattr(self, 'enable_camera_control', False)}") | |
| self._ts_print_count += 1 | |
| # Positions with FPS normalization (critical for temporal consistency) | |
| scale_factors = SpatioTemporalScaleFactors(time=8, width=32, height=32) | |
| latent_coords = patchifier.get_patch_grid_bounds(output_shape=target_shape, device=device) | |
| positions = get_pixel_coords( | |
| latent_coords=latent_coords, | |
| scale_factors=scale_factors, | |
| causal_fix=True | |
| ).float() | |
| # Normalize time dimension by FPS (important for correct motion) | |
| if getattr(self, 'normalize_time_by_fps', True): | |
| positions[:, 0, ...] = positions[:, 0, ...] / fps | |
| positions = positions.to(dtype) | |
| # Context: stack if list | |
| if isinstance(context, list): | |
| # Pad to max length and stack | |
| max_len = max(c.shape[0] for c in context) | |
| context_padded = [] | |
| for c in context: | |
| if c.shape[0] < max_len: | |
| pad = torch.zeros(max_len - c.shape[0], c.shape[1], device=c.device, dtype=c.dtype) | |
| c = torch.cat([c, pad], dim=0) | |
| context_padded.append(c) | |
| context_tensor = torch.stack(context_padded).to(dtype) # [B, L, C] | |
| else: | |
| context_tensor = context.to(dtype) | |
| # Patch embedding | |
| hidden = self.patchify_proj(latent) | |
| tokens_per_frame = height * width | |
| cond_tokens = cond_latent_frames * tokens_per_frame if cond_latent_frames > 0 else 0 | |
| if cond_latent_frames > 0 and cond_tokens < num_tokens: | |
| sigma_gen = timesteps[0, -1, 0:1] # [1] sigma of the generated frames | |
| unique_sigmas = torch.stack([torch.zeros_like(sigma_gen), sigma_gen]) # [2, 1] | |
| unique_scaled = unique_sigmas * self.timestep_scale_multiplier | |
| unique_emb, unique_embedded = self.adaln_single( | |
| unique_scaled.squeeze(-1), hidden_dtype=hidden.dtype, | |
| ) # unique_emb: [2, 6*dim], unique_embedded: [2, dim] | |
| token_indices = torch.ones(num_tokens, dtype=torch.long, device=device) | |
| token_indices[:cond_tokens] = 0 | |
| timestep_emb = unique_emb[token_indices].unsqueeze(0) # [1, num_tokens, 6*dim] | |
| embedded_timestep = unique_embedded[token_indices].unsqueeze(0) # [1, num_tokens, dim] | |
| else: | |
| single_sigma = timesteps[0, 0, 0:1] # [1] | |
| single_scaled = single_sigma * self.timestep_scale_multiplier | |
| single_emb, single_embedded = self.adaln_single( | |
| single_scaled, hidden_dtype=hidden.dtype, | |
| ) # single_emb: [1, 6*dim], single_embedded: [1, dim] | |
| timestep_emb = single_emb.unsqueeze(0).expand(batch_size, num_tokens, -1) # [B, num_tokens, 6*dim] | |
| embedded_timestep = single_embedded.unsqueeze(0).expand(batch_size, num_tokens, -1) # [B, num_tokens, dim] | |
| # ===== Discrete action injection (add to timestep_emb via adaln_single) ===== | |
| # FSDP requires ALL ranks to access the same parameters in forward pass. | |
| # Must always access action_embedder params even when action_labels is None. | |
| if (getattr(self, 'enable_camera_control', False) and | |
| getattr(self, 'enable_discrete_camera', False)): | |
| if action_labels is not None: | |
| action_emb = self.action_embedder(action_labels) # [B, dim] or [B, T, dim] | |
| if action_emb.dim() == 3: | |
| action_emb = action_emb.repeat_interleave(tokens_per_frame, dim=1) # [B, T*H*W, dim] | |
| if action_emb.shape[1] != num_tokens: | |
| action_emb = action_emb.permute(0, 2, 1) | |
| action_emb = F.interpolate(action_emb.float(), size=num_tokens, mode='linear', align_corners=True).to(dtype=action_emb.dtype) | |
| action_emb = action_emb.permute(0, 2, 1) | |
| action_emb = self.discrete_action_dropout(action_emb) | |
| action_emb_mod = self.adaln_single.linear( | |
| self.adaln_single.silu(action_emb) | |
| ) # [B, num_tokens, 6*dim] or [B, 6*dim] | |
| if not hasattr(self, '_act_diag_cnt'): | |
| self._act_diag_cnt = 0 | |
| if self._act_diag_cnt < 2: | |
| _ratio = action_emb_mod.abs().mean() / (timestep_emb.abs().mean() + 1e-12) | |
| print(f"[ActionDiag] action_emb_mod: shape={action_emb_mod.shape}, " | |
| f"timestep_emb: shape={timestep_emb.shape}, " | |
| f"mean_abs_ratio={_ratio:.4f}") | |
| self._act_diag_cnt += 1 | |
| if action_emb_mod.dim() == 2: | |
| timestep_emb = timestep_emb + action_emb_mod.unsqueeze(1) | |
| else: | |
| timestep_emb = timestep_emb + action_emb_mod | |
| else: | |
| # Dummy forward to keep FSDP params in sync across ranks | |
| dummy_labels = torch.zeros(1, dtype=torch.long, device=device) | |
| _ = self.action_embedder(dummy_labels) | |
| _ = self.discrete_action_dropout(torch.zeros(1, self.inner_dim, device=device, dtype=dtype)) | |
| # Caption projection | |
| context_proj = self.caption_projection(context_tensor) | |
| context_proj = context_proj.view(batch_size, -1, hidden.shape[-1]) | |
| # ===== Continuous camera signal processing ===== | |
| # FSDP requires ALL ranks to access the same parameters in forward pass. | |
| # Must always access camera adapter params even when plucker_coords is None. | |
| plucker_emb = None | |
| if (getattr(self, 'enable_camera_control', False) and | |
| getattr(self, 'enable_continuous_camera', False)): | |
| if plucker_coords is not None: | |
| # All modes now receive pixel-resolution Plucker [B, 6, F, H_pixel, W_pixel] | |
| # WanCameraAdapter handles spatial downsample internally (PixelUnshuffle + strided Conv) | |
| if self.camera_injection_mode == "scale_shift": | |
| # plucker_processor is now a WanCameraAdapter | |
| plucker_emb = self.plucker_processor( | |
| plucker_coords, num_frames, height, width) # [B, L, dim] | |
| # Store raw adapter output for aux rotation loss (before dropout) | |
| if self.training and getattr(self, '_enable_cam_aux', False): | |
| self._cam_features_for_aux = plucker_emb | |
| self._cam_aux_shape = (num_frames, height, width) | |
| plucker_emb = self.camera_dropout(plucker_emb) | |
| # Sequence length safety check | |
| if plucker_emb.shape[1] != num_tokens: | |
| plucker_emb = plucker_emb.permute(0, 2, 1) | |
| plucker_emb = torch.nn.functional.interpolate( | |
| plucker_emb.float(), size=num_tokens, mode='linear', align_corners=True, | |
| ).to(dtype=plucker_emb.dtype) | |
| plucker_emb = plucker_emb.permute(0, 2, 1) | |
| # Cross Normalization: align plucker_emb distribution to hidden | |
| if getattr(self, '_camera_use_cross_norm', False): | |
| _mean_h = hidden.detach().mean(dim=-1, keepdim=True) | |
| _std_h = hidden.detach().std(dim=-1, keepdim=True) | |
| _mean_c = plucker_emb.mean(dim=-1, keepdim=True) | |
| _std_c = plucker_emb.std(dim=-1, keepdim=True) | |
| plucker_emb = (plucker_emb - _mean_c) * (_std_h / (_std_c + 1e-12)) + _mean_h | |
| plucker_emb = plucker_emb * self._cross_norm_scale | |
| elif self.camera_injection_mode == "additive": | |
| # additive_camera_adapter is now a WanCameraAdapter | |
| cam_features = self.additive_camera_adapter( | |
| plucker_coords, num_frames, height, width) # [B, L, C] | |
| # Store raw adapter output for aux rotation loss (before dropout) | |
| if self.training and getattr(self, '_enable_cam_aux', False): | |
| self._cam_features_for_aux = cam_features | |
| self._cam_aux_shape = (num_frames, height, width) | |
| cam_features = self.camera_dropout(cam_features) | |
| if cam_features.shape[1] != num_tokens: | |
| cam_features = cam_features.permute(0, 2, 1) | |
| cam_features = torch.nn.functional.interpolate( | |
| cam_features.float(), size=num_tokens, mode='linear', align_corners=True, | |
| ).to(dtype=cam_features.dtype) | |
| cam_features = cam_features.permute(0, 2, 1) | |
| if not hasattr(self, '_cam_diag_cnt'): | |
| self._cam_diag_cnt = 0 | |
| if self._cam_diag_cnt < 2: | |
| _ratio = cam_features.abs().mean() / (hidden.abs().mean() + 1e-12) | |
| print(f"[CameraDiag] cam_features: mean_abs={cam_features.abs().mean():.6f}, " | |
| f"hidden: mean_abs={hidden.abs().mean():.6f}, " | |
| f"ratio={_ratio:.4f} (needs to exceed ~0.05 to have a visible effect)") | |
| self._cam_diag_cnt += 1 | |
| hidden = hidden + cam_features | |
| del cam_features # free immediately | |
| elif self.camera_injection_mode == "wan_inject": | |
| # wan_inject: plucker_coords is at pixel resolution [B, 6, F_pixel, H_pixel, W_pixel] | |
| # WanCameraAdapter handles downsample + conv2d internally | |
| cam_features = self.wan_camera_adapter( | |
| plucker_coords, num_frames, height, width) # [B, L, C] | |
| # Store raw adapter output for aux rotation loss (before dropout) | |
| if self.training and getattr(self, '_enable_cam_aux', False): | |
| self._cam_features_for_aux = cam_features | |
| self._cam_aux_shape = (num_frames, height, width) | |
| cam_features = self.camera_dropout(cam_features) | |
| if cam_features.shape[1] != num_tokens: | |
| cam_features = cam_features.permute(0, 2, 1) | |
| cam_features = torch.nn.functional.interpolate( | |
| cam_features.float(), size=num_tokens, mode='linear', align_corners=True, | |
| ).to(dtype=cam_features.dtype) | |
| cam_features = cam_features.permute(0, 2, 1) | |
| if not hasattr(self, '_cam_diag_cnt'): | |
| self._cam_diag_cnt = 0 | |
| if self._cam_diag_cnt < 2: | |
| _ratio = cam_features.abs().mean() / (hidden.abs().mean() + 1e-12) | |
| print(f"[CameraDiag-WanInject] cam_features: mean_abs={cam_features.abs().mean():.6f}, " | |
| f"hidden: mean_abs={hidden.abs().mean():.6f}, " | |
| f"ratio={_ratio:.4f} (needs to exceed ~0.05 to have a visible effect)") | |
| self._cam_diag_cnt += 1 | |
| hidden = hidden + cam_features | |
| del cam_features # free immediately | |
| else: | |
| # No camera data: dummy forward to keep FSDP params in sync across ranks | |
| # All modes use WanCameraAdapter now, min spatial = 32x32 | |
| # F=1 pixel frame → temporal unshuffle repeats to S frames → F_lat=1 | |
| if self.camera_injection_mode == "scale_shift": | |
| dummy_pl = torch.zeros(1, 6, 1, 32, 32, device=device, dtype=dtype) | |
| _ = self.plucker_processor(dummy_pl, 1, 1, 1) | |
| _ = self.camera_dropout(torch.zeros(1, 1, self.inner_dim, device=device, dtype=dtype)) | |
| elif self.camera_injection_mode == "additive": | |
| dummy_pl = torch.zeros(1, 6, 1, 32, 32, device=device, dtype=dtype) | |
| _ = self.additive_camera_adapter(dummy_pl, 1, 1, 1) | |
| _ = self.camera_dropout(torch.zeros(1, 1, self.inner_dim, device=device, dtype=dtype)) | |
| elif self.camera_injection_mode == "wan_inject": | |
| dummy_pl = torch.zeros(1, 6, 1, 32, 32, device=device, dtype=dtype) | |
| _ = self.wan_camera_adapter(dummy_pl, 1, 1, 1) | |
| _ = self.camera_dropout(torch.zeros(1, 1, self.inner_dim, device=device, dtype=dtype)) | |
| # Clear aux cache for dummy path | |
| if getattr(self, '_enable_cam_aux', False): | |
| self._cam_features_for_aux = None | |
| self._cam_aux_shape = None | |
| # Positional embeddings (RoPE) | |
| freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch | |
| freqs = precompute_freqs_cis( | |
| indices_grid=positions, | |
| dim=self.inner_dim, | |
| out_dtype=latent.dtype, | |
| theta=self.positional_embedding_theta, | |
| max_pos=self.positional_embedding_max_pos, | |
| use_middle_indices_grid=self.use_middle_indices_grid, | |
| num_attention_heads=self.num_attention_heads, | |
| rope_type=self.rope_type, | |
| freq_grid_generator=freq_grid_generator, | |
| normalize_positions=self.normalize_rope_positions, | |
| ) | |
| # ===== Context Parallel: pad + scatter sequence-dimension tensors ===== | |
| _cp_orig_seq_len = None | |
| if is_cp_enabled(): | |
| cp_size = get_cp_world_size() | |
| # Pad sequence to be divisible by cp_size (handles arbitrary T*H*W) | |
| hidden, _cp_orig_seq_len = pad_to_cp_divisible(hidden, dim=1) | |
| timestep_emb, _ = pad_to_cp_divisible(timestep_emb, dim=1) | |
| freqs_cos, _ = pad_to_cp_divisible(freqs[0], dim=2) | |
| freqs_sin, _ = pad_to_cp_divisible(freqs[1], dim=2) | |
| freqs = (freqs_cos, freqs_sin) | |
| if plucker_emb is not None: | |
| plucker_emb, _ = pad_to_cp_divisible(plucker_emb, dim=1) | |
| # Scatter | |
| hidden = scatter_sequence(hidden, dim=1) | |
| timestep_emb = scatter_sequence(timestep_emb, dim=1) | |
| freqs = (scatter_sequence(freqs[0], dim=2), scatter_sequence(freqs[1], dim=2)) | |
| if plucker_emb is not None: | |
| plucker_emb = scatter_sequence(plucker_emb, dim=1) | |
| # Update video_shape T dimension for sparse attention | |
| if video_shape is not None: | |
| T_vs, H_vs, W_vs = video_shape | |
| # Use padded T: hidden.shape[1] = padded_seq / cp_size, T_padded = padded_seq / (H*W) | |
| padded_seq = hidden.shape[1] * cp_size | |
| T_padded = padded_seq // (H_vs * W_vs) | |
| video_shape = (T_padded // cp_size, H_vs, W_vs) | |
| # Transformer blocks with gradient checkpointing support (WAN-style) | |
| def create_custom_forward(module): | |
| def custom_forward(*inputs, **kwargs): | |
| return module(*inputs, **kwargs) | |
| return custom_forward | |
| for block in self.blocks: | |
| block_kwargs = { | |
| "timesteps": timestep_emb, | |
| "freqs": freqs, | |
| "context": context_proj, | |
| "context_mask": None, | |
| "perturbations": None, | |
| "plucker_emb": plucker_emb, | |
| "video_shape": video_shape, # For sparse attention | |
| } | |
| if torch.is_grad_enabled() and self.gradient_checkpointing: | |
| hidden = torch.utils.checkpoint.checkpoint( | |
| create_custom_forward(block), | |
| hidden, | |
| **block_kwargs, | |
| use_reentrant=False, | |
| ) | |
| else: | |
| hidden = block(x=hidden, **block_kwargs) | |
| # ===== Context Parallel: gather sequence back + unpad ===== | |
| if is_cp_enabled(): | |
| hidden = gather_for_loss(hidden, dim=1) | |
| if _cp_orig_seq_len is not None: | |
| hidden = unpad_from_cp(hidden, _cp_orig_seq_len, dim=1) | |
| # Output processing | |
| out = self._process_output(self.scale_shift_table, self.norm_out, self.proj_out, hidden, embedded_timestep) | |
| # Unpatchify: [B, num_tokens, C] -> [B, C, F, H, W] | |
| out_bcfhw = patchifier.unpatchify(out, target_shape) | |
| return out_bcfhw | |
| def init_weights(self): | |
| """ | |
| Initialize model parameters using Xavier initialization. | |
| """ | |
| # basic init | |
| for m in self.modules(): | |
| if isinstance(m, nn.Linear): | |
| nn.init.xavier_uniform_(m.weight) | |
| if m.bias is not None: | |
| nn.init.zeros_(m.bias) | |
| # init video embeddings | |
| if self.model_type.is_video_enabled(): | |
| # patchify_proj is the equivalent of patch_embedding | |
| nn.init.xavier_uniform_(self.patchify_proj.weight) | |
| # caption_projection is the text embedding | |
| for m in self.caption_projection.modules(): | |
| if isinstance(m, nn.Linear): | |
| nn.init.normal_(m.weight, std=.02) | |
| # adaln_single contains time embedding | |
| for m in self.adaln_single.modules(): | |
| if isinstance(m, nn.Linear): | |
| nn.init.normal_(m.weight, std=.02) | |
| # proj_out is the output head | |
| nn.init.zeros_(self.proj_out.weight) | |
| # ============================================================================= | |
| # Utility Functions | |
| # ============================================================================= | |
| def to_denoised( | |
| sample: torch.Tensor, | |
| velocity: torch.Tensor, | |
| sigma: float | torch.Tensor, | |
| calc_dtype: torch.dtype = torch.float32, | |
| ) -> torch.Tensor: | |
| """ | |
| Convert the sample and its denoising velocity to denoised sample. | |
| Formula: x0 = sample - velocity * sigma | |
| """ | |
| if isinstance(sigma, torch.Tensor): | |
| sigma = sigma.to(calc_dtype) | |
| return (sample.to(calc_dtype) - velocity.to(calc_dtype) * sigma).to(sample.dtype) | |
| # ============================================================================= | |
| # X0Model Wrapper | |
| # ============================================================================= | |
| class X0Model(nn.Module): | |
| """X0 model wrapper: expose a velocity model as an x0 model.""" | |
| def __init__(self, velocity_model: LTX2Model): | |
| super().__init__() | |
| self.velocity_model = velocity_model | |
| def forward( | |
| self, x, t, context, seq_len, **kwargs | |
| ) -> torch.Tensor: | |
| """ | |
| Forward pass: outputs x0 (denoised samples) instead of velocity. | |
| Args: | |
| x: List of input video tensors [C, F, H, W] | |
| t: Timesteps [B] | |
| context: List of text embeddings [L, C] | |
| seq_len: Maximum sequence length | |
| Returns: | |
| x0 prediction [B, C, F, H, W] | |
| """ | |
| velocity = self.velocity_model(x, t, context, seq_len, **kwargs) | |
| # Convert velocity to x0: x0 = xt - sigma * velocity | |
| sigma = t.float() / 1000.0 | |
| while sigma.dim() < velocity.dim(): | |
| sigma = sigma.unsqueeze(-1) | |
| # x is List[[C, F, H, W]], stack to [B, C, F, H, W] | |
| if isinstance(x, list): | |
| x_stacked = torch.stack(x) | |
| else: | |
| x_stacked = x | |
| x0 = x_stacked - sigma * velocity | |
| return x0 | |