Add opt-in native fp16 and regional torch.compile, optimize RoFormer/MPS inference, and move to a validated PyTorch 2.13 baseline - #298
Conversation
- eliminate redundant RoFormer tail chunks and reuse consecutive model loads - keep supported MPS spectral work and bounded accumulators on-device - add observable precision modes and regional compilation with safe fallbacks - preserve float32 numerical islands and scaled BS-RoFormer attention
- publish platform-aware requirements through Poetry 2 and PEP 621 metadata - require PyTorch 2.13 on Apple arm64 while preserving the existing 2.8 lock on other Python <3.14 platforms - use the first torch and torchvision pair with CPython 3.14 wheels and mirror torchvision's Python 3.14.1 exclusion
- document Apple Silicon MPS spectral paths, bounded buffers, and the PyTorch baseline - explain precision and regional compilation capabilities and fallbacks - describe effective-mode reporting and consecutive model reuse
- align the contributor CUDA environment with the validated runtime\n- preserve the existing published range and Windows development lock
- Forward linear_transformer_depth through the normalized loader path. - Preserve zero-depth behavior for existing BS-RoFormer configurations. - Cover string normalization and constructor forwarding with unit tests.
- Name the pinned rotary-embedding-torch 0.6.5 behavior precisely\n- Link the still-open upstream device-hardcoding issue\n- Document why audio-separator keeps rotary angle construction in float32
- Explain when reused model weights remain allocated or are replaced. - Document the intentionally per-separation Demucs network lifecycle.
- declare packaging as a direct runtime dependency - document the CUDA 13 driver floor for the contributor lock - clarify the locked rotary dependency and fallback warning
- Derive the budget from the free Metal working set instead of a constant - Keep the 1 GiB floor when Metal cannot report a working-set size - Add AUDIO_SEPARATOR_MPS_BUFFER_BUDGET_GIB to override the heuristic - Name the buffers that move to CPU in the fallback logs and the README Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
- Replace the MPS buffer budget internals with the threshold and its override - Drop the rotary-embedding-torch pinning rationale and the compile retry mechanics - Merge the duplicate VR/Demucs rows in the verified-combination table Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (2)
WalkthroughThe PR adds execution-policy resolution for autocast, native FP16, and regional Torch compilation. It adds device capability probes, CPU fallbacks, memory-aware accumulation, model reuse, cleanup handling, CLI options, packaging updates, documentation, and tests. ChangesExecution and device optimization
CLI, packaging, and documentation
Estimated code review effort: 4 (Complex) | ~75 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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: 3
🧹 Nitpick comments (14)
tests/unit/test_model_reuse.py (2)
314-314: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReplace the assigned lambda with a
def.Ruff reports E731 for this line. Use a named function so the lint passes.
♻️ Proposed change
def test_vr_model_retries_after_weight_loading_failure(): - placeholder = lambda: None + def placeholder(): + return None + separator = _make_vr_separator(placeholder)As per coding guidelines: "Use ruff for code linting and formatting checks".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/test_model_reuse.py` at line 314, Replace the lambda assigned to placeholder with a named def function named placeholder, preserving its no-argument, no-op behavior so Ruff E731 is resolved.Sources: Coding guidelines, Linters/SAST tools
269-269: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueHardcoded
/tmpmodel paths trigger ruff S108 in both new test files. The shared root cause is the use of literal/tmp/...strings as stand-in model paths. Replace them with the pytesttmp_pathfixture, which also removes the platform assumption.
tests/unit/test_model_reuse.py#L269-L269: accepttmp_pathin_make_vr_separator, setseparator.model_pathfrom it, and update the matching assertion at line 307. Apply the same change to the/tmp/second.ckptand/tmp/model.ckptliterals at lines 53, 132, 172, 199, and 223.tests/unit/test_demucs_cleanup.py#L14-L14: accepttmp_pathin the test and setseparator.model_pathfrom it.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/test_model_reuse.py` at line 269, Replace hardcoded /tmp model paths with pytest tmp_path fixtures to remove ruff S108 violations and platform assumptions. In tests/unit/test_model_reuse.py lines 269-269, update _make_vr_separator to accept tmp_path, derive separator.model_path from it, update the matching assertion at line 307, and apply the same conversion to literals at lines 53, 132, 172, 199, and 223. In tests/unit/test_demucs_cleanup.py lines 14-14, accept tmp_path in the test and derive separator.model_path from it.Sources: Coding guidelines, Linters/SAST tools
tests/unit/test_demucs_import.py (1)
5-16: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the imported
demucsmodules fromsys.modulesafter the test.
monkeypatch.syspath_prependrestoressys.pathat teardown, but it does not remove entries fromsys.modules. The top-leveldemucs,demucs.hdemucs,demucs.htdemucs, anddemucs.specmodules stay cached for the rest of the session. The same source files are also imported asaudio_separator.separator.uvr_lib_v5.demucs.*, so two distinct class objects forHDemucsandHTDemucsremain loaded. Any laterisinstanceor identity check across the two import paths can then fail depending on test order.♻️ Proposed change
import importlib +import sys from pathlib import Path def test_checkpoint_compatible_top_level_demucs_import(monkeypatch): """Demucs modules remain importable under checkpoint-compatible top-level names.""" uvr_lib_path = Path(__file__).resolve().parents[2] / "audio_separator" / "separator" / "uvr_lib_v5" monkeypatch.syspath_prepend(str(uvr_lib_path)) + for name in list(sys.modules): + if name == "demucs" or name.startswith("demucs."): + monkeypatch.delitem(sys.modules, name) hdemucs = importlib.import_module("demucs.hdemucs") htdemucs = importlib.import_module("demucs.htdemucs") spec = importlib.import_module("demucs.spec")
monkeypatch.delitemrestores the previoussys.modulesstate at teardown, which also discards the modules imported inside the test.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/test_demucs_import.py` around lines 5 - 16, Update test_checkpoint_compatible_top_level_demucs_import to remove the imported top-level demucs modules from sys.modules via monkeypatch.delitem after importing them, including demucs, demucs.hdemucs, demucs.htdemucs, and demucs.spec, so teardown restores the prior module state.audio_separator/separator/uvr_lib_v5/roformer/mel_band_roformer.py (1)
436-441: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueLine 441 is now redundant.
Lines 436-437 align
masksto thestft_reprreal dtype before both tensors become complex. After line 439,masksandstft_reprtherefore already share the same complex dtype, somasks.type(stft_repr.dtype)on line 441 is a no-op. Remove it to keep one dtype-alignment point.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@audio_separator/separator/uvr_lib_v5/roformer/mel_band_roformer.py` around lines 436 - 441, Remove the redundant masks.type(stft_repr.dtype) call after the torch.view_as_complex conversions in the mask-processing flow, keeping the earlier dtype alignment before conversion as the single normalization point.audio_separator/separator/execution_policy.py (1)
77-97: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueLog the device type that the capability lookup used.
Line 78 keys the capability lookup on
capability_device_type, but the warning at lines 82-86 reportsdevice_type. For DirectML, the two values can differ, so the warning can name a device that was not checked. Usecapability_device_typein the native FP16 warning and in the compile warning at lines 103-108 for consistent diagnostics.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@audio_separator/separator/execution_policy.py` around lines 77 - 97, The native FP16 unsupported warning and the compile warning should report the device identifier used for capability lookup. Update the relevant logger calls in the precision-selection flow, including the block around use_native_fp16 and the compile warning, to use capability_device_type instead of device_type while preserving all other behavior.audio_separator/separator/uvr_lib_v5/device_utils.py (1)
85-98: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRecord why the probe fails, and silence the lint rule explicitly.
The probe must catch any backend error, so the broad
except Exceptionis correct here. Two improvements apply:
- The result is cached by
lru_cache. A transient failure, for example a temporary allocation failure, permanently forces the CPU path for that device. A debug log makes that outcome diagnosable.- Ruff reports BLE001 on line 97. A
# noqa: BLE001with a reason documents the intent.♻️ Proposed change
- except Exception: + except Exception as error: # noqa: BLE001 - any backend error means the op is unusable + logger.debug("Complex spectral probe failed for %s: %s", device_type, error) return FalseAdd a module-level logger:
import logging logger = logging.getLogger(__name__)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@audio_separator/separator/uvr_lib_v5/device_utils.py` around lines 85 - 98, Update the probe’s broad exception handler in the cached device-probing function to log the caught backend error at debug level before returning False, preserving the catch-all behavior. Add the module-level logger using logging.getLogger(__name__), and annotate the broad except with a reasoned # noqa: BLE001 suppression.Source: Linters/SAST tools
audio_separator/separator/uvr_lib_v5/roformer/bs_roformer.py (1)
478-483: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueThe MPS fallback path copies the spectrum across devices twice.
When
x_is_mpsis true, line 479 computes the STFT on CPU, line 480 movesstft_reprback to the model device, and line 537 moves it to CPU again for the complex multiply. The intermediate move is only needed sorearrangeruns on the device. Keepingstft_repron CPU until line 543 removes one full-spectrum copy in each direction. This is a performance improvement only; the numerical result does not change.Also applies to: 536-542
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@audio_separator/separator/uvr_lib_v5/roformer/bs_roformer.py` around lines 478 - 483, Update the x_is_mps/x_is_dml STFT path so stft_repr remains on CPU after torch.view_as_real instead of being moved to device. Adjust the corresponding rearrange and complex-multiply flow around stft_repr to keep it CPU-resident until the existing final transfer, preserving numerical behavior while removing the redundant device copies.tests/unit/test_bs_roformer_fp16.py (1)
41-41: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd
strict=Trueto the zip.Ruff reports B905 here. The two lists come from the same
model.modules()traversal, so their lengths always match.strict=Truerecords that invariant and clears the lint finding.As per coding guidelines: "Use ruff for code linting and formatting checks".
♻️ Proposed change
- for rotary, frequencies in zip(rotary_modules, rotary_frequencies): + for rotary, frequencies in zip(rotary_modules, rotary_frequencies, strict=True):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/test_bs_roformer_fp16.py` at line 41, Update the zip call in the rotary/frequency iteration to pass strict=True, recording that rotary_modules and rotary_frequencies must have matching lengths and resolving Ruff B905 without changing the loop behavior.Sources: Coding guidelines, Linters/SAST tools
audio_separator/separator/architectures/demucs_separator.py (1)
136-138: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReset
demucs_model_instanceto None instead of deleting the attribute.
__init__setsself.demucs_model_instance = Noneat line 86. Thedelat line 138 removes that attribute from the instance, so after the firstseparate()call any read outsideseparate()raisesAttributeError. Assigning None releases the model reference just as effectively and keeps the attribute contract stable across separations. Updatetests/unit/test_demucs_cleanup.pyto assertseparator.demucs_model_instance is Noneif you accept this.♻️ Proposed change
finally: - if hasattr(self, "demucs_model_instance"): - del self.demucs_model_instance + self.demucs_model_instance = None🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@audio_separator/separator/architectures/demucs_separator.py` around lines 136 - 138, Update the cleanup in the finally block of the separator flow to assign None to self.demucs_model_instance instead of deleting the attribute, preserving the attribute initialized by __init__ across repeated separations. Update tests/unit/test_demucs_cleanup.py to assert demucs_model_instance is None after cleanup.audio_separator/separator/architectures/mdxc_separator.py (1)
218-254: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueClarify the private compile-state guard and narrow the lazy retry.
_configure_model_compilationsaves and restorestransformer._compiled_call_impl, a private attribute. Add a short comment near the guard that PyTorch has no public API to de-compileModuleback to its original eager implementation; this prevents a Python 3.11+ upgrade from hiding why the private-path fallback exists.
_run_roformer_modelretries the chunk on anyExceptionwhenis_torch_compiledis true. Catch only the failures Dynamo might produce, or chain the retry failure withraise ... from excso the original traceback is not replaced.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@audio_separator/separator/architectures/mdxc_separator.py` around lines 218 - 254, Add a brief comment beside the _compiled_call_impl capability guard in _configure_model_compilation explaining that PyTorch lacks a public API to restore a Module’s original eager implementation. In _run_roformer_model, narrow the retry handler to Dynamo/torch.compile-related failures, or preserve the original exception by chaining any retry failure with raise-from while retaining the existing eager fallback behavior.tests/unit/test_execution_policy.py (1)
10-10: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename the
compileparameter to avoid shadowing the builtin.Ruff reports A002 for this argument. Rename it to
torch_compileand update the call sites in this file.♻️ Proposed rename
-def _resolve(*, device="mps", requested_device=None, model="mel_band_roformer", autocast=False, native=False, compile=False, pytorch=True): +def _resolve(*, device="mps", requested_device=None, model="mel_band_roformer", autocast=False, native=False, torch_compile=False, pytorch=True): logger = Mock() policy = resolve_execution_policy( device=torch.device(device), requested_device=torch.device(requested_device) if requested_device else None, model_family=model, use_autocast=autocast, use_native_fp16=native, - use_torch_compile=compile, + use_torch_compile=torch_compile,As per coding guidelines: "Use ruff for code linting and formatting checks".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/test_execution_policy.py` at line 10, Rename the compile parameter in _resolve to torch_compile to avoid shadowing the built-in, and update every call site in tests/unit/test_execution_policy.py to use the new keyword while preserving the existing behavior.Sources: Coding guidelines, Linters/SAST tools
tests/unit/test_mps_native_fp16.py (1)
197-207: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd
strict=Trueto thezipcall.Ruff reports B905. The two sequences are built from the same filtered
model.modules()scan, so a length mismatch signals a real defect.strict=Trueturns that into an explicit error instead of a silent truncation. The other new test file in this PR already usesstrict=True.♻️ Proposed fix
rotary_modules = [module for module in model.modules() if isinstance(module, RotaryEmbedding)] - for rotary, frequencies in zip(rotary_modules, rotary_frequencies): + for rotary, frequencies in zip(rotary_modules, rotary_frequencies, strict=True):As per coding guidelines: "Use ruff for code linting and formatting checks".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/test_mps_native_fp16.py` around lines 197 - 207, Update the zip call in _half_preserving_rotary_frequencies to use strict=True, preserving the existing pairing and assignment behavior while raising an error if the rotary module and saved-frequency sequences differ in length.Sources: Coding guidelines, Linters/SAST tools
audio_separator/separator/separator.py (1)
1117-1134: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winCleanup errors after a successful separation discard the output files.
The
finallyblock raisescleanup_errorwhen separation succeeded. The caller then losesoutput_fileseven though the stems were written to disk.clear_gpu_cacheandclear_file_specific_pathsare housekeeping steps, so a failure there is not equivalent to a separation failure. Consider logging the cleanup error and returning the output files, or document the current contract explicitly.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@audio_separator/separator/separator.py` around lines 1117 - 1134, The finally block in the separation flow must not raise cleanup_error after successful separation, because this discards valid output_files. Update the cleanup handling around clear_gpu_cache and clear_file_specific_paths to log housekeeping failures and preserve the successful return of output files; continue retaining the existing failure-path behavior for separation errors.tests/unit/test_mps_torch_compile.py (1)
25-38: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winBuild the transformers inside the test instead of at parametrize time.
The three modules are constructed when pytest collects this file. They are built even when the test skips on PyTorch below 2.6, and the same instances persist for the whole session.
torch._dynamo.explaintraces them, so shared instances can carry compilation state between runs. Pass factories and call them inside the test body.♻️ Proposed refactor
`@pytest.mark.parametrize`( - "transformer", + "build_transformer", [ - MelBandTransformer(dim=16, depth=1, dim_head=4, heads=2, rotary_embed=RotaryEmbedding(dim=4), flash_attn=True), - BSTransformer(dim=16, depth=1, dim_head=4, heads=2, rotary_embed=RotaryEmbedding(dim=4), flash_attn=True), - BSTransformer(dim=16, depth=1, dim_head=4, heads=2, flash_attn=True, linear_attn=True), + lambda: MelBandTransformer(dim=16, depth=1, dim_head=4, heads=2, rotary_embed=RotaryEmbedding(dim=4), flash_attn=True), + lambda: BSTransformer(dim=16, depth=1, dim_head=4, heads=2, rotary_embed=RotaryEmbedding(dim=4), flash_attn=True), + lambda: BSTransformer(dim=16, depth=1, dim_head=4, heads=2, flash_attn=True, linear_attn=True), ], ids=["mel-band", "bs-rotary", "bs-linear"], ) -def test_regional_transformer_is_captured_as_one_dynamo_graph(transformer): +def test_regional_transformer_is_captured_as_one_dynamo_graph(build_transformer): if version.parse(torch.__version__.split("+")[0]) < version.parse("2.6"): pytest.skip("Regional compilation requires PyTorch 2.6 or newer") + transformer = build_transformer() explanation = torch._dynamo.explain(transformer.eval())(torch.randn(2, 8, 16))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/test_mps_torch_compile.py` around lines 25 - 38, Replace the parametrized transformer instances with factory callables, preserving the existing three configurations and test IDs. In test_regional_transformer_is_captured_as_one_dynamo_graph, perform the PyTorch version skip before invoking the selected factory, then construct a fresh transformer and pass it to torch._dynamo.explain.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@audio_separator/separator/architectures/mdxc_separator.py`:
- Around line 205-214: Update the rotary cache invalidation in the loop over
RotaryEmbedding modules after restoring frequencies: clear both cached_freqs and
cached_freqs_seq_len so subsequent lookups cannot reuse stale angles or cache
metadata.
In `@audio_separator/separator/uvr_lib_v5/roformer/rotary.py`:
- Around line 12-48: Pin the rotary-embedding-torch dependency to the 0.6.5
implementation required by _float32_frequencies and rotate_queries_or_keys,
rather than allowing arbitrary 0.6.x patches; update the rotate_queries_or_keys
docstring to document that these helpers rely on internal rotary-embedding-torch
attributes and the pinned dependency behavior.
In `@tests/unit/test_roformer_rotary.py`:
- Line 43: Update the exact-equality torch.testing.assert_close assertions
comparing rotate_queries_or_keys with _float32_reference at the referenced
locations to use a small nonzero tolerance, including both rtol and atol as
appropriate. Apply the same tolerance consistently at all three assertion sites
while preserving the existing comparisons.
---
Nitpick comments:
In `@audio_separator/separator/architectures/demucs_separator.py`:
- Around line 136-138: Update the cleanup in the finally block of the separator
flow to assign None to self.demucs_model_instance instead of deleting the
attribute, preserving the attribute initialized by __init__ across repeated
separations. Update tests/unit/test_demucs_cleanup.py to assert
demucs_model_instance is None after cleanup.
In `@audio_separator/separator/architectures/mdxc_separator.py`:
- Around line 218-254: Add a brief comment beside the _compiled_call_impl
capability guard in _configure_model_compilation explaining that PyTorch lacks a
public API to restore a Module’s original eager implementation. In
_run_roformer_model, narrow the retry handler to Dynamo/torch.compile-related
failures, or preserve the original exception by chaining any retry failure with
raise-from while retaining the existing eager fallback behavior.
In `@audio_separator/separator/execution_policy.py`:
- Around line 77-97: The native FP16 unsupported warning and the compile warning
should report the device identifier used for capability lookup. Update the
relevant logger calls in the precision-selection flow, including the block
around use_native_fp16 and the compile warning, to use capability_device_type
instead of device_type while preserving all other behavior.
In `@audio_separator/separator/separator.py`:
- Around line 1117-1134: The finally block in the separation flow must not raise
cleanup_error after successful separation, because this discards valid
output_files. Update the cleanup handling around clear_gpu_cache and
clear_file_specific_paths to log housekeeping failures and preserve the
successful return of output files; continue retaining the existing failure-path
behavior for separation errors.
In `@audio_separator/separator/uvr_lib_v5/device_utils.py`:
- Around line 85-98: Update the probe’s broad exception handler in the cached
device-probing function to log the caught backend error at debug level before
returning False, preserving the catch-all behavior. Add the module-level logger
using logging.getLogger(__name__), and annotate the broad except with a reasoned
# noqa: BLE001 suppression.
In `@audio_separator/separator/uvr_lib_v5/roformer/bs_roformer.py`:
- Around line 478-483: Update the x_is_mps/x_is_dml STFT path so stft_repr
remains on CPU after torch.view_as_real instead of being moved to device. Adjust
the corresponding rearrange and complex-multiply flow around stft_repr to keep
it CPU-resident until the existing final transfer, preserving numerical behavior
while removing the redundant device copies.
In `@audio_separator/separator/uvr_lib_v5/roformer/mel_band_roformer.py`:
- Around line 436-441: Remove the redundant masks.type(stft_repr.dtype) call
after the torch.view_as_complex conversions in the mask-processing flow, keeping
the earlier dtype alignment before conversion as the single normalization point.
In `@tests/unit/test_bs_roformer_fp16.py`:
- Line 41: Update the zip call in the rotary/frequency iteration to pass
strict=True, recording that rotary_modules and rotary_frequencies must have
matching lengths and resolving Ruff B905 without changing the loop behavior.
In `@tests/unit/test_demucs_import.py`:
- Around line 5-16: Update test_checkpoint_compatible_top_level_demucs_import to
remove the imported top-level demucs modules from sys.modules via
monkeypatch.delitem after importing them, including demucs, demucs.hdemucs,
demucs.htdemucs, and demucs.spec, so teardown restores the prior module state.
In `@tests/unit/test_execution_policy.py`:
- Line 10: Rename the compile parameter in _resolve to torch_compile to avoid
shadowing the built-in, and update every call site in
tests/unit/test_execution_policy.py to use the new keyword while preserving the
existing behavior.
In `@tests/unit/test_model_reuse.py`:
- Line 314: Replace the lambda assigned to placeholder with a named def function
named placeholder, preserving its no-argument, no-op behavior so Ruff E731 is
resolved.
- Line 269: Replace hardcoded /tmp model paths with pytest tmp_path fixtures to
remove ruff S108 violations and platform assumptions. In
tests/unit/test_model_reuse.py lines 269-269, update _make_vr_separator to
accept tmp_path, derive separator.model_path from it, update the matching
assertion at line 307, and apply the same conversion to literals at lines 53,
132, 172, 199, and 223. In tests/unit/test_demucs_cleanup.py lines 14-14, accept
tmp_path in the test and derive separator.model_path from it.
In `@tests/unit/test_mps_native_fp16.py`:
- Around line 197-207: Update the zip call in
_half_preserving_rotary_frequencies to use strict=True, preserving the existing
pairing and assignment behavior while raising an error if the rotary module and
saved-frequency sequences differ in length.
In `@tests/unit/test_mps_torch_compile.py`:
- Around line 25-38: Replace the parametrized transformer instances with factory
callables, preserving the existing three configurations and test IDs. In
test_regional_transformer_is_captured_as_one_dynamo_graph, perform the PyTorch
version skip before invoking the selected factory, then construct a fresh
transformer and pass it to torch._dynamo.explain.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: b9c42467-e3e0-4b7f-bd08-d540580549ad
⛔ Files ignored due to path filters (1)
poetry.lockis excluded by!**/*.lock
📒 Files selected for processing (38)
README.mdaudio_separator/separator/architectures/demucs_separator.pyaudio_separator/separator/architectures/mdx_separator.pyaudio_separator/separator/architectures/mdxc_separator.pyaudio_separator/separator/architectures/vr_separator.pyaudio_separator/separator/common_separator.pyaudio_separator/separator/execution_policy.pyaudio_separator/separator/roformer/configuration_normalizer.pyaudio_separator/separator/roformer/roformer_loader.pyaudio_separator/separator/separator.pyaudio_separator/separator/uvr_lib_v5/demucs/hdemucs.pyaudio_separator/separator/uvr_lib_v5/demucs/htdemucs.pyaudio_separator/separator/uvr_lib_v5/demucs/spec.pyaudio_separator/separator/uvr_lib_v5/device_utils.pyaudio_separator/separator/uvr_lib_v5/roformer/attend.pyaudio_separator/separator/uvr_lib_v5/roformer/bs_roformer.pyaudio_separator/separator/uvr_lib_v5/roformer/mel_band_roformer.pyaudio_separator/separator/uvr_lib_v5/roformer/rotary.pyaudio_separator/separator/uvr_lib_v5/stft.pyaudio_separator/separator/uvr_lib_v5/tfc_tdf_v3.pyaudio_separator/utils/cli.pypyproject.tomltests/unit/test_bs_roformer_fp16.pytests/unit/test_cli.pytests/unit/test_configuration_normalizer.pytests/unit/test_demucs_cleanup.pytests/unit/test_demucs_import.pytests/unit/test_device_utils.pytests/unit/test_execution_policy.pytests/unit/test_mdxc_roformer_chunk_starts.pytests/unit/test_model_reuse.pytests/unit/test_mps_device_accumulation.pytests/unit/test_mps_native_fp16.pytests/unit/test_mps_stft_helpers.pytests/unit/test_mps_torch_compile.pytests/unit/test_roformer_dml_forward.pytests/unit/test_roformer_rotary.pytests/unit/test_separator_api_compatibility.py
- Hosted CI Macs expose a paravirtual Metal device (VirtualMac*) whose half-precision accumulation cannot meet the 30 dB gate - Detect virtualization via hw.model so real Apple GPUs keep the gates Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 94011c06-5dfe-4b06-b046-006289dd82ea
📒 Files selected for processing (3)
audio_separator/separator/uvr_lib_v5/roformer/rotary.pytests/unit/test_bs_roformer_fp16.pytests/unit/test_mps_native_fp16.py
🚧 Files skipped from review as they are similar to previous changes (1)
- audio_separator/separator/uvr_lib_v5/roformer/rotary.py
- Catch subprocess.TimeoutExpired, which is not an OSError - Treat a non-zero sysctl exit as not virtualized - Unknown environments keep the SNR gates active Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unit/test_bs_roformer_fp16.py`:
- Around line 21-25: Pin the Darwin helper’s sysctl subprocess invocation to the
absolute system executable path instead of relying on PATH lookup. Update the
subprocess.run calls in tests/unit/test_bs_roformer_fp16.py lines 21-25 and
tests/unit/test_mps_native_fp16.py lines 23-27; both sites require the same
direct change while preserving their existing arguments and error handling.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 816275b9-407a-4b15-b4c6-396a8bdb699c
📒 Files selected for processing (2)
tests/unit/test_bs_roformer_fp16.pytests/unit/test_mps_native_fp16.py
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- Hosted macOS runners dropped the same large wheel download mid-transfer in two consecutive runs, cancelling the whole matrix through fail-fast - Completed downloads are reused from the poetry cache between attempts Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Note for reviewers: this PR now includes one CI config change (commit 0b6771c), plus a test-side change that alters what CI displays. Flagging both here so they don't hide in the diff.
The four fp16 SNR gates now skip on virtualized Apple GPUs ( |
1 participant
Interactive walkthrough: https://claude.ai/code/artifact/af4c08ee-e129-424f-a84e-43101b1b9e58
The same content as this description, as a page you can drive: a resolver that shows what any device/model/precision/compile request actually activates, a chunk schedule you can replay at any input length, a spill calculator for the MPS buffer budget, and the benchmark tables as charts.
Summary
This PR makes PyTorch stem separation faster and more robust while preserving default outputs. It adds two opt-in, independent execution controls — precision (
--use_autocast/ new--use_native_fp16) and regional compilation (new--use_torch_compile) — fixes several RoFormer correctness issues (rotary precision, tail/short-input chunk scheduling, linear-attention layouts), makes model loading reusable, and moves the validated baseline to PyTorch 2.13. Every unsupported combination logs a warning and safely continues with today's float32/eager behavior.Headline results, measured against
v0.44.5on the exact same 99.000 s stereo 44.1 kHz float32 WAV (4,365,900 frames) for every MPS and CUDA timing cell, with each side built from its own lock (so the numbers reflect the combined effect of this PR — code changes plus the Torch 2.8 → 2.13 dependency move — not a code-only attribution):v0.44.5— 3.26–18.63 % (fp32) and 14.41–25.65 % (autocast).v0.44.5at fp32; with autocast, the three RoFormers and VR are 4.25–8.01 % faster, while HTDemucs autocast is 2.63 % slower.torch.compile(opt-in) on the released RoFormers cuts warm time further at equal precision: up to 33.93 % on MPS (fp32) and up to 43.18 % on CUDA (native fp16).What changes for users
Independent precision and compilation axes
--use_autocastand the new--use_native_fp16are mutually exclusive precision modes (enforced in the CLI and theSeparatorconstructor). The new--use_torch_compileis orthogonal and combines with any supported precision —autocast + compileis a valid pair.autocast_disabled(device)suppresses autocast on any backend, and fixing degraded CPU/MPS autocast angle precision is its whole purpose. STFT/ISTFT need no guard at all — they sit outside the low-precision region, before the cast intoband_splitand after the mask is cast back.Verified combinations are intentionally conservative:
torch.compilefp32,autocast,native_fp16fp32,autocastfp32,autocast(as today)fp32(as today)Observable fallbacks
load_model(), the read-only propertiesSeparator.effective_precision("fp32" | "autocast" | "native_fp16") andSeparator.effective_torch_compilereport what was actually activated, so warning-based fallbacks are visible to callers.effective_torch_compilereportsFalse.Model reuse and lifecycle (deliberate, documented behavior change)
load_model()now reuses the loaded instance when the same single model is requested again. Loading copiesSeparatorconfiguration into the architecture instance — output directory and format, normalization settings, architecture parameters, and the requested precision/compile settings — soload_model(..., force_reload=True)exists for the one case where a caller mutates such configuration after the first load and wants the same model rebuilt with the new values. Ordinary fixed-configuration use never needs it. A failed (re)load keeps the previously working model and its metadata intact.separate()internal load/release lifecycle (its lightweight wrapper is reusable, memory behavior unchanged), now with exception-safe cleanup.Correctness and robustness fixes
RoFormer chunk scheduler —
v0.44.5re-anchors an overrunning chunk tomix[:, -chunk_size:]and writes it atresult.shape[-1] - chunk_size. When the last two start positions on the step grid both overrun the end of the input, that produces two forwards over the byte-identical slice, written to the identical offset. The overlap-add is a weighted average (result / counter), so the duplicate adds the same Hamming window tocountertwice and the tail chunk ends up double-weighted. On the 99 s input (chunk 485,100 samples = 11.000 s, step 352,800 = 8.000 s) forwards per run drop 13 → 12 with unchanged coverage, and across the 3.0 s the tail chunk shares with its predecessor (88.0–91.0 s) its weight goes from 2:1 back to the intended 1:1 — at the midpoint of that region, from 66.67 % to 50.00 % of the blend.This changes the output versus
v0.44.5in that overlap, deliberately:v0.44.5's weighting was the bug. Outside the overlap the tail chunk is the only contributor, so2wy/2w = wy/wand the samples are identical. The numerical-parity section below is a within-branch comparison across execution modes and does not coverv0.44.5-vs-PR output equality, so a reviewer diffing againstv0.44.5should expect a difference confined to that tail window.The saving is input-length dependent, not universal. The duplicate only appears when two grid positions overrun the end, which happens for
chunk/step - 1= 37.5 % of input lengths at these settings. At 98 s it is 13 → 12 like 99 s; at 99.5 s, 100 s and 110 s both revisions produce identical schedules, identical forward counts, and identical output. The 99 s benchmark input happens to fall on the saving side, so every RoFormer timing cell in this campaign includes it.Inputs shorter than one chunk now work:
L - chunk_sizegoes negative onv0.44.5whilelengthis forced tochunk_size, so a 5 s input crashes withThe size of tensor a (176400) must match the size of tensor b (220500). The same crash reproduces withv0.44.5's code on Torch 2.13, so it is a code bug, not a Torch difference; this PR clamps the tail start to 0 and returns the full 220,500 frames from a single forward. The automatic short-audio segment override also no longer mutates persistent separator state.Rotary embeddings stay float32 on every backend —
rotary-embedding-torch0.6.x disables autocast only for CUDA (still true in 0.9.1; tracked in Avoid hard-coding autocast device parameter in rotary_embedding_torch.py lucidrains/rotary-embedding-torch#46), so CPU/MPS autocast could degrade angle precision. Rotation now runs inside a device-generic autocast-disabled float32 region, replaces any low-precision cached angles with float32, and skips cache mutation while Dynamo traces (avoiding per-instance recompilations).Linear-attention BS-RoFormer layouts now load — configs with
linear_transformer_depth > 0previously failed to construct (Attend.__init__() got an unexpected keyword argument 'scale').Attendnow honorsscaleon both the SDPA and einsum paths, and the loader/normalizer forwardlinear_transformer_depth.SDPA context migrated from the deprecated
torch.backends.cuda.sdp_kerneltotorch.nn.attention.sdpa_kernelwith the same effective backend set; this is also what lets Dynamo trace attention without graph breaks.MPS complex ops are probed at runtime — STFT/ISTFT, complex multiply, and the scatter op are probed once per device; supported spectral work stays on-device, otherwise the legacy CPU hop is preserved.
AUDIO_SEPARATOR_FORCE_CPU_COMPLEX=1forces the legacy path for diagnosis. Non-CaC Demucs Wiener masking deliberately stays on CPU on MPS.Bounded MPS accumulation, sized per device — duration-scaled overlap-add/accumulator buffers stay on MPS while their estimated footprint fits a budget, and fall back to CPU beyond it, so long inputs cannot exhaust the Metal working set. The budget is half of the free working set —
recommended_max_memory() - driver_allocated_memory()— floored at 1 GiB. Model weights are already resident when the decision is made, so the buffers are measured against what is actually left, and can never take more room than they leave behind for activations.driver_allocated_memory()counts the allocator's cached blocks, so free room is understated rather than overstated, and the budget varies with what the process has already allocated.AUDIO_SEPARATOR_MPS_BUFFER_BUDGET_GIBoverrides it, and every fallback logs the estimate alongside the budget it was compared against.Model inference always runs on MPS. Only the duration-scaled buffers move: for RoFormer the overlap-add
result/counterbuffers, the Hammingwindow, and each chunk's output as it is accumulated; for non-RoFormer MDXC the padded mix, its chunk view, andaccumulated_outputs; for Demucs the full-track mix and the returned sources. VR and ONNX MDX never allocate these buffers, so the budget does not apply to them.Measured on Apple M4 Pro / 24 GB, macOS 15.3.1: Metal reports a 16 GiB working set, and roughly 1 GiB of model weights are resident at the decision point, giving a budget near 7.5 GiB. Input duration at which each path spills, at 44.1 kHz stereo:
htdemucs_ftNote that the spilled path remains unmeasured — the 99 s benchmark input estimates 0.13 GiB for a 2-stem RoFormer and 0.37 GiB for HTDemucs, so no timing cell in this campaign crossed the budget. The change is covered by unit tests (budget scaling, the 1 GiB floor, env override, and a failing or absent Metal query falling back to the floor), not by a benchmark.
Dependencies and packaging
Published package metadata (what pip users get):
torch>=2.13,<3on macOS arm64 only (Torch 2.13 Apple Silicon wheels target macOS 14+); all other platforms keeptorch>=2.3,<3.requires-python = ">=3.10,!=3.14.1"(3.14.1 is excluded by the Python metadata of torchvision 0.28, which the Python 3.14 wheel set needs).packagingis now a declared dependency (it was already imported and always present transitively).cpu/gpu/dml) are unchanged.Contributor lock (what
poetry installgets):lock-version 2.1), so contributors need Poetry ≥ 2.0; CI already installs current Poetry via pipx.nvidia-smion the integration runners before merging. This applies to the contributor lock only; published metadata still allows Torch 2.3+ on Linux.Measured results
Method. Warm steady state per cell: one excluded warm-up separation, then the median of three timed
separate()calls. Timed work includes input decode, all architecture-internal work insideseparate()(for HTDemucs that includes its per-call network build, checkpoint read, and release), inference, WAV output, and device synchronization. Runner setup,Separatorconstruction, and top-levelload_model()are excluded — cold-start latency (including compile warm-up) is not measured. 76 formal cells (38 per accelerator) all used the identical input file; every percentage below compares two cells with identical device, input, model, precision, Python version, and cooldown protocol, and cells from different cooldown protocols are never combined or ranked. Absolute seconds must not be compared between MPS and CUDA — the hardware differs.v0.44.5Each side was installed from its own lock, so
v0.44.5-vs-PR numbers are the combined code + dependency effect. (The separately-run linear-attention fixture cells are the one exception to the MPS Python version; that split is explained where the fixture is introduced, and no comparison crosses Python versions.)Five released models, treated as peers:
mel_band_roformer_kim_ft2_bleedless_unwa.ckptmel_band_roformer_karaoke_gabox_v2.ckptbs_roformer_vocals_revive_unwa.ckpthtdemucs.yamlUVR-DeEcho-DeReverb.pthv0.44.5vs this PR, warm eager (median seconds; change vsv0.44.5)MPS (Apple M4 Pro)
v0.44.5fp32v0.44.5autocastCUDA (Google Colab Tesla T4)
v0.44.5fp32v0.44.5autocastHTDemucs and VR are not targets of native fp16 or regional compilation; their deltas here are the eager-path + dependency effect only.
Within this PR: released RoFormer precision × compile matrix
Every cell below runs this branch — this table compares execution modes within the PR, not
v0.44.5vs PR. Values are median seconds; parenthesized deltas compare compile against same-precision eager.MPS
CUDA (Google Colab Tesla T4)
The fastest measured condition (bold) differs by accelerator: on MPS it was fp32 + compile for all three RoFormers; on CUDA the fp16-family (autocast or native fp16) combined with compile won, with the exact winner model-dependent. All 36 PR RoFormer cells (2 devices × 3 models × 3 precisions × 2 execution modes) reported effective settings identical to the requested ones, and all 24 compile logs show zero graph breaks, zero regional-compile failures, and zero eager fallbacks.
Linear-attention architecture fixture (not a released model)
To exercise the
linear_transformer_depth > 0code path, a depth-1 linear-attention variant was derived from the released BS-RoFormer checkpoint. It is not a released or trained model, so it carries no separation-quality claim and is kept out of the released-model tables.v0.44.5fails to load the fixture on both accelerators with thescaleTypeError above; this PR runs all 12 cells:Native fp16 and memory
After warm eager runs on MPS, retained RoFormer model tensors roughly halve versus autocast, and post-run MPS allocator usage shrinks accordingly:
Whole-process peak RSS did not decrease in these isolated runs (allocator caches, compiler/runtime areas, and temporaries dominate), so the claim is limited to retained model tensors and post-run device allocation.
Numerical parity
Waveform comparisons across execution modes of the same checkpoint were all valid: 32 comparisons on MPS (minimum finite SNR 43.25 dB, minimum correlation 0.99998) and 26 selected comparisons on CUDA (minimum finite SNR 53.90 dB, minimum correlation 0.999998). CUDA coverage is representative rather than exhaustive. This is execution-mode numerical parity only — it is not ground-truth SDR or listening quality.
Scope notes
Verification
poetry check --lockpass.Known limitations
v0.44.5-vs-PR numbers are own-lock comparisons: combined code + Torch 2.8→2.13 effect, deliberately not attributed per factor.Summary by CodeRabbit