# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import math from functools import lru_cache from typing import Any import torch import torch.nn as nn from sglang.multimodal_gen.configs.models.dits import WanVideoConfig from sglang.multimodal_gen.configs.sample.wan import WanTeaCacheParams from sglang.multimodal_gen.runtime.distributed import ( divide, get_sp_group, get_sp_world_size, get_tp_world_size, sequence_model_parallel_all_gather, ) from sglang.multimodal_gen.runtime.layers.attention import ( MinimalA2AAttnOp, UlyssesAttention_VSA, USPAttention, ) from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd from sglang.multimodal_gen.runtime.layers.layernorm import ( FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidualLayerNormScaleShift, tensor_parallel_rms_norm, ) from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( NDRotaryEmbedding, _apply_rotary_emb, apply_flashinfer_rope_qk_inplace, ) from sglang.multimodal_gen.runtime.layers.visual_embedding import ( ModulateProjection, PatchEmbed, TimestepEmbedder, ) from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context 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.server_args import get_global_server_args from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) _is_cuda = current_platform.is_cuda() class WanImageEmbedding(torch.nn.Module): def __init__(self, in_features: int, out_features: int): super().__init__() self.norm1 = FP32LayerNorm(in_features) self.ff = MLP(in_features, in_features, out_features, act_type="gelu") self.norm2 = FP32LayerNorm(out_features) def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: dtype = encoder_hidden_states_image.dtype hidden_states = self.norm1(encoder_hidden_states_image) hidden_states = self.ff(hidden_states) hidden_states = self.norm2(hidden_states).to(dtype) return hidden_states class WanTimeTextImageEmbedding(nn.Module): def __init__( self, dim: int, time_freq_dim: int, text_embed_dim: int, image_embed_dim: int | None = None, ): super().__init__() self.time_embedder = TimestepEmbedder( dim, frequency_embedding_size=time_freq_dim, act_layer="silu" ) self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu") self.text_embedder = MLP( text_embed_dim, dim, dim, bias=True, act_type="gelu_pytorch_tanh" ) self.image_embedder = None if image_embed_dim is not None: self.image_embedder = WanImageEmbedding(image_embed_dim, dim) def forward( self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, encoder_hidden_states_image: torch.Tensor | None = None, timestep_seq_len: int | None = None, ): temb = self.time_embedder(timestep, timestep_seq_len) timestep_proj = self.time_modulation(temb) encoder_hidden_states = self.text_embedder(encoder_hidden_states) if encoder_hidden_states_image is not None: assert self.image_embedder is not None encoder_hidden_states_image = self.image_embedder( encoder_hidden_states_image ) return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image class WanSelfAttention(nn.Module): def __init__( self, dim: int, num_heads: int, window_size=(-1, -1), qk_norm=True, eps=1e-6, parallel_attention=False, supported_attention_backends: set[AttentionBackendEnum] | None = None, ) -> None: assert dim % num_heads == 0 super().__init__() self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads self.window_size = window_size self.qk_norm = qk_norm self.eps = eps self.parallel_attention = parallel_attention tp_size = get_tp_world_size() # layers self.to_q = ColumnParallelLinear(dim, dim, gather_output=False) self.to_k = ColumnParallelLinear(dim, dim, gather_output=False) self.to_v = ColumnParallelLinear(dim, dim, gather_output=False) self.to_out = RowParallelLinear(dim, dim, input_is_parallel=True) self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() self.tp_rmsnorm = tp_size > 1 and qk_norm self.local_num_heads = divide(num_heads, tp_size) # Scaled dot product attention self.attn = USPAttention( num_heads=self.local_num_heads, head_size=self.head_dim, dropout_rate=0, softmax_scale=None, causal=False, supported_attention_backends=supported_attention_backends, ) def forward(self, x: torch.Tensor, context: torch.Tensor, context_lens: int): r""" Args: x(Tensor): Shape [B, L, num_heads, C / num_heads] seq_lens(Tensor): Shape [B] grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] """ pass class WanT2VCrossAttention(WanSelfAttention): def forward(self, x, context, context_lens): r""" Args: x(Tensor): Shape [B, L1, C] context(Tensor): Shape [B, L2, C] context_lens(Tensor): Shape [B] """ q, _ = self.to_q(x) if self.tp_rmsnorm: q = tensor_parallel_rms_norm(q, self.norm_q) else: q = self.norm_q(q) q = q.unflatten(2, (self.local_num_heads, self.head_dim)) k, _ = self.to_k(context) if self.tp_rmsnorm: k = tensor_parallel_rms_norm(k, self.norm_k) else: k = self.norm_k(k) k = k.unflatten(2, (self.local_num_heads, self.head_dim)) v, _ = self.to_v(context) v = v.unflatten(2, (self.local_num_heads, self.head_dim)) # compute attention x = self.attn(q, k, v) # output x = x.flatten(2) x, _ = self.to_out(x) return x class WanI2VCrossAttention(WanSelfAttention): def __init__( self, dim: int, num_heads: int, window_size=(-1, -1), qk_norm=True, eps=1e-6, supported_attention_backends: set[AttentionBackendEnum] | None = None, ) -> None: super().__init__( dim, num_heads, window_size, qk_norm, eps, supported_attention_backends=supported_attention_backends, ) self.add_k_proj = ColumnParallelLinear(dim, dim, gather_output=False) self.add_v_proj = ColumnParallelLinear(dim, dim, gather_output=False) self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() def forward(self, x, context, context_lens): r""" Args: x(Tensor): Shape [B, L1, C] context(Tensor): Shape [B, L2, C] context_lens(Tensor): Shape [B] """ context_img = context[:, :257] context = context[:, 257:] q, _ = self.to_q(x) if self.tp_rmsnorm: q = tensor_parallel_rms_norm(q, self.norm_q) else: q = self.norm_q(q) q = q.unflatten(2, (self.local_num_heads, self.head_dim)) k, _ = self.to_k(context) if self.tp_rmsnorm: k = tensor_parallel_rms_norm(k, self.norm_k) else: k = self.norm_k(k) k = k.unflatten(2, (self.local_num_heads, self.head_dim)) v, _ = self.to_v(context) v = v.unflatten(2, (self.local_num_heads, self.head_dim)) k_img, _ = self.add_k_proj(context_img) if self.tp_rmsnorm: k_img = tensor_parallel_rms_norm(k_img, self.norm_added_k) else: k_img = self.norm_added_k(k_img) k_img = k_img.unflatten(2, (self.local_num_heads, self.head_dim)) v_img, _ = self.add_v_proj(context_img) v_img = v_img.unflatten(2, (self.local_num_heads, self.head_dim)) img_x = self.attn(q, k_img, v_img) x = self.attn(q, k, v) # output x = x.flatten(2) img_x = img_x.flatten(2) x = x + img_x x, _ = self.to_out(x) return x class WanTransformerBlock(nn.Module): def __init__( self, dim: int, ffn_dim: int, num_heads: int, qk_norm: str = "rms_norm_across_heads", cross_attn_norm: bool = False, eps: float = 1e-6, added_kv_proj_dim: int | None = None, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", attention_type: str = "original", sla_topk: float = 0.1, ): super().__init__() # 1. Self-attention self.norm1 = LayerNormScaleShift( dim, eps=eps, elementwise_affine=False, dtype=torch.float32, ) self.to_q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False) self.to_k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False) self.to_v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False) self.to_out = RowParallelLinear(dim, dim, bias=True, reduce_results=True) tp_size = get_tp_world_size() self.local_num_heads = divide(num_heads, tp_size) self_attn_backends = supported_attention_backends if attention_type in ("sla", "sagesla"): self.attn1 = MinimalA2AAttnOp( num_heads=self.local_num_heads, head_size=dim // num_heads, attention_type=attention_type, topk=sla_topk, supported_attention_backends={ AttentionBackendEnum.SLA_ATTN, AttentionBackendEnum.SAGE_SLA_ATTN, }, ) else: self.attn1 = USPAttention( num_heads=self.local_num_heads, head_size=dim // num_heads, causal=False, supported_attention_backends=self_attn_backends, prefix=f"{prefix}.attn1", ) self.hidden_dim = dim self.num_attention_heads = num_heads self.dim_head = dim // num_heads if qk_norm == "rms_norm": self.norm_q = RMSNorm(self.dim_head, eps=eps) self.norm_k = RMSNorm(self.dim_head, eps=eps) elif qk_norm == "rms_norm_across_heads": # LTX applies qk norm across all heads self.norm_q = RMSNorm(dim, eps=eps) self.norm_k = RMSNorm(dim, eps=eps) else: logger.error("QK Norm type not supported") raise Exception assert cross_attn_norm is True self.qk_norm = qk_norm self.tp_rmsnorm = qk_norm == "rms_norm_across_heads" and tp_size > 1 self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, eps=eps, elementwise_affine=True, dtype=torch.float32, ) # 2. Cross-attention cross_attn_backends = { b for b in supported_attention_backends if not b.is_sparse } if added_kv_proj_dim is not None: # I2V self.attn2 = WanI2VCrossAttention( dim, num_heads, qk_norm=qk_norm, eps=eps, supported_attention_backends=cross_attn_backends, ) else: # T2V self.attn2 = WanT2VCrossAttention( dim, num_heads, qk_norm=qk_norm, eps=eps, supported_attention_backends=cross_attn_backends, ) self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, eps=eps, elementwise_affine=False, dtype=torch.float32, ) # 3. Feed-forward self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") self.mlp_residual = MulAdd() self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype if temb.dim() == 4: # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( self.scale_shift_table.unsqueeze(0) + temb.float() ).chunk(6, dim=2) # batch_size, seq_len, 1, inner_dim shift_msa = shift_msa.squeeze(2) scale_msa = scale_msa.squeeze(2) gate_msa = gate_msa.squeeze(2) c_shift_msa = c_shift_msa.squeeze(2) c_scale_msa = c_scale_msa.squeeze(2) c_gate_msa = c_gate_msa.squeeze(2) else: # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B) e = self.scale_shift_table + temb.float() shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( e.chunk(6, dim=1) ) assert shift_msa.dtype == torch.float32 # 1. Self-attention norm_hidden_states = self.norm1(hidden_states, shift_msa, scale_msa) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) if self.norm_q is not None: if self.tp_rmsnorm: query = tensor_parallel_rms_norm(query, self.norm_q) else: query = self.norm_q(query) if self.norm_k is not None: if self.tp_rmsnorm: key = tensor_parallel_rms_norm(key, self.norm_k) else: key = self.norm_k(key) query = query.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head)) key = key.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head)) value = value.squeeze(1).unflatten(2, (self.local_num_heads, self.dim_head)) # Apply rotary embeddings cos, sin = freqs_cis if _is_cuda and query.shape == key.shape: cos_sin_cache = torch.cat( [ cos.to(dtype=torch.float32).contiguous(), sin.to(dtype=torch.float32).contiguous(), ], dim=-1, ) query, key = apply_flashinfer_rope_qk_inplace( query, key, cos_sin_cache, is_neox=False ) else: query, key = _apply_rotary_emb( query, cos, sin, is_neox_style=False ), _apply_rotary_emb(key, cos, sin, is_neox_style=False) attn_output = self.attn1(query, key, value) attn_output = attn_output.flatten(2) attn_output, _ = self.to_out(attn_output) attn_output = attn_output.squeeze(1) null_shift = null_scale = torch.zeros( (1,), device=hidden_states.device, dtype=hidden_states.dtype ) norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale ) norm_hidden_states, hidden_states = norm_hidden_states.to( orig_dtype ), hidden_states.to(orig_dtype) # 2. Cross-attention attn_output = self.attn2( norm_hidden_states, context=encoder_hidden_states, context_lens=None ) norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa ) norm_hidden_states, hidden_states = norm_hidden_states.to( orig_dtype ), hidden_states.to(orig_dtype) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) hidden_states = self.mlp_residual(ff_output, c_gate_msa, hidden_states) hidden_states = hidden_states.to(orig_dtype) return hidden_states class WanTransformerBlock_VSA(nn.Module): def __init__( self, dim: int, ffn_dim: int, num_heads: int, qk_norm: str = "rms_norm_across_heads", cross_attn_norm: bool = False, eps: float = 1e-6, added_kv_proj_dim: int | None = None, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", ): super().__init__() # 1. Self-attention self.norm1 = LayerNormScaleShift( dim, eps=eps, elementwise_affine=False, dtype=torch.float32, ) self.to_q = ColumnParallelLinear(dim, dim, bias=True, gather_output=True) self.to_k = ColumnParallelLinear(dim, dim, bias=True, gather_output=True) self.to_v = ColumnParallelLinear(dim, dim, bias=True, gather_output=True) self.to_gate_compress = ColumnParallelLinear( dim, dim, bias=True, gather_output=True ) self.to_out = ColumnParallelLinear(dim, dim, bias=True, gather_output=True) self.attn1 = UlyssesAttention_VSA( num_heads=num_heads, head_size=dim // num_heads, causal=False, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn1", ) self.hidden_dim = dim self.num_attention_heads = num_heads dim_head = dim // num_heads if qk_norm == "rms_norm": self.norm_q = RMSNorm(dim_head, eps=eps) self.norm_k = RMSNorm(dim_head, eps=eps) elif qk_norm == "rms_norm_across_heads": # LTX applies qk norm across all heads self.norm_q = RMSNorm(dim, eps=eps) self.norm_k = RMSNorm(dim, eps=eps) else: logger.error("QK Norm type not supported") raise Exception assert cross_attn_norm is True self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, eps=eps, elementwise_affine=True, dtype=torch.float32, ) # 2. Cross-attention cross_attn_backends = { b for b in supported_attention_backends if not b.is_sparse } if added_kv_proj_dim is not None: # I2V self.attn2 = WanI2VCrossAttention( dim, num_heads, qk_norm=qk_norm, eps=eps, supported_attention_backends=cross_attn_backends, ) else: # T2V self.attn2 = WanT2VCrossAttention( dim, num_heads, qk_norm=qk_norm, eps=eps, supported_attention_backends=cross_attn_backends, ) self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( dim, eps=eps, elementwise_affine=False, dtype=torch.float32, ) # 3. Feed-forward self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") self.mlp_residual = MulAdd() self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype # assert orig_dtype != torch.float32 e = self.scale_shift_table + temb.float() shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( 6, dim=1 ) assert shift_msa.dtype == torch.float32 # 1. Self-attention norm_hidden_states = self.norm1(hidden_states, shift_msa, scale_msa) query, _ = self.to_q(norm_hidden_states) key, _ = self.to_k(norm_hidden_states) value, _ = self.to_v(norm_hidden_states) gate_compress, _ = self.to_gate_compress(norm_hidden_states) if self.norm_q is not None: query = self.norm_q(query) if self.norm_k is not None: key = self.norm_k(key) query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) gate_compress = gate_compress.squeeze(1).unflatten( 2, (self.num_attention_heads, -1) ) # Apply rotary embeddings cos, sin = freqs_cis if _is_cuda and query.shape == key.shape: cos_sin_cache = torch.cat( [ cos.to(dtype=torch.float32).contiguous(), sin.to(dtype=torch.float32).contiguous(), ], dim=-1, ) query, key = apply_flashinfer_rope_qk_inplace( query, key, cos_sin_cache, is_neox=False ) else: query, key = _apply_rotary_emb( query, cos, sin, is_neox_style=False ), _apply_rotary_emb(key, cos, sin, is_neox_style=False) attn_output = self.attn1(query, key, value, gate_compress=gate_compress) attn_output = attn_output.flatten(2) attn_output, _ = self.to_out(attn_output) attn_output = attn_output.squeeze(1) null_shift = null_scale = torch.zeros((1,), device=hidden_states.device) norm_hidden_states, hidden_states = self.self_attn_residual_norm( hidden_states, attn_output, gate_msa, null_shift, null_scale ) norm_hidden_states, hidden_states = norm_hidden_states.to( orig_dtype ), hidden_states.to(orig_dtype) # 2. Cross-attention attn_output = self.attn2( norm_hidden_states, context=encoder_hidden_states, context_lens=None ) norm_hidden_states, hidden_states = self.cross_attn_residual_norm( hidden_states, attn_output, 1, c_shift_msa, c_scale_msa ) norm_hidden_states, hidden_states = norm_hidden_states.to( orig_dtype ), hidden_states.to(orig_dtype) # 3. Feed-forward ff_output = self.ffn(norm_hidden_states) hidden_states = self.mlp_residual(ff_output, c_gate_msa, hidden_states) hidden_states = hidden_states.to(orig_dtype) return hidden_states class WanTransformer3DModel(CachableDiT, OffloadableDiTMixin): _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions _compile_conditions = WanVideoConfig()._compile_conditions _supported_attention_backends = WanVideoConfig()._supported_attention_backends param_names_mapping = WanVideoConfig().param_names_mapping reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping def __init__(self, config: WanVideoConfig, hf_config: dict[str, Any]) -> None: super().__init__(config=config, hf_config=hf_config) inner_dim = config.num_attention_heads * config.attention_head_dim self.hidden_size = config.hidden_size self.num_attention_heads = config.num_attention_heads self.in_channels = config.in_channels self.out_channels = config.out_channels self.num_channels_latents = config.num_channels_latents self.patch_size = config.patch_size self.text_len = config.text_len # 1. Patch & position embedding self.patch_embedding = PatchEmbed( in_chans=config.in_channels, embed_dim=inner_dim, patch_size=config.patch_size, flatten=False, ) # 2. Condition embeddings self.condition_embedder = WanTimeTextImageEmbedding( dim=inner_dim, time_freq_dim=config.freq_dim, text_embed_dim=config.text_dim, image_embed_dim=config.image_dim, ) # 3. Transformer blocks attn_backend = get_global_server_args().attention_backend transformer_block = ( WanTransformerBlock_VSA if (attn_backend and attn_backend.lower() == "video_sparse_attn") else WanTransformerBlock ) self.blocks = nn.ModuleList( [ transformer_block( inner_dim, config.ffn_dim, config.num_attention_heads, config.qk_norm, config.cross_attn_norm, config.eps, config.added_kv_proj_dim, self._supported_attention_backends | {AttentionBackendEnum.VIDEO_SPARSE_ATTN}, prefix=f"{config.prefix}.blocks.{i}", attention_type=config.attention_type, sla_topk=config.sla_topk, ) for i in range(config.num_layers) ] ) # 4. Output norm & projection self.norm_out = LayerNormScaleShift( inner_dim, eps=config.eps, elementwise_affine=False, dtype=torch.float32, ) self.proj_out = ColumnParallelLinear( inner_dim, config.out_channels * math.prod(config.patch_size), bias=True, gather_output=True, ) self.scale_shift_table = nn.Parameter( torch.randn(1, 2, inner_dim) / inner_dim**0.5 ) # For type checking self.cnt = 0 self.__post_init__() # misc self.sp_size = get_sp_world_size() # Get rotary embeddings d = self.hidden_size // self.num_attention_heads self.rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] self.rotary_emb = NDRotaryEmbedding( rope_dim_list=self.rope_dim_list, rope_theta=10000, dtype=( torch.float32 if current_platform.is_mps() or current_platform.is_musa() else torch.float64 ), ) self.layer_names = ["blocks"] @lru_cache(maxsize=1) def _compute_rope_for_sequence_shard( self, local_len: int, rank: int, frame_stride_local: int, width_local: int, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor]: token_start = rank * local_len token_indices = torch.arange( token_start, token_start + local_len, device=device, dtype=torch.long, ) t_idx = token_indices // frame_stride_local rem = token_indices % frame_stride_local h_idx = rem // width_local w_idx = rem % width_local positions = torch.stack((t_idx, h_idx, w_idx), dim=1) return self.rotary_emb.forward_uncached(positions) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, guidance=None, **kwargs, ) -> torch.Tensor: forward_batch = get_forward_context().forward_batch if forward_batch is not None: sequence_shard_enabled = ( forward_batch.enable_sequence_shard and self.sp_size > 1 ) else: sequence_shard_enabled = False self.enable_teacache = ( forward_batch is not None and forward_batch.enable_teacache ) orig_dtype = hidden_states.dtype if not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] if ( isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0 ): encoder_hidden_states_image = encoder_hidden_states_image[0] else: encoder_hidden_states_image = None batch_size, num_channels, num_frames, height, width = hidden_states.shape p_t, p_h, p_w = self.patch_size post_patch_num_frames = num_frames // p_t post_patch_height = height // p_h post_patch_width = width // p_w if not sequence_shard_enabled: # The rotary embedding layer correctly handles SP offsets internally. freqs_cos, freqs_sin = self.rotary_emb.forward_from_grid( ( post_patch_num_frames * self.sp_size, post_patch_height, post_patch_width, ), shard_dim=0, start_frame=0, device=hidden_states.device, ) assert freqs_cos.dtype == torch.float32 assert freqs_cos.device == hidden_states.device freqs_cis = ( (freqs_cos.float(), freqs_sin.float()) if freqs_cos is not None else None ) hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) # shape is [B, T' * H' * W', C] seq_len_orig = hidden_states.shape[1] seq_shard_pad = 0 if sequence_shard_enabled: if seq_len_orig % self.sp_size != 0: seq_shard_pad = self.sp_size - (seq_len_orig % self.sp_size) pad = torch.zeros( (batch_size, seq_shard_pad, hidden_states.shape[2]), dtype=hidden_states.dtype, device=hidden_states.device, ) hidden_states = torch.cat([hidden_states, pad], dim=1) sp_rank = get_sp_group().rank_in_group local_seq_len = hidden_states.shape[1] // self.sp_size hidden_states = hidden_states.view( batch_size, self.sp_size, local_seq_len, hidden_states.shape[2] ) hidden_states = hidden_states[:, sp_rank, :, :] frame_stride = post_patch_height * post_patch_width freqs_cos, freqs_sin = self._compute_rope_for_sequence_shard( local_seq_len, sp_rank, frame_stride, post_patch_width, hidden_states.device, ) freqs_cis = (freqs_cos.float(), freqs_sin.float()) # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v) if timestep.dim() == 2: # ti2v ts_seq_len = timestep.shape[1] timestep = timestep.flatten() # batch_size * seq_len else: ts_seq_len = None temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = ( self.condition_embedder( timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len, ) ) if ts_seq_len is not None: # batch_size, seq_len, 6, inner_dim timestep_proj = timestep_proj.unflatten(2, (6, -1)) else: # batch_size, 6, inner_dim timestep_proj = timestep_proj.unflatten(1, (6, -1)) if encoder_hidden_states_image is not None: encoder_hidden_states = torch.concat( [encoder_hidden_states_image, encoder_hidden_states], dim=1 ) encoder_hidden_states = ( encoder_hidden_states.to(orig_dtype) if not current_platform.is_amp_supported() else encoder_hidden_states ) # cast to orig_dtype for MPS assert encoder_hidden_states.dtype == orig_dtype # 4. Transformer blocks # if caching is enabled, we might be able to skip the forward pass should_skip_forward = self.should_skip_forward_for_cached_states( timestep_proj=timestep_proj, temb=temb ) if should_skip_forward: hidden_states = self.retrieve_cached_states(hidden_states) else: # if teacache is enabled, we need to cache the original hidden states if self.enable_teacache: original_hidden_states = hidden_states.clone() for block in self.blocks: hidden_states = block( hidden_states, encoder_hidden_states, timestep_proj, freqs_cis ) # if teacache is enabled, we need to cache the original hidden states if self.enable_teacache: self.maybe_cache_states(hidden_states, original_hidden_states) self.cnt += 1 if sequence_shard_enabled: hidden_states = hidden_states.contiguous() hidden_states = sequence_model_parallel_all_gather(hidden_states, dim=1) if seq_shard_pad > 0: hidden_states = hidden_states[:, :seq_len_orig, :] # 5. Output norm, projection & unpatchify if temb.dim() == 3: # batch_size, seq_len, inner_dim (wan 2.2 ti2v) shift, scale = ( self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2) ).chunk(2, dim=2) shift = shift.squeeze(2) scale = scale.squeeze(2) else: # batch_size, inner_dim shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) hidden_states = self.norm_out(hidden_states, shift, scale) hidden_states, _ = self.proj_out(hidden_states) hidden_states = hidden_states.reshape( batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1, ) hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) return output def maybe_cache_states( self, hidden_states: torch.Tensor, original_hidden_states: torch.Tensor ) -> None: """Cache residual with CFG positive/negative separation.""" residual = hidden_states.squeeze(0) - original_hidden_states if not self.is_cfg_negative: self.previous_residual = residual else: self.previous_residual_negative = residual def should_skip_forward_for_cached_states(self, **kwargs) -> bool: if not self.enable_teacache: return False ctx = self._get_teacache_context() if ctx is None: return False # Wan uses WanTeaCacheParams with additional fields teacache_params = ctx.teacache_params assert isinstance( teacache_params, WanTeaCacheParams ), "teacache_params is not a WanTeaCacheParams" # Initialize Wan-specific parameters use_ret_steps = teacache_params.use_ret_steps cutoff_steps = teacache_params.get_cutoff_steps(ctx.num_inference_steps) ret_steps = teacache_params.ret_steps # Adjust ret_steps and cutoff_steps for non-CFG mode # (WanTeaCacheParams uses *2 factor assuming CFG) if not ctx.do_cfg: ret_steps = ret_steps // 2 cutoff_steps = cutoff_steps // 2 timestep_proj = kwargs["timestep_proj"] temb = kwargs["temb"] modulated_inp = timestep_proj if use_ret_steps else temb self.is_cfg_negative = ctx.is_cfg_negative # Wan uses ret_steps/cutoff_steps for boundary detection is_boundary_step = self.cnt < ret_steps or self.cnt >= cutoff_steps # Use shared helper to compute cache decision should_calc = self._compute_teacache_decision( modulated_inp=modulated_inp, is_boundary_step=is_boundary_step, coefficients=ctx.coefficients, teacache_thresh=ctx.teacache_thresh, ) return not should_calc def retrieve_cached_states(self, hidden_states: torch.Tensor) -> torch.Tensor: """Retrieve cached residual with CFG positive/negative separation.""" if not self.is_cfg_negative: return hidden_states + self.previous_residual else: return hidden_states + self.previous_residual_negative EntryClass = WanTransformer3DModel