DFAGWE/infinitetalk2
0
1# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.2import math3 4import torch5import torch.cuda.amp as amp6import torch.nn as nn7from diffusers.configuration_utils import ConfigMixin, register_to_config8from diffusers.models.modeling_utils import ModelMixin9 10from .attention import flash_attention11 12__all__ = ['WanModel']13 14T5_CONTEXT_TOKEN_NUMBER = 51215FIRST_LAST_FRAME_CONTEXT_TOKEN_NUMBER = 257 * 216 17 18def sinusoidal_embedding_1d(dim, position):19 # preprocess20 assert dim % 2 == 021 half = dim // 222 position = position.type(torch.float64)23 24 # calculation25 sinusoid = torch.outer(26 position, torch.pow(10000, -torch.arange(half).to(position).div(half)))27 x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)28 return x29 30 31@amp.autocast(enabled=False)32def rope_params(max_seq_len, dim, theta=10000):33 assert dim % 2 == 034 freqs = torch.outer(35 torch.arange(max_seq_len),36 1.0 / torch.pow(theta,37 torch.arange(0, dim, 2).to(torch.float64).div(dim)))38 freqs = torch.polar(torch.ones_like(freqs), freqs)39 return freqs40 41 42@amp.autocast(enabled=False)43def rope_apply(x, grid_sizes, freqs):44 n, c = x.size(2), x.size(3) // 245 46 # split freqs47 freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)48 49 # loop over samples50 output = []51 for i, (f, h, w) in enumerate(grid_sizes.tolist()):52 seq_len = f * h * w53 54 # precompute multipliers55 x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(56 seq_len, n, -1, 2))57 freqs_i = torch.cat([58 freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),59 freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),60 freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)61 ],62 dim=-1).reshape(seq_len, 1, -1)63 64 # apply rotary embedding65 x_i = torch.view_as_real(x_i * freqs_i).flatten(2)66 x_i = torch.cat([x_i, x[i, seq_len:]])67 68 # append to collection69 output.append(x_i)70 return torch.stack(output).float()71 72 73class WanRMSNorm(nn.Module):74 75 def __init__(self, dim, eps=1e-5):76 super().__init__()77 self.dim = dim78 self.eps = eps79 self.weight = nn.Parameter(torch.ones(dim))80 81 def forward(self, x):82 r"""83 Args:84 x(Tensor): Shape [B, L, C]85 """86 return self._norm(x.float()).type_as(x) * self.weight87 88 def _norm(self, x):89 return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)90 91 92class WanLayerNorm(nn.LayerNorm):93 94 def __init__(self, dim, eps=1e-6, elementwise_affine=False):95 super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)96 97 def forward(self, x):98 r"""99 Args:100 x(Tensor): Shape [B, L, C]101 """102 return super().forward(x.float()).type_as(x)103 104 105class WanSelfAttention(nn.Module):106 107 def __init__(self,108 dim,109 num_heads,110 window_size=(-1, -1),111 qk_norm=True,112 eps=1e-6):113 assert dim % num_heads == 0114 super().__init__()115 self.dim = dim116 self.num_heads = num_heads117 self.head_dim = dim // num_heads118 self.window_size = window_size119 self.qk_norm = qk_norm120 self.eps = eps121 122 # layers123 self.q = nn.Linear(dim, dim)124 self.k = nn.Linear(dim, dim)125 self.v = nn.Linear(dim, dim)126 self.o = nn.Linear(dim, dim)127 self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()128 self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()129 130 def forward(self, x, seq_lens, grid_sizes, freqs):131 r"""132 Args:133 x(Tensor): Shape [B, L, num_heads, C / num_heads]134 seq_lens(Tensor): Shape [B]135 grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)136 freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]137 """138 b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim139 140 # query, key, value function141 def qkv_fn(x):142 q = self.norm_q(self.q(x)).view(b, s, n, d)143 k = self.norm_k(self.k(x)).view(b, s, n, d)144 v = self.v(x).view(b, s, n, d)145 return q, k, v146 147 q, k, v = qkv_fn(x)148 149 x = flash_attention(150 q=rope_apply(q, grid_sizes, freqs),151 k=rope_apply(k, grid_sizes, freqs),152 v=v,153 k_lens=seq_lens,154 window_size=self.window_size)155 156 # output157 x = x.flatten(2)158 x = self.o(x)159 return x160 161 162class WanT2VCrossAttention(WanSelfAttention):163 164 def forward(self, x, context, context_lens):165 r"""166 Args:167 x(Tensor): Shape [B, L1, C]168 context(Tensor): Shape [B, L2, C]169 context_lens(Tensor): Shape [B]170 """171 b, n, d = x.size(0), self.num_heads, self.head_dim172 173 # compute query, key, value174 q = self.norm_q(self.q(x)).view(b, -1, n, d)175 k = self.norm_k(self.k(context)).view(b, -1, n, d)176 v = self.v(context).view(b, -1, n, d)177 178 # compute attention179 x = flash_attention(q, k, v, k_lens=context_lens)180 181 # output182 x = x.flatten(2)183 x = self.o(x)184 return x185 186 187class WanI2VCrossAttention(WanSelfAttention):188 189 def __init__(self,190 dim,191 num_heads,192 window_size=(-1, -1),193 qk_norm=True,194 eps=1e-6):195 super().__init__(dim, num_heads, window_size, qk_norm, eps)196 197 self.k_img = nn.Linear(dim, dim)198 self.v_img = nn.Linear(dim, dim)199 # self.alpha = nn.Parameter(torch.zeros((1, )))200 self.norm_k_img = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()201 202 def forward(self, x, context, context_lens):203 r"""204 Args:205 x(Tensor): Shape [B, L1, C]206 context(Tensor): Shape [B, L2, C]207 context_lens(Tensor): Shape [B]208 """209 image_context_length = context.shape[1] - T5_CONTEXT_TOKEN_NUMBER210 context_img = context[:, :image_context_length]211 context = context[:, image_context_length:]212 b, n, d = x.size(0), self.num_heads, self.head_dim213 214 # compute query, key, value215 q = self.norm_q(self.q(x)).view(b, -1, n, d)216 k = self.norm_k(self.k(context)).view(b, -1, n, d)217 v = self.v(context).view(b, -1, n, d)218 k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d)219 v_img = self.v_img(context_img).view(b, -1, n, d)220 img_x = flash_attention(q, k_img, v_img, k_lens=None)221 # compute attention222 x = flash_attention(q, k, v, k_lens=context_lens)223 224 # output225 x = x.flatten(2)226 img_x = img_x.flatten(2)227 x = x + img_x228 x = self.o(x)229 return x230 231 232WAN_CROSSATTENTION_CLASSES = {233 't2v_cross_attn': WanT2VCrossAttention,234 'i2v_cross_attn': WanI2VCrossAttention,235}236 237 238class WanAttentionBlock(nn.Module):239 240 def __init__(self,241 cross_attn_type,242 dim,243 ffn_dim,244 num_heads,245 window_size=(-1, -1),246 qk_norm=True,247 cross_attn_norm=False,248 eps=1e-6):249 super().__init__()250 self.dim = dim251 self.ffn_dim = ffn_dim252 self.num_heads = num_heads253 self.window_size = window_size254 self.qk_norm = qk_norm255 self.cross_attn_norm = cross_attn_norm256 self.eps = eps257 258 # layers259 self.norm1 = WanLayerNorm(dim, eps)260 self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm,261 eps)262 self.norm3 = WanLayerNorm(263 dim, eps,264 elementwise_affine=True) if cross_attn_norm else nn.Identity()265 self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim,266 num_heads,267 (-1, -1),268 qk_norm,269 eps)270 self.norm2 = WanLayerNorm(dim, eps)271 self.ffn = nn.Sequential(272 nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),273 nn.Linear(ffn_dim, dim))274 275 # modulation276 self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)277 278 def forward(279 self,280 x,281 e,282 seq_lens,283 grid_sizes,284 freqs,285 context,286 context_lens,287 ):288 r"""289 Args:290 x(Tensor): Shape [B, L, C]291 e(Tensor): Shape [B, 6, C]292 seq_lens(Tensor): Shape [B], length of each sequence in batch293 grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)294 freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]295 """296 assert e.dtype == torch.float32297 with amp.autocast(dtype=torch.float32):298 e = (self.modulation.to(e.device) + e).chunk(6, dim=1)299 assert e[0].dtype == torch.float32300 301 # self-attention302 y = self.self_attn(303 self.norm1(x).float() * (1 + e[1]) + e[0], seq_lens, grid_sizes,304 freqs)305 with amp.autocast(dtype=torch.float32):306 x = x + y * e[2]307 308 # cross-attention & ffn function309 def cross_attn_ffn(x, context, context_lens, e):310 x = x + self.cross_attn(self.norm3(x), context, context_lens)311 y = self.ffn(self.norm2(x).float() * (1 + e[4]) + e[3])312 with amp.autocast(dtype=torch.float32):313 x = x + y * e[5]314 return x315 316 x = cross_attn_ffn(x, context, context_lens, e)317 return x318 319 320class Head(nn.Module):321 322 def __init__(self, dim, out_dim, patch_size, eps=1e-6):323 super().__init__()324 self.dim = dim325 self.out_dim = out_dim326 self.patch_size = patch_size327 self.eps = eps328 329 # layers330 out_dim = math.prod(patch_size) * out_dim331 self.norm = WanLayerNorm(dim, eps)332 self.head = nn.Linear(dim, out_dim)333 334 # modulation335 self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)336 337 def forward(self, x, e):338 r"""339 Args:340 x(Tensor): Shape [B, L1, C]341 e(Tensor): Shape [B, C]342 """343 assert e.dtype == torch.float32344 with amp.autocast(dtype=torch.float32):345 e = (self.modulation.to(e.device) + e.unsqueeze(1)).chunk(2, dim=1)346 x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))347 return x348 349 350class MLPProj(torch.nn.Module):351 352 def __init__(self, in_dim, out_dim, flf_pos_emb=False):353 super().__init__()354 355 self.proj = torch.nn.Sequential(356 torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim),357 torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim),358 torch.nn.LayerNorm(out_dim))359 if flf_pos_emb: # NOTE: we only use this for `flf2v`360 self.emb_pos = nn.Parameter(361 torch.zeros(1, FIRST_LAST_FRAME_CONTEXT_TOKEN_NUMBER, 1280))362 363 def forward(self, image_embeds):364 if hasattr(self, 'emb_pos'):365 bs, n, d = image_embeds.shape366 image_embeds = image_embeds.view(-1, 2 * n, d)367 image_embeds = image_embeds + self.emb_pos368 clip_extra_context_tokens = self.proj(image_embeds)369 return clip_extra_context_tokens370 371 372class WanModel(ModelMixin, ConfigMixin):373 r"""374 Wan diffusion backbone supporting both text-to-video and image-to-video.375 """376 377 ignore_for_config = [378 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size'379 ]380 _no_split_modules = ['WanAttentionBlock']381 382 @register_to_config383 def __init__(self,384 model_type='t2v',385 patch_size=(1, 2, 2),386 text_len=512,387 in_dim=16,388 dim=2048,389 ffn_dim=8192,390 freq_dim=256,391 text_dim=4096,392 out_dim=16,393 num_heads=16,394 num_layers=32,395 window_size=(-1, -1),396 qk_norm=True,397 cross_attn_norm=True,398 eps=1e-6):399 r"""400 Initialize the diffusion model backbone.401 402 Args:403 model_type (`str`, *optional*, defaults to 't2v'):404 Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video) or 'flf2v' (first-last-frame-to-video) or 'vace'405 patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):406 3D patch dimensions for video embedding (t_patch, h_patch, w_patch)407 text_len (`int`, *optional*, defaults to 512):408 Fixed length for text embeddings409 in_dim (`int`, *optional*, defaults to 16):410 Input video channels (C_in)411 dim (`int`, *optional*, defaults to 2048):412 Hidden dimension of the transformer413 ffn_dim (`int`, *optional*, defaults to 8192):414 Intermediate dimension in feed-forward network415 freq_dim (`int`, *optional*, defaults to 256):416 Dimension for sinusoidal time embeddings417 text_dim (`int`, *optional*, defaults to 4096):418 Input dimension for text embeddings419 out_dim (`int`, *optional*, defaults to 16):420 Output video channels (C_out)421 num_heads (`int`, *optional*, defaults to 16):422 Number of attention heads423 num_layers (`int`, *optional*, defaults to 32):424 Number of transformer blocks425 window_size (`tuple`, *optional*, defaults to (-1, -1)):426 Window size for local attention (-1 indicates global attention)427 qk_norm (`bool`, *optional*, defaults to True):428 Enable query/key normalization429 cross_attn_norm (`bool`, *optional*, defaults to False):430 Enable cross-attention normalization431 eps (`float`, *optional*, defaults to 1e-6):432 Epsilon value for normalization layers433 """434 435 super().__init__()436 437 assert model_type in ['t2v', 'i2v', 'flf2v', 'vace']438 self.model_type = model_type439 440 self.patch_size = patch_size441 self.text_len = text_len442 self.in_dim = in_dim443 self.dim = dim444 self.ffn_dim = ffn_dim445 self.freq_dim = freq_dim446 self.text_dim = text_dim447 self.out_dim = out_dim448 self.num_heads = num_heads449 self.num_layers = num_layers450 self.window_size = window_size451 self.qk_norm = qk_norm452 self.cross_attn_norm = cross_attn_norm453 self.eps = eps454 455 # embeddings456 self.patch_embedding = nn.Conv3d(457 in_dim, dim, kernel_size=patch_size, stride=patch_size)458 self.text_embedding = nn.Sequential(459 nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),460 nn.Linear(dim, dim))461 462 self.time_embedding = nn.Sequential(463 nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))464 self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))465 466 # blocks467 cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'468 self.blocks = nn.ModuleList([469 WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,470 window_size, qk_norm, cross_attn_norm, eps)471 for _ in range(num_layers)472 ])473 474 # head475 self.head = Head(dim, out_dim, patch_size, eps)476 477 # buffers (don't use register_buffer otherwise dtype will be changed in to())478 assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0479 d = dim // num_heads480 self.freqs = torch.cat([481 rope_params(1024, d - 4 * (d // 6)),482 rope_params(1024, 2 * (d // 6)),483 rope_params(1024, 2 * (d // 6))484 ],485 dim=1)486 487 if model_type == 'i2v' or model_type == 'flf2v':488 self.img_emb = MLPProj(1280, dim, flf_pos_emb=model_type == 'flf2v')489 490 # initialize weights491 self.init_weights()492 493 def forward(494 self,495 x,496 t,497 context,498 seq_len,499 clip_fea=None,500 y=None,501 ):502 r"""503 Forward pass through the diffusion model504 505 Args:506 x (List[Tensor]):507 List of input video tensors, each with shape [C_in, F, H, W]508 t (Tensor):509 Diffusion timesteps tensor of shape [B]510 context (List[Tensor]):511 List of text embeddings each with shape [L, C]512 seq_len (`int`):513 Maximum sequence length for positional encoding514 clip_fea (Tensor, *optional*):515 CLIP image features for image-to-video mode or first-last-frame-to-video mode516 y (List[Tensor], *optional*):517 Conditional video inputs for image-to-video mode, same shape as x518 519 Returns:520 List[Tensor]:521 List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]522 """523 if self.model_type == 'i2v' or self.model_type == 'flf2v':524 assert clip_fea is not None and y is not None525 # params526 device = self.patch_embedding.weight.device527 if self.freqs.device != device:528 self.freqs = self.freqs.to(device)529 530 if y is not None:531 x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]532 533 # embeddings534 x = [self.patch_embedding(u.unsqueeze(0)) for u in x]535 grid_sizes = torch.stack(536 [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])537 x = [u.flatten(2).transpose(1, 2) for u in x]538 seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)539 assert seq_lens.max() <= seq_len540 x = torch.cat([541 torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],542 dim=1) for u in x543 ])544 545 # time embeddings546 with amp.autocast(dtype=torch.float32):547 e = self.time_embedding(548 sinusoidal_embedding_1d(self.freq_dim, t).float())549 e0 = self.time_projection(e).unflatten(1, (6, self.dim))550 assert e.dtype == torch.float32 and e0.dtype == torch.float32551 552 # context553 context_lens = None554 context = self.text_embedding(555 torch.stack([556 torch.cat(557 [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])558 for u in context559 ]))560 561 if clip_fea is not None:562 context_clip = self.img_emb(clip_fea) # bs x 257 (x2) x dim563 context = torch.concat([context_clip, context], dim=1)564 565 # arguments566 kwargs = dict(567 e=e0,568 seq_lens=seq_lens,569 grid_sizes=grid_sizes,570 freqs=self.freqs,571 context=context,572 context_lens=context_lens)573 574 for block in self.blocks:575 x = block(x, **kwargs)576 577 # head578 x = self.head(x, e)579 580 # unpatchify581 x = self.unpatchify(x, grid_sizes)582 return [u.float() for u in x]583 584 def unpatchify(self, x, grid_sizes):585 r"""586 Reconstruct video tensors from patch embeddings.587 588 Args:589 x (List[Tensor]):590 List of patchified features, each with shape [L, C_out * prod(patch_size)]591 grid_sizes (Tensor):592 Original spatial-temporal grid dimensions before patching,593 shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)594 595 Returns:596 List[Tensor]:597 Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]598 """599 600 c = self.out_dim601 out = []602 for u, v in zip(x, grid_sizes.tolist()):603 u = u[:math.prod(v)].view(*v, *self.patch_size, c)604 u = torch.einsum('fhwpqrc->cfphqwr', u)605 u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])606 out.append(u)607 return out608 609 def init_weights(self):610 r"""611 Initialize model parameters using Xavier initialization.612 """613 614 # basic init615 for m in self.modules():616 if isinstance(m, nn.Linear):617 nn.init.xavier_uniform_(m.weight)618 if m.bias is not None:619 nn.init.zeros_(m.bias)620 621 # init embeddings622 nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))623 for m in self.text_embedding.modules():624 if isinstance(m, nn.Linear):625 nn.init.normal_(m.weight, std=.02)626 for m in self.time_embedding.modules():627 if isinstance(m, nn.Linear):628 nn.init.normal_(m.weight, std=.02)629 630 # init output layer631 nn.init.zeros_(self.head.head.weight)632