Multiple talkers

Initial commit, works but needs more utility for the mask creation.

Based mostly on Rudra-ai-coder's modifications.

Co-Authored-By: Rudra-ai-coder <177262225+rudra-ai-coder@users.noreply.github.com>

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

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"
 }

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

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
                 )