From e976583cf47e1e89d6eb67253add2d2e293a1587 Mon Sep 17 00:00:00 2001 From: "Zhou, Wei" Date: Fri, 14 Mar 2025 20:00:01 -0400 Subject: [PATCH] Add wan dit --- fastvideo/v1/layers/rotary_embedding.py | 2 +- fastvideo/v1/layers/visual_embedding.py | 6 +- fastvideo/v1/models/dits/wanvideo.py | 456 ++++++++++++++++++++++++ fastvideo/v1/models/registry.py | 2 + 4 files changed, 462 insertions(+), 4 deletions(-) create mode 100644 fastvideo/v1/models/dits/wanvideo.py diff --git a/fastvideo/v1/layers/rotary_embedding.py b/fastvideo/v1/layers/rotary_embedding.py index f557ba1ae..252f29f7d 100644 --- a/fastvideo/v1/layers/rotary_embedding.py +++ b/fastvideo/v1/layers/rotary_embedding.py @@ -349,7 +349,7 @@ def get_1d_rotary_pos_embed( if theta_rescale_factor != 1.0: theta *= theta_rescale_factor**(dim / (dim - 2)) - freqs = 1.0 / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].float() / dim)) # [D/2] + freqs = 1.0 / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].to(torch.float64) / dim)) # [D/2] freqs = torch.outer(pos * interpolation_factor, freqs) # [S, D/2] freqs_cos = freqs.cos() # [S, D/2] freqs_sin = freqs.sin() # [S, D/2] diff --git a/fastvideo/v1/layers/visual_embedding.py b/fastvideo/v1/layers/visual_embedding.py index 13a61aef1..31324644c 100644 --- a/fastvideo/v1/layers/visual_embedding.py +++ b/fastvideo/v1/layers/visual_embedding.py @@ -87,8 +87,8 @@ class TimestepEmbedder(nn.Module): ) def forward(self, t): - t_freq = timestep_embedding(t, self.frequency_embedding_size, self.max_period) - t_freq = t_freq.to(self.mlp.fc_in.weight.dtype) + t_freq = timestep_embedding(t, self.frequency_embedding_size, self.max_period).float() + # t_freq = t_freq.to(self.mlp.fc_in.weight.dtype) t_emb = self.mlp(t_freq) return t_emb @@ -106,7 +106,7 @@ def timestep_embedding(t, dim, max_period=10000): Tensor of shape [B, dim] with embeddings """ half = dim // 2 - freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(device=t.device) + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float64) / half).to(device=t.device) args = t[:, None].float() * freqs[None] embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) if dim % 2: diff --git a/fastvideo/v1/models/dits/wanvideo.py b/fastvideo/v1/models/dits/wanvideo.py new file mode 100644 index 000000000..aef8f0445 --- /dev/null +++ b/fastvideo/v1/models/dits/wanvideo.py @@ -0,0 +1,456 @@ +import torch +import torch.nn as nn +import math +from typing import Optional, Tuple, List, Union, Dict, Any +from fastvideo.v1.attention.flash_attn import DistributedAttention, LocalAttention +from fastvideo.v1.layers.linear import ReplicatedLinear +from fastvideo.v1.layers.layernorm import LayerNormScaleShift, ScaleResidual, ScaleResidualLayerNormScaleShift, RMSNorm +from fastvideo.v1.layers.visual_embedding import PatchEmbed, TimestepEmbedder, ModulateProjection +from fastvideo.v1.layers.rotary_embedding import _apply_rotary_emb, get_rotary_pos_embed +from fastvideo.v1.distributed.parallel_state import get_sequence_model_parallel_world_size +# from torch.nn import RMSNorm +# TODO: RMSNorm .... +from fastvideo.v1.layers.mlp import MLP +from fastvideo.v1.models.dits.base import BaseDiT + +class WanImageEmbedding(torch.nn.Module): + def __init__(self, in_features: int, out_features: int): + super().__init__() + + self.norm1 = nn.LayerNorm(in_features) + self.ff = MLP(in_features, in_features, out_features, act_type="gelu") + self.norm2 = nn.LayerNorm(out_features) + + def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: + hidden_states = self.norm1(encoder_hidden_states_image) + hidden_states = self.ff(hidden_states) + hidden_states = self.norm2(hidden_states) + return hidden_states + +class WanTimeTextImageEmbedding(nn.Module): + def __init__( + self, + dim: int, + time_freq_dim: int, + text_embed_dim: int, + image_embed_dim: Optional[int] = 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: Optional[torch.Tensor] = None, + ): + with torch.cuda.amp.autocast(dtype=torch.float32): + temb = self.time_embedder(timestep.float()) + timestep_proj = self.time_modulation(temb) + + encoder_hidden_states = self.text_embedder(encoder_hidden_states) + if encoder_hidden_states_image 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, + num_heads, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6, + parallel_attention=False): + 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 + + # layers + self.to_q = ReplicatedLinear(dim, dim) + self.to_k = ReplicatedLinear(dim, dim) + self.to_v = ReplicatedLinear(dim, dim) + self.to_out = ReplicatedLinear(dim, dim) + 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() + + # Scaled dot product attention + self.attn = LocalAttention(dropout_rate=0, softmax_scale=None, causal=False) + + def forward(self, x, seq_lens, grid_sizes, freqs): + 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] + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q.forward_native(self.to_q(x)[0]).view(b, -1, n, d) + k = self.norm_k.forward_native(self.to_k(context)[0]).view(b, -1, n, d) + v = self.to_v(context)[0].view(b, -1, n, d) + + # 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, + num_heads, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6): + super().__init__(dim, num_heads, window_size, qk_norm, eps) + + self.add_k_proj = ReplicatedLinear(dim, dim) + self.add_v_proj = ReplicatedLinear(dim, dim) + self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_added_q = 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:] + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q.forward_native(self.to_q(x)[0]).view(b, -1, n, d) + k = self.norm_k.forward_native(self.to_k(context)[0]).view(b, -1, n, d) + v = self.to_v(context)[0].view(b, -1, n, d) + k_img = self.norm_added_k.forward_native(self.add_k_proj(context_img)[0]).view(b, -1, n, d) + v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d) + img_x = self.attn(q, k_img, v_img) + # compute attention + 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: Optional[int] = None, + ): + super().__init__() + + # 1. Self-attention + self.norm1 = nn.LayerNorm(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 = DistributedAttention( + dropout_rate=0.0, + causal=False + ) + 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: + print("QK Norm type not supported") + raise Exception + assert cross_attn_norm is True + self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, norm_type="layer", eps=eps, elementwise_affine=True, dtype=torch.float32) + + # 2. Cross-attention + if added_kv_proj_dim is not None: + # I2V + self.attn2 = WanI2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps) + else: + # T2V + self.attn2 = WanT2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps) + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim, norm_type="layer", 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 = ScaleResidual() + + 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] = None, + ) -> 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 temb.dtype == torch.float32 + with torch.cuda.amp.autocast(dtype=torch.float32): + e = self.scale_shift_table + temb + 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.float()).to(dtype=orig_dtype) * (1 + scale_msa) + shift_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: + query = self.norm_q.forward_native(query) + if self.norm_k is not None: + key = self.norm_k.forward_native(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)) + # Apply rotary embeddings + cos, sin = freqs_cis + 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.tensor([0], device=hidden_states.device) + norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa, null_shift, null_scale) + + # 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) + + # 3. Feed-forward + ff_output = self.ffn(norm_hidden_states) + hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) + + return hidden_states + +class WanTransformer3DModel(BaseDiT): + _fsdp_shard_conditions = [ + lambda n, m: "blocks" in n and str.isdigit(n.split(".")[-1]), + ] + _param_names_mapping = { + r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1", + + r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1", + r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1", + r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1", + r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1", + r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1", + r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$": r"condition_embedder.image_embedder.ff.fc_in.\1", + r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$": r"condition_embedder.image_embedder.ff.fc_out.\1", + + r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$": r"blocks.\1.to_q.\2", + r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$": r"blocks.\1.to_k.\2", + r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$": r"blocks.\1.to_v.\2", + r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": r"blocks.\1.to_out.\2", + r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2", + r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2", + + r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"blocks.\1.attn2.to_out.\2", + + r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": r"blocks.\1.ffn.fc_in.\2", + r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", + + r"blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2", + } + def __init__( + self, + patch_size: Tuple[int] = (1, 2, 2), + text_len = 512, + num_attention_heads: int = 40, + attention_head_dim: int = 128, + in_channels: int = 16, + out_channels: int = 16, + text_dim: int = 4096, + freq_dim: int = 256, + ffn_dim: int = 13824, + num_layers: int = 40, + cross_attn_norm: bool = True, + qk_norm: Optional[str] = "rms_norm_across_heads", + eps: float = 1e-6, + image_dim: Optional[int] = None, + added_kv_proj_dim: Optional[int] = None, + rope_max_seq_len: int = 1024, + ) -> None: + super().__init__() + + inner_dim = num_attention_heads * attention_head_dim + self.inner_dim = inner_dim + self.num_attention_heads = num_attention_heads + self.out_channels = out_channels or in_channels + self.patch_size = patch_size + self.text_len = text_len + + # 1. Patch & position embedding + self.patch_embedding = PatchEmbed(in_chans=in_channels, embed_dim=inner_dim, patch_size=patch_size, flatten=False) + + # 2. Condition embeddings + self.condition_embedder = WanTimeTextImageEmbedding( + dim=inner_dim, + time_freq_dim=freq_dim, + text_embed_dim=text_dim, + image_embed_dim=image_dim, + ) + + # 3. Transformer blocks + self.blocks = nn.ModuleList( + [ + WanTransformerBlock( + inner_dim, ffn_dim, num_attention_heads, qk_norm, cross_attn_norm, eps, added_kv_proj_dim + ) + for _ in range(num_layers) + ] + ) + + # 4. Output norm & projection + self.norm_out = LayerNormScaleShift(inner_dim, norm_type="layer", eps=eps, elementwise_affine=False, dtype=torch.float32) + self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.LongTensor, + encoder_hidden_states: torch.Tensor, + seq_len: Optional[int] = None, + encoder_hidden_states_image: Optional[torch.Tensor] = None, + y: Optional[torch.Tensor] = None, + return_dict: bool = True, + attention_kwargs: Optional[Dict[str, Any]] = None, + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + if y is not None: + hidden_states = torch.cat([hidden_states, y], dim=1) + + 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.inner_dim // 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_sequence_model_parallel_world_size(), post_patch_height, post_patch_width), self.inner_dim, self.num_attention_heads, rope_dim_list, rope_theta=10000) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + + hidden_states = self.patch_embedding(hidden_states) + grid_sizes = torch.stack( + [torch.tensor(hidden_states[0].shape[1:], dtype=torch.long)]) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + if seq_len is None: + seq_len = hidden_states.size(1) + hidden_states = torch.cat([hidden_states, hidden_states.new_zeros(1, seq_len - hidden_states.size(1), hidden_states.size(2))], dim=1) + + encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1) + + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( + timestep, encoder_hidden_states, encoder_hidden_states_image + ) + 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) + + # 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 + ) + else: + for block in self.blocks: + hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis) + + # 5. Output norm, projection & unpatchify + with torch.cuda.amp.autocast(dtype=torch.float32): + shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states.float(), shift, scale) + hidden_states = self.proj_out(hidden_states) + + output = self.unpatchify(hidden_states, grid_sizes) + + return output.float() + + def unpatchify(self, x, grid_sizes): + r""" + Reconstruct video tensors from patch embeddings. + + Args: + x (List[Tensor]): + List of patchified features, each with shape [L, C_out * prod(patch_size)] + grid_sizes (Tensor): + Original spatial-temporal grid dimensions before patching, + shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches) + + Returns: + Tensor: + Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8] + """ + + c = self.out_channels + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[:math.prod(v)].view(*v, *self.patch_size, c) + u = u.permute(6, 0, 3, 1, 4, 2, 5) + # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous()) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + out = torch.cat(out, dim=0) + return out \ No newline at end of file diff --git a/fastvideo/v1/models/registry.py b/fastvideo/v1/models/registry.py index b665a6ccc..4537a7de2 100644 --- a/fastvideo/v1/models/registry.py +++ b/fastvideo/v1/models/registry.py @@ -16,10 +16,12 @@ from fastvideo.v1.logger import logger # huggingface class name: (component_name, fastvideo module name, fastvideo class name) _TEXT_TO_VIDEO_DIT_MODELS = { "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"), + "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"), } _IMAGE_TO_VIDEO_DIT_MODELS = { # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"), + "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"), } _TEXT_ENCODER_MODELS = {