import torch import torch.nn as nn from ovi.modules.model import WanLayerNorm, WanModel, WanRMSNorm, gradient_checkpointing, rope_apply from ovi.modules.attention import flash_attention from ovi.distributed_comms.communications import all_gather, all_to_all_4D from ovi.distributed_comms.parallel_states import nccl_info, get_sequence_parallel_state from einops import rearrange class FusionModel(nn.Module): def __init__(self, video_config=None, audio_config=None, av2av_edit=False, concat_edit_source_latents=True, use_siga=False): super().__init__() has_video = True has_audio = True if video_config is not None: self.video_model = WanModel( av2av_edit=av2av_edit, concat_edit_source_latents=concat_edit_source_latents, use_siga=use_siga, **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( av2av_edit=av2av_edit, concat_edit_source_latents=concat_edit_source_latents, use_siga=use_siga, **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 = get_sequence_parallel_state() if self.use_sp: self.sp_size = nccl_info.sp_size self.sp_rank = nccl_info.rank_within_group self.inject_cross_attention_kv_projections() # self.gradient_checkpointing = True # 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 trainable_modules(self): trainable_modules = filter(lambda p: p.requires_grad, self.parameters()) return trainable_modules def enable_siga_gate_visualization( self, enabled=True, record_passes=("pos",), step_interval=10, ): if self.video_model is not None: self.video_model.enable_siga_gate_visualization( enabled=enabled, record_passes=record_passes, step_interval=step_interval, ) if self.audio_model is not None: self.audio_model.enable_siga_gate_visualization( enabled=enabled, record_passes=record_passes, step_interval=step_interval, ) def set_siga_gate_visualization_context(self, step_index=None, pass_name=None): if self.video_model is not None: self.video_model.set_siga_gate_visualization_context( step_index=step_index, pass_name=pass_name, ) if self.audio_model is not None: self.audio_model.set_siga_gate_visualization_context( step_index=step_index, pass_name=pass_name, ) def reset_siga_gate_visualization(self): if self.video_model is not None: self.video_model.reset_siga_gate_visualization() if self.audio_model is not None: self.audio_model.reset_siga_gate_visualization() def get_siga_gate_visualization(self): if self.video_model is None: return None return self.video_model.get_siga_gate_visualization() 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, attn_block=None, siga_attn_block=None, source_context=None, source_context_lens=None, is_audio_type=None, ): # import ipdb; ipdb.set_trace() 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 if self.use_sp: q = all_to_all_4D(q, scatter_dim=2, gather_dim=1) k = torch.chunk(k, self.sp_size, dim=2)[self.sp_rank] v = torch.chunk(v, self.sp_size, dim=2)[self.sp_rank] if k_img is not None: k_img = torch.chunk(k_img, self.sp_size, dim=2)[self.sp_rank] if v_img is not None: v_img = torch.chunk(v_img, self.sp_size, dim=2)[self.sp_rank] x = flash_attention(q, k, v, k_lens=context_lens) if attn_block is not None and attn_block.use_siga and source_context is not None: if (is_audio_type): q_siga, k_siga, v_siga = siga_attn_block.qkv_fn(src_seq, source_context) if self.use_sp: q_siga = all_to_all_4D(q_siga, scatter_dim=2, gather_dim=1) k_siga = torch.chunk(k_siga, self.sp_size, dim=2)[self.sp_rank] v_siga = torch.chunk(v_siga, self.sp_size, dim=2)[self.sp_rank] attn_siga = flash_attention(q_siga, k_siga, v_siga, k_lens=source_context_lens) x_flat = x.flatten(2) attn_siga_flat = attn_siga.flatten(2) gate_input = torch.cat([src_seq, attn_siga_flat, x_flat], dim=-1) source_gate = torch.sigmoid(attn_block.siga_gate(gate_input)) # print(source_gate) mixed = (1 - source_gate) * x_flat + source_gate * attn_siga_flat x = mixed.view_as(x) else: t = src_grid_sizes[0][0] h = src_grid_sizes[0][1] w = src_grid_sizes[0][2] src_seq_frames = rearrange(src_seq, 'b (t h w) d -> (b t) (h w) d', t=t, h=h, w=w) source_context_frames = rearrange(source_context, 'b (t h w) d -> (b t) (h w) d', t=t, h=h, w=w) source_context_lens_frames = (source_context_lens // t).repeat_interleave(t) q_siga, k_siga, v_siga = siga_attn_block.qkv_fn(src_seq_frames, source_context_frames) attn_siga = flash_attention(q_siga, k_siga, v_siga, k_lens=source_context_lens_frames) attn_siga_flat = attn_siga.flatten(2) attn_siga_flat = rearrange(attn_siga_flat, '(b t) (h w) d -> b (t h w) d', t=t, h=h, w=w) x_flat = x.flatten(2) gate_input = torch.cat([src_seq, attn_siga_flat, x_flat], dim=-1) source_gate = torch.sigmoid(attn_block.siga_gate(gate_input)) recorder = getattr(attn_block, "siga_gate_recorder", None) if recorder is not None: recorder( source_gate=source_gate, grid_sizes=src_grid_sizes, block_index=getattr(attn_block, "siga_block_index", None), ) mixed = (1 - source_gate) * x_flat + source_gate * attn_siga_flat x = mixed.view_as(x) if k_img is not None: img_x = flash_attention(q, k_img, v_img, k_lens=None) x = 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) if self.use_sp: k_target = all_to_all_4D(k_target, scatter_dim=2, gather_dim=1) # [B, L, H/P, C/H] v_target = all_to_all_4D(v_target, scatter_dim=2, gather_dim=1) # [B, L, H/P, C/H] q = rope_apply(q, src_grid_sizes, src_freqs) k_target = rope_apply(k_target, target_grid_sizes, target_freqs) target_x = flash_attention(q, k_target, v_target, k_lens=target_seq_lens) x = x + target_x if self.use_sp: x = all_to_all_4D(x, scatter_dim=1, gather_dim=2) # [B, L/P, H, C/H] 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, source_context=None, source_context_lens=None, is_audio_type=None): src_norm = attn_block.norm3(src_seq) if attn_block.use_siga and source_context is not None: attn_out = self.single_fusion_cross_attention_forward(attn_block.cross_attn, src_norm, 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, attn_block=attn_block, siga_attn_block=attn_block.siga_source_attn, source_context=source_context, source_context_lens=source_context_lens, is_audio_type=is_audio_type ) else: attn_out = self.single_fusion_cross_attention_forward(attn_block.cross_attn, src_norm, 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, is_audio_type=is_audio_type ) src_seq = src_seq + attn_out 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, vid_source_context, vid_source_context_lens, audio_e, audio_seq_lens, audio_grid_sizes, audio_freqs, audio_context, audio_context_lens, audio_source_context, audio_source_context_lens ): ## audio modulation assert audio_e.dtype == torch.bfloat16 assert len(audio_e.shape) == 4 and audio_e.size(2) == 6 and audio_e.shape[1] == audio.shape[1], f"{audio_e.shape}, {audio.shape}" with torch.amp.autocast('cuda', dtype=torch.bfloat16): audio_e = audio_block.modulation(audio_e).chunk(6, dim=2) assert audio_e[0].dtype == torch.bfloat16 # audio self-attention audio_y = audio_block.self_attn( audio_block.norm1(audio).bfloat16() * (1 + audio_e[1].squeeze(2)) + audio_e[0].squeeze(2), audio_seq_lens, audio_grid_sizes, audio_freqs) with torch.amp.autocast('cuda', dtype=torch.bfloat16): audio = audio + audio_y * audio_e[2].squeeze(2) ## video modulation assert len(vid_e.shape) == 4 and vid_e.size(2) == 6 and vid_e.shape[1] == vid.shape[1], f"{vid_e.shape}, {vid.shape}" with torch.amp.autocast('cuda', dtype=torch.bfloat16): vid_e = vid_block.modulation(vid_e).chunk(6, dim=2) # video self-attention vid_y = vid_block.self_attn( vid_block.norm1(vid).bfloat16() * (1 + vid_e[1].squeeze(2)) + vid_e[0].squeeze(2), vid_seq_lens, vid_grid_sizes, vid_freqs) with torch.amp.autocast('cuda', dtype=torch.bfloat16): vid = vid + vid_y * vid_e[2].squeeze(2) 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, source_context=audio_source_context, source_context_lens=audio_source_context_lens, is_audio_type=True, ) 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, og_audio, audio_seq_lens, audio_grid_sizes, audio_freqs, vid_context, vid_context_lens, vid_e, source_context=vid_source_context, source_context_lens=vid_source_context_lens, is_audio_type=False, ) 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, slg_layer=False, vid_source=None, audio_source=None ): # import pdb; pdb.set_trace() assert clip_fea is None assert y is None if vid is None or all([x is None for x in vid]): assert vid_context is None assert vid_seq_len is None assert self.audio_model is not None return None, self.audio_model(x=audio, t=t, context=audio_context, seq_len=audio_seq_len, clip_fea=clip_fea_audio, y=None, source_x=audio_source) if audio is None or all([x is None for x in audio]): assert clip_fea_audio is None assert audio_context is None assert audio_seq_len is None assert self.video_model is not None return self.video_model(x=vid, t=t, context=vid_context, seq_len=vid_seq_len, clip_fea=clip_fea, y=y, first_frame_is_clean=first_frame_is_clean, source_x=vid_source), None vid, vid_e, vid_kwargs = self.video_model.prepare_transformer_block_kwargs( x=vid, t=t, context=vid_context, seq_len=vid_seq_len, clip_fea=clip_fea, y=y, first_frame_is_clean=first_frame_is_clean, source_x=vid_source ) audio, audio_e, audio_kwargs = self.audio_model.prepare_transformer_block_kwargs( x=audio, t=t, context=audio_context, seq_len=audio_seq_len, clip_fea=clip_fea_audio, y=None, first_frame_is_clean=False, source_x=audio_source ) kwargs = self.merge_kwargs(vid_kwargs, audio_kwargs) # import pdb; pdb.set_trace() for i in range(self.num_blocks): """ 1 fusion block refers to 1 audio block with 1 video block. """ if slg_layer > 0 and i == slg_layer: continue # vid_block = self.video_model.blocks[i] audio_block = self.audio_model.blocks[i] # with torch.autograd.graph.save_on_cpu(): vid, audio = gradient_checkpointing( # enabled=(self.training and self.gradient_checkpointing), enabled=True, # enabled = False, module=self.single_fusion_block_forward, vid_block=vid_block, audio_block=audio_block, vid=vid, audio=audio, **kwargs ) vid = self.video_model.post_transformer_block_out(vid, vid_kwargs['grid_sizes'], vid_e) audio = self.audio_model.post_transformer_block_out(audio, audio_kwargs['grid_sizes'], audio_e) return vid, audio 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 set_rope_params(self, has_video = True, has_audio = True): if has_video: self.video_model.set_rope_params() if has_audio: self.audio_model.set_rope_params()