Fix/infinitetalk-multi#1258
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request generalizes the InfiniteTalk model and runner to support multi-person inputs (more than two people) by updating the RoPE calculation, enabling directory-based configuration loading, supporting multi-track audio mixing, and improving frame mapping and audio muxing. The review feedback highlights several critical bug fixes and optimizations, including preventing None audio paths from being treated as strings, slicing mask and bounding box inputs to avoid shape mismatches, simplifying RoPE range calculations, ensuring robust rounding in torch.linspace, and optimizing memory usage during audio summation.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| cond_audio, mask_files, bbox = self._load_directory_audio_input(audio_path) | ||
| else: | ||
| cond_audio = {} | ||
| audio_paths = [item.strip() for item in str(audio_path).split(",") if item.strip()] |
There was a problem hiding this comment.
If audio_path is None, calling str(audio_path) converts it to the string "None". This will result in audio_paths being ["None"], which bypasses the empty check in _load_input_data and eventually causes a FileNotFoundError when trying to load the audio file. It is safer to default audio_path to an empty string if it is None.
| audio_paths = [item.strip() for item in str(audio_path).split(",") if item.strip()] | |
| audio_path_str = audio_path or "" | |
| audio_paths = [item.strip() for item in audio_path_str.split(",") if item.strip()] |
| raise ValueError("InfiniteTalk multi-person input requires one mask file for each person audio.") | ||
| background_mask = torch.zeros([self.src_h, self.src_w]) | ||
| for _, person_bbox in self.input_data["bbox"].items(): | ||
| for _, mask_file in self._sorted_person_items(self.input_data["mask_files"]): |
There was a problem hiding this comment.
If the number of mask files in self.input_data["mask_files"] is greater than human_num, looping over all of them will append too many masks to human_masks. This will cause a shape mismatch in the model (which expects exactly human_num + 1 masks). We should slice the sorted list to [:human_num] to ensure we only process the active speakers' masks.
| for _, mask_file in self._sorted_person_items(self.input_data["mask_files"]): | |
| for _, mask_file in self._sorted_person_items(self.input_data["mask_files"])[:human_num]: |
| human_masks.append(human_mask) | ||
| elif "bbox" in self.input_data: | ||
| background_mask = torch.zeros([self.src_h, self.src_w]) | ||
| for _, person_bbox in self._sorted_person_items(self.input_data["bbox"]): |
There was a problem hiding this comment.
If the number of bboxes in self.input_data["bbox"] is greater than human_num, looping over all of them will append too many masks to human_masks. This will cause a shape mismatch in the model (which expects exactly human_num + 1 masks). We should slice the sorted list to [:human_num] to ensure we only process the active speakers' bboxes.
| for _, person_bbox in self._sorted_person_items(self.input_data["bbox"]): | |
| for _, person_bbox in self._sorted_person_items(self.input_data["bbox"])[:human_num]: |
| h1 = torch.tensor(self.rope_h1, dtype=dtype, device=device) | ||
| h2 = torch.tensor(self.rope_h2, dtype=dtype, device=device) | ||
| if human_num == 2: | ||
| return torch.stack([h1, h2], dim=0) | ||
| starts = torch.linspace(h1[0], h2[0], human_num, dtype=dtype, device=device) | ||
| ends = torch.linspace(h1[1], h2[1], human_num, dtype=dtype, device=device) | ||
| return torch.stack([starts, ends], dim=1) |
There was a problem hiding this comment.
The _multi_human_rope_ranges method can be simplified by directly using the scalar values from self.rope_h1 and self.rope_h2 in torch.linspace. This avoids creating temporary tensors h1 and h2 on the device and eliminates the redundant if human_num == 2 branch, as torch.linspace with steps=2 naturally produces the same output.
starts = torch.linspace(self.rope_h1[0], self.rope_h2[0], human_num, dtype=dtype, device=device)
ends = torch.linspace(self.rope_h1[1], self.rope_h2[1], human_num, dtype=dtype, device=device)
return torch.stack([starts, ends], dim=1)| per_frame = torch.zeros(audio_tokens, dtype=encoder_k.dtype, device=encoder_k.device) | ||
| per_frame[: audio_tokens // 2] = (self.rope_h1[0] + self.rope_h1[1]) / 2 | ||
| per_frame[audio_tokens // 2 :] = (self.rope_h2[0] + self.rope_h2[1]) / 2 | ||
| token_edges = torch.linspace(0, audio_tokens, human_map_count + 1, dtype=torch.int64, device=encoder_k.device) |
There was a problem hiding this comment.
Using torch.linspace with an integer dtype can lead to inconsistent rounding or truncation behaviors across different PyTorch versions. It is safer and more robust to generate the linspace as float, round it, and then cast to int64.
| token_edges = torch.linspace(0, audio_tokens, human_map_count + 1, dtype=torch.int64, device=encoder_k.device) | |
| token_edges = torch.linspace(0, audio_tokens, human_map_count + 1, device=encoder_k.device).round().to(torch.int64) |
| else: | ||
| raise ValueError(f"Unsupported InfiniteTalk audio_type: {audio_type}") | ||
|
|
||
| return prepared, np.sum(np.stack(prepared, axis=0), axis=0) |
There was a problem hiding this comment.
Using np.stack on the list of prepared audio arrays creates a large intermediate 2D array in memory before summing it. Since all arrays in prepared have the same shape, we can use np.add.reduce(prepared) to perform element-wise addition directly, which is more memory-efficient.
| return prepared, np.sum(np.stack(prepared, axis=0), axis=0) | |
| return prepared, np.add.reduce(prepared) |
2 participants
把infinitetalk的多人逻辑和sekotalk对齐了一下，之后可以上线这个多人版本，实测效果比目前的多人好