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>
WalkthroughThe change 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 extensive tests. ChangesExecution and accelerator support
Estimated code review effort: 5 (Critical) | ~120 minutes Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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
| # Preserve full-precision rotary angles before converting the rest of the model. | ||
| rotary_frequencies = [ | ||
| (module, module.freqs.detach().float().clone()) | ||
| for module in self.model_run.modules() | ||
| if isinstance(module, RotaryEmbedding) | ||
| ] | ||
| self.model_run.half() | ||
| for rotary_embedding, frequencies in rotary_frequencies: | ||
| rotary_embedding.freqs.data = frequencies.to(rotary_embedding.freqs.device) | ||
| rotary_embedding.cached_freqs = None |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
rotary-embedding-torch 0.6.1 RotaryEmbedding freqs cached_freqs cache_if_possible attributes
💡 Result:
In the rotary-embedding-torch library, RotaryEmbedding is a class designed to implement rotary positional embeddings for PyTorch transformers [1]. The attributes cache_if_possible, freqs, and cached_freqs are central to the library's performance optimizations for sequence processing and inference caching [2][3]. Attributes and Caching Mechanics cache_if_possible: This is a boolean initialization argument (defaulting to True) that determines whether the model should attempt to cache calculated frequency values to improve efficiency [2][3]. freqs: This attribute stores the frequency values used for rotation [2]. Depending on the configuration, it may be a learned parameter (nn.Parameter) or a fixed tensor [2][3]. When learned_freq is set to True during initialization, these frequencies are optimized during training, and the library logic typically disables caching for these values to ensure accuracy [4]. cached_freqs: This is a registered buffer (managed via torch.register_buffer) that stores the precomputed rotary frequencies [2][3]. By storing these in a persistent-free buffer (persistent=False), the library avoids recomputing the embeddings for every forward pass during inference, provided that the sequence length is within the defined cache_max_seq_len [2][5]. Interaction The library uses these attributes to manage state during the forward pass. When cache_if_possible is enabled and the learned_freq parameter is False, the module checks if the required frequencies are available in the cache (often tracked by a cached_freqs_seq_len buffer) before calculating them [2][5][4]. If a cache miss occurs or the sequence length is new, it performs the computation and updates the cached_freqs buffer [2][5]. This approach is particularly beneficial for generative tasks where rotary embeddings are applied repeatedly to the same keys and queries in an autoregressive fashion [1][6].
Citations:
- 1: https://github.com/lucidrains/rotary-embedding-torch
- 2: https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
- 3: https://github.com/DecartAI/open-oasis-tnr/blob/master/rotary_embedding_torch.py
- 4: lucidrains/rotary-embedding-torch@2cab6c4
- 5: lucidrains/rotary-embedding-torch@0.8.1...0.8.2
- 6: https://gitlab.com/lucidrains/rotary-embedding-torch
🌐 Web query:
lucidrains rotary-embedding-torch source cached_freqs_seq_len freqs.data set
💡 Result:
In the lucidrains rotary-embedding-torch library, cached_freqs_seq_len is a buffer used to track the current number of cached frequency positions [1][2]. Specifically, within the RotaryEmbedding class: 1. Initialization: It is registered as a buffer named cached_freqs_seq_len (initialized to 0) to maintain state during inference or caching operations [1][2]. A corresponding buffer named cached_freqs is also registered as a tensor of shape (cache_max_seq_len, dim) [1][2]. 2. Function: It acts as a counter for how many positions have been pre-computed and stored in the cached_freqs buffer [1]. When generating frequencies, the library checks this length to determine if it needs to compute new frequencies or if it can utilize the existing cache [3]. 3. Data Handling: The actual frequency data is stored in the cached_freqs buffer [1][2]. When new frequencies are computed and added to the cache (typically during a forward pass), cached_freqs_seq_len is updated to reflect the new number of stored positions, often using.copy_ for in-place buffer updates [2]. This mechanism allows the model to avoid recomputing trigonometric frequency values for already-seen sequence positions during iterative inference [3][2].
Citations:
- 1: https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
- 2: lucidrains/rotary-embedding-torch@0.8.1...0.8.2
- 3: Slower than absolute positional embeddings? lucidrains/rotary-embedding-torch#36
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Find rotary_embedding_torch package files in sandbox"
python3 - <<'PY'
import importlib.util, pathlib, json, os, sys
for spec_name in ["rotary_embedding_torch", "rotary_embedding_torch.rotary_embedding_torch"]:
spec = importlib.util.find_spec(spec_name)
print(f"{spec_name}: origin={spec.origin!r}")
if spec and spec.origin:
root = pathlib.Path(spec.origin).parent
for p in root.glob("**/*.py"):
print(p)
PY
echo
echo "Find package source under installed pip cache if present"
python3 - <<'PY'
import site, pathlib, os
from venv import EnvBuilder
# avoid creating; just print common user/site-packages roots
for root in site.getsitepackages() + [site.getusersitepackages()]:
if root:
path = pathlib.Path(root)
if path.exists():
for p in list(path.glob("rotary_embedding_torch*")) + list(path.glob("rotary_embedding*")):
print(p)
PY
echo
echo "Search installed package source for cached_freqs_seq_len"
python3 - <<'PY'
import importlib.util, pathlib, glob
for spec_name in ["rotary_embedding_torch", "rotary_embedding_torch.rotary_embedding_torch"]:
spec = importlib.util.find_spec(spec_name)
if not spec or not spec.origin:
continue
root = pathlib.Path(spec.origin).parent
for py in glob.glob(str(root / "**" / "*.py"), recursive=True):
text = pathlib.Path(py).read_text(errors="ignore")
if "cached_freqs_seq_len" in text or "cached_freqs" in text:
print(f"\n--- {py} ---")
for i,line in enumerate(text.splitlines(),1):
if "cached_freqs_seq_len" in line or "cached_freqs" in line:
print(f"{i}: {line}")
PY
echo
echo "Search repo usage for cached_freqs_seq_len"
rg -n "cached_freqs_seq_len|cached_freqs|RotaryEmbedding" . || trueRepository: nomadkaraoke/python-audio-separator
Length of output: 359
🌐 Web query:
"rotary-embedding-torch" "cached_freqs_seq_len" "cache_len" "if cache_len"
💡 Result:
In the context of the rotary-embedding-torch library, these variables are part of an optimization mechanism used to cache precomputed frequency (cos/sin) tensors, avoiding redundant calculations during model forward passes [1][2]. Key details regarding their roles: 1. cache_max_seq_len: This is a configuration parameter (defaulting to 8192 in recent versions) that defines the maximum number of sequence positions the module is allowed to pre-allocate in its internal buffer [2]. 2. cached_freqs_seq_len: This is a buffer (registered via register_buffer, and thus not a learnable parameter) that tracks the current length of the sequences actually stored within the cache [1][2]. 3. Logic: When a forward pass occurs, the library checks if caching is enabled, if the frequencies are not learned, and if the requested sequence length (including any provided offset) fits within the cache limits: (offset + seq_len) <= self.cache_max_seq_len [3][2]. If these conditions are met and the requested range is already within the bounds of the current cached data, it retrieves the cached frequencies [3][2]. Otherwise, it computes the frequencies and updates the cached buffer and the cached_freqs_seq_len value [3][2]. This caching logic is designed to improve performance during inference or training on fixed-length sequences, though some users have reported that the overhead of checking cache bounds can occasionally impact performance compared to highly optimized custom kernels [4]. Sources: [4][1][3][2]
Citations:
- 1: https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
- 2: lucidrains/rotary-embedding-torch@0.8.1...0.8.2
- 3: lucidrains/rotary-embedding-torch@93e57f0
- 4: Slower than absolute positional embeddings? lucidrains/rotary-embedding-torch#36
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import importlib.util, sys, pkgutil, importlib
for name in ["rotary_embedding_torch", "rotary_embedding_torch.rotary_embedding_torch"]:
spec = importlib.util.find_spec(name)
print(f"{name}: {spec}")
PY
python3 - <<'PY'
import pkg_resources
try:
for r in pkg_resources.working_set:
if "rotary" in r.project_name.lower():
print(r)
try:
for path in r._get_metadata_path("top_level.txt"):
print("metadata:", path)
except Exception as e:
print("metadata error:", e)
except Exception as e:
print("pkg_resources error:", e)
PYRepository: nomadkaraoke/python-audio-separator
Length of output: 468
🌐 Web query:
github lucidrains rotary-embedding-torch cached_freqs_seq_len source
💡 Result:
In the lucidrains/rotary-embedding-torch library, cached_freqs_seq_len is an internal attribute used to track the number of precomputed rotary embedding frequencies currently stored in the model's buffer [1]. When the RotaryEmbedding module is initialized with caching enabled (via cache_if_possible), it allocates a buffer for frequencies up to a specified maximum length (cache_max_seq_len) [1]. The cached_freqs_seq_len attribute maintains the count of how many positions have been cached, allowing the implementation to efficiently manage and reuse these computations during inference or forward passes [1][2]. This mechanism helps optimize performance by avoiding redundant calculations of sine and cosine values for rotary embeddings, particularly in scenarios involving sequence generation [2].
Citations:
- 1: https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
- 2: Slower than absolute positional embeddings? lucidrains/rotary-embedding-torch#36
Reset the rotary cache length after changing freqs.
RotaryEmbedding uses cached_freqs plus cached_freqs_seq_len to reuse precomputed angles. Reassigning freqs.data and clearing only cached_freqs leaves cached_freqs_seq_len with its old length, so subsequent lookups can reuse stale cached angles in the new precision. Reset or invalidate the cache length for each rotary module after updating freqs.
🤖 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 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.
| def _float32_frequencies(rotary_embed, *, seq_len: int, device: torch.device) -> torch.Tensor: | ||
| """Build or retrieve rotary angles without allowing autocast to reduce precision.""" | ||
| # A compiled graph is shared by many regional Transformer instances. Their | ||
| # time and frequency embeddings have different cache shapes, so reading or | ||
| # mutating cached_freqs here creates instance-state guards and repeated | ||
| # Dynamo recompilations. Angle construction is cheap relative to attention | ||
| # and becomes part of the compiled graph, so skip the mutable cache while | ||
| # Dynamo is tracing. | ||
| is_compiling = torch.compiler.is_compiling() | ||
| should_cache = ( | ||
| not is_compiling | ||
| and rotary_embed.cache_if_possible | ||
| and not rotary_embed.learned_freq | ||
| and rotary_embed.freqs_for != "pixel" | ||
| ) | ||
| cached_freqs = rotary_embed.cached_freqs if should_cache else None | ||
|
|
||
| if ( | ||
| should_cache | ||
| and cached_freqs is not None | ||
| and cached_freqs.dtype == torch.float32 | ||
| and cached_freqs.device == device | ||
| and seq_len <= cached_freqs.shape[0] | ||
| ): | ||
| return cached_freqs[:seq_len].detach() | ||
|
|
||
| positions = rotary_embed.get_seq_pos(seq_len, device=device, dtype=torch.float32) | ||
| base_frequencies = rotary_embed.freqs.to(dtype=torch.float32) | ||
| frequencies = torch.einsum("..., f -> ... f", positions, base_frequencies) | ||
| frequencies = torch.repeat_interleave(frequencies, 2, dim=-1) | ||
|
|
||
| if should_cache: | ||
| # Replace an older low-precision cache instead of allowing it to be | ||
| # reused after leaving an autocast region. | ||
| rotary_embed.cached_freqs = frequencies.detach() | ||
|
|
||
| return frequencies |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm the pinned rotary-embedding-torch range and the attributes this helper relies on.
fd -t f 'pyproject.toml' --exec rg -n -C 2 'rotary'
rg -nP -C 2 'rotary-embedding-torch|rotary_embedding_torch' --type=toml --type=txtRepository: nomadkaraoke/python-audio-separator
Length of output: 324
🌐 Web query:
rotary-embedding-torch RotaryEmbedding cached_freqs buffer get_seq_pos API changes 0.6
💡 Result:
The changes you are inquiring about primarily concern the rotary-embedding-torch library, a common implementation of Rotary Positional Embeddings (RoPE). In recent updates to this library, the get_seq_pos method was updated to improve its robustness regarding device and data type selection [1][2]. Specifically, previous versions required explicit arguments for device and dtype. More recent iterations (e.g., seen in commits around July 2025) updated the API signature to: def get_seq_pos(self, seq_len, device = None, dtype = None, offset = 0): device = default(device, self.device) dtype = default(dtype, self.cached_freqs.dtype) return (torch.arange(seq_len, device = device, dtype = dtype) + offset) / self.interpolate_factor This change allows these parameters to default to the module's internal device and the cached_freqs buffer's data type, respectively, simplifying calls within the rotate_queries_or_keys and related methods [1][2]. Additionally, the library's caching mechanism—specifically for cached_freqs and cached_scales—has been refined [3]. Updates replaced the older tmp_store mechanism with direct buffer manipulation [3]. Instead of simply checking the shape, the library now tracks the valid cache length using explicit buffers (e.g., cached_freqs_seq_len) [3][4]. This ensures that when cache_if_possible is enabled, the code properly checks and updates the cached tensors only when the requested sequence length and offset are within the pre-allocated cache_max_seq_len buffer [3]. It is important to note that while version numbers like "0.6" may appear in various projects (e.g., vLLM or Open WebUI releases) [5][6], the specific API changes described above are characteristic of the evolution of the lucidrains/rotary-embedding-torch repository rather than a single unified "0.6" version change across the entire ecosystem [1][2][3]. Always check the specific commit history or README of the library version you are using if you encounter compatibility issues.
Citations:
- 1: https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
- 2: lucidrains/rotary-embedding-torch@0.8.7...0.8.8
- 3: lucidrains/rotary-embedding-torch@93e57f0
- 4: https://github.com/DecartAI/open-oasis-tnr/blob/master/rotary_embedding_torch.py
- 5: https://raw.githubusercontent.com/vllm-project/vllm/refs/tags/v0.6.6.post1/vllm/model_executor/layers/rotary_embedding.py
- 6: issue: Hybrid search fails since version 0.6.0 open-webui/open-webui#12347
🏁 Script executed:
#!/bin/bash
# Inspect the helper and nearby ROTARYEMBEDDER imports/usages, plus dependency file context.
set -e
printf '--- file outline ---\n'
ast-grep outline audio_separator/separator/uvr_lib_v5/roformer/rotaly.py --view expanded 2>/dev/null || true
ast-grep outline audio_separator/separator/uvr_lib_v5/roformer/rotary.py --view expanded 2>/dev/null || true
printf '\n--- rotary.py lines 1-130 ---\n'
sed -n '1,130p' audio_separator/separator/uvr_lib_v5/roformer/rotary.py | nl -ba
printf '\n--- dependency matches ---\n'
fd -t f 'pyproject.toml|requirements*.txt|setup.cfg|setup.py' --exec sh -c 'echo "### $1"; rg -n -C 3 "rotary-embedding-torch|rotary_embedding_torch" "$1" || true' sh {}Repository: nomadkaraoke/python-audio-separator
Length of output: 264
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- file outline ---'
ast-grep outline audio_separator/separator/uvr_lib_v5/roformer/rotary.py --view expanded || true
printf '%s\n' ''
printf '%s\n' '--- rotary.py lines 1-130 ---'
sed -n '1,130p' audio_separator/separator/uvr_lib_v5/roformer/rotary.py | nl -ba
printf '%s\n' ''
printf '%s\n' '--- dependency matches ---'
for f in $(fd -t f 'pyproject.toml|requirements.*txt|setup.cfg|setup.py'); do
printf '%s\n' "### $f"
rg -n -C 3 'rotary-embedding-torch|rotary_embedding_torch' "$f" || true
done
printf '%s\n' ''
printf '%s\n' '--- ROTARYEMBEDDER assignments/usages ---'
rg -n -C 3 'ROTARYEMBEDDER|cache_if_possible|cached_freqs|default_seq_dim|learned_freq|freqs_for' audio_separator/separator/uvr_lib_v5/roformer -t py || true
printf '%s\n' ''
printf '%s\n' '--- locked metadata mentioning rotary version ---'
fd -t f '.*lock|Pipfile.lock|poetry.lock|uv.lock|requirements.*txt|environment.yml' --exec sh -c 'echo "### $1"; rg -n -C 2 "rotary-embedding-torch|rotary_embedding_torch" "$1" | head -80 || true' sh {} || trueRepository: nomadkaraoke/python-audio-separator
Length of output: 558
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- rotary.py lines 1-130 ---'
sed -n '1,130p' audio_separator/separator/uvr_lib_v5/roformer/rotary.py | awk '{printf "%4d\t%s\n", NR, $0}'
printf '%s\n' ''
printf '%s\n' '--- dependency matches ---'
while IFS= read -r f; do
printf '%s\n' "### $f"
rg -n -C 3 'rotary-embedding-torch|rotary_embedding_torch' "$f" || true
done < <(fd -t f 'pyproject.toml|requirements.*txt|setup.cfg|setup.py')
printf '%s\n' ''
printf '%s\n' '--- ROTARYEMBEDDER assignments/usages ---'
rg -n -C 4 'ROTARYEMBEDDER|cache_if_possible|cached_freqs|default_seq_dim|learned_freq|freqs_for' audio_separator/separator/uvr_lib_v5/roformer -g '*.py' || true
printf '%s\n' ''
printf '%s\n' '--- locked metadata mentioning rotary version ---'
while IFS= read -r f; do
printf '%s\n' "### $f"
rg -n -C 2 'rotary-embedding-torch|rotary_embedding_torch' "$f" | head -120 || true
done < <(fd -t f '.*lock|Pipfile.lock|poetry.lock|uv.lock|requirements.*txt|environment.yml')Repository: nomadkaraoke/python-audio-separator
Length of output: 3861
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- dependency matches ---'
while IFS= read -r f; do
printf '%s\n' "### $f"
rg -n -C 3 'rotary-embedding-torch|rotary_embedding_torch' "$f" || true
done < <(fd -t f 'pyproject.toml|requirements.*txt|setup.cfg|setup.py')
printf '%s\n' ''
printf '%s\n' '--- ROTARYEMBEDDER assignments/usages ---'
rg -n -C 4 'ROTARYEMBEDDER|cache_if_possible|cached_freqs|default_seq_dim|learned_freq|freqs_for' audio_separator/separator/uvr_lib_v5/roformer -g '*.py' || true
printf '%s\n' ''
printf '%s\n' '--- locked metadata mentioning rotary version ---'
while IFS= read -r f; do
printf '%s\n' "### $f"
rg -n -C 2 'rotary-embedding-torch|rotary_embedding_torch' "$f" | head -120 || true
done < <(fd -t f '.*lock|Pipfile.lock|poetry.lock|uv.lock|requirements.*txt|environment.yml')
printf '%s\n' ''
printf '%s\n' '--- source dependency files ---'
fd -t f 'pyproject.toml|requirements.*txt|setup.cfg|setup.py'Repository: nomadkaraoke/python-audio-separator
Length of output: 252
🏁 Script executed:
#!/bin/bash
set -u
printf '%s\n' '--- dependency matches ---'
for f in $(fd -t f 'pyproject.toml|requirements.*txt|setup.cfg|setup.py'); do
printf '%s\n' "### $f"
rg -n -C 3 'rotary-embedding-torch|rotary_embedding_torch' "$f" || true
done
printf '%s\n' ''
printf '%s\n' '--- ROTARYEMBEDDER assignments/usages ---'
rg -n -C 4 'ROTARYEMBEDDER|cache_if_possible|cached_freqs|default_seq_dim|learned_freq|freqs_for' audio_separator/separator/uvr_lib_v5/roformer -g '*.py' || true
printf '%s\n' ''
printf '%s\n' '--- locked metadata mentioning rotary version ---'
while IFS= read -r f; do
printf '%s\n' "### $f"
rg -n -C 2 'rotary-embedding-torch|rotary_embedding_torch' "$f" | head -120 || true
done < <(find . -maxdepth 3 \( -iname '*lock' -o -name 'Pipfile.lock' -o -name 'poetry.lock' -o -name 'uv.lock' -o -name 'requirements*.txt' -o -name 'environment.yml' \) -type f)Repository: nomadkaraoke/python-audio-separator
Length of output: 5414
🌐 Web query:
github lucidrains rotary embedding torch 0.6.1 RotaryEmbedding cached_freqs default_seq_dim
💡 Result:
In the rotary-embedding-torch library, cached_freqs and default_seq_dim are core components of the RotaryEmbedding class [1]. Cached Frequencies (cached_freqs) The cached_freqs buffer is used to store precomputed trigonometric frequency values for efficient retrieval during the forward pass [1]. By registering it as a buffer (using self.register_buffer), the library ensures it is moved alongside the model when calling .to(device) [1][2]. - It is initialized during init with shape (cache_max_seq_len, dim) [1][2]. - When cache_if_possible is enabled, the model attempts to reuse these cached values to reduce redundant computations during inference or repeated forward passes [1][3]. - Updates to cached_freqs occur during the forward pass if the requested sequence length exceeds the currently cached length, provided the frequencies are not being learned (learned_freq=False) [4]. Default Sequence Dimension (default_seq_dim) The default_seq_dim is an internal attribute used to determine which dimension of the input tensor contains the sequence positions [1]. - It is calculated based on the seq_before_head_dim parameter: if seq_before_head_dim is True, default_seq_dim is set to -3; otherwise, it defaults to -2 [1]. - This allows the library to handle various tensor layouts (e.g., where the head dimension might be at the end or positioned differently relative to the sequence length) [1]. Regarding version 0.6.1: While specific code for version 0.6.1 is dated [5], these attributes have been standard parts of the RotaryEmbedding architecture in this library for several years [1][2]. Earlier versions used different methods for registering these buffers (such as tmp_store) [2], but the logical roles of cached_freqs and default_seq_dim remain fundamentally the same [1][2].
Citations:
- 1: https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
- 2: lucidrains/rotary-embedding-torch@0.8.1...0.8.2
- 3: Slower than absolute positional embeddings? lucidrains/rotary-embedding-torch#36
- 4: lucidrains/rotary-embedding-torch@2cab6c4
- 5: https://gitlab.com/lucidrains/rotary-embedding-torch
🌐 Web query:
github lucidrains rotary embedding torch versions tags 0.6.1 cached_freqs register_buffer
💡 Result:
In the lucidrains/rotary-embedding-torch library, the use of register_buffer for cached_freqs and related parameters was introduced to manage memory and persistence for positional embedding frequency caching [1]. Specifically, starting around version 0.8.1, the library moved away from a temporary storage method (tmp_store) to explicitly using register_buffer with persistent=False for several tensors [1]. This implementation ensures that: 1. cached_freqs and cached_scales are registered as buffers, allowing them to be moved to the appropriate device (e.g., GPU) automatically with the module [2][1]. 2. The persistent=False flag ensures these buffers are not included in the model's state_dict, which is appropriate for transient cache data [2][1]. 3. The module tracks cache usage through buffers like cached_freqs_seq_len, which are updated during forward passes to efficiently reuse precomputed rotary embedding frequencies [2][3][1]. This design facilitates the library's ability to cache frequencies and scale factors up to a user-defined cache_max_seq_len, improving performance during inference by avoiding redundant computations [2][1]. Versions prior to this transition (such as 0.6.1) managed caching through different mechanisms, often relying on internal dictionaries or simpler registration patterns that were later refactored for better PyTorch compatibility and performance [4][5][1].
Citations:
- 1: lucidrains/rotary-embedding-torch@0.8.0...0.8.1
- 2: https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
- 3: Slower than absolute positional embeddings? lucidrains/rotary-embedding-torch#36
- 4: lucidrains/rotary-embedding-torch@2cab6c4
- 5: lucidrains/rotary-embedding-torch@22cac59
🌐 Web query:
lucidrains rotary-embedding-torch 0.6.1 cached_freqs registered buffer get_seq_pos source
💡 Result:
In the rotary-embedding-torch library, cached_freqs is a registered buffer used to store precomputed rotary embedding frequencies, which improves performance by avoiding redundant calculations [1][2]. The implementation details involve: 1. Registered Buffer: It is registered using self.register_buffer('cached_freqs', ..., persistent=False) [1][2]. This ensures it is treated as part of the module's state but is not saved in the state_dict (it is not a parameter and typically recomputed or ignored during serialization) [1][2]. 2. Functionality: During the forward pass, the model checks if the required sequence lengths can be retrieved from cached_freqs [1][2]. If conditions are met (e.g., cache_if_possible is enabled, frequencies are not being learned), the cached values are returned [1][2][3]. 3. get_seq_pos: The get_seq_pos method is used to calculate the sequence positions, often incorporating an offset (useful for inference with KV caching) and interpolate_factor (for length extrapolation) [1][4][5][6]. It generally returns (torch.arange(seq_len, ...) + offset) / self.interpolate_factor [4][5][6]. While your query specifically mentions version 0.6.1, the cache-related implementation has evolved significantly through later versions (e.g., introduction of cached_freqs_seq_len in commit 93e57f0 to manage cache validity more effectively) [2]. Older versions may have used different mechanisms, such as tmp_store or basic dictionaries, before moving to the register_buffer approach [2][3].
Citations:
- 1: https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
- 2: lucidrains/rotary-embedding-torch@93e57f0
- 3: lucidrains/rotary-embedding-torch@2cab6c4
- 4: lucidrains/rotary-embedding-torch@0.8.7...0.8.8
- 5: https://github.com/DecartAI/open-oasis-tnr/blob/master/rotary_embedding_torch.py
- 6: canerturkmen/autogluon@1323a66
🌐 Web query:
raw.githubusercontent.com lucidrains rotary-embedding-torch 0.6.1 rotary_embedding_torch.py cache_if_possible cached_freqs register_buffer
💡 Result:
In the rotary-embedding-torch library by lucidrains, the register_buffer method is used to manage cached frequency tensors, ensuring they are properly registered with the PyTorch module without being considered trainable parameters [1][2]. Specifically, in versions starting from 0.8.1, the library utilizes register_buffer to handle caching for rotary frequency embeddings [1][2]. The cached_freqs tensor, which stores precomputed frequency values, is registered with persistent=False, meaning it will not be saved as part of the model's state_dict [1][2]. This registration allows the module to track the cache state—specifically the cached_freqs and a cached_freqs_seq_len tensor—efficiently [1][2]. The mechanism functions as follows: 1. When initialized, the module uses register_buffer to allocate space for cached_freqs (a tensor of shape [cache_max_seq_len, dim]) and a scalar cached_freqs_seq_len [1][2]. 2. During the forward pass, the code evaluates the condition cache_if_possible. If this flag is enabled, and other criteria (such as not learning the frequencies and staying within the pre-defined cache length) are met, the module checks if the requested sequence length can be satisfied by the existing cache [1][2]. 3. If the data is available in the cache, it is returned; otherwise, the frequencies are computed and, if caching is applicable, the cache is effectively updated [1][2]. This approach replaced older, more manual caching methods (like the use of a temporary store or custom dictionaries) used in earlier versions of the repository [3][4].
Citations:
- 1: lucidrains/rotary-embedding-torch@0.8.1...0.8.2
- 2: lucidrains/rotary-embedding-torch@0.8.0...0.8.1
- 3: https://github.com/ISEE213/Current-Diffusion-Model/blob/v1.6/list_rotary.py
- 4: lucidrains/rotary-embedding-torch@2cab6c4
🌐 Web query:
"rotary-embedding-torch" "0.6.0" "cached_freqs" "register_buffer"
💡 Result:
In the rotary-embedding-torch library, cached_freqs and register_buffer are used to manage precomputed rotary embedding frequencies for efficiency, particularly to avoid redundant calculations during forward passes [1][2]. The register_buffer method is used to register cached_freqs as a buffer in the PyTorch module [1][2]. This ensures the tensor is part of the module's state—meaning it is moved to the correct device alongside the module—but it is explicitly marked as non-persistent (persistent=False) so it is not included in the model's state_dict during saving, as these frequencies can be recomputed [1][2]. Key implementation details include: - Buffer Initialization: self.register_buffer('cached_freqs', torch.zeros(cache_max_seq_len, dim), persistent=False) creates a placeholder buffer for the frequencies [1][2]. - Dynamic Updates: The library monitors the length of the cached frequencies using another registered buffer, cached_freqs_seq_len [2]. During the forward pass, the implementation checks if the required sequence length is within the precomputed cache [2]. If not, or if the frequencies are being learned (which prevents caching), the library computes them dynamically [3][2]. - Purpose: By caching these values, the library reduces the computational overhead of generating frequency tensors repeatedly [2][4]. As of the current development state (post-0.6.0), the library moved away from earlier temporary storage methods (self.tmp_store) toward this more idiomatic register_buffer approach to handle device placement and state management more reliably [2].
Citations:
- 1: https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
- 2: lucidrains/rotary-embedding-torch@93e57f0
- 3: lucidrains/rotary-embedding-torch@2cab6c4
- 4: Slower than absolute positional embeddings? lucidrains/rotary-embedding-torch#36
🌐 Web query:
cache_if_possible freqs_for 0.6.1 rotary embedding torch
💡 Result:
In the rotary-embedding-torch library, cache_if_possible and freqs_for are configuration parameters for the RotaryEmbedding class [1][2]. The freqs_for parameter defines how frequency bases are generated for the rotary embedding. It accepts values such as 'lang' (for language models, using theta-based frequencies), 'pixel' (for image/spatial data), and 'constant' (for fixed frequencies) [3][1][4]. The cache_if_possible parameter (defaulting to True) controls whether the computed frequency tensors should be cached for efficiency [1][2]. The library implements a mechanism to verify cache eligibility, which checks if: - Caching is enabled (cache_if_possible is True) [1][5]. - The frequencies are not learned (learned_freq is False) [1][5]. - A sequence length is provided [1][5]. - The mode is not 'pixel' (as caching logic often differs or is restricted for non-1D sequences) [1][5]. - The requested sequence position plus length is within the defined cache_max_seq_len [1]. Version 0.6.1 is a specific release of the rotary-embedding-torch library [6]. While the fundamental architecture involving these parameters has remained consistent across versions, users should consult the documentation or source code specific to version 0.6.1 if they observe variations in behavior or available configuration options compared to the latest main branch [3][6].
Citations:
- 1: https://context7.com/lucidrains/rotary-embedding-torch/llms.txt
- 2: https://github.com/DecartAI/open-oasis-tnr/blob/master/rotary_embedding_torch.py
- 3: https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
- 4: https://github.com/ISEE213/Current-Diffusion-Model/blob/v1.6/list_rotary.py
- 5: https://huggingface.co/tuandunghcmut/vlm_clone_2/commit/46977c9b9c90c81776066aabca5bdbb40a98bc41
- 6: https://gitlab.com/lucidrains/rotary-embedding-torch
🌐 Web query:
raw.githubusercontent.com lucidrains rotary-embedding-torch v0.6.1 rotary_embedding_torch.py
💡 Result:
The library rotary-embedding-torch, authored by Phil Wang (lucidrains), is a PyTorch implementation of rotary positional embeddings (RoPE) [1][2]. While the specific file requested (rotary_embedding_torch.py from version 0.6.1) is an older iteration of the library's core module, the project is actively maintained, with current versions exceeding 0.9.x as of August 2026 [3][4]. The library is designed to provide efficient and flexible rotary embeddings, supporting features such as axial rotary embeddings, length extrapolation, and integration with fused flash attention [4]. For users looking to work with the library, it is recommended to use the latest version via pip install rotary-embedding-torch or conda-forge [1][3][4]. The source code for the current version can be accessed via the official GitHub repository, where historical versions like 0.6.1 can also be inspected by checking out the relevant tags or commit history [5][1]. The primary class, RotaryEmbedding, allows for straightforward instantiation and application to attention mechanisms [1]. You can instantiate it using RotaryEmbedding(dim =...) and apply it to queries and keys using the.rotate_queries_or_keys method, which simplifies the rotation process [1][6].
Citations:
- 1: https://github.com/lucidrains/rotary-embedding-torch
- 2: https://github.com/lucidrains/rotary-embedding-torch/blob/dc109d43621a6d0dacdacfa55e56afa146681bc0/setup.py
- 3: https://anaconda.org/conda-forge/rotary-embedding-torch
- 4: https://pypi.org/project/rotary-embedding-torch/
- 5: https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
- 6: lucidrains/rotary-embedding-torch@2887009
Keep the rotary dependency pinned and document the internal-attribute use.
_float32_frequencies and rotate_queries_or_keys access cache_if_possible, cached_freqs, learned_freq, freqs_for, default_seq_dim, and get_seq_pos directly on rotary-embedding-torch. These are internal attributes, not a public API. The dependency is limited only to >=0.6.1 in pyproject.toml, so a 0.6.x patch can change them with wrong rotary angles instead of an import error.
Update the rotary-embedding-torch 0.6.5 docstring at rotate_queries_or_keys to match the dependency context, and narrow/annotate the dependency range if this helper needs a specific implementation behavior.
🤖 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/rotary.py` around lines 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.
| expected = _float32_reference(rotary, tensor) | ||
| assert actual.dtype == tensor.dtype | ||
| assert rotary.cached_freqs.dtype == torch.float32 | ||
| torch.testing.assert_close(actual, expected, rtol=0, atol=0) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Exact-equality assertions can make these tests flaky.
rotate_queries_or_keys uses apply_rotary_emb from rotary-embedding-torch on the non-DirectML path. _float32_reference uses a hand-written t * cos + rotate_half(t) * sin. The two express the same mathematics, but they do not guarantee bit-identical results across PyTorch versions, backends, or half-precision rounding order. rtol=0, atol=0 therefore asserts more than the code contract.
Use a small tolerance so the tests verify precision, not bit patterns.
💚 Proposed change
- torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+ torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3)Apply the same change on lines 50 and 94.
Also applies to: 50-50, 94-94
🤖 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_roformer_rotary.py` at 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.
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