diff --git a/multitalk/nodes.py b/multitalk/nodes.py index 83606568..183d413c 100644 --- a/multitalk/nodes.py +++ b/multitalk/nodes.py @@ -352,6 +352,7 @@ class MultiTalkSilentEmbeds: def INPUT_TYPES(s): return {"required": { "num_frames": ("INT", {"default": 81, "min": 1, "max": 10000, "step": 1, "tooltip": "The total frame count to generate."}), + "audio_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01, "tooltip": "Strength of the audio conditioning"}), }, } @@ -360,17 +361,17 @@ def INPUT_TYPES(s): FUNCTION = "process" CATEGORY = "WanVideoWrapper" - def process(self, num_frames): + def process(self, num_frames, audio_scale): silence_path = os.path.join(script_directory, "encoded_silence.safetensors") encoded_silence = load_torch_file(silence_path)["audio_emb"] - + target_frames = num_frames repeats = (target_frames + encoded_silence.shape[0] - 1) // encoded_silence.shape[0] repeated = encoded_silence.repeat(repeats, 1, 1) repeated = repeated[:target_frames] multitalk_embeds = { "audio_features": repeated, - "audio_scale": 1.0, + "audio_scale": audio_scale, "audio_cfg_scale": 1.0, "ref_target_masks": None } diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py index ce101d8a..9417f9f6 100644 --- a/wanvideo/modules/model.py +++ b/wanvideo/modules/model.py @@ -1343,8 +1343,26 @@ def forward( if audio_output_cond is not None: x_audio = torch.cat([audio_output_cond, x_audio], dim=1).contiguous() else: - x_audio = self.audio_cross_attn(self.norm_x(x.to(self.norm_x.weight.dtype)).to(input_dtype), encoder_hidden_states=multitalk_audio_embedding, - shape=grid_sizes[0], x_ref_attn_map=x_ref_attn_map, human_num=human_num) + # audio temporal frames may differ from grid_sizes when external conditioning + # (e.g. SCAIL pose ref_latent) modifies the spatial-temporal token layout + x_normed = self.norm_x(x.to(self.norm_x.weight.dtype)).to(input_dtype) + N_h, N_w = grid_sizes[0][1].item(), grid_sizes[0][2].item() + S = N_h * N_w + audio_N_t = multitalk_audio_embedding.shape[0] + audio_tokens = audio_N_t * S + if grid_sizes[0][0].item() != audio_N_t or x_normed.shape[1] != audio_tokens: + x_for_audio = x_normed[:, :audio_tokens] + audio_shape = grid_sizes[0].clone() + audio_shape[0] = audio_N_t + else: + x_for_audio = x_normed + audio_shape = grid_sizes[0] + x_audio = self.audio_cross_attn(x_for_audio, encoder_hidden_states=multitalk_audio_embedding, + shape=audio_shape, x_ref_attn_map=x_ref_attn_map, human_num=human_num) + if x_normed.shape[1] > audio_tokens: + x_audio_padded = torch.zeros(x.shape, dtype=x_audio.dtype, device=x_audio.device) + x_audio_padded[:, :audio_tokens] = x_audio + x_audio = x_audio_padded x.add_(x_audio, alpha=audio_scale) del x_audio