durgappc/infinitetalk2
0
1# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.2import math3import numpy as np4import os5import torch6import torch.cuda.amp as amp7import torch.nn as nn8import torch.nn.functional as F9 10from einops import rearrange11from diffusers import ModelMixin12from diffusers.configuration_utils import ConfigMixin, register_to_config13 14from .attention import flash_attention, SingleStreamMutiAttention15from ..utils.multitalk_utils import get_attn_map_with_target16import logging17try:18 from sageattention import sageattn19 USE_SAGEATTN = True20 logging.info("Using sageattn")21except:22 USE_SAGEATTN = False23 24__all__ = ['WanModel']25 26 27 28def sinusoidal_embedding_1d(dim, position):29 # preprocess30 assert dim % 2 == 031 half = dim // 232 position = position.type(torch.float64)33 34 # calculation35 sinusoid = torch.outer(36 position, torch.pow(10000, -torch.arange(half).to(position).div(half)))37 x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)38 return x39 40 41@amp.autocast(enabled=False)42def rope_params(max_seq_len, dim, theta=10000):43 44 assert dim % 2 == 045 freqs = torch.outer(46 torch.arange(max_seq_len),47 1.0 / torch.pow(theta,48 torch.arange(0, dim, 2).to(torch.float64).div(dim)))49 freqs = torch.polar(torch.ones_like(freqs), freqs)50 return freqs51 52 53@amp.autocast(enabled=False)54def rope_apply(x, grid_sizes, freqs):55 s, n, c = x.size(1), x.size(2), x.size(3) // 256 57 freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)58 59 output = []60 for i, (f, h, w) in enumerate(grid_sizes.tolist()):61 seq_len = f * h * w62 63 x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(64 s, n, -1, 2))65 freqs_i = torch.cat([66 freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),67 freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),68 freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)69 ],70 dim=-1).reshape(seq_len, 1, -1)71 freqs_i = freqs_i.to(device=x_i.device)72 x_i = torch.view_as_real(x_i * freqs_i).flatten(2)73 x_i = torch.cat([x_i, x[i, seq_len:]])74 75 output.append(x_i)76 return torch.stack(output).float()77 78 79class WanRMSNorm(nn.Module):80 81 def __init__(self, dim, eps=1e-5):82 super().__init__()83 self.dim = dim84 self.eps = eps85 self.weight = nn.Parameter(torch.ones(dim))86 87 def forward(self, x):88 r"""89 Args:90 x(Tensor): Shape [B, L, C]91 """92 return self._norm(x.float()).type_as(x) * self.weight93 94 def _norm(self, x):95 return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)96 97 98class WanLayerNorm(nn.LayerNorm):99 100 def __init__(self, dim, eps=1e-6, elementwise_affine=False):101 super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)102 103 def forward(self, inputs: torch.Tensor) -> torch.Tensor:104 origin_dtype = inputs.dtype105 out = F.layer_norm(106 inputs.float(), 107 self.normalized_shape, 108 None if self.weight is None else self.weight.float(), 109 None if self.bias is None else self.bias.float() ,110 self.eps111 ).to(origin_dtype)112 return out113 114 115class WanSelfAttention(nn.Module):116 117 def __init__(self,118 dim,119 num_heads,120 window_size=(-1, -1),121 qk_norm=True,122 eps=1e-6):123 assert dim % num_heads == 0124 super().__init__()125 self.dim = dim126 self.num_heads = num_heads127 self.head_dim = dim // num_heads128 self.window_size = window_size129 self.qk_norm = qk_norm130 self.eps = eps131 132 # layers133 self.q = nn.Linear(dim, dim)134 self.k = nn.Linear(dim, dim)135 self.v = nn.Linear(dim, dim)136 self.o = nn.Linear(dim, dim)137 self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()138 self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()139 140 def forward(self, x, seq_lens, grid_sizes, freqs, ref_target_masks=None):141 b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim142 143 # query, key, value function144 def qkv_fn(x):145 q = self.norm_q(self.q(x)).view(b, s, n, d)146 k = self.norm_k(self.k(x)).view(b, s, n, d)147 v = self.v(x).view(b, s, n, d)148 return q, k, v149 q, k, v = qkv_fn(x)150 151 q = rope_apply(q, grid_sizes, freqs)152 k = rope_apply(k, grid_sizes, freqs)153 154 if USE_SAGEATTN:155 x = sageattn(q.to(torch.bfloat16), k.to(torch.bfloat16), v, tensor_layout='NHD')156 else:157 x = flash_attention(158 q=q,159 k=k,160 v=v,161 k_lens=seq_lens,162 window_size=self.window_size163 ).type_as(x)164 165 # output166 x = x.flatten(2)167 x = self.o(x)168 with torch.no_grad():169 x_ref_attn_map = get_attn_map_with_target(q.type_as(x), k.type_as(x), grid_sizes[0], 170 ref_target_masks=ref_target_masks)171 172 return x, x_ref_attn_map173 174 175class WanI2VCrossAttention(WanSelfAttention):176 177 def __init__(self,178 dim,179 num_heads,180 window_size=(-1, -1),181 qk_norm=True,182 eps=1e-6):183 super().__init__(dim, num_heads, window_size, qk_norm, eps)184 185 self.k_img = nn.Linear(dim, dim)186 self.v_img = nn.Linear(dim, dim)187 self.norm_k_img = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()188 189 def forward(self, x, context, context_lens):190 context_img = context[:, :257]191 context = context[:, 257:]192 b, n, d = x.size(0), self.num_heads, self.head_dim193 194 # compute query, key, value195 q = self.norm_q(self.q(x)).view(b, -1, n, d)196 k = self.norm_k(self.k(context)).view(b, -1, n, d)197 v = self.v(context).view(b, -1, n, d)198 k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d)199 v_img = self.v_img(context_img).view(b, -1, n, d)200 if USE_SAGEATTN:201 img_x = sageattn(q, k_img, v_img, tensor_layout='NHD')202 x = sageattn(q, k, v, tensor_layout='NHD')203 else: 204 img_x = flash_attention(q, k_img, v_img, k_lens=None)205 # compute attention206 x = flash_attention(q, k, v, k_lens=context_lens)207 208 # output209 x = x.flatten(2)210 img_x = img_x.flatten(2)211 x = x + img_x212 x = self.o(x)213 return x214 215 216class WanAttentionBlock(nn.Module):217 218 def __init__(self,219 cross_attn_type,220 dim,221 ffn_dim,222 num_heads,223 window_size=(-1, -1),224 qk_norm=True,225 cross_attn_norm=False,226 eps=1e-6,227 output_dim=768,228 norm_input_visual=True,229 class_range=24,230 class_interval=4):231 super().__init__()232 self.dim = dim233 self.ffn_dim = ffn_dim234 self.num_heads = num_heads235 self.window_size = window_size236 self.qk_norm = qk_norm237 self.cross_attn_norm = cross_attn_norm238 self.eps = eps239 240 # layers241 self.norm1 = WanLayerNorm(dim, eps)242 self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, eps)243 self.norm3 = WanLayerNorm(244 dim, eps,245 elementwise_affine=True) if cross_attn_norm else nn.Identity()246 self.cross_attn = WanI2VCrossAttention(dim,247 num_heads,248 (-1, -1),249 qk_norm,250 eps)251 self.norm2 = WanLayerNorm(dim, eps)252 self.ffn = nn.Sequential(253 nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),254 nn.Linear(ffn_dim, dim))255 256 # modulation257 self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)258 259 # init audio module260 self.audio_cross_attn = SingleStreamMutiAttention(261 dim=dim,262 encoder_hidden_states_dim=output_dim,263 num_heads=num_heads,264 qk_norm=False,265 qkv_bias=True,266 eps=eps,267 norm_layer=WanRMSNorm,268 class_range=class_range,269 class_interval=class_interval270 )271 self.norm_x = WanLayerNorm(dim, eps, elementwise_affine=True) if norm_input_visual else nn.Identity()272 273 274 def forward(275 self,276 x,277 e,278 seq_lens,279 grid_sizes,280 freqs,281 context,282 context_lens,283 audio_embedding=None,284 ref_target_masks=None,285 human_num=None,286 ):287 288 dtype = x.dtype289 assert e.dtype == torch.float32290 with amp.autocast(dtype=torch.float32):291 e = (self.modulation.to(e.device) + e).chunk(6, dim=1)292 assert e[0].dtype == torch.float32293 294 # self-attention295 y, x_ref_attn_map = self.self_attn(296 (self.norm1(x).float() * (1 + e[1]) + e[0]).type_as(x), seq_lens, grid_sizes,297 freqs, ref_target_masks=ref_target_masks)298 with amp.autocast(dtype=torch.float32):299 x = x + y * e[2]300 301 x = x.to(dtype)302 303 # cross-attention of text304 x = x + self.cross_attn(self.norm3(x), context, context_lens)305 306 # cross attn of audio307 x_a = self.audio_cross_attn(self.norm_x(x), encoder_hidden_states=audio_embedding,308 shape=grid_sizes[0], x_ref_attn_map=x_ref_attn_map, human_num=human_num)309 x = x + x_a310 311 y = self.ffn((self.norm2(x).float() * (1 + e[4]) + e[3]).to(dtype))312 with amp.autocast(dtype=torch.float32):313 x = x + y * e[5]314 315 316 x = x.to(dtype)317 318 return x319 320 321class Head(nn.Module):322 323 def __init__(self, dim, out_dim, patch_size, eps=1e-6):324 super().__init__()325 self.dim = dim326 self.out_dim = out_dim327 self.patch_size = patch_size328 self.eps = eps329 330 # layers331 out_dim = math.prod(patch_size) * out_dim332 self.norm = WanLayerNorm(dim, eps)333 self.head = nn.Linear(dim, out_dim)334 335 # modulation336 self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)337 338 def forward(self, x, e):339 r"""340 Args:341 x(Tensor): Shape [B, L1, C]342 e(Tensor): Shape [B, C]343 """344 assert e.dtype == torch.float32345 with amp.autocast(dtype=torch.float32):346 e = (self.modulation.to(e.device) + e.unsqueeze(1)).chunk(2, dim=1)347 x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))348 return x349 350 351class MLPProj(torch.nn.Module):352 353 def __init__(self, in_dim, out_dim):354 super().__init__()355 356 self.proj = torch.nn.Sequential(357 torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim),358 torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim),359 torch.nn.LayerNorm(out_dim))360 361 def forward(self, image_embeds):362 clip_extra_context_tokens = self.proj(image_embeds)363 return clip_extra_context_tokens364 365 366class AudioProjModel(ModelMixin, ConfigMixin):367 def __init__(368 self,369 seq_len=5,370 seq_len_vf=12,371 blocks=12, 372 channels=768, 373 intermediate_dim=512,374 output_dim=768,375 context_tokens=32,376 norm_output_audio=False,377 ):378 super().__init__()379 380 self.seq_len = seq_len381 self.blocks = blocks382 self.channels = channels383 self.input_dim = seq_len * blocks * channels 384 self.input_dim_vf = seq_len_vf * blocks * channels385 self.intermediate_dim = intermediate_dim386 self.context_tokens = context_tokens387 self.output_dim = output_dim388 389 # define multiple linear layers390 self.proj1 = nn.Linear(self.input_dim, intermediate_dim)391 self.proj1_vf = nn.Linear(self.input_dim_vf, intermediate_dim)392 self.proj2 = nn.Linear(intermediate_dim, intermediate_dim)393 self.proj3 = nn.Linear(intermediate_dim, context_tokens * output_dim)394 self.norm = nn.LayerNorm(output_dim) if norm_output_audio else nn.Identity()395 396 def forward(self, audio_embeds, audio_embeds_vf):397 video_length = audio_embeds.shape[1] + audio_embeds_vf.shape[1]398 B, _, _, S, C = audio_embeds.shape399 400 # process audio of first frame401 audio_embeds = rearrange(audio_embeds, "bz f w b c -> (bz f) w b c")402 batch_size, window_size, blocks, channels = audio_embeds.shape403 audio_embeds = audio_embeds.view(batch_size, window_size * blocks * channels)404 405 # process audio of latter frame406 audio_embeds_vf = rearrange(audio_embeds_vf, "bz f w b c -> (bz f) w b c")407 batch_size_vf, window_size_vf, blocks_vf, channels_vf = audio_embeds_vf.shape408 audio_embeds_vf = audio_embeds_vf.view(batch_size_vf, window_size_vf * blocks_vf * channels_vf)409 410 # first projection411 audio_embeds = torch.relu(self.proj1(audio_embeds)) 412 audio_embeds_vf = torch.relu(self.proj1_vf(audio_embeds_vf)) 413 audio_embeds = rearrange(audio_embeds, "(bz f) c -> bz f c", bz=B)414 audio_embeds_vf = rearrange(audio_embeds_vf, "(bz f) c -> bz f c", bz=B)415 audio_embeds_c = torch.concat([audio_embeds, audio_embeds_vf], dim=1) 416 batch_size_c, N_t, C_a = audio_embeds_c.shape417 audio_embeds_c = audio_embeds_c.view(batch_size_c*N_t, C_a)418 419 # second projection420 audio_embeds_c = torch.relu(self.proj2(audio_embeds_c))421 422 context_tokens = self.proj3(audio_embeds_c).reshape(batch_size_c*N_t, self.context_tokens, self.output_dim)423 424 # normalization and reshape425 with amp.autocast(dtype=torch.float32):426 context_tokens = self.norm(context_tokens)427 context_tokens = rearrange(context_tokens, "(bz f) m c -> bz f m c", f=video_length)428 429 return context_tokens430 431 432class WanModel(ModelMixin, ConfigMixin):433 r"""434 Wan diffusion backbone supporting both text-to-video and image-to-video.435 """436 437 ignore_for_config = [438 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size'439 ]440 _no_split_modules = ['WanAttentionBlock']441 442 @register_to_config443 def __init__(self,444 model_type='i2v',445 patch_size=(1, 2, 2),446 text_len=512,447 in_dim=16,448 dim=2048,449 ffn_dim=8192,450 freq_dim=256,451 text_dim=4096,452 out_dim=16,453 num_heads=16,454 num_layers=32,455 window_size=(-1, -1),456 qk_norm=True,457 cross_attn_norm=True,458 eps=1e-6,459 # audio params460 audio_window=5,461 intermediate_dim=512,462 output_dim=768,463 context_tokens=32,464 vae_scale=4, # vae timedownsample scale465 466 norm_input_visual=True,467 norm_output_audio=True,468 weight_init=True):469 super().__init__()470 471 assert model_type == 'i2v', 'MultiTalk model requires your model_type is i2v.'472 self.model_type = model_type473 474 self.patch_size = patch_size475 self.text_len = text_len476 self.in_dim = in_dim477 self.dim = dim478 self.ffn_dim = ffn_dim479 self.freq_dim = freq_dim480 self.text_dim = text_dim481 self.out_dim = out_dim482 self.num_heads = num_heads483 self.num_layers = num_layers484 self.window_size = window_size485 self.qk_norm = qk_norm486 self.cross_attn_norm = cross_attn_norm487 self.eps = eps488 489 490 self.norm_output_audio = norm_output_audio491 self.audio_window = audio_window492 self.intermediate_dim = intermediate_dim493 self.vae_scale = vae_scale494 495 496 # embeddings497 self.patch_embedding = nn.Conv3d(498 in_dim, dim, kernel_size=patch_size, stride=patch_size)499 self.text_embedding = nn.Sequential(500 nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),501 nn.Linear(dim, dim))502 503 self.time_embedding = nn.Sequential(504 nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))505 self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))506 507 # blocks508 cross_attn_type = 'i2v_cross_attn'509 self.blocks = nn.ModuleList([510 WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,511 window_size, qk_norm, cross_attn_norm, eps, 512 output_dim=output_dim, norm_input_visual=norm_input_visual)513 for _ in range(num_layers)514 ])515 516 # head517 self.head = Head(dim, out_dim, patch_size, eps)518 519 assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0520 d = dim // num_heads521 self.freqs = torch.cat([522 rope_params(1024, d - 4 * (d // 6)),523 rope_params(1024, 2 * (d // 6)),524 rope_params(1024, 2 * (d // 6))525 ],526 dim=1)527 528 if model_type == 'i2v':529 self.img_emb = MLPProj(1280, dim)530 else:531 raise NotImplementedError('Not supported model type.')532 533 # init audio adapter534 self.audio_proj = AudioProjModel(535 seq_len=audio_window,536 seq_len_vf=audio_window+vae_scale-1,537 intermediate_dim=intermediate_dim,538 output_dim=output_dim,539 context_tokens=context_tokens,540 norm_output_audio=norm_output_audio,541 )542 543 544 # initialize weights545 if weight_init:546 self.init_weights()547 548 def init_freqs(self):549 d = self.dim // self.num_heads550 self.freqs = torch.cat([551 rope_params(1024, d - 4 * (d // 6)),552 rope_params(1024, 2 * (d // 6)),553 rope_params(1024, 2 * (d // 6))554 ],555 dim=1)556 557 def teacache_init(558 self,559 use_ret_steps=True,560 teacache_thresh=0.2,561 sample_steps=40,562 model_scale='infinitetalk-480',563 ):564 print("teacache_init")565 self.enable_teacache = True566 567 self.__class__.cnt = 0568 self.__class__.num_steps = sample_steps*3569 self.__class__.teacache_thresh = teacache_thresh570 self.__class__.accumulated_rel_l1_distance_even = 0571 self.__class__.accumulated_rel_l1_distance_odd = 0572 self.__class__.previous_e0_even = None573 self.__class__.previous_e0_odd = None574 self.__class__.previous_residual_even = None575 self.__class__.previous_residual_odd = None576 self.__class__.use_ret_steps = use_ret_steps577 578 if use_ret_steps:579 if model_scale == 'infinitetalk-480':580 self.__class__.coefficients = [ 2.57151496e+05, -3.54229917e+04, 1.40286849e+03, -1.35890334e+01, 1.32517977e-01]581 if model_scale == 'infinitetalk-720':582 self.__class__.coefficients = [ 8.10705460e+03, 2.13393892e+03, -3.72934672e+02, 1.66203073e+01, -4.17769401e-02]583 self.__class__.ret_steps = 5*3584 self.__class__.cutoff_steps = sample_steps*3585 else:586 if model_scale == 'infinitetalk-480':587 self.__class__.coefficients = [-3.02331670e+02, 2.23948934e+02, -5.25463970e+01, 5.87348440e+00, -2.01973289e-01]588 589 if model_scale == 'infinitetalk-720':590 self.__class__.coefficients = [-114.36346466, 65.26524496, -18.82220707, 4.91518089, -0.23412683]591 self.__class__.ret_steps = 1*3592 self.__class__.cutoff_steps = sample_steps*3 - 3593 print("teacache_init done")594 595 def disable_teacache(self):596 self.enable_teacache = False597 598 def forward(599 self,600 x,601 t,602 context,603 seq_len,604 clip_fea=None,605 y=None,606 audio=None,607 ref_target_masks=None,608 ):609 assert clip_fea is not None and y is not None610 611 _, T, H, W = x[0].shape612 N_t = T // self.patch_size[0]613 N_h = H // self.patch_size[1]614 N_w = W // self.patch_size[2]615 616 if y is not None:617 x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]618 x[0] = x[0].to(context[0].dtype)619 620 # embeddings621 x = [self.patch_embedding(u.unsqueeze(0)) for u in x]622 grid_sizes = torch.stack(623 [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])624 x = [u.flatten(2).transpose(1, 2) for u in x]625 seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)626 assert seq_lens.max() <= seq_len627 x = torch.cat([628 torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],629 dim=1) for u in x630 ])631 632 # time embeddings633 with amp.autocast(dtype=torch.float32):634 e = self.time_embedding(635 sinusoidal_embedding_1d(self.freq_dim, t).float())636 e0 = self.time_projection(e).unflatten(1, (6, self.dim))637 assert e.dtype == torch.float32 and e0.dtype == torch.float32638 639 # text embedding640 context_lens = None641 context = self.text_embedding(642 torch.stack([643 torch.cat(644 [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])645 for u in context646 ]))647 648 # clip embedding649 if clip_fea is not None:650 context_clip = self.img_emb(clip_fea) 651 context = torch.concat([context_clip, context], dim=1).to(x.dtype)652 653 654 audio_cond = audio.to(device=x.device, dtype=x.dtype)655 first_frame_audio_emb_s = audio_cond[:, :1, ...] 656 latter_frame_audio_emb = audio_cond[:, 1:, ...] 657 latter_frame_audio_emb = rearrange(latter_frame_audio_emb, "b (n_t n) w s c -> b n_t n w s c", n=self.vae_scale) 658 middle_index = self.audio_window // 2659 latter_first_frame_audio_emb = latter_frame_audio_emb[:, :, :1, :middle_index+1, ...] 660 latter_first_frame_audio_emb = rearrange(latter_first_frame_audio_emb, "b n_t n w s c -> b n_t (n w) s c") 661 latter_last_frame_audio_emb = latter_frame_audio_emb[:, :, -1:, middle_index:, ...] 662 latter_last_frame_audio_emb = rearrange(latter_last_frame_audio_emb, "b n_t n w s c -> b n_t (n w) s c") 663 latter_middle_frame_audio_emb = latter_frame_audio_emb[:, :, 1:-1, middle_index:middle_index+1, ...] 664 latter_middle_frame_audio_emb = rearrange(latter_middle_frame_audio_emb, "b n_t n w s c -> b n_t (n w) s c") 665 latter_frame_audio_emb_s = torch.concat([latter_first_frame_audio_emb, latter_middle_frame_audio_emb, latter_last_frame_audio_emb], dim=2) 666 audio_embedding = self.audio_proj(first_frame_audio_emb_s, latter_frame_audio_emb_s) 667 human_num = len(audio_embedding)668 audio_embedding = torch.concat(audio_embedding.split(1), dim=2).to(x.dtype)669 670 671 # convert ref_target_masks to token_ref_target_masks672 if ref_target_masks is not None:673 ref_target_masks = ref_target_masks.unsqueeze(0).to(torch.float32) 674 token_ref_target_masks = nn.functional.interpolate(ref_target_masks, size=(N_h, N_w), mode='nearest') 675 token_ref_target_masks = token_ref_target_masks.squeeze(0)676 token_ref_target_masks = (token_ref_target_masks > 0)677 token_ref_target_masks = token_ref_target_masks.view(token_ref_target_masks.shape[0], -1) 678 token_ref_target_masks = token_ref_target_masks.to(x.dtype)679 680 # teacache681 if self.enable_teacache:682 modulated_inp = e0 if self.use_ret_steps else e683 if self.cnt%3==0: # cond684 if self.cnt < self.ret_steps or self.cnt >= self.cutoff_steps:685 should_calc_cond = True686 self.accumulated_rel_l1_distance_cond = 0687 else:688 rescale_func = np.poly1d(self.coefficients)689 self.accumulated_rel_l1_distance_cond += rescale_func(((modulated_inp-self.previous_e0_cond).abs().mean() / self.previous_e0_cond.abs().mean()).cpu().item())690 if self.accumulated_rel_l1_distance_cond < self.teacache_thresh:691 should_calc_cond = False692 else:693 should_calc_cond = True694 self.accumulated_rel_l1_distance_cond = 0695 self.previous_e0_cond = modulated_inp.clone()696 elif self.cnt%3==1: # drop_text697 if self.cnt < self.ret_steps or self.cnt >= self.cutoff_steps:698 should_calc_drop_text = True699 self.accumulated_rel_l1_distance_drop_text = 0700 else:701 rescale_func = np.poly1d(self.coefficients)702 self.accumulated_rel_l1_distance_drop_text += rescale_func(((modulated_inp-self.previous_e0_drop_text).abs().mean() / self.previous_e0_drop_text.abs().mean()).cpu().item())703 if self.accumulated_rel_l1_distance_drop_text < self.teacache_thresh:704 should_calc_drop_text = False705 else:706 should_calc_drop_text = True707 self.accumulated_rel_l1_distance_drop_text = 0708 self.previous_e0_drop_text = modulated_inp.clone()709 else: # uncond710 if self.cnt < self.ret_steps or self.cnt >= self.cutoff_steps:711 should_calc_uncond = True712 self.accumulated_rel_l1_distance_uncond = 0713 else:714 rescale_func = np.poly1d(self.coefficients)715 self.accumulated_rel_l1_distance_uncond += rescale_func(((modulated_inp-self.previous_e0_uncond).abs().mean() / self.previous_e0_uncond.abs().mean()).cpu().item())716 if self.accumulated_rel_l1_distance_uncond < self.teacache_thresh:717 should_calc_uncond = False718 else:719 should_calc_uncond = True720 self.accumulated_rel_l1_distance_uncond = 0721 self.previous_e0_uncond = modulated_inp.clone()722 723 # arguments724 kwargs = dict(725 e=e0,726 seq_lens=seq_lens,727 grid_sizes=grid_sizes,728 freqs=self.freqs,729 context=context,730 context_lens=context_lens,731 audio_embedding=audio_embedding,732 ref_target_masks=token_ref_target_masks,733 human_num=human_num,734 )735 if self.enable_teacache:736 if self.cnt%3==0:737 if not should_calc_cond:738 x += self.previous_residual_cond739 else:740 ori_x = x.clone()741 for block in self.blocks:742 x = block(x, **kwargs)743 self.previous_residual_cond = x - ori_x744 elif self.cnt%3==1:745 if not should_calc_drop_text:746 x += self.previous_residual_drop_text747 else:748 ori_x = x.clone()749 for block in self.blocks:750 x = block(x, **kwargs)751 self.previous_residual_drop_text = x - ori_x752 else:753 if not should_calc_uncond:754 x += self.previous_residual_uncond755 else:756 ori_x = x.clone()757 for block in self.blocks:758 x = block(x, **kwargs)759 self.previous_residual_uncond = x - ori_x760 else:761 for block in self.blocks:762 x = block(x, **kwargs)763 764 # head765 x = self.head(x, e)766 767 # unpatchify768 x = self.unpatchify(x, grid_sizes)769 if self.enable_teacache:770 self.cnt += 1771 if self.cnt >= self.num_steps:772 self.cnt = 0773 774 return torch.stack(x).float()775 776 777 def unpatchify(self, x, grid_sizes):778 r"""779 Reconstruct video tensors from patch embeddings.780 781 Args:782 x (List[Tensor]):783 List of patchified features, each with shape [L, C_out * prod(patch_size)]784 grid_sizes (Tensor):785 Original spatial-temporal grid dimensions before patching,786 shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)787 788 Returns:789 List[Tensor]:790 Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]791 """792 793 c = self.out_dim794 out = []795 for u, v in zip(x, grid_sizes.tolist()):796 u = u[:math.prod(v)].view(*v, *self.patch_size, c)797 u = torch.einsum('fhwpqrc->cfphqwr', u)798 u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])799 out.append(u)800 return out801 802 def init_weights(self):803 r"""804 Initialize model parameters using Xavier initialization.805 """806 807 # basic init808 for m in self.modules():809 if isinstance(m, nn.Linear):810 nn.init.xavier_uniform_(m.weight)811 if m.bias is not None:812 nn.init.zeros_(m.bias)813 814 # init embeddings815 nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))816 for m in self.text_embedding.modules():817 if isinstance(m, nn.Linear):818 nn.init.normal_(m.weight, std=.02)819 for m in self.time_embedding.modules():820 if isinstance(m, nn.Linear):821 nn.init.normal_(m.weight, std=.02)822 823 # init output layer824 nn.init.zeros_(self.head.head.weight)