# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 import math from typing import Any import torch import torch.nn as nn from torch.nn.attention.flex_attention import ( BlockMask, create_block_mask, flex_attention, ) from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin # wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention # see https://github.com/pytorch/pytorch/issues/133254 # change to default for other models flex_attention = torch.compile( flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs" ) import torch.distributed as dist from sglang.multimodal_gen.configs.models.dits import WanVideoConfig from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size from sglang.multimodal_gen.runtime.layers.attention import LocalAttention from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd from sglang.multimodal_gen.runtime.layers.layernorm import ( FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidualLayerNormScaleShift, ) from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear from sglang.multimodal_gen.runtime.layers.mlp import MLP from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( _apply_rotary_emb, get_rotary_pos_embed, ) from sglang.multimodal_gen.runtime.layers.visual_embedding import PatchEmbed from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT from sglang.multimodal_gen.runtime.models.dits.wanvideo import ( WanT2VCrossAttention, WanTimeTextImageEmbedding, ) from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) class CausalWanSelfAttention(nn.Module): def __init__( self, dim: int, num_heads: int, local_attn_size: int = -1, sink_size: int = 0, qk_norm=True, eps=1e-6, parallel_attention=False, ) -> None: assert dim % num_heads == 0 super().__init__() self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads self.local_attn_size = local_attn_size self.sink_size = sink_size self.qk_norm = qk_norm self.eps = eps self.parallel_attention = parallel_attention self.max_attention_size = ( 32760 if local_attn_size == -1 else local_attn_size * 1560 ) # Scaled dot product attention self.attn = LocalAttention( num_heads=num_heads, head_size=self.head_dim, dropout_rate=0, softmax_scale=None, causal=False, supported_attention_backends=( AttentionBackendEnum.FA, AttentionBackendEnum.AITER, AttentionBackendEnum.TORCH_SDPA, ), ) def forward( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], block_mask: BlockMask, kv_cache: dict | None = None, current_start: int = 0, cache_start: int | None = None, ): 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] """ if cache_start is None: cache_start = current_start cos, sin = freqs_cis roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v) roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v) if kv_cache is None: # Padding for flex attention padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] padded_roped_query = torch.cat( [ roped_query, torch.zeros( [q.shape[0], padded_length, q.shape[2], q.shape[3]], device=q.device, dtype=v.dtype, ), ], dim=1, ) padded_roped_key = torch.cat( [ roped_key, torch.zeros( [k.shape[0], padded_length, k.shape[2], k.shape[3]], device=k.device, dtype=v.dtype, ), ], dim=1, ) padded_v = torch.cat( [ v, torch.zeros( [v.shape[0], padded_length, v.shape[2], v.shape[3]], device=v.device, dtype=v.dtype, ), ], dim=1, ) x = flex_attention( query=padded_roped_query.transpose(2, 1), key=padded_roped_key.transpose(2, 1), value=padded_v.transpose(2, 1), block_mask=block_mask, )[:, :, :-padded_length].transpose(2, 1) else: frame_seqlen = q.shape[1] current_end = current_start + roped_query.shape[1] sink_tokens = self.sink_size * frame_seqlen # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache kv_cache_size = kv_cache["k"].shape[1] num_new_tokens = roped_query.shape[1] if ( self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and ( num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size ) ): # Calculate the number of new tokens added in this step # Shift existing cache content left to discard oldest tokens # Clone the source slice to avoid overlapping memory error num_evicted_tokens = ( num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size ) num_rolled_tokens = ( kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens ) kv_cache["k"][ :, sink_tokens : sink_tokens + num_rolled_tokens ] = kv_cache["k"][ :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens, ].clone() kv_cache["v"][ :, sink_tokens : sink_tokens + num_rolled_tokens ] = kv_cache["v"][ :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens, ].clone() # Insert the new keys/values at the end local_end_index = ( kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() - num_evicted_tokens ) local_start_index = local_end_index - num_new_tokens kv_cache["k"][:, local_start_index:local_end_index] = roped_key kv_cache["v"][:, local_start_index:local_end_index] = v else: # Assign new keys/values directly up to current_end local_end_index = ( kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() ) local_start_index = local_end_index - num_new_tokens kv_cache["k"] = kv_cache["k"].detach() kv_cache["v"] = kv_cache["v"].detach() # logger.info("kv_cache['k'] is in comp graph: %s", kv_cache["k"].requires_grad or kv_cache["k"].grad_fn is not None) kv_cache["k"][:, local_start_index:local_end_index] = roped_key kv_cache["v"][:, local_start_index:local_end_index] = v x = self.attn( roped_query, kv_cache["k"][ :, max(0, local_end_index - self.max_attention_size) : local_end_index, ], kv_cache["v"][ :, max(0, local_end_index - self.max_attention_size) : local_end_index, ], ) kv_cache["global_end_index"].fill_(current_end) kv_cache["local_end_index"].fill_(local_end_index) return x class CausalWanTransformerBlock(nn.Module): def __init__( self, dim: int, ffn_dim: int, num_heads: int, local_attn_size: int = -1, sink_size: int = 0, 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 = FP32LayerNorm(dim, eps, elementwise_affine=False) self.to_q = ReplicatedLinear(dim, dim, bias=True) self.to_k = ReplicatedLinear(dim, dim, bias=True) self.to_v = ReplicatedLinear(dim, dim, bias=True) self.to_out = ReplicatedLinear(dim, dim, bias=True) self.attn1 = CausalWanSelfAttention( dim, num_heads, local_attn_size=local_attn_size, sink_size=sink_size, qk_norm=qk_norm, eps=eps, ) self.hidden_dim = dim self.num_attention_heads = num_heads self.local_attn_size = local_attn_size 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: print("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 # Only T2V for now cross_attn_backends = { b for b in supported_attention_backends if not b.is_sparse } 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], block_mask: BlockMask, kv_cache: dict | None = None, crossattn_cache: dict | None = None, current_start: int = 0, cache_start: int | None = None, ) -> torch.Tensor: # hidden_states.shape: [batch_size, seq_length, inner_dim] # temb.shape: [batch_size, num_frames, 6, inner_dim] if hidden_states.dim() == 4: hidden_states = hidden_states.squeeze(1) num_frames = temb.shape[1] frame_seqlen = hidden_states.shape[1] // num_frames bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype # assert orig_dtype != torch.float32 e = self.scale_shift_table + temb.float() # e.shape: [batch_size, num_frames, 6, inner_dim] assert e.shape == (bs, num_frames, 6, self.hidden_dim) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk( 6, dim=2 ) # *_msa.shape: [batch_size, num_frames, 1, inner_dim] assert shift_msa.dtype == torch.float32 # 1. Self-attention norm_hidden_states = ( ( self.norm1(hidden_states.float()).unflatten( dim=1, sizes=(num_frames, frame_seqlen) ) * (1 + scale_msa) + shift_msa ) .flatten(1, 2) .to(orig_dtype) ) 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: 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)) attn_output = self.attn1( query, key, value, freqs_cis, block_mask, kv_cache, current_start, cache_start, ) attn_output = attn_output.flatten(2) attn_output, _ = self.to_out(attn_output) attn_output = attn_output.squeeze(1) null_shift = null_scale = torch.zeroes( (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, crossattn_cache=crossattn_cache, ) 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 CausalWanTransformer3DModel(BaseDiT, 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.attention_head_dim = config.attention_head_dim 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 self.local_attn_size = config.local_attn_size # 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 self.blocks = nn.ModuleList( [ CausalWanTransformerBlock( inner_dim, config.ffn_dim, config.num_attention_heads, config.local_attn_size, config.sink_size, config.qk_norm, config.cross_attn_norm, config.eps, config.added_kv_proj_dim, self._supported_attention_backends, prefix=f"{config.prefix}.blocks.{i}", ) 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 = nn.Linear( inner_dim, config.out_channels * math.prod(config.patch_size) ) self.scale_shift_table = nn.Parameter( torch.randn(1, 2, inner_dim) / inner_dim**0.5 ) self.gradient_checkpointing = False # Causal-specific self.block_mask = None self.num_frame_per_block = config.arch_config.num_frames_per_block assert self.num_frame_per_block <= 3 self.independent_first_frame = False self.__post_init__() self.layer_names = [ "blocks", ] @staticmethod def _prepare_blockwise_causal_attn_mask( device: torch.device | str, num_frames: int = 21, frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1, ) -> BlockMask: """ we will divide the token sequence into the following format [1 latent frame] [1 latent frame] ... [1 latent frame] We use flexattention to construct the attention mask """ total_length = num_frames * frame_seqlen # we do right padding to get to a multiple of 128 padded_length = math.ceil(total_length / 128) * 128 - total_length ends = torch.zeros( total_length + padded_length, device=device, dtype=torch.long ) # Block-wise causal mask will attend to all elements that are before the end of the current chunk frame_indices = torch.arange( start=0, end=total_length, step=frame_seqlen * num_frame_per_block, device=device, ) for tmp in frame_indices: ends[tmp : tmp + frame_seqlen * num_frame_per_block] = ( tmp + frame_seqlen * num_frame_per_block ) def attention_mask(b, h, q_idx, kv_idx): if local_attn_size == -1: return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) else: return ( (kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen)) ) | (q_idx == kv_idx) # return ((kv_idx < total_length) & (q_idx < total_length)) | (q_idx == kv_idx) # bidirectional mask block_mask = create_block_mask( attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, KV_LEN=total_length + padded_length, _compile=False, device=device, ) if not dist.is_initialized() or dist.get_rank() == 0: print( f" cache a block wise causal mask with block size of {num_frame_per_block} frames" ) print(block_mask) # import imageio # import numpy as np # from torch.nn.attention.flex_attention import create_mask # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length + # padded_length, KV_LEN=total_length + padded_length, device=device) # import cv2 # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask)) return block_mask def _forward_inference( 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, kv_cache: dict = None, crossattn_cache: dict = None, current_start: int = 0, cache_start: int = 0, start_frame: int = 0, **kwargs, ) -> torch.Tensor: r""" Run the diffusion model with kv caching. See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details. This function will be run for num_frame times. Process the latent frames one by one (1560 tokens each) """ 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 # Get rotary embeddings d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] freqs_cos, freqs_sin = get_rotary_pos_embed( ( post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width, ), self.hidden_size, self.num_attention_heads, rope_dim_list, dtype=( torch.float32 if current_platform.is_mps() or current_platform.is_musa() else torch.float64 ), rope_theta=10000, start_frame=start_frame, # Assume that start_frame is 0 when kv_cache is None ) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(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) temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = ( self.condition_embedder( timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image ) ) timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten( dim=0, sizes=timestep.shape ) 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 current_platform.is_mps() else encoder_hidden_states ) # cast to orig_dtype for MPS assert encoder_hidden_states.dtype == orig_dtype # 4. Transformer blocks for block_index, block in enumerate(self.blocks): if torch.is_grad_enabled() and self.gradient_checkpointing: causal_kwargs = { "kv_cache": kv_cache[block_index], "current_start": current_start, "cache_start": cache_start, "block_mask": self.block_mask, } hidden_states = self._gradient_checkpointing_func( block, hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, **causal_kwargs, ) else: causal_kwargs = { "kv_cache": kv_cache[block_index], "crossattn_cache": crossattn_cache[block_index], "current_start": current_start, "cache_start": cache_start, "block_mask": self.block_mask, } hidden_states = block( hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, **causal_kwargs, ) # 5. Output norm, projection & unpatchify temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) 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 _forward_train( 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, start_frame: int = 0, **kwargs, ) -> torch.Tensor: 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 # Get rotary embeddings d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] freqs_cos, freqs_sin = get_rotary_pos_embed( ( post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width, ), self.hidden_size, self.num_attention_heads, rope_dim_list, dtype=( torch.float32 if current_platform.is_mps() or current_platform.is_musa() else torch.float64 ), rope_theta=10000, start_frame=start_frame, ) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) freqs_cis = ( (freqs_cos.float(), freqs_sin.float()) if freqs_cos is not None else None ) # Construct blockwise causal attn mask if self.block_mask is None: self.block_mask = self._prepare_blockwise_causal_attn_mask( device=hidden_states.device, num_frames=num_frames, frame_seqlen=post_patch_height * post_patch_width, num_frame_per_block=self.num_frame_per_block, local_attn_size=self.local_attn_size, ) hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = ( self.condition_embedder( timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image ) ) timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten( dim=0, sizes=timestep.shape ) 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 current_platform.is_mps() else encoder_hidden_states ) # cast to orig_dtype for MPS assert encoder_hidden_states.dtype == orig_dtype # 4. Transformer blocks if torch.is_grad_enabled() and self.gradient_checkpointing: for block in self.blocks: hidden_states = self._gradient_checkpointing_func( block, hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, block_mask=self.block_mask, ) else: for block in self.blocks: hidden_states = block( hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, block_mask=self.block_mask, ) # 5. Output norm, projection & unpatchify temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) 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 forward(self, *args, **kwargs): if kwargs.get("kv_cache") is not None: return self._forward_inference(*args, **kwargs) else: return self._forward_train(*args, **kwargs) EntryClass = CausalWanTransformer3DModel