feat: Implement batch processing for the MDXC separator#262
Closed
pedroalmeida415 wants to merge 6 commits intonomadkaraoke:mainfrom
Closed
feat: Implement batch processing for the MDXC separator#262pedroalmeida415 wants to merge 6 commits intonomadkaraoke:mainfrom
pedroalmeida415 wants to merge 6 commits intonomadkaraoke:mainfrom
Conversation
WalkthroughReplaces per-step Roformer demux with a DataLoader-based batching path: adds Changes
Sequence DiagramsequenceDiagram
participant Sep as Separator
participant DS as RoformerDataset
participant DL as DataLoader
participant Model as RoformerModel
participant OA as OverlapAdd
Sep->>DS: Create(mix, chunk_size, step)
Sep->>DL: Create(DS, batch_size, num_workers, pin_memory)
loop per-batch
DL->>DS: __getitem__ (batch fetch)
DS-->>DL: (chunks, start_indices, lengths)
DL-->>Sep: Batch of chunks
Sep->>Model: Move batch to device and forward
Model-->>Sep: Predictions (batch)
Sep->>OA: Overlap-add using start_indices and lengths (accumulate on CPU)
end
OA-->>Sep: Final demixed outputs
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
📝 Coding Plan
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
audio_separator/separator/architectures/mdxc_separator.py (1)
347-364: Good implementation of batched inference with DataLoader.The DataLoader integration with
pin_memoryfor CUDA and the batch-wise overlap-add logic is well structured.Minor suggestion: Consider using
xs.shape[0]instead oflen(xs)at line 354 for clearer tensor semantics.♻️ Optional: Use tensor shape accessor
- for b in range(len(xs)): + for b in range(xs.shape[0]):🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@audio_separator/separator/architectures/mdxc_separator.py` around lines 347 - 364, Replace the Python built-in len(xs) with the tensor shape accessor xs.shape[0] in the batch loop to use tensor semantics; in mdxc_separator.py within the batched inference loop that calls self.model_run(parts) and then iterates over outputs, change "for b in range(len(xs)):" to "for b in range(xs.shape[0]):" so the loop uses the tensor's first-dimension size reliably when xs is a tensor returned by self.model_run.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@audio_separator/separator/architectures/mdxc_separator.py`:
- Around line 17-40: The RoformerDataset.__getitem__ last-chunk handling can
produce a negative start_idx and an incorrect length when mix.shape[1] <
chunk_size; update the block to clamp start_idx with start_idx = max(0,
self.mix.shape[1] - self.chunk_size) and set length from the actual slice
(length = part.shape[-1]) after computing part, so that for very short audio the
returned part, start_idx, and length are consistent and non-negative (refer to
RoformerDataset, __getitem__, self.mix, chunk_size, and length).
---
Nitpick comments:
In `@audio_separator/separator/architectures/mdxc_separator.py`:
- Around line 347-364: Replace the Python built-in len(xs) with the tensor shape
accessor xs.shape[0] in the batch loop to use tensor semantics; in
mdxc_separator.py within the batched inference loop that calls
self.model_run(parts) and then iterates over outputs, change "for b in
range(len(xs)):" to "for b in range(xs.shape[0]):" so the loop uses the tensor's
first-dimension size reliably when xs is a tensor returned by self.model_run.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: fc0d3223-134e-403d-8fbc-c4caa174747a
📒 Files selected for processing (3)
audio_separator/separator/architectures/mdxc_separator.pyaudio_separator/separator/separator.pyaudio_separator/utils/cli.py
There was a problem hiding this comment.
🧹 Nitpick comments (3)
tests/unit/test_cli.py (1)
31-31: Add one non-default--mdxc_num_workerstest.Line 31 only covers the default
0. If the CLI accepted--mdxc_num_workersbut still always forwarded0, this suite would still pass. A small targeted test with a non-default value would close that gap.def test_cli_mdxc_num_workers_argument(common_expected_args): test_args = ["cli.py", "test_audio.mp3", "--mdxc_num_workers=2"] with patch("sys.argv", test_args): with patch("audio_separator.separator.Separator") as mock_separator: mock_separator.return_value.separate.return_value = ["output_file.mp3"] main() expected_args = common_expected_args.copy() expected_args["mdxc_params"] = { **common_expected_args["mdxc_params"], "num_workers": 2, } mock_separator.assert_called_once_with(**expected_args)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unit/test_cli.py` at line 31, Add a new unit test that verifies the CLI forwards a non-default --mdxc_num_workers value into the mdxc_params passed to Separator: call main() with argv containing "--mdxc_num_workers=2", patch audio_separator.separator.Separator, assert Separator was called once and that its mdxc_params dict has "num_workers": 2 (keeping other mdxc_params from common_expected_args); reference main(), Separator, and the mdxc_params key when implementing the test.audio_separator/separator/architectures/mdxc_separator.py (2)
35-35: Avoid scheduling the same tail window twice.Line 35 schedules starts all the way to the raw tail, and Lines 60-65 remap any short tail back to
mix.shape[1] - chunk_size. When that last full-window start is already on a step boundary, the same end chunk gets inferred twice.♻️ Proposed fix
- self.indices = list(range(0, mix.shape[1], step)) + if mix.shape[1] <= chunk_size: + self.indices = [0] + else: + last_start = mix.shape[1] - chunk_size + self.indices = list(range(0, last_start + 1, step)) + if self.indices[-1] != last_start: + self.indices.append(last_start)- # We need to handle the last chunk where part is smaller than chunk_size - if length < self.chunk_size and self.mix.shape[1] >= self.chunk_size: - # Take the last chunk_size from the end - part = self.mix[:, -self.chunk_size :] - length = self.chunk_size - start_idx = self.mix.shape[1] - self.chunk_size - # If mix is shorter than chunk_size, keep original part and length - return part, start_idx, lengthAlso applies to: 60-65
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@audio_separator/separator/architectures/mdxc_separator.py` at line 35, The current scheduling can append a tail start that duplicates an existing start when mix.shape[1]-chunk_size falls on a step boundary; update the logic that builds self.indices so that the final list never contains duplicate starts: after creating indices = list(range(0, mix.shape[1], step)) and before applying the tail remap (the code that references chunk_size and mix.shape[1]), compute last_start = mix.shape[1] - chunk_size and ensure you only append or remap to last_start if it is not already present (or alternatively deduplicate self.indices preserving order), so the tail window is not scheduled twice.
377-385: Keep the host/device transfers batched.Line 377 pins the batch, but Line 380 still uses a blocking
.to(device). Line 384 then copies each item back separately with.cpu(), which reintroduces per-sample D2H transfers immediately after batching.⚡ Proposed fix
- for parts, start_idxs, lengths in tqdm(dataloader): - parts = parts.to(device) - xs = self.model_run(parts) + for parts, start_idxs, lengths in tqdm(dataloader): + parts = parts.to(device, non_blocking=(device.type == "cuda")) + xs = self.model_run(parts).cpu() for b in range(xs.shape[0]): - x = xs[b].cpu() + x = xs[b] start_idx = start_idxs[b].item() length = lengths[b].item()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@audio_separator/separator/architectures/mdxc_separator.py` around lines 377 - 385, The per-sample device-to-host transfers happen because you call .to(device) on the whole batch (parts) but then call .cpu() inside the loop for each sample (xs[b].cpu()); fix by doing the host copy once per batch: after xs = self.model_run(parts) perform a single batch-level transfer (e.g., xs = xs.detach().cpu()) and, if needed, ensure start_idxs is on CPU (e.g., start_idxs = start_idxs.cpu()) before the inner loop, then iterate over xs[b] and start_idxs[b].item() without per-sample .cpu() calls. This keeps transfers batched and avoids repeated D2H copies.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@audio_separator/separator/architectures/mdxc_separator.py`:
- Line 35: The current scheduling can append a tail start that duplicates an
existing start when mix.shape[1]-chunk_size falls on a step boundary; update the
logic that builds self.indices so that the final list never contains duplicate
starts: after creating indices = list(range(0, mix.shape[1], step)) and before
applying the tail remap (the code that references chunk_size and mix.shape[1]),
compute last_start = mix.shape[1] - chunk_size and ensure you only append or
remap to last_start if it is not already present (or alternatively deduplicate
self.indices preserving order), so the tail window is not scheduled twice.
- Around line 377-385: The per-sample device-to-host transfers happen because
you call .to(device) on the whole batch (parts) but then call .cpu() inside the
loop for each sample (xs[b].cpu()); fix by doing the host copy once per batch:
after xs = self.model_run(parts) perform a single batch-level transfer (e.g., xs
= xs.detach().cpu()) and, if needed, ensure start_idxs is on CPU (e.g.,
start_idxs = start_idxs.cpu()) before the inner loop, then iterate over xs[b]
and start_idxs[b].item() without per-sample .cpu() calls. This keeps transfers
batched and avoids repeated D2H copies.
In `@tests/unit/test_cli.py`:
- Line 31: Add a new unit test that verifies the CLI forwards a non-default
--mdxc_num_workers value into the mdxc_params passed to Separator: call main()
with argv containing "--mdxc_num_workers=2", patch
audio_separator.separator.Separator, assert Separator was called once and that
its mdxc_params dict has "num_workers": 2 (keeping other mdxc_params from
common_expected_args); reference main(), Separator, and the mdxc_params key when
implementing the test.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 1a40c725-b4a4-4590-89c1-8271d945ff54
📒 Files selected for processing (2)
audio_separator/separator/architectures/mdxc_separator.pytests/unit/test_cli.py
- Refactor RoformerDataset to calculate tail-window remapping during initialization. - Fix an issue where the final audio chunk was duplicated when the tail start fell exactly on a step boundary. - Optimize MDXCSeparator inference by moving device-to-host transfers (.cpu()) outside the per-sample loop. - Batch the D2H copies to reduce synchronization overhead and improve processing speed. - Simplify RoformerDataset.__getitem__ to remove redundant re-calculation logic.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
audio_separator/separator/separator.py (1)
100-104: Consider usingNoneas default for mutable dict parameters (pre-existing pattern).The static analysis tool flags the use of mutable default arguments (dicts) for
mdx_params,vr_params,demucs_params, andmdxc_params. While this is a pre-existing pattern and not introduced by this PR, mutable defaults can cause subtle bugs if the dict is mutated during runtime.The standard Python idiom is to use
Noneand initialize within the function:♻️ Proposed refactor (optional, for future consideration)
def __init__( self, ... - mdx_params={"hop_length": 1024, "segment_size": 256, "overlap": 0.25, "batch_size": 1, "enable_denoise": False}, - vr_params={"batch_size": 1, "window_size": 512, "aggression": 5, "enable_tta": False, "enable_post_process": False, "post_process_threshold": 0.2, "high_end_process": False}, - demucs_params={"segment_size": "Default", "shifts": 2, "overlap": 0.25, "segments_enabled": True}, - mdxc_params={"segment_size": 256, "override_model_segment_size": False, "batch_size": 1, "overlap": 8, "pitch_shift": 0, "num_workers": 0}, + mdx_params=None, + vr_params=None, + demucs_params=None, + mdxc_params=None, info_only=False, ): """Initialize the separator.""" + if mdx_params is None: + mdx_params = {"hop_length": 1024, "segment_size": 256, "overlap": 0.25, "batch_size": 1, "enable_denoise": False} + if vr_params is None: + vr_params = {"batch_size": 1, "window_size": 512, "aggression": 5, "enable_tta": False, "enable_post_process": False, "post_process_threshold": 0.2, "high_end_process": False} + if demucs_params is None: + demucs_params = {"segment_size": "Default", "shifts": 2, "overlap": 0.25, "segments_enabled": True} + if mdxc_params is None: + mdxc_params = {"segment_size": 256, "override_model_segment_size": False, "batch_size": 1, "overlap": 8, "pitch_shift": 0, "num_workers": 0}This is flagged by Ruff B006. Since this pattern is pre-existing and the current PR only adds a key, this can be addressed in a separate cleanup PR if desired.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@audio_separator/separator/separator.py` around lines 100 - 104, The function currently uses mutable dict defaults for parameters mdx_params, vr_params, demucs_params, and mdxc_params which triggers Ruff B006; change each default to None and inside the function (where these params are used/merged) initialize them with the existing literal dicts only when the argument is None (e.g., if mdx_params is None: mdx_params = {...}); reference the parameter names mdx_params, vr_params, demucs_params, and mdxc_params and ensure any downstream usage expects those newly-initialized dicts so behavior remains unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@audio_separator/separator/separator.py`:
- Around line 100-104: The function currently uses mutable dict defaults for
parameters mdx_params, vr_params, demucs_params, and mdxc_params which triggers
Ruff B006; change each default to None and inside the function (where these
params are used/merged) initialize them with the existing literal dicts only
when the argument is None (e.g., if mdx_params is None: mdx_params = {...});
reference the parameter names mdx_params, vr_params, demucs_params, and
mdxc_params and ensure any downstream usage expects those newly-initialized
dicts so behavior remains unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: dcc3eebe-8d78-4f60-a0bc-83e4e62237ab
📒 Files selected for processing (2)
README.mdaudio_separator/separator/separator.py
Author
|
@beveradb let me know what you think. |
Author
|
After careful analysis and exhaustive benchmarks, I noticed that this approach was actually causing the process to slowdown overall, with marginal gains on big files (+-1 second). I'll be closing the PR, please disregard this changes. Benchmark results (original implementation using the published PyPi package is just faster out of the box):
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
This PR properly implements batch processing for the MDXC Arch models, addressing a massive processing bottleneck. Previously, setting
batch_sizefor MDXC separator did not actually influence the execution loop for Roformer models, causing the GPU to be severely underutilized on high-overlap segment sliding window operations.By restructuring the iteration to leverage
torch.utils.data.DataLoaderover a customRoformerDataset, this implementation lazily calculates coordinate slice bounds instead of redundantly slicing massive sections of audio into RAM.Why this approach?
Datasethandles indices behind the scenes, ensuring deterministic memory bounds regardless of track size.SlidingWindowInfereror Asteroid'sLambdaOverlapAdd, overlapping regions are gracefully interpolated through ahammingwindow weighting buffer to completely nullify any clipping or "seams" at segment boundaries.Benchmarks
File Details: ~22 minutes of PCM_16 audio
Model Details:
mel_band_roformer_kim_ft_unwa.ckptUsing standard inference parameters
batch_size=4Speedup is ~2.75x
Inference
use_autocast=Truebatch_size=4Speedup is ~2.2x
System Environment
Summary by CodeRabbit
New Features
Refactor
Tests
Documentation