diff --git a/multitalk/multitalk.py b/multitalk/multitalk.py index e89a399..9749ea0 100644 --- a/multitalk/multitalk.py +++ b/multitalk/multitalk.py @@ -297,6 +297,15 @@ class SingleStreamAttention(nn.Module): return x class SingleStreamMultiAttention(SingleStreamAttention): + """Multi-speaker rotary-position cross-attention. + + This implementation generalises the original 2-speaker logic to an arbitrary + number of voices. Each speaker is allocated a contiguous *class_interval* + segment inside a shared *class_range* rotary bucket. The centre of each + bucket is applied to that speaker's KV tokens while queries are modulated + per-token according to which speaker dominates the pixel. + """ + def __init__( self, dim: int, @@ -324,87 +333,122 @@ class SingleStreamMultiAttention(SingleStreamAttention): eps=eps, attention_mode=attention_mode, ) + + # Rotary-embedding layout parameters self.class_interval = class_interval self.class_range = class_range - self.rope_h1 = (0, self.class_interval) - self.rope_h2 = (self.class_range - self.class_interval, self.class_range) + self.max_humans = self.class_range // self.class_interval + + # Constant bucket used for background tokens self.rope_bak = int(self.class_range // 2) self.rope_1d = RotaryPositionalEmbedding1D(self.head_dim) self.attention_mode = attention_mode - def forward(self, + def forward( + self, x: torch.Tensor, encoder_hidden_states: torch.Tensor, shape=None, x_ref_attn_map=None, - human_num=None) -> torch.Tensor: + human_num=None, + ) -> torch.Tensor: encoder_hidden_states = encoder_hidden_states.squeeze(0) - if human_num == 1: + + # Single-speaker fall-through + if human_num is None or human_num <= 1: return super().forward(x, encoder_hidden_states, shape) + # Safety check: do we have enough buckets? + assert human_num <= self.max_humans, ( + f"Configured for at most {self.max_humans} speakers but got {human_num}.") + N_t, _, _ = shape x = rearrange(x, "B (N_t S) C -> (B N_t) S C", N_t=N_t) - # get q for hidden_state + # Query projection B, N, C = x.shape q = self.q_linear(x) - q_shape = (B, N, self.num_heads, self.head_dim) - q = q.view(q_shape).permute((0, 2, 1, 3)) - + q = q.view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3) if self.qk_norm: q = self.q_norm(q) - - max_values = x_ref_attn_map.max(1).values[:, None, None] - min_values = x_ref_attn_map.min(1).values[:, None, None] - max_min_values = torch.cat([max_values, min_values], dim=2) + # 1) Build per-speaker rotary ranges + rope_ranges = [ + (i * self.class_interval, (i + 1) * self.class_interval) + for i in range(human_num) + ] - human1_max_value, human1_min_value = max_min_values[0, :, 0].max(), max_min_values[0, :, 1].min() - human2_max_value, human2_min_value = max_min_values[1, :, 0].max(), max_min_values[1, :, 1].min() + # 2) Normalise each speaker's attention map into its own bucket + human_norm_list = [] + for idx in range(human_num): + attn_map = x_ref_attn_map[idx] + att_min, att_max = attn_map.min(), attn_map.max() + human_norm = normalize_and_scale( + attn_map, (att_min, att_max), rope_ranges[idx] + ) + human_norm_list.append(human_norm) - human1 = normalize_and_scale(x_ref_attn_map[0], (human1_min_value, human1_max_value), (self.rope_h1[0], self.rope_h1[1])) - human2 = normalize_and_scale(x_ref_attn_map[1], (human2_min_value, human2_max_value), (self.rope_h2[0], self.rope_h2[1])) - back = torch.full((x_ref_attn_map.size(1),), self.rope_bak, dtype=human1.dtype).to(human1.device) + # Background constant bucket + back = torch.full( + (x_ref_attn_map.size(1),), + self.rope_bak, + dtype=x_ref_attn_map.dtype, + device=x_ref_attn_map.device, + ) + + # Token-wise speaker dominance max_indices = x_ref_attn_map.argmax(dim=0) - normalized_map = torch.stack([human1, human2, back], dim=1) - normalized_pos = normalized_map[range(x_ref_attn_map.size(1)), max_indices] # N + normalized_map = torch.stack(human_norm_list + [back], dim=1) + normalized_pos = normalized_map[torch.arange(x_ref_attn_map.size(1)), max_indices] + # Apply rotary to Q q = rearrange(q, "(B N_t) H S C -> B H (N_t S) C", N_t=N_t) q = self.rope_1d(q, normalized_pos) q = rearrange(q, "B H (N_t S) C -> (B N_t) H S C", N_t=N_t) + # Keys / Values _, N_a, _ = encoder_hidden_states.shape encoder_kv = self.kv_linear(encoder_hidden_states) - encoder_kv_shape = (B, N_a, 2, self.num_heads, self.head_dim) - encoder_kv = encoder_kv.view(encoder_kv_shape).permute((2, 0, 3, 1, 4)) + encoder_kv = encoder_kv.view(B, N_a, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) encoder_k, encoder_v = encoder_kv.unbind(0) - if self.qk_norm: encoder_k = self.add_k_norm(encoder_k) - - per_frame = torch.zeros(N_a, dtype=encoder_k.dtype).to(encoder_k.device) - per_frame[:per_frame.size(0)//2] = (self.rope_h1[0] + self.rope_h1[1]) / 2 - per_frame[per_frame.size(0)//2:] = (self.rope_h2[0] + self.rope_h2[1]) / 2 - encoder_pos = torch.concat([per_frame]*N_t, dim=0) + # Rotary for keys – assign centre of each speaker bucket to its context tokens + tokens_per_human = N_a // human_num + encoder_pos_list = [] + for i in range(human_num): + start, end = rope_ranges[i] + centre = (start + end) / 2 + encoder_pos_list.append( + torch.full( + (tokens_per_human,), centre, dtype=encoder_k.dtype, device=encoder_k.device + ) + ) + encoder_pos = torch.cat(encoder_pos_list * N_t, dim=0) + encoder_k = rearrange(encoder_k, "(B N_t) H S C -> B H (N_t S) C", N_t=N_t) encoder_k = self.rope_1d(encoder_k, encoder_pos) encoder_k = rearrange(encoder_k, "B H (N_t S) C -> (B N_t) H S C", N_t=N_t) - x = torch.nn.functional.scaled_dot_product_attention( - q, encoder_k, encoder_v, attn_mask=None, is_causal=False, dropout_p=0.0) + # Final attention + q = rearrange(q, "B H M K -> B M H K") + encoder_k = rearrange(encoder_k, "B H M K -> B M H K") + encoder_v = rearrange(encoder_v, "B H M K -> B M H K") + x = attention( + q, encoder_k, encoder_v, attention_mode=self.attention_mode + ) + # x is already in B M H K format from attention function - # linear transform - x_output_shape = (B, N, C) - x = x.transpose(1, 2) - x = x.reshape(x_output_shape) + # Linear projection + x = x.reshape(B, N, C) x = self.proj(x) x = self.proj_drop(x) - # reshape x to origin shape + # Restore original layout x = rearrange(x, "(B N_t) S C -> B (N_t S) C", N_t=N_t) return x \ No newline at end of file diff --git a/multitalk/nodes.py b/multitalk/nodes.py index 23b9897..84af6a1 100644 --- a/multitalk/nodes.py +++ b/multitalk/nodes.py @@ -4,6 +4,7 @@ from comfy.utils import load_torch_file, common_upscale from accelerate import init_empty_weights from accelerate.utils import set_module_tensor_to_device import torch +from ..utils import log class MultiTalkModelLoader: @@ -81,13 +82,20 @@ class MultiTalkWav2VecEmbeds: def INPUT_TYPES(s): return {"required": { "wav2vec_model": ("WAV2VECMODEL",), - "audio": ("AUDIO",), + "audio_1": ("AUDIO",), "normalize_loudness": ("BOOLEAN", {"default": True}), "num_frames": ("INT", {"default": 81, "min": 1, "max": 1000, "step": 1}), - "fps": ("FLOAT", {"default": 23.0, "min": 1.0, "max": 60.0, "step": 0.1}), + "fps": ("FLOAT", {"default": 25.0, "min": 1.0, "max": 60.0, "step": 0.1}), "audio_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "tooltip": "Strength of the audio conditioning"}), "audio_cfg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "tooltip": "When not 1.0, an extra model pass without audio conditioning is done: slower inference but more motion is allowed"}), - }, + "multi_audio_type": (["para", "add"], {"default": "para", "tooltip": "'para' overlay speakers in parallel, 'add' concatenate sequentially"}), + }, + "optional" : { + "audio_2": ("AUDIO",), + "audio_3": ("AUDIO",), + "audio_4": ("AUDIO",), + "ref_target_masks": ("MASK", {"tooltip": "Per-speaker semantic mask(s) in pixel space. Supply one mask per speaker (plus optional background) to guide mouth assignment"}), + } } RETURN_TYPES = ("MULTITALK_EMBEDS", "AUDIO", ) @@ -95,7 +103,7 @@ class MultiTalkWav2VecEmbeds: FUNCTION = "process" CATEGORY = "WanVideoWrapper" - def process(self, wav2vec_model, normalize_loudness, fps, num_frames, audio, audio_scale, audio_cfg_scale): + def process(self, wav2vec_model, normalize_loudness, fps, num_frames, audio_1, audio_scale, audio_cfg_scale, multi_audio_type, audio_2=None, audio_3=None, audio_4=None, ref_target_masks=None): import torchaudio import numpy as np from einops import rearrange @@ -108,60 +116,208 @@ class MultiTalkWav2VecEmbeds: sr = 16000 - audio_input = audio["waveform"] - sample_rate = audio["sample_rate"] - if sample_rate != sr: - audio_input = torchaudio.functional.resample(audio_input, sample_rate, sr) - audio_input = audio_input[0][0] + audio_inputs = [audio_1, audio_2, audio_3, audio_4] + audio_inputs = [a for a in audio_inputs if a is not None] - start_time = 0 - end_time = num_frames / fps + multitalk_audio_features = [] + seq_lengths = [] + audio_outputs = [] # for debugging / optional saving – choose first as return - start_sample = int(start_time * sr) - end_sample = int(end_time * sr) + for audio in audio_inputs: + audio_input = audio["waveform"] + sample_rate = audio["sample_rate"] - try: - audio_segment = audio_input[start_sample:end_sample] - except: - audio_segment = audio_input + if sample_rate != 16000: + audio_input = torchaudio.functional.resample(audio_input, sample_rate, sr) + audio_input = audio_input[0][0] - audio_segment = audio_segment.numpy() + start_time = 0 + end_time = num_frames / fps - if normalize_loudness: - audio_segment = loudness_norm(audio_segment, sr=sr) + start_sample = int(start_time * sr) + end_sample = int(end_time * sr) - audio_feature = np.squeeze( - wav2vec_feature_extractor(audio_segment, sampling_rate=sr).input_values - ) + try: + audio_segment = audio_input[start_sample:end_sample] + except Exception: + audio_segment = audio_input - audio_feature = torch.from_numpy(audio_feature).float().to(device=device) - audio_feature = audio_feature.unsqueeze(0) + audio_segment = audio_segment.numpy() - # audio encoder - audio_duration = len(audio_segment) / sr - video_length = audio_duration * fps - print("Audio duration:", audio_duration, "Video length:", video_length) - embeddings = wav2vec(audio_feature.to(dtype), seq_len=int(video_length), output_hidden_states=True) + if normalize_loudness: + audio_segment = loudness_norm(audio_segment, sr=sr) - if len(embeddings) == 0: - print("Fail to extract audio embedding") - return None + audio_feature = np.squeeze( + wav2vec_feature_extractor(audio_segment, sampling_rate=sr).input_values + ) - audio_emb = torch.stack(embeddings.hidden_states[1:], dim=1).squeeze(0) - audio_emb = rearrange(audio_emb, "b s d -> s b d") + audio_feature = torch.from_numpy(audio_feature).float().to(device=device) + audio_feature = audio_feature.unsqueeze(0) + + # audio encoder + audio_duration = len(audio_segment) / sr + video_length = audio_duration * fps + + embeddings = wav2vec(audio_feature.to(dtype), seq_len=int(video_length), output_hidden_states=True) + + if len(embeddings) == 0: + print("Fail to extract audio embedding for one speaker") + continue + + audio_emb = torch.stack(embeddings.hidden_states[1:], dim=1).squeeze(0) + audio_emb = rearrange(audio_emb, "b s d -> s b d") + + multitalk_audio_features.append(audio_emb.cpu().detach()) + seq_lengths.append(audio_emb.shape[0]) + + waveform_tensor = torch.from_numpy(audio_segment).float().cpu().unsqueeze(0).unsqueeze(0) # (B, C, N) + audio_outputs.append({"waveform": waveform_tensor, "sample_rate": sr}) + + log.info("[MultiTalk] --- Raw speaker lengths (samples) ---") + for idx, ao in enumerate(audio_outputs): + log.info(f" speaker {idx+1}: {ao['waveform'].shape[-1]} samples (shape: {ao['waveform'].shape})") + + # Pad / combine depending on multi_audio_type + if len(multitalk_audio_features) > 1: + if multi_audio_type == "para": + max_len = max(seq_lengths) + padded = [] + for emb in multitalk_audio_features: + if emb.shape[0] < max_len: + pad = torch.zeros(max_len - emb.shape[0], *emb.shape[1:], dtype=emb.dtype) + emb = torch.cat([emb, pad], dim=0) + padded.append(emb) + multitalk_audio_features = padded + elif multi_audio_type == "add": + total_len = sum(seq_lengths) + full_list = [] + offset = 0 + for emb, length in zip(multitalk_audio_features, seq_lengths): + full = torch.zeros(total_len, *emb.shape[1:], dtype=emb.dtype) + full[offset:offset+length] = emb + full_list.append(full) + offset += length + multitalk_audio_features = full_list + + # fallback + if len(multitalk_audio_features) == 0: + raise RuntimeError("No valid audio embeddings extracted, please check inputs") multitalk_embeds = { - "audio_features": audio_emb, + "audio_features": multitalk_audio_features, "audio_scale": audio_scale, - "audio_cfg_scale": audio_cfg_scale + "audio_cfg_scale": audio_cfg_scale, + "ref_target_masks": ref_target_masks } - audio_output = { - "waveform": audio_feature.unsqueeze(0).cpu(), - "sample_rate": sr - } + if len(audio_outputs) == 1: # single speaker + out_audio = audio_outputs[0] + else: # multi speaker + if multi_audio_type == "para": + # Overlay speakers in parallel – mix waveforms to same length (max len) + max_len = max([a["waveform"].shape[-1] for a in audio_outputs]) + mixed = torch.zeros(1, 1, max_len, dtype=audio_outputs[0]["waveform"].dtype) + for a in audio_outputs: + w = a["waveform"] + if w.shape[-1] < max_len: + w = torch.nn.functional.pad(w, (0, max_len - w.shape[-1])) + mixed += w + out_audio = {"waveform": mixed, "sample_rate": sr} + else: # "add" – sequential concatenate with silent padding for other speakers + total_len = sum([a["waveform"].shape[-1] for a in audio_outputs]) + mixed = torch.zeros(1, 1, total_len, dtype=audio_outputs[0]["waveform"].dtype) + offset = 0 + for a in audio_outputs: + w = a["waveform"] + mixed[:, :, offset:offset + w.shape[-1]] += w + offset += w.shape[-1] + out_audio = {"waveform": mixed, "sample_rate": sr} + + # Debug: log final mixed audio length and mode + total_samples_raw = sum([ao["waveform"].shape[-1] for ao in audio_outputs]) + log.info(f"[MultiTalk] total raw duration = {total_samples_raw/sr:.3f}s") + log.info(f"[MultiTalk] multi_audio_type={multi_audio_type} | final waveform shape={out_audio['waveform'].shape} | length={out_audio['waveform'].shape[-1]} samples | seconds={out_audio['waveform'].shape[-1]/sr:.3f}s (expected {'sum' if multi_audio_type=='add' else 'max'} of raw)") + + return (multitalk_embeds, out_audio) - return (multitalk_embeds, audio_output) +class MultiTalkReferenceMasks: + @classmethod + def INPUT_TYPES(s): + return{ + "required" : { + "width": ("INT", {"default": 832, "min": 64, "max": 4096, "step": 16}), + "height": ("INT", {"default": 480, "min": 64, "max": 4096, "step": 16}), + "human_number": ("INT", {"default": 2, "min": 1, "max": 4, "step": 1, "tooltip": "Number of speakers (1-4)"}), + }, + + "optional" : { + # Bounding boxes as comma-separated string: "x_min,y_min,x_max,y_max" in pixel coordinates + "bbox_person1": ("STRING", {"default": "", "multiline": False, "tooltip": "Bounding box for speaker 1 (x_min,y_min,x_max,y_max). Leave empty to auto-split."}), + "bbox_person2": ("STRING", {"default": "", "multiline": False, "tooltip": "Bounding box for speaker 2."}), + "bbox_person3": ("STRING", {"default": "", "multiline": False, "tooltip": "Bounding box for speaker 3."}), + "bbox_person4": ("STRING", {"default": "", "multiline": False, "tooltip": "Bounding box for speaker 4."}), + "face_scale": ("FLOAT", {"default": 0.05, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Used when bboxes are not provided, defines central face band height."}), + } + } + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("ref_target_masks",) + FUNCTION = "generate" + CATEGORY = "WanVideoWrapper" + + def _parse_bbox(self, bbox_str): + try: + parts = [int(float(x.strip())) for x in bbox_str.split(',')] + if len(parts) == 4: + return parts # x_min, y_min, x_max, y_max + except Exception: + pass + return None + + def _build_mask_from_bbox(self, h, w, bbox): + x_min, y_min, x_max, y_max = bbox + mask = torch.zeros(h, w) + mask[x_min:x_max, y_min:y_max] = 1.0 + return mask + + def generate(self, width, height, human_number, bbox_person1="", bbox_person2="", bbox_person3="", bbox_person4="", face_scale=0.05): + device = mm.get_torch_device() + human_masks = [] + + # Build human masks based on inputs + if human_number == 1: + # Single speaker covers whole frame + human_masks.append(torch.ones(height, width)) + elif 2 <= human_number <= 4: + # Gather bbox strings list up to human_number + bbox_strings = [bbox_person1, bbox_person2, bbox_person3, bbox_person4][:human_number] + + # Pre-compute default vertical splits for fallback + segment_w = width // human_number + x_min_def = int(height * face_scale) + x_max_def = int(height * (1.0 - face_scale)) + + for idx in range(human_number): + bbox = self._parse_bbox(bbox_strings[idx]) + if bbox is None: + # create default bbox in segment idx + y_start = idx * segment_w + y_end = (idx + 1) * segment_w if idx < human_number - 1 else width + y_min_def = int(y_start + segment_w * face_scale) + y_max_def = int(y_end - segment_w * face_scale) + bbox = [x_min_def, y_min_def, x_max_def, y_max_def] + human_masks.append(self._build_mask_from_bbox(height, width, bbox)) + else: + raise ValueError("human_number must be between 1 and 4 for this node.") + + # Background mask – 1 where no speaker mask, 0 where speaker present + combined = torch.stack(human_masks, 0).sum(dim=0).clamp_max(1) + bg_mask = (1.0 - combined) + human_masks.append(bg_mask) + + ref_target_masks = torch.stack(human_masks, 0).float().to(device) # (N, H, W) + return (ref_target_masks,) + class WanVideoImageToVideoMultiTalk: @classmethod @@ -236,12 +392,14 @@ class WanVideoImageToVideoMultiTalk: NODE_CLASS_MAPPINGS = { "MultiTalkModelLoader": MultiTalkModelLoader, "MultiTalkWav2VecEmbeds": MultiTalkWav2VecEmbeds, - "WanVideoImageToVideoMultiTalk": WanVideoImageToVideoMultiTalk + "WanVideoImageToVideoMultiTalk": WanVideoImageToVideoMultiTalk, + "MultiTalkReferenceMasks": MultiTalkReferenceMasks } NODE_DISPLAY_NAME_MAPPINGS = { "MultiTalkModelLoader": "MultiTalk Model Loader", "MultiTalkWav2VecEmbeds": "MultiTalk Wav2Vec Embeds", - "WanVideoImageToVideoMultiTalk": "WanVideo Image To Video MultiTalk" + "WanVideoImageToVideoMultiTalk": "WanVideo Image To Video MultiTalk", + "MultiTalkReferenceMasks": "MultiTalk Reference Masks" } \ No newline at end of file diff --git a/nodes.py b/nodes.py index 2368b98..e238fe3 100644 --- a/nodes.py +++ b/nodes.py @@ -2960,12 +2960,25 @@ class WanVideoSampler: audio_cfg_scale = [audio_cfg_scale] * (steps +1) log.info(f"Audio proj shape: {audio_proj.shape}, audio context lens: {audio_context_lens}") elif multitalk_embeds is not None: - multitalk_audio_embedding = multitalk_embeds["audio_features"].to(device, dtype) - audio_scale = multitalk_embeds["audio_scale"] - audio_cfg_scale = multitalk_embeds["audio_cfg_scale"] + # Handle single or multiple speaker embeddings + audio_features_in = multitalk_embeds.get("audio_features", None) + if audio_features_in is None: + multitalk_audio_embedding = None + else: + if isinstance(audio_features_in, list): + multitalk_audio_embedding = [emb.to(device, dtype) for emb in audio_features_in] + else: + # keep backward-compatibility with single tensor input + multitalk_audio_embedding = [audio_features_in.to(device, dtype)] + + audio_scale = multitalk_embeds.get("audio_scale", 1.0) + audio_cfg_scale = multitalk_embeds.get("audio_cfg_scale", 1.0) + ref_target_masks = multitalk_embeds.get("ref_target_masks", None) if not isinstance(audio_cfg_scale, list): - audio_cfg_scale = [audio_cfg_scale] * (steps +1) - log.info(f"Multitalk audio features shape: {multitalk_audio_embedding.shape}") + audio_cfg_scale = [audio_cfg_scale] * (steps + 1) + + shapes = [tuple(e.shape) for e in multitalk_audio_embedding] + log.info(f"Multitalk audio features shapes (per speaker): {shapes}") minimax_latents = minimax_mask_latents = None @@ -3346,12 +3359,13 @@ class WanVideoSampler: z_pos = z_neg = torch.cat([z, minimax_latents, minimax_mask_latents], dim=0) if not multitalk_sampling and multitalk_audio_embedding is not None: - audio_embedding = [multitalk_audio_embedding] + audio_embedding = multitalk_audio_embedding audio_embs = [] - indices = (torch.arange(4 + 1) - 2) * 1 + indices = (torch.arange(4 + 1) - 2) * 1 + human_num = len(audio_embedding) # split audio with window size if context_window is None: - for human_idx in range(1): + for human_idx in range(human_num): center_indices = torch.arange( 0, latent_video_length * 4 + 1 if add_cond is not None else (latent_video_length-1) * 4 + 1, @@ -3363,7 +3377,7 @@ class WanVideoSampler: audio_emb = audio_embedding[human_idx][center_indices].unsqueeze(0).to(device) audio_embs.append(audio_emb) else: - for human_idx in range(1): + for human_idx in range(human_num): audio_start = context_window[0] * 4 audio_end = context_window[-1] * 4 + 1 print("audio_start: ", audio_start, "audio_end: ", audio_end) @@ -3405,6 +3419,7 @@ class WanVideoSampler: "nag_params": text_embeds.get("nag_params", {}), "nag_context": text_embeds.get("nag_prompt_embeds", None), "multitalk_audio": multitalk_audio_input if multitalk_audio_embedding is not None else None, + "ref_target_masks": ref_target_masks if multitalk_audio_embedding is not None else None } batch_size = 1 diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py index 395d654..166d8a6 100644 --- a/wanvideo/modules/model.py +++ b/wanvideo/modules/model.py @@ -1472,14 +1472,14 @@ class WanModel(ModelMixin, ConfigMixin): multitalk_audio_embedding = torch.concat(multitalk_audio_embedding.split(1), dim=2).to(x.dtype) # convert ref_target_masks to token_ref_target_masks - # !not implemented! + token_ref_target_masks = None if ref_target_masks is not None: ref_target_masks = ref_target_masks.unsqueeze(0).to(torch.float32) token_ref_target_masks = nn.functional.interpolate(ref_target_masks, size=(H // 2, W // 2), mode='nearest') token_ref_target_masks = token_ref_target_masks.squeeze(0) token_ref_target_masks = (token_ref_target_masks > 0) token_ref_target_masks = token_ref_target_masks.view(token_ref_target_masks.shape[0], -1) - token_ref_target_masks = token_ref_target_masks.to(x.dtype) + token_ref_target_masks = token_ref_target_masks.to(x.dtype).to(device) should_calc = True accumulated_rel_l1_distance = torch.tensor(0.0, dtype=torch.float32, device=device) @@ -1590,7 +1590,7 @@ class WanModel(ModelMixin, ConfigMixin): nag_context=nag_context, is_uncond = is_uncond, multitalk_audio_embedding=multitalk_audio_embedding if multitalk_audio is not None else None, - ref_target_masks=ref_target_masks if multitalk_audio is not None else None, + ref_target_masks=token_ref_target_masks if multitalk_audio is not None else None, human_num=human_num if multitalk_audio is not None else 0 )