ALSv/self-forcing
0
1# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.2import math3 4import torch5import torch.nn as nn6from diffusers.configuration_utils import ConfigMixin, register_to_config7from diffusers.models.modeling_utils import ModelMixin8from einops import repeat9 10from .attention import flash_attention11 12__all__ = ['WanModel']13 14 15def sinusoidal_embedding_1d(dim, position):16 # preprocess17 assert dim % 2 == 018 half = dim // 219 position = position.type(torch.float64)20 21 # calculation22 sinusoid = torch.outer(23 position, torch.pow(10000, -torch.arange(half).to(position).div(half)))24 x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)25 return x26 27 28# @amp.autocast(enabled=False)29def rope_params(max_seq_len, dim, theta=10000):30 assert dim % 2 == 031 freqs = torch.outer(32 torch.arange(max_seq_len),33 1.0 / torch.pow(theta,34 torch.arange(0, dim, 2).to(torch.float64).div(dim)))35 freqs = torch.polar(torch.ones_like(freqs), freqs)36 return freqs37 38 39# @amp.autocast(enabled=False)40def rope_apply(x, grid_sizes, freqs):41 n, c = x.size(2), x.size(3) // 242 43 # split freqs44 freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)45 46 # loop over samples47 output = []48 for i, (f, h, w) in enumerate(grid_sizes.tolist()):49 seq_len = f * h * w50 51 # precompute multipliers52 x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(53 seq_len, n, -1, 2))54 freqs_i = torch.cat([55 freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),56 freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),57 freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)58 ],59 dim=-1).reshape(seq_len, 1, -1)60 61 # apply rotary embedding62 x_i = torch.view_as_real(x_i * freqs_i).flatten(2)63 x_i = torch.cat([x_i, x[i, seq_len:]])64 65 # append to collection66 output.append(x_i)67 return torch.stack(output).type_as(x)68 69 70class WanRMSNorm(nn.Module):71 72 def __init__(self, dim, eps=1e-5):73 super().__init__()74 self.dim = dim75 self.eps = eps76 self.weight = nn.Parameter(torch.ones(dim))77 78 def forward(self, x):79 r"""80 Args:81 x(Tensor): Shape [B, L, C]82 """83 return self._norm(x.float()).type_as(x) * self.weight84 85 def _norm(self, x):86 return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)87 88 89class WanLayerNorm(nn.LayerNorm):90 91 def __init__(self, dim, eps=1e-6, elementwise_affine=False):92 super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)93 94 def forward(self, x):95 r"""96 Args:97 x(Tensor): Shape [B, L, C]98 """99 return super().forward(x).type_as(x)100 101 102class WanSelfAttention(nn.Module):103 104 def __init__(self,105 dim,106 num_heads,107 window_size=(-1, -1),108 qk_norm=True,109 eps=1e-6):110 assert dim % num_heads == 0111 super().__init__()112 self.dim = dim113 self.num_heads = num_heads114 self.head_dim = dim // num_heads115 self.window_size = window_size116 self.qk_norm = qk_norm117 self.eps = eps118 119 # layers120 self.q = nn.Linear(dim, dim)121 self.k = nn.Linear(dim, dim)122 self.v = nn.Linear(dim, dim)123 self.o = nn.Linear(dim, dim)124 self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()125 self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()126 127 def forward(self, x, seq_lens, grid_sizes, freqs):128 r"""129 Args:130 x(Tensor): Shape [B, L, num_heads, C / num_heads]131 seq_lens(Tensor): Shape [B]132 grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)133 freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]134 """135 b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim136 137 # query, key, value function138 def qkv_fn(x):139 q = self.norm_q(self.q(x)).view(b, s, n, d)140 k = self.norm_k(self.k(x)).view(b, s, n, d)141 v = self.v(x).view(b, s, n, d)142 return q, k, v143 144 q, k, v = qkv_fn(x)145 146 x = flash_attention(147 q=rope_apply(q, grid_sizes, freqs),148 k=rope_apply(k, grid_sizes, freqs),149 v=v,150 k_lens=seq_lens,151 window_size=self.window_size)152 153 # output154 x = x.flatten(2)155 x = self.o(x)156 return x157 158 159class WanT2VCrossAttention(WanSelfAttention):160 161 def forward(self, x, context, context_lens, crossattn_cache=None):162 r"""163 Args:164 x(Tensor): Shape [B, L1, C]165 context(Tensor): Shape [B, L2, C]166 context_lens(Tensor): Shape [B]167 crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding.168 """169 b, n, d = x.size(0), self.num_heads, self.head_dim170 171 # compute query, key, value172 q = self.norm_q(self.q(x)).view(b, -1, n, d)173 174 if crossattn_cache is not None:175 if not crossattn_cache["is_init"]:176 crossattn_cache["is_init"] = True177 k = self.norm_k(self.k(context)).view(b, -1, n, d)178 v = self.v(context).view(b, -1, n, d)179 crossattn_cache["k"] = k180 crossattn_cache["v"] = v181 else:182 k = crossattn_cache["k"]183 v = crossattn_cache["v"]184 else:185 k = self.norm_k(self.k(context)).view(b, -1, n, d)186 v = self.v(context).view(b, -1, n, d)187 188 # compute attention189 x = flash_attention(q, k, v, k_lens=context_lens)190 191 # output192 x = x.flatten(2)193 x = self.o(x)194 return x195 196 197class WanGanCrossAttention(WanSelfAttention):198 199 def forward(self, x, context, crossattn_cache=None):200 r"""201 Args:202 x(Tensor): Shape [B, L1, C]203 context(Tensor): Shape [B, L2, C]204 context_lens(Tensor): Shape [B]205 crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding.206 """207 b, n, d = x.size(0), self.num_heads, self.head_dim208 209 # compute query, key, value210 qq = self.norm_q(self.q(context)).view(b, 1, -1, d)211 212 kk = self.norm_k(self.k(x)).view(b, -1, n, d)213 vv = self.v(x).view(b, -1, n, d)214 215 # compute attention216 x = flash_attention(qq, kk, vv)217 218 # output219 x = x.flatten(2)220 x = self.o(x)221 return x222 223 224class WanI2VCrossAttention(WanSelfAttention):225 226 def __init__(self,227 dim,228 num_heads,229 window_size=(-1, -1),230 qk_norm=True,231 eps=1e-6):232 super().__init__(dim, num_heads, window_size, qk_norm, eps)233 234 self.k_img = nn.Linear(dim, dim)235 self.v_img = nn.Linear(dim, dim)236 # self.alpha = nn.Parameter(torch.zeros((1, )))237 self.norm_k_img = WanRMSNorm(238 dim, eps=eps) if qk_norm else nn.Identity()239 240 def forward(self, x, context, context_lens):241 r"""242 Args:243 x(Tensor): Shape [B, L1, C]244 context(Tensor): Shape [B, L2, C]245 context_lens(Tensor): Shape [B]246 """247 context_img = context[:, :257]248 context = context[:, 257:]249 b, n, d = x.size(0), self.num_heads, self.head_dim250 251 # compute query, key, value252 q = self.norm_q(self.q(x)).view(b, -1, n, d)253 k = self.norm_k(self.k(context)).view(b, -1, n, d)254 v = self.v(context).view(b, -1, n, d)255 k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d)256 v_img = self.v_img(context_img).view(b, -1, n, d)257 img_x = flash_attention(q, k_img, v_img, k_lens=None)258 # compute attention259 x = flash_attention(q, k, v, k_lens=context_lens)260 261 # output262 x = x.flatten(2)263 img_x = img_x.flatten(2)264 x = x + img_x265 x = self.o(x)266 return x267 268 269WAN_CROSSATTENTION_CLASSES = {270 't2v_cross_attn': WanT2VCrossAttention,271 'i2v_cross_attn': WanI2VCrossAttention,272}273 274 275class WanAttentionBlock(nn.Module):276 277 def __init__(self,278 cross_attn_type,279 dim,280 ffn_dim,281 num_heads,282 window_size=(-1, -1),283 qk_norm=True,284 cross_attn_norm=False,285 eps=1e-6):286 super().__init__()287 self.dim = dim288 self.ffn_dim = ffn_dim289 self.num_heads = num_heads290 self.window_size = window_size291 self.qk_norm = qk_norm292 self.cross_attn_norm = cross_attn_norm293 self.eps = eps294 295 # layers296 self.norm1 = WanLayerNorm(dim, eps)297 self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm,298 eps)299 self.norm3 = WanLayerNorm(300 dim, eps,301 elementwise_affine=True) if cross_attn_norm else nn.Identity()302 self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim,303 num_heads,304 (-1, -1),305 qk_norm,306 eps)307 self.norm2 = WanLayerNorm(dim, eps)308 self.ffn = nn.Sequential(309 nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),310 nn.Linear(ffn_dim, dim))311 312 # modulation313 self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)314 315 def forward(316 self,317 x,318 e,319 seq_lens,320 grid_sizes,321 freqs,322 context,323 context_lens,324 ):325 r"""326 Args:327 x(Tensor): Shape [B, L, C]328 e(Tensor): Shape [B, 6, C]329 seq_lens(Tensor): Shape [B], length of each sequence in batch330 grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)331 freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]332 """333 # assert e.dtype == torch.float32334 # with amp.autocast(dtype=torch.float32):335 e = (self.modulation + e).chunk(6, dim=1)336 # assert e[0].dtype == torch.float32337 338 # self-attention339 y = self.self_attn(340 self.norm1(x) * (1 + e[1]) + e[0], seq_lens, grid_sizes,341 freqs)342 # with amp.autocast(dtype=torch.float32):343 x = x + y * e[2]344 345 # cross-attention & ffn function346 def cross_attn_ffn(x, context, context_lens, e):347 x = x + self.cross_attn(self.norm3(x), context, context_lens)348 y = self.ffn(self.norm2(x) * (1 + e[4]) + e[3])349 # with amp.autocast(dtype=torch.float32):350 x = x + y * e[5]351 return x352 353 x = cross_attn_ffn(x, context, context_lens, e)354 return x355 356 357class GanAttentionBlock(nn.Module):358 359 def __init__(self,360 dim=1536,361 ffn_dim=8192,362 num_heads=12,363 window_size=(-1, -1),364 qk_norm=True,365 cross_attn_norm=True,366 eps=1e-6):367 super().__init__()368 self.dim = dim369 self.ffn_dim = ffn_dim370 self.num_heads = num_heads371 self.window_size = window_size372 self.qk_norm = qk_norm373 self.cross_attn_norm = cross_attn_norm374 self.eps = eps375 376 # layers377 # self.norm1 = WanLayerNorm(dim, eps)378 # self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm,379 # eps)380 self.norm3 = WanLayerNorm(381 dim, eps,382 elementwise_affine=True) if cross_attn_norm else nn.Identity()383 384 self.norm2 = WanLayerNorm(dim, eps)385 self.ffn = nn.Sequential(386 nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),387 nn.Linear(ffn_dim, dim))388 389 self.cross_attn = WanGanCrossAttention(dim, num_heads,390 (-1, -1),391 qk_norm,392 eps)393 394 # modulation395 # self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)396 397 def forward(398 self,399 x,400 context,401 # seq_lens,402 # grid_sizes,403 # freqs,404 # context,405 # context_lens,406 ):407 r"""408 Args:409 x(Tensor): Shape [B, L, C]410 e(Tensor): Shape [B, 6, C]411 seq_lens(Tensor): Shape [B], length of each sequence in batch412 grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)413 freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]414 """415 # assert e.dtype == torch.float32416 # with amp.autocast(dtype=torch.float32):417 # e = (self.modulation + e).chunk(6, dim=1)418 # assert e[0].dtype == torch.float32419 420 # # self-attention421 # y = self.self_attn(422 # self.norm1(x) * (1 + e[1]) + e[0], seq_lens, grid_sizes,423 # freqs)424 # # with amp.autocast(dtype=torch.float32):425 # x = x + y * e[2]426 427 # cross-attention & ffn function428 def cross_attn_ffn(x, context):429 token = context + self.cross_attn(self.norm3(x), context)430 y = self.ffn(self.norm2(token)) + token # * (1 + e[4]) + e[3])431 # with amp.autocast(dtype=torch.float32):432 # x = x + y * e[5]433 return y434 435 x = cross_attn_ffn(x, context)436 return x437 438 439class Head(nn.Module):440 441 def __init__(self, dim, out_dim, patch_size, eps=1e-6):442 super().__init__()443 self.dim = dim444 self.out_dim = out_dim445 self.patch_size = patch_size446 self.eps = eps447 448 # layers449 out_dim = math.prod(patch_size) * out_dim450 self.norm = WanLayerNorm(dim, eps)451 self.head = nn.Linear(dim, out_dim)452 453 # modulation454 self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)455 456 def forward(self, x, e):457 r"""458 Args:459 x(Tensor): Shape [B, L1, C]460 e(Tensor): Shape [B, C]461 """462 # assert e.dtype == torch.float32463 # with amp.autocast(dtype=torch.float32):464 e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)465 x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))466 return x467 468 469class MLPProj(torch.nn.Module):470 471 def __init__(self, in_dim, out_dim):472 super().__init__()473 474 self.proj = torch.nn.Sequential(475 torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim),476 torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim),477 torch.nn.LayerNorm(out_dim))478 479 def forward(self, image_embeds):480 clip_extra_context_tokens = self.proj(image_embeds)481 return clip_extra_context_tokens482 483 484class RegisterTokens(nn.Module):485 def __init__(self, num_registers: int, dim: int):486 super().__init__()487 self.register_tokens = nn.Parameter(torch.randn(num_registers, dim) * 0.02)488 self.rms_norm = WanRMSNorm(dim, eps=1e-6)489 490 def forward(self):491 return self.rms_norm(self.register_tokens)492 493 def reset_parameters(self):494 nn.init.normal_(self.register_tokens, std=0.02)495 496 497class WanModel(ModelMixin, ConfigMixin):498 r"""499 Wan diffusion backbone supporting both text-to-video and image-to-video.500 """501 502 ignore_for_config = [503 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size'504 ]505 _no_split_modules = ['WanAttentionBlock']506 _supports_gradient_checkpointing = True507 508 @register_to_config509 def __init__(self,510 model_type='t2v',511 patch_size=(1, 2, 2),512 text_len=512,513 in_dim=16,514 dim=2048,515 ffn_dim=8192,516 freq_dim=256,517 text_dim=4096,518 out_dim=16,519 num_heads=16,520 num_layers=32,521 window_size=(-1, -1),522 qk_norm=True,523 cross_attn_norm=True,524 eps=1e-6):525 r"""526 Initialize the diffusion model backbone.527 528 Args:529 model_type (`str`, *optional*, defaults to 't2v'):530 Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)531 patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):532 3D patch dimensions for video embedding (t_patch, h_patch, w_patch)533 text_len (`int`, *optional*, defaults to 512):534 Fixed length for text embeddings535 in_dim (`int`, *optional*, defaults to 16):536 Input video channels (C_in)537 dim (`int`, *optional*, defaults to 2048):538 Hidden dimension of the transformer539 ffn_dim (`int`, *optional*, defaults to 8192):540 Intermediate dimension in feed-forward network541 freq_dim (`int`, *optional*, defaults to 256):542 Dimension for sinusoidal time embeddings543 text_dim (`int`, *optional*, defaults to 4096):544 Input dimension for text embeddings545 out_dim (`int`, *optional*, defaults to 16):546 Output video channels (C_out)547 num_heads (`int`, *optional*, defaults to 16):548 Number of attention heads549 num_layers (`int`, *optional*, defaults to 32):550 Number of transformer blocks551 window_size (`tuple`, *optional*, defaults to (-1, -1)):552 Window size for local attention (-1 indicates global attention)553 qk_norm (`bool`, *optional*, defaults to True):554 Enable query/key normalization555 cross_attn_norm (`bool`, *optional*, defaults to False):556 Enable cross-attention normalization557 eps (`float`, *optional*, defaults to 1e-6):558 Epsilon value for normalization layers559 """560 561 super().__init__()562 563 assert model_type in ['t2v', 'i2v']564 self.model_type = model_type565 566 self.patch_size = patch_size567 self.text_len = text_len568 self.in_dim = in_dim569 self.dim = dim570 self.ffn_dim = ffn_dim571 self.freq_dim = freq_dim572 self.text_dim = text_dim573 self.out_dim = out_dim574 self.num_heads = num_heads575 self.num_layers = num_layers576 self.window_size = window_size577 self.qk_norm = qk_norm578 self.cross_attn_norm = cross_attn_norm579 self.eps = eps580 self.local_attn_size = 21581 582 # embeddings583 self.patch_embedding = nn.Conv3d(584 in_dim, dim, kernel_size=patch_size, stride=patch_size)585 self.text_embedding = nn.Sequential(586 nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),587 nn.Linear(dim, dim))588 589 self.time_embedding = nn.Sequential(590 nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))591 self.time_projection = nn.Sequential(592 nn.SiLU(), nn.Linear(dim, dim * 6))593 594 # blocks595 cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'596 self.blocks = nn.ModuleList([597 WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,598 window_size, qk_norm, cross_attn_norm, eps)599 for _ in range(num_layers)600 ])601 602 # head603 self.head = Head(dim, out_dim, patch_size, eps)604 605 # buffers (don't use register_buffer otherwise dtype will be changed in to())606 assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0607 d = dim // num_heads608 self.freqs = torch.cat([609 rope_params(1024, d - 4 * (d // 6)),610 rope_params(1024, 2 * (d // 6)),611 rope_params(1024, 2 * (d // 6))612 ],613 dim=1)614 615 if model_type == 'i2v':616 self.img_emb = MLPProj(1280, dim)617 618 # initialize weights619 self.init_weights()620 621 self.gradient_checkpointing = False622 623 def _set_gradient_checkpointing(self, module, value=False):624 self.gradient_checkpointing = value625 626 def forward(627 self,628 *args,629 **kwargs630 ):631 # if kwargs.get('classify_mode', False) is True:632 # kwargs.pop('classify_mode')633 # return self._forward_classify(*args, **kwargs)634 # else:635 return self._forward(*args, **kwargs)636 637 def _forward(638 self,639 x,640 t,641 context,642 seq_len,643 classify_mode=False,644 concat_time_embeddings=False,645 register_tokens=None,646 cls_pred_branch=None,647 gan_ca_blocks=None,648 clip_fea=None,649 y=None,650 ):651 r"""652 Forward pass through the diffusion model653 654 Args:655 x (List[Tensor]):656 List of input video tensors, each with shape [C_in, F, H, W]657 t (Tensor):658 Diffusion timesteps tensor of shape [B]659 context (List[Tensor]):660 List of text embeddings each with shape [L, C]661 seq_len (`int`):662 Maximum sequence length for positional encoding663 clip_fea (Tensor, *optional*):664 CLIP image features for image-to-video mode665 y (List[Tensor], *optional*):666 Conditional video inputs for image-to-video mode, same shape as x667 668 Returns:669 List[Tensor]:670 List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]671 """672 if self.model_type == 'i2v':673 assert clip_fea is not None and y is not None674 # params675 device = self.patch_embedding.weight.device676 if self.freqs.device != device:677 self.freqs = self.freqs.to(device)678 679 if y is not None:680 x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]681 682 # embeddings683 x = [self.patch_embedding(u.unsqueeze(0)) for u in x]684 grid_sizes = torch.stack(685 [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])686 x = [u.flatten(2).transpose(1, 2) for u in x]687 seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)688 assert seq_lens.max() <= seq_len689 x = torch.cat([690 torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],691 dim=1) for u in x692 ])693 694 # time embeddings695 # with amp.autocast(dtype=torch.float32):696 e = self.time_embedding(697 sinusoidal_embedding_1d(self.freq_dim, t).type_as(x))698 e0 = self.time_projection(e).unflatten(1, (6, self.dim))699 # assert e.dtype == torch.float32 and e0.dtype == torch.float32700 701 # context702 context_lens = None703 context = self.text_embedding(704 torch.stack([705 torch.cat(706 [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])707 for u in context708 ]))709 710 if clip_fea is not None:711 context_clip = self.img_emb(clip_fea) # bs x 257 x dim712 context = torch.concat([context_clip, context], dim=1)713 714 # arguments715 kwargs = dict(716 e=e0,717 seq_lens=seq_lens,718 grid_sizes=grid_sizes,719 freqs=self.freqs,720 context=context,721 context_lens=context_lens)722 723 def create_custom_forward(module):724 def custom_forward(*inputs, **kwargs):725 return module(*inputs, **kwargs)726 return custom_forward727 728 # TODO: Tune the number of blocks for feature extraction729 final_x = None730 if classify_mode:731 assert register_tokens is not None732 assert gan_ca_blocks is not None733 assert cls_pred_branch is not None734 735 final_x = []736 registers = repeat(register_tokens(), "n d -> b n d", b=x.shape[0])737 # x = torch.cat([registers, x], dim=1)738 739 gan_idx = 0740 for ii, block in enumerate(self.blocks):741 if torch.is_grad_enabled() and self.gradient_checkpointing:742 x = torch.utils.checkpoint.checkpoint(743 create_custom_forward(block),744 x, **kwargs,745 use_reentrant=False,746 )747 else:748 x = block(x, **kwargs)749 750 if classify_mode and ii in [13, 21, 29]:751 gan_token = registers[:, gan_idx: gan_idx + 1]752 final_x.append(gan_ca_blocks[gan_idx](x, gan_token))753 gan_idx += 1754 755 if classify_mode:756 final_x = torch.cat(final_x, dim=1)757 if concat_time_embeddings:758 final_x = cls_pred_branch(torch.cat([final_x, 10 * e[:, None, :]], dim=1).view(final_x.shape[0], -1))759 else:760 final_x = cls_pred_branch(final_x.view(final_x.shape[0], -1))761 762 # head763 x = self.head(x, e)764 765 # unpatchify766 x = self.unpatchify(x, grid_sizes)767 768 if classify_mode:769 return torch.stack(x), final_x770 771 return torch.stack(x)772 773 def _forward_classify(774 self,775 x,776 t,777 context,778 seq_len,779 register_tokens,780 cls_pred_branch,781 clip_fea=None,782 y=None,783 ):784 r"""785 Feature extraction through the diffusion model786 787 Args:788 x (List[Tensor]):789 List of input video tensors, each with shape [C_in, F, H, W]790 t (Tensor):791 Diffusion timesteps tensor of shape [B]792 context (List[Tensor]):793 List of text embeddings each with shape [L, C]794 seq_len (`int`):795 Maximum sequence length for positional encoding796 clip_fea (Tensor, *optional*):797 CLIP image features for image-to-video mode798 y (List[Tensor], *optional*):799 Conditional video inputs for image-to-video mode, same shape as x800 801 Returns:802 List[Tensor]:803 List of video features with original input shapes [C_block, F, H / 8, W / 8]804 """805 if self.model_type == 'i2v':806 assert clip_fea is not None and y is not None807 # params808 device = self.patch_embedding.weight.device809 if self.freqs.device != device:810 self.freqs = self.freqs.to(device)811 812 if y is not None:813 x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]814 815 # embeddings816 x = [self.patch_embedding(u.unsqueeze(0)) for u in x]817 grid_sizes = torch.stack(818 [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])819 x = [u.flatten(2).transpose(1, 2) for u in x]820 seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)821 assert seq_lens.max() <= seq_len822 x = torch.cat([823 torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],824 dim=1) for u in x825 ])826 827 # time embeddings828 # with amp.autocast(dtype=torch.float32):829 e = self.time_embedding(830 sinusoidal_embedding_1d(self.freq_dim, t).type_as(x))831 e0 = self.time_projection(e).unflatten(1, (6, self.dim))832 # assert e.dtype == torch.float32 and e0.dtype == torch.float32833 834 # context835 context_lens = None836 context = self.text_embedding(837 torch.stack([838 torch.cat(839 [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])840 for u in context841 ]))842 843 if clip_fea is not None:844 context_clip = self.img_emb(clip_fea) # bs x 257 x dim845 context = torch.concat([context_clip, context], dim=1)846 847 # arguments848 kwargs = dict(849 e=e0,850 seq_lens=seq_lens,851 grid_sizes=grid_sizes,852 freqs=self.freqs,853 context=context,854 context_lens=context_lens)855 856 def create_custom_forward(module):857 def custom_forward(*inputs, **kwargs):858 return module(*inputs, **kwargs)859 return custom_forward860 861 # TODO: Tune the number of blocks for feature extraction862 for block in self.blocks[:16]:863 if torch.is_grad_enabled() and self.gradient_checkpointing:864 x = torch.utils.checkpoint.checkpoint(865 create_custom_forward(block),866 x, **kwargs,867 use_reentrant=False,868 )869 else:870 x = block(x, **kwargs)871 872 # unpatchify873 x = self.unpatchify(x, grid_sizes, c=self.dim // 4)874 return torch.stack(x)875 876 def unpatchify(self, x, grid_sizes, c=None):877 r"""878 Reconstruct video tensors from patch embeddings.879 880 Args:881 x (List[Tensor]):882 List of patchified features, each with shape [L, C_out * prod(patch_size)]883 grid_sizes (Tensor):884 Original spatial-temporal grid dimensions before patching,885 shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)886 887 Returns:888 List[Tensor]:889 Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]890 """891 892 c = self.out_dim if c is None else c893 out = []894 for u, v in zip(x, grid_sizes.tolist()):895 u = u[:math.prod(v)].view(*v, *self.patch_size, c)896 u = torch.einsum('fhwpqrc->cfphqwr', u)897 u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])898 out.append(u)899 return out900 901 def init_weights(self):902 r"""903 Initialize model parameters using Xavier initialization.904 """905 906 # basic init907 for m in self.modules():908 if isinstance(m, nn.Linear):909 nn.init.xavier_uniform_(m.weight)910 if m.bias is not None:911 nn.init.zeros_(m.bias)912 913 # init embeddings914 nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))915 for m in self.text_embedding.modules():916 if isinstance(m, nn.Linear):917 nn.init.normal_(m.weight, std=.02)918 for m in self.time_embedding.modules():919 if isinstance(m, nn.Linear):920 nn.init.normal_(m.weight, std=.02)921 922 # init output layer923 nn.init.zeros_(self.head.head.weight)924