import torch import torch.nn as nn from .model import WanLayerNorm, WanModel, WanRMSNorm, rope_apply from shared.attention import pay_attention ##### Enjoy this spagheti VRAM optimizations done by DeepBeepMeep ! # I am sure you are a nice person and as you copy this code, you will give me officially proper credits: # Please link to https://github.com/deepbeepmeep/Wan2GP and @deepbeepmeep on twitter def reshape_latent(latent, latent_frames): return latent.reshape(latent.shape[0], latent_frames, -1, latent.shape[-1] ) def restore_latent_shape(latent): return latent.reshape(latent.shape[0], -1, latent.shape[-1] ) class FusionModel(nn.Module): def __init__(self, video_config=None, audio_config=None): super().__init__() has_video = True has_audio = True if video_config is not None: self.video_model = WanModel(**video_config) else: has_video = False self.video_model = None print("Warning: No video model is provided!") if audio_config is not None: self.audio_model = WanModel(**audio_config) else: has_audio = False self.audio_model = None print("Warning: No audio model is provided!") if has_video and has_audio: assert len(self.video_model.blocks) == len(self.audio_model.blocks) self.num_blocks = len(self.video_model.blocks) self.use_sp = False self.inject_cross_attention_kv_projections() self.init_weights() def inject_cross_attention_kv_projections(self): for vid_block in self.video_model.blocks: vid_block.cross_attn.k_fusion = nn.Linear(vid_block.dim, vid_block.dim) vid_block.cross_attn.v_fusion = nn.Linear(vid_block.dim, vid_block.dim) vid_block.cross_attn.pre_attn_norm_fusion = WanLayerNorm(vid_block.dim, elementwise_affine=True) vid_block.cross_attn.norm_k_fusion = WanRMSNorm(vid_block.dim, eps=1e-6) if vid_block.qk_norm else nn.Identity() for audio_block in self.audio_model.blocks: audio_block.cross_attn.k_fusion = nn.Linear(audio_block.dim, audio_block.dim) audio_block.cross_attn.v_fusion = nn.Linear(audio_block.dim, audio_block.dim) audio_block.cross_attn.pre_attn_norm_fusion = WanLayerNorm(audio_block.dim, elementwise_affine=True) audio_block.cross_attn.norm_k_fusion = WanRMSNorm(audio_block.dim, eps=1e-6) if audio_block.qk_norm else nn.Identity() def merge_kwargs(self, vid_kwargs, audio_kwargs): """ keys in each kwarg: e seq_lens grid_sizes freqs context context_lens """ merged_kwargs = {} for key in vid_kwargs: merged_kwargs[f"vid_{key}"] = vid_kwargs[key] for key in audio_kwargs: merged_kwargs[f"audio_{key}"] = audio_kwargs[key] return merged_kwargs def single_fusion_cross_attention_forward(self, cross_attn_block, src_seq, src_grid_sizes, src_freqs, target_seq, target_seq_lens, target_grid_sizes, target_freqs, context, context_lens ): b, n, d = src_seq.size(0), cross_attn_block.num_heads, cross_attn_block.head_dim if hasattr(cross_attn_block, "k_img"): ## means is i2v block q, k, v, k_img, v_img = cross_attn_block.qkv_fn(src_seq, context) else: ## means is t2v block q, k, v = cross_attn_block.qkv_fn(src_seq, context) k_img = v_img = None qkv_list =[q, k, v] del k, v x = pay_attention(qkv_list) if k_img is not None: qkv_list =[q, k_img, v_img] del k_img, v_img img_x = pay_attention(qkv_list) # img_x = flash_attention(q, k_img, v_img, k_lens=None) x += img_x is_vid = src_grid_sizes.shape[1] > 1 # compute target attention target_seq = cross_attn_block.pre_attn_norm_fusion(target_seq) k_target = cross_attn_block.norm_k_fusion(cross_attn_block.k_fusion(target_seq)).view(b, -1, n, d) v_target = cross_attn_block.v_fusion(target_seq).view(b, -1, n, d) q = rope_apply(q, src_grid_sizes, src_freqs) k_target = rope_apply(k_target, target_grid_sizes, target_freqs) qkv_list =[q, k_target, v_target] del q, k_target, v_target target_x = pay_attention(qkv_list) # target_x = flash_attention(q, k_target, v_target, k_lens=target_seq_lens) x += target_x x = x.flatten(2) # [B, L/P, C] x = cross_attn_block.o(x) return x def single_fusion_cross_attention_ffn_forward(self, attn_block, src_seq, src_grid_sizes, src_freqs, target_seq, target_seq_lens, target_grid_sizes, target_freqs, context, context_lens, src_e): src_seq += self.single_fusion_cross_attention_forward(attn_block.cross_attn, attn_block.norm3(src_seq), src_grid_sizes=src_grid_sizes, src_freqs=src_freqs, target_seq=target_seq, target_seq_lens=target_seq_lens, target_grid_sizes=target_grid_sizes, target_freqs=target_freqs, context=context, context_lens=context_lens ) latent_frames = src_e[0].shape[0] y = attn_block.norm2(src_seq).to(torch.bfloat16) y = reshape_latent(y , latent_frames) y *= (1 + src_e[4].squeeze(2)) y += src_e[3].squeeze(2) y = restore_latent_shape(y) # y = attn_block.ffn(y) ffn = attn_block.ffn[0] gelu = attn_block.ffn[1] ffn2= attn_block.ffn[2] y_shape = y.shape y = y.view(-1, y_shape[-1]) chunk_size = int(y.shape[0]/2.7) chunks =torch.split(y, chunk_size) for y_chunk in chunks: mlp_chunk = ffn(y_chunk) mlp_chunk = gelu(mlp_chunk) y_chunk[...] = ffn2(mlp_chunk) del mlp_chunk y = y.view(y_shape) src_seq, y = reshape_latent(src_seq , latent_frames), reshape_latent(y , latent_frames) src_seq.addcmul_(y, src_e[5].squeeze(2)) src_seq = restore_latent_shape(src_seq) del y # # y = attn_block.ffn(attn_block.norm2(src_seq).bfloat16() * (1 + src_e[4].squeeze(2)) + src_e[3].squeeze(2)) # with torch.amp.autocast('cuda', dtype=torch.bfloat16): # src_seq = src_seq + y * src_e[5].squeeze(2) return src_seq def single_fusion_block_forward(self, vid_block, audio_block, vid, audio, vid_e, vid_seq_lens, vid_grid_sizes, vid_freqs, vid_context, vid_context_lens, audio_e, audio_seq_lens, audio_grid_sizes, audio_freqs, audio_context, audio_context_lens, ): ## audio modulation audio_e = audio_block.modulation(audio_e).chunk(6, dim=1) # audio self-attention audio_y = audio_block.norm1(audio).to(torch.bfloat16) audio_y *= (1 + audio_e[1].squeeze(2)) audio_y += audio_e[0].squeeze(2) audio_y = audio_block.self_attn(audio_y, audio_seq_lens, audio_grid_sizes, audio_freqs) audio.addcmul_(audio_y, audio_e[2].squeeze(2)) del audio_y latent_frames = vid_e.shape[0] ## video modulation vid_e = vid_block.modulation(vid_e).chunk(6, dim=1) # video self-attention vid_y = vid_block.norm1(vid).to(torch.bfloat16) vid_y = reshape_latent(vid_y , latent_frames) vid_y *= (1 + vid_e[1].squeeze(2)) vid_y += vid_e[0].squeeze(2) vid_y = restore_latent_shape(vid_y) vid_y = vid_block.self_attn(vid_y, vid_seq_lens, vid_grid_sizes, vid_freqs) vid, vid_y = reshape_latent(vid , latent_frames), reshape_latent(vid_y , latent_frames) vid.addcmul_(vid_y, vid_e[2].squeeze(2)) vid = restore_latent_shape(vid) del vid_y # og_audio = audio # audio cross-attention audio = self.single_fusion_cross_attention_ffn_forward( audio_block, audio, audio_grid_sizes, audio_freqs, vid, vid_seq_lens, vid_grid_sizes, vid_freqs, audio_context, audio_context_lens, audio_e, ) if audio is None: return None, None # if torch.equal(og_audio, audio): # print("Audio should be changed after cross-attention!") # assert not torch.equal(og_audio, audio), "Audio should be changed after cross-attention!" # video cross-attention vid = self.single_fusion_cross_attention_ffn_forward( vid_block, vid, vid_grid_sizes, vid_freqs, audio, audio_seq_lens, audio_grid_sizes, audio_freqs, vid_context, vid_context_lens, vid_e, ) if vid is None: return None, None return vid, audio def forward( self, vid, audio, t, vid_context, audio_context, vid_seq_len, audio_seq_len, clip_fea=None, clip_fea_audio=None, y=None, first_frame_is_clean=False, computed_perturbation_layers=None, callback = None, pipeline= None, x_id_list = 0, video_freqs = None, audio_freqs = None, ): vid_list = [] vid_e_list = [] audio_list = [] audio_e_list = [] kwargs_list = [] for one_vid_context, one_audio_context in zip(vid_context, audio_context): one_vid, one_vid_e, vid_kwargs = self.video_model.prepare_transformer_block_kwargs( x=[vid], t=t, context=[one_vid_context], seq_len=vid_seq_len, clip_fea=clip_fea, y=y, first_frame_is_clean=first_frame_is_clean, freqs=video_freqs ) vid_list.append(one_vid) vid_e_list.append(one_vid_e) one_vid = one_vid_e = None one_audio, one_audio_e, audio_kwargs = self.audio_model.prepare_transformer_block_kwargs( x=[audio], t=t, context=[one_audio_context], seq_len=audio_seq_len, clip_fea=clip_fea_audio, y=None, first_frame_is_clean=False, freqs=audio_freqs ) audio_list.append(one_audio) audio_e_list.append(one_audio_e) one_audio = one_audio_e = None kwargs_list.append(self.merge_kwargs(vid_kwargs, audio_kwargs)) for i in range(self.num_blocks): """ 1 fusion block refers to 1 audio block with 1 video block. """ if callback != None: callback(-1, None, False, True) vid_block = self.video_model.blocks[i] audio_block = self.audio_model.blocks[i] for x_id, one_vid, one_vid_e, one_audio, one_audio_e, one_kwargs in zip(x_id_list, vid_list, vid_e_list, audio_list, audio_e_list, kwargs_list): if pipeline._interrupt: return None, None if x_id == 1 and computed_perturbation_layers is not None and i in computed_perturbation_layers: continue # one_vid[...], one_audio[...] = self.single_fusion_block_forward( a, b = self.single_fusion_block_forward( vid_block=vid_block, audio_block=audio_block, vid=one_vid, audio=one_audio, **one_kwargs ) one_vid[...], one_audio[...] = a, b for i, (x_id, one_vid, one_vid_e, one_audio, one_audio_e) in enumerate(zip(x_id_list, vid_list, vid_e_list, audio_list, audio_e_list)): one_vid = self.video_model.post_transformer_block_out(one_vid, vid_kwargs['grid_sizes'], one_vid_e) vid_list[i] = one_vid one_vid = None one_audio = self.audio_model.post_transformer_block_out(one_audio, audio_kwargs['grid_sizes'], one_audio_e) audio_list[i] = one_audio one_audio = None if len(vid_list) == 1: return vid_list[0], audio_list[0] return vid_list, audio_list def init_weights(self): if self.audio_model is not None: self.audio_model.init_weights() if self.video_model is not None: self.video_model.init_weights() for name, mod in self.video_model.named_modules(): if "fusion" in name and isinstance(mod, nn.Linear): with torch.no_grad(): mod.weight.div_(10.0) def custom_compile(self, compile_kwargs): torch.compile(self.single_fusion_block_forward, **compile_kwargs)