From 7048cab45268f680d8fd91886a82bc960fe21778 Mon Sep 17 00:00:00 2001 From: Andrew Beveridge Date: Mon, 20 Jul 2026 14:43:14 -0400 Subject: [PATCH 01/11] ci: add windows-directml integration job on self-hosted Windows GPU runners (#292) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Runs RoFormer (incl. the exact model from the issue), MDX, and VR separations with --use_directml on the gha-runner-gpu-windows ephemeral fleet (T4, WDDM). Asserts per model: - DirectML actually engaged (torch device + DmlExecutionProvider) - reference-image quality (same thresholds as the CUDA jobs) - non-silent, finite output (RMS floor — silent garbage must fail) - RoFormer loads via the NEW implementation with no legacy fallback continue-on-error until stabilized, then promoted to the gate + ruleset. torch-directml is layered via pip in-job: current releases pin torch==2.4.1, so poetry's lock resolves the dml extra to torch-directml 0.1.13 (torch-1.13 era) — unusable. pip resolves the modern wheel in this job's venv only; no other platform's torch changes. Co-Authored-By: Claude Fable 5 --- .github/workflows/run-integration-tests.yaml | 52 ++++++- docs/CI-GPU-RUNNERS.md | 14 +- tests/integration/test_windows_directml.py | 138 +++++++++++++++++++ 3 files changed, 198 insertions(+), 6 deletions(-) create mode 100644 tests/integration/test_windows_directml.py diff --git a/.github/workflows/run-integration-tests.yaml b/.github/workflows/run-integration-tests.yaml index a407d07..ee4209a 100644 --- a/.github/workflows/run-integration-tests.yaml +++ b/.github/workflows/run-integration-tests.yaml @@ -218,10 +218,59 @@ jobs: *.flac tests/*.flac + # ── Windows DirectML integration (self-hosted ephemeral Windows GPU) ── + # + # Runs on the gha-runner-gpu-windows ephemeral fleet (T4 in WDDM mode; + # see docs/CI-GPU-RUNNERS.md). Experimental: continue-on-error until + # promoted to a required check. + # + # torch-directml is layered via pip, NOT the poetry `dml` extra: current + # torch-directml releases pin torch==2.4.1, so poetry's lock resolved the + # extra to torch-directml 0.1.13 (a torch-1.13-era relic that must not be + # used). pip resolves the modern 0.2.5.dev wheel + a compatible torch in + # this job's venv only, leaving every other platform's torch untouched. + + windows-directml: + needs: changes + if: needs.changes.outputs.should_run == 'true' + runs-on: [self-hosted, windows, gpu] + timeout-minutes: 40 + continue-on-error: true + defaults: + run: + shell: powershell + env: + AUDIO_SEPARATOR_MODEL_DIR: C:\audio-separator-models + steps: + - uses: actions/checkout@v4 + - name: Verify GPU is in WDDM mode (DirectML requirement) + run: | + nvidia-smi --query-gpu=driver_version,name,driver_model.current --format=csv,noheader + $model = nvidia-smi --query-gpu=driver_model.current --format=csv,noheader + if ($model -notmatch "WDDM") { throw "GPU not in WDDM mode: $model" } + - name: Install Poetry dependencies (base) + run: poetry install + - name: Layer DirectML packages via pip + run: poetry run pip install onnxruntime-directml torch-directml + - name: Verify DirectML availability + run: | + poetry run python -c "import torch_directml; assert torch_directml.is_available(); print('torch_directml OK:', torch_directml.device_name(0))" + poetry run python -c "import onnxruntime as ort; ps = ort.get_available_providers(); print(ps); assert 'DmlExecutionProvider' in ps" + - name: "Run: DirectML separations (RoFormer, MDX, VR)" + run: poetry run pytest -sv tests/integration/test_windows_directml.py + - name: Upload test artifacts + if: always() + uses: actions/upload-artifact@v4 + with: + name: windows-directml-results + path: | + *.flac + tests/*.flac + # ── Gate job for branch protection ──────────────────────────────── integration-test: - needs: [changes, ensemble-presets, core-models, stems-and-quality, windows-cpu-integration] + needs: [changes, ensemble-presets, core-models, stems-and-quality, windows-cpu-integration, windows-directml] if: always() runs-on: ubuntu-latest steps: @@ -236,6 +285,7 @@ jobs: echo "core-models: ${{ needs.core-models.result }}" echo "stems-and-quality: ${{ needs.stems-and-quality.result }}" echo "windows-cpu-integration: ${{ needs.windows-cpu-integration.result }}" + echo "windows-directml: ${{ needs.windows-directml.result }} (continue-on-error; see job for real outcome)" if [[ "${{ needs.ensemble-presets.result }}" == "failure" ]] || \ [[ "${{ needs.core-models.result }}" == "failure" ]] || \ diff --git a/docs/CI-GPU-RUNNERS.md b/docs/CI-GPU-RUNNERS.md index 4c90470..c8f984b 100644 --- a/docs/CI-GPU-RUNNERS.md +++ b/docs/CI-GPU-RUNNERS.md @@ -82,11 +82,15 @@ Two tiers (added 2026-07 for RoFormer/DirectML support, issue #292): runs one model per DirectML-relevant architecture (RoFormer, MDX, VR) end-to-end on CPU. Models are cached via `actions/cache`. Python 3.12 (torch-directml has no 3.13 wheels, and the DML jobs must match). -2. **`windows-directml`** (planned) — self-hosted ephemeral Windows Server + - T4 VM (`gha-runner-gpu-windows` family), runs separation with - `--use_directml` and compares output quality against CPU results. The - image uses the NVIDIA **GRID** driver (WDDM mode) — the datacenter driver - puts the T4 in TCC mode, which has no DirectX support and breaks DirectML. +2. **`windows-directml`** (experimental, `continue-on-error`) — self-hosted + ephemeral Windows Server + T4 VM (`gha-runner-gpu-windows` family), runs + RoFormer/MDX/VR separations with `--use_directml`, asserting reference + quality, non-silent finite output, and new-implementation RoFormer loads + (no silent legacy fallback). The image uses the NVIDIA **GRID** driver + (WDDM mode) — the datacenter driver puts the T4 in TCC mode, which has no + DirectX support and breaks DirectML. torch-directml is layered via pip in + the job (poetry's lock resolves the `dml` extra to an unusable + torch-1.13-era torch-directml because modern releases pin torch==2.4.1). ## Required GitHub branch protection checks diff --git a/tests/integration/test_windows_directml.py b/tests/integration/test_windows_directml.py new file mode 100644 index 0000000..a3ad976 --- /dev/null +++ b/tests/integration/test_windows_directml.py @@ -0,0 +1,138 @@ +"""Windows DirectML integration tests (issue #292). + +Runs real separations with --use_directml on a Windows machine with a +DirectX 12 GPU (CI: the self-hosted gha-runner-gpu-windows ephemeral fleet, +NVIDIA T4 in WDDM mode). Asserts, per DML-supported architecture: + + * the separation completes and produces non-empty, non-silent, finite audio + (DirectML producing silent garbage must fail, even if it "succeeds"), + * output quality matches the committed reference images (same waveform/ + spectrogram similarity thresholds as the Linux CUDA integration tests), + * RoFormer models load via the NEW implementation with no silent fallback + to legacy (the loader map_location fix regression guard). + +Skipped automatically when torch-directml isn't installed, so it never runs +in the Linux/macOS/Windows-CPU jobs. +""" + +import os +import subprocess +import sys + +import numpy as np +import pytest +import soundfile as sf + +# Skip BEFORE importing test_cli_integration (it pulls in skimage/matplotlib +# via tests/utils) so DML-less environments skip cleanly at collection time. +torch_directml = pytest.importorskip("torch_directml") + +sys.path.append(os.path.dirname(os.path.dirname(__file__))) # tests/ (for utils) +sys.path.append(os.path.dirname(__file__)) # tests/integration/ +from test_cli_integration import resolve_cli_executable, validate_audio_output + +pytestmark = pytest.mark.skipif( + not torch_directml.is_available(), reason="torch-directml reports no DirectML device" +) + +INPUT_FILE = "tests/inputs/mardy20s.flac" +REFERENCE_DIR = "tests/inputs/reference" + +# (model, expected output files, is_roformer) +DML_MODEL_PARAMS = [ + ( + # The exact model from the issue #292 report + "model_bs_roformer_ep_317_sdr_12.9755.ckpt", + [ + "mardy20s_(Instrumental)_model_bs_roformer_ep_317_sdr_12.flac", + "mardy20s_(Vocals)_model_bs_roformer_ep_317_sdr_12.flac", + ], + True, + ), + ( + "mel_band_roformer_karaoke_aufr33_viperx_sdr_10.1956.ckpt", + [ + "mardy20s_(Instrumental)_mel_band_roformer_karaoke_aufr33_viperx_sdr_10.flac", + "mardy20s_(Vocals)_mel_band_roformer_karaoke_aufr33_viperx_sdr_10.flac", + ], + True, + ), + ( + # MDX — regression guard: already worked on DirectML before #292 + "UVR-MDX-NET-Inst_HQ_4.onnx", + [ + "mardy20s_(Instrumental)_UVR-MDX-NET-Inst_HQ_4.flac", + "mardy20s_(Vocals)_UVR-MDX-NET-Inst_HQ_4.flac", + ], + False, + ), + ( + # VR — regression guard: already worked on DirectML before #292 + "2_HP-UVR.pth", + [ + "mardy20s_(Instrumental)_2_HP-UVR.flac", + "mardy20s_(Vocals)_2_HP-UVR.flac", + ], + False, + ), +] + +# Same defaults as test_cli_integration; DML numerics may drift slightly more +# than CUDA vs the CPU-generated references, revisit per-model if needed. +WAVEFORM_THRESHOLD = 0.90 +SPECTROGRAM_THRESHOLD = 0.80 + +# Guards against "ran fine, produced silence" — well below any real stem, +# well above float noise. +RMS_FLOOR = 1e-4 + + +def _assert_audible_and_finite(path): + data, _sr = sf.read(path) + assert np.isfinite(data).all(), f"{path} contains NaN/Inf samples" + rms = float(np.sqrt(np.mean(np.square(data)))) + assert rms > RMS_FLOOR, f"{path} is (near-)silent: rms={rms:.2e}" + + +@pytest.mark.parametrize("model,expected_files,is_roformer", DML_MODEL_PARAMS) +def test_dml_separation(model, expected_files, is_roformer): + for f in expected_files: + if os.path.exists(f): + os.remove(f) + + result = subprocess.run( + [resolve_cli_executable(), "--use_directml", "--log_level", "debug", "-m", model, INPUT_FILE], + capture_output=True, + text=True, + check=False, + ) + log_text = (result.stdout or "") + (result.stderr or "") + + assert result.returncode == 0, f"CLI failed for {model}:\n{log_text[-4000:]}" + + # DirectML must actually be engaged — not silently fallen back to CPU. + assert "DirectML is available in Torch, setting Torch device to DirectML" in log_text + assert "ONNXruntime has DmlExecutionProvider available, enabling acceleration" in log_text + + if is_roformer: + # Loader regression guard: the map_location fix means the NEW + # implementation must load — a silent legacy fallback would keep CI + # green while shipping the unfixed path. + assert "Fell back to legacy" not in log_text, f"{model} silently fell back to legacy implementation" + assert "with new implementation" in log_text, f"{model} did not report new-implementation load" + + for output_file in expected_files: + assert os.path.exists(output_file), f"Output file {output_file} was not created" + assert os.path.getsize(output_file) > 0, f"Output file {output_file} is empty" + _assert_audible_and_finite(output_file) + + if os.environ.get("SKIP_AUDIO_VALIDATION") != "1": + waveform_match, spectrogram_match = validate_audio_output( + output_file, REFERENCE_DIR, WAVEFORM_THRESHOLD, SPECTROGRAM_THRESHOLD + ) + assert waveform_match, f"Waveform similarity below threshold for {output_file}" + assert spectrogram_match, f"Spectrogram similarity below threshold for {output_file}" + + for f in expected_files: + if os.path.exists(f): + os.remove(f) From a5e99a9f7e50730350cc3f7bb4b359e5098aaec3 Mon Sep 17 00:00:00 2001 From: Andrew Beveridge Date: Mon, 20 Jul 2026 14:50:30 -0400 Subject: [PATCH 02/11] fix(ci): avoid Windows MAX_PATH failure in windows-directml poetry install First live run failed with [WinError 206]: the SYSTEM-profile pypoetry venv path plus onnx's long test-data paths exceeds 260 chars. Use a short venv base (C:\venvs) and enable the LongPathsEnabled registry key (the runner executes as SYSTEM, and Python is longPathAware). Run also PROVED the stack: Windows ephemeral runner dispatched, booted, registered, ran the job; nvidia-smi reports '596.36, Tesla T4, WDDM'. Co-Authored-By: Claude Fable 5 --- .github/workflows/run-integration-tests.yaml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/workflows/run-integration-tests.yaml b/.github/workflows/run-integration-tests.yaml index ee4209a..9ae6d41 100644 --- a/.github/workflows/run-integration-tests.yaml +++ b/.github/workflows/run-integration-tests.yaml @@ -241,8 +241,14 @@ jobs: shell: powershell env: AUDIO_SEPARATOR_MODEL_DIR: C:\audio-separator-models + # Short venv base path: the default SYSTEM-profile pypoetry cache dir + # plus onnx's very long test-data paths exceeds Windows' 260-char + # MAX_PATH and breaks pip ([WinError 206]). + POETRY_VIRTUALENVS_PATH: C:\venvs steps: - uses: actions/checkout@v4 + - name: Enable Windows long paths + run: Set-ItemProperty -Path 'HKLM:\SYSTEM\CurrentControlSet\Control\FileSystem' -Name LongPathsEnabled -Value 1 -Type DWord - name: Verify GPU is in WDDM mode (DirectML requirement) run: | nvidia-smi --query-gpu=driver_version,name,driver_model.current --format=csv,noheader From 44f5c45c0d9d4418f12223c51befe5f47762196d Mon Sep 17 00:00:00 2001 From: Andrew Beveridge Date: Mon, 20 Jul 2026 14:57:43 -0400 Subject: [PATCH 03/11] fix(ci): install VC++ runtime for windows-directml job MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit torch (c10.dll) and onnxruntime both fail with [WinError 126] on a fresh Windows Server image — they need the MSVC 2015-2022 runtime. Co-Authored-By: Claude Fable 5 --- .github/workflows/run-integration-tests.yaml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/run-integration-tests.yaml b/.github/workflows/run-integration-tests.yaml index 9ae6d41..cebac03 100644 --- a/.github/workflows/run-integration-tests.yaml +++ b/.github/workflows/run-integration-tests.yaml @@ -258,6 +258,10 @@ jobs: run: poetry install - name: Layer DirectML packages via pip run: poetry run pip install onnxruntime-directml torch-directml + - name: Install VC++ runtime (torch/onnxruntime native DLLs need it) + run: | + curl.exe -fsSL -o "$env:TEMP\vc_redist.x64.exe" https://aka.ms/vs/17/release/vc_redist.x64.exe + Start-Process -FilePath "$env:TEMP\vc_redist.x64.exe" -ArgumentList "/install", "/quiet", "/norestart" -Wait - name: Verify DirectML availability run: | poetry run python -c "import torch_directml; assert torch_directml.is_available(); print('torch_directml OK:', torch_directml.device_name(0))" From d72436c0c3c979d3941cdda1b762f9dbebb43258 Mon Sep 17 00:00:00 2001 From: Andrew Beveridge Date: Mon, 20 Jul 2026 15:26:55 -0400 Subject: [PATCH 04/11] fix: einsum attention fallback on DirectML + CLI exits nonzero on failed separation Real-DML CI (first windows-directml runs) confirmed the last predicted gap: F.scaled_dot_product_attention is unimplemented on torch-directml and fails mid-forward with D3D12 'The parameter is incorrect.'. Attend now routes privateuseone tensors to the existing einsum math path, gated so every other device keeps flash behavior (equivalence test: math path matches SDPA output on CPU). MDX and VR already passed on real DirectML in the same run. The failure also exposed silent-failure bugs: Separator.separate() swallows per-file exceptions without tracebacks, and the CLI then printed 'Separation complete!' and exited 0 having produced nothing. Errors now log with exc_info, and the CLI exits 1 when no output files were produced. Bumps to 0.44.5 (code changes ship in the wheel). Co-Authored-By: Claude Fable 5 --- audio_separator/separator/separator.py | 4 +-- .../separator/uvr_lib_v5/roformer/attend.py | 14 ++++++++- audio_separator/utils/cli.py | 9 ++++++ pyproject.toml | 2 +- tests/unit/test_roformer_dml_forward.py | 30 +++++++++++++++++++ 5 files changed, 55 insertions(+), 4 deletions(-) diff --git a/audio_separator/separator/separator.py b/audio_separator/separator/separator.py index 6b1259d..4561e0c 100644 --- a/audio_separator/separator/separator.py +++ b/audio_separator/separator/separator.py @@ -984,7 +984,7 @@ def separate(self, audio_file_path, custom_output_names=None): files_output = self._separate_file(full_path, custom_output_names) output_files.extend(files_output) except Exception as e: - self.logger.error(f"Failed to process file {full_path}: {e}") + self.logger.error(f"Failed to process file {full_path}: {e}", exc_info=True) else: # If the path is a file, process it directly self.logger.info(f"Processing file: {path}") @@ -992,7 +992,7 @@ def separate(self, audio_file_path, custom_output_names=None): files_output = self._separate_file(path, custom_output_names) output_files.extend(files_output) except Exception as e: - self.logger.error(f"Failed to process file {path}: {e}") + self.logger.error(f"Failed to process file {path}: {e}", exc_info=True) return output_files diff --git a/audio_separator/separator/uvr_lib_v5/roformer/attend.py b/audio_separator/separator/uvr_lib_v5/roformer/attend.py index cc01093..e5ffef5 100644 --- a/audio_separator/separator/uvr_lib_v5/roformer/attend.py +++ b/audio_separator/separator/uvr_lib_v5/roformer/attend.py @@ -15,6 +15,16 @@ # helpers +def _is_dml_device(device) -> bool: + """torch-directml devices use torch's out-of-tree backend slot (privateuseone). + + F.scaled_dot_product_attention is not implemented by torch-directml (fails + with a D3D12 'The parameter is incorrect.' error), so DML tensors must take + the plain einsum attention path. Module-level so tests can patch it. + """ + return device.type == "privateuseone" + + def exists(val): return val is not None @@ -93,7 +103,9 @@ def forward(self, q, k, v): scale = q.shape[-1] ** -0.5 - if self.flash: + # DML has no SDPA — fall through to the einsum path. Gated so every + # other device keeps its exact existing behavior. (Issue #292) + if self.flash and not _is_dml_device(device): return self.flash_attn(q, k, v) # similarity diff --git a/audio_separator/utils/cli.py b/audio_separator/utils/cli.py index 6c33fae..23bbc90 100755 --- a/audio_separator/utils/cli.py +++ b/audio_separator/utils/cli.py @@ -296,4 +296,13 @@ def main(): separator.load_model(model_filename=model_filenames) output_files = separator.separate(audio_files, custom_output_names=args.custom_output_names) + + if not output_files: + # Separator.separate logs and swallows per-file errors so a batch can + # continue past one bad file — but if NOTHING was produced, the run + # failed and the CLI must not report success (this previously exited 0 + # with "Separation complete!" after logging an error). + logger.error("Separation produced no output files — see errors above.") + sys.exit(1) + logger.info(f"Separation complete! Output file(s): {' '.join(output_files)}") diff --git a/pyproject.toml b/pyproject.toml index 401dda3..c6f3ad9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api" [tool.poetry] name = "audio-separator" -version = "0.44.4" +version = "0.44.5" description = "Easy to use audio stem separation, using various models from UVR trained primarily by @Anjok07" authors = ["Andrew Beveridge "] license = "MIT" diff --git a/tests/unit/test_roformer_dml_forward.py b/tests/unit/test_roformer_dml_forward.py index ba6d55c..f9b73c4 100644 --- a/tests/unit/test_roformer_dml_forward.py +++ b/tests/unit/test_roformer_dml_forward.py @@ -88,3 +88,33 @@ def test_forced_dml_branch_matches_normal_cpu_output(self): assert hopped.device == audio.device assert hopped.shape == normal.shape assert torch.allclose(normal, hopped, atol=1e-6), "DML CPU-hop branch changed the output" + + +class TestAttendDmlFallback: + """SDPA is unimplemented on torch-directml — Attend must route DML tensors + to the einsum math path. Forcing the gate on CPU proves the fallback path + produces the same attention output as SDPA.""" + + def test_is_dml_device(self): + from audio_separator.separator.uvr_lib_v5.roformer import attend as attend_mod + + assert not attend_mod._is_dml_device(torch.device("cpu")) + assert attend_mod._is_dml_device(torch.device("privateuseone", 0)) + + def test_math_fallback_matches_sdpa_output(self): + from audio_separator.separator.uvr_lib_v5.roformer import attend as attend_mod + + torch.manual_seed(0) + att = attend_mod.Attend(flash=True) + att.eval() + q = torch.randn(2, 4, 16, 32) + k = torch.randn(2, 4, 16, 32) + v = torch.randn(2, 4, 16, 32) + + with torch.no_grad(): + flash_out = att(q, k, v) + with patch.object(attend_mod, "_is_dml_device", return_value=True): + math_out = att(q, k, v) + + assert math_out.shape == flash_out.shape + assert torch.allclose(flash_out, math_out, atol=1e-5), "einsum fallback diverges from SDPA" From 3a70443a9d1455429fe34e7ac88ffc9ee95131f8 Mon Sep 17 00:00:00 2001 From: Andrew Beveridge Date: Mon, 20 Jul 2026 15:43:14 -0400 Subject: [PATCH 05/11] fix: compute rotary embeddings without zero-width concat on DirectML MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Real-DML CI traceback (enabled by the previous commit's exc_info fix) pinpointed the RoFormer forward failure: rotary_embedding_torch's apply_rotary_emb concatenates empty unrotated edge slices, and torch-directml rejects zero-sized tensor ops ('The parameter is incorrect.'). These models rotate the full head dim, so the concat is a mathematical no-op — compute t*cos + rotate_half(t)*sin directly on DML. Gated on privateuseone; unit test proves exact equivalence with the library implementation, and the forced-DML full-model equivalence tests now exercise this path too. Co-Authored-By: Claude Fable 5 --- .../uvr_lib_v5/roformer/bs_roformer.py | 26 +++++++++++++++++-- .../uvr_lib_v5/roformer/mel_band_roformer.py | 26 +++++++++++++++++-- tests/unit/test_roformer_dml_forward.py | 22 ++++++++++++++++ 3 files changed, 70 insertions(+), 4 deletions(-) diff --git a/audio_separator/separator/uvr_lib_v5/roformer/bs_roformer.py b/audio_separator/separator/uvr_lib_v5/roformer/bs_roformer.py index 0a9adad..4f4543c 100644 --- a/audio_separator/separator/uvr_lib_v5/roformer/bs_roformer.py +++ b/audio_separator/separator/uvr_lib_v5/roformer/bs_roformer.py @@ -11,6 +11,7 @@ from beartype import beartype from rotary_embedding_torch import RotaryEmbedding +from rotary_embedding_torch.rotary_embedding_torch import rotate_half as _rotate_half_no_cat from einops import rearrange, pack, unpack from einops.layers.torch import Rearrange @@ -27,6 +28,27 @@ def _is_dml_device(device) -> bool: """ return device.type == "privateuseone" + +def _rotate_queries_or_keys(rotary_embed, t): + """Apply rotary position embedding, avoiding zero-width tensor ops on DML. + + rotary_embedding_torch's apply_rotary_emb concatenates (possibly empty) + unrotated edge slices around the rotated block; torch-directml rejects + zero-sized tensor ops with 'The parameter is incorrect.'. These models + always rotate the full head dimension, so the edge slices are empty and + the concat is a no-op — compute the rotation directly instead. Verified + equivalent to the library implementation by unit test. (Issue #292) + """ + if not _is_dml_device(t.device): + return rotary_embed.rotate_queries_or_keys(t) + seq_len = t.shape[-2] + freqs = rotary_embed.forward(rotary_embed.get_seq_pos(seq_len, device=t.device, dtype=t.dtype), seq_len=seq_len) + if freqs.shape[-1] != t.shape[-1]: + # Partial-dim rotation would need the edge concat — unreachable here + # (RotaryEmbedding(dim=dim_head) rotates the full head dim). + return rotary_embed.rotate_queries_or_keys(t) + return t * freqs.cos() + _rotate_half_no_cat(t) * freqs.sin() + def exists(val): return val is not None @@ -98,8 +120,8 @@ def forward(self, x): q, k, v = rearrange(self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads) if exists(self.rotary_embed): - q = self.rotary_embed.rotate_queries_or_keys(q) - k = self.rotary_embed.rotate_queries_or_keys(k) + q = _rotate_queries_or_keys(self.rotary_embed, q) + k = _rotate_queries_or_keys(self.rotary_embed, k) out = self.attend(q, k, v) diff --git a/audio_separator/separator/uvr_lib_v5/roformer/mel_band_roformer.py b/audio_separator/separator/uvr_lib_v5/roformer/mel_band_roformer.py index a4ecda1..5fa0ce3 100644 --- a/audio_separator/separator/uvr_lib_v5/roformer/mel_band_roformer.py +++ b/audio_separator/separator/uvr_lib_v5/roformer/mel_band_roformer.py @@ -11,6 +11,7 @@ from beartype import beartype from rotary_embedding_torch import RotaryEmbedding +from rotary_embedding_torch.rotary_embedding_torch import rotate_half as _rotate_half_no_cat from einops import rearrange, pack, unpack, reduce, repeat @@ -26,6 +27,27 @@ def _is_dml_device(device) -> bool: """ return device.type == "privateuseone" + +def _rotate_queries_or_keys(rotary_embed, t): + """Apply rotary position embedding, avoiding zero-width tensor ops on DML. + + rotary_embedding_torch's apply_rotary_emb concatenates (possibly empty) + unrotated edge slices around the rotated block; torch-directml rejects + zero-sized tensor ops with 'The parameter is incorrect.'. These models + always rotate the full head dimension, so the edge slices are empty and + the concat is a no-op — compute the rotation directly instead. Verified + equivalent to the library implementation by unit test. (Issue #292) + """ + if not _is_dml_device(t.device): + return rotary_embed.rotate_queries_or_keys(t) + seq_len = t.shape[-2] + freqs = rotary_embed.forward(rotary_embed.get_seq_pos(seq_len, device=t.device, dtype=t.dtype), seq_len=seq_len) + if freqs.shape[-1] != t.shape[-1]: + # Partial-dim rotation would need the edge concat — unreachable here + # (RotaryEmbedding(dim=dim_head) rotates the full head dim). + return rotary_embed.rotate_queries_or_keys(t) + return t * freqs.cos() + _rotate_half_no_cat(t) * freqs.sin() + def exists(val): return val is not None @@ -93,8 +115,8 @@ def forward(self, x): q, k, v = rearrange(self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads) if exists(self.rotary_embed): - q = self.rotary_embed.rotate_queries_or_keys(q) - k = self.rotary_embed.rotate_queries_or_keys(k) + q = _rotate_queries_or_keys(self.rotary_embed, q) + k = _rotate_queries_or_keys(self.rotary_embed, k) out = self.attend(q, k, v) diff --git a/tests/unit/test_roformer_dml_forward.py b/tests/unit/test_roformer_dml_forward.py index f9b73c4..382df66 100644 --- a/tests/unit/test_roformer_dml_forward.py +++ b/tests/unit/test_roformer_dml_forward.py @@ -118,3 +118,25 @@ def test_math_fallback_matches_sdpa_output(self): assert math_out.shape == flash_out.shape assert torch.allclose(flash_out, math_out, atol=1e-5), "einsum fallback diverges from SDPA" + + +class TestRotaryNoCatEquivalence: + """The DML rotary path computes t*cos + rotate_half(t)*sin directly, + skipping the (empty) edge concat torch-directml rejects. It must match + the library implementation exactly.""" + + def test_manual_rotation_matches_library(self): + from rotary_embedding_torch import RotaryEmbedding + + torch.manual_seed(0) + rotary = RotaryEmbedding(dim=64) + t = torch.randn(2, 8, 16, 64) + + library = rotary.rotate_queries_or_keys(t) + with patch.object(bs_mod, "_is_dml_device", return_value=True): + manual_bs = bs_mod._rotate_queries_or_keys(rotary, t) + with patch.object(mel_mod, "_is_dml_device", return_value=True): + manual_mel = mel_mod._rotate_queries_or_keys(rotary, t) + + assert torch.allclose(library, manual_bs, atol=1e-6) + assert torch.allclose(library, manual_mel, atol=1e-6) From 1462b41edd6131d2615a85a65ae3647067bf48e3 Mon Sep 17 00:00:00 2001 From: Andrew Beveridge Date: Mon, 20 Jul 2026 16:01:49 -0400 Subject: [PATCH 06/11] fix: bound attention memory on DirectML with batch-sliced einsum MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Real-DML run confirmed the new implementation loads (loading stats: new_implementation_success=1) and the rotary fix works, but inference then died with 'DML allocator out of memory': the einsum attention path materializes the full (b h i j) similarity tensor — ~1.3GB per layer at segment 801 across 62 bands — which torch-directml's allocator can't handle. Slice the batch dim (step 8) on DML only; equivalence test covers the ragged final slice. Co-Authored-By: Claude Fable 5 --- .../separator/uvr_lib_v5/roformer/attend.py | 14 +++++++++++++ tests/unit/test_roformer_dml_forward.py | 20 +++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/audio_separator/separator/uvr_lib_v5/roformer/attend.py b/audio_separator/separator/uvr_lib_v5/roformer/attend.py index e5ffef5..02d84b4 100644 --- a/audio_separator/separator/uvr_lib_v5/roformer/attend.py +++ b/audio_separator/separator/uvr_lib_v5/roformer/attend.py @@ -108,6 +108,20 @@ def forward(self, q, k, v): if self.flash and not _is_dml_device(device): return self.flash_attn(q, k, v) + if _is_dml_device(device): + # Materializing the full (b h i j) similarity tensor at once + # (~1.3GB/layer at segment 801 × 62 bands) OOMs torch-directml's + # allocator. Slice over the batch dim to bound peak memory; the + # result is mathematically identical. + outs = [] + step = 8 + for i in range(0, q.shape[0], step): + sim = einsum(f"b h i d, b h j d -> b h i j", q[i : i + step], k[i : i + step]) * scale + attn = sim.softmax(dim=-1) + attn = self.attn_dropout(attn) + outs.append(einsum(f"b h i j, b h j d -> b h i d", attn, v[i : i + step])) + return torch.cat(outs, dim=0) + # similarity sim = einsum(f"b h i d, b h j d -> b h i j", q, k) * scale diff --git a/tests/unit/test_roformer_dml_forward.py b/tests/unit/test_roformer_dml_forward.py index 382df66..b33e1d3 100644 --- a/tests/unit/test_roformer_dml_forward.py +++ b/tests/unit/test_roformer_dml_forward.py @@ -119,6 +119,26 @@ def test_math_fallback_matches_sdpa_output(self): assert math_out.shape == flash_out.shape assert torch.allclose(flash_out, math_out, atol=1e-5), "einsum fallback diverges from SDPA" + def test_sliced_attention_matches_unsliced(self): + # The DML path slices the batch dim (step=8) to bound peak memory; + # use a batch that is neither a multiple of the step nor smaller than + # it, to cover the ragged final slice. + from audio_separator.separator.uvr_lib_v5.roformer import attend as attend_mod + + torch.manual_seed(0) + att = attend_mod.Attend(flash=False) + att.eval() + q = torch.randn(19, 4, 16, 32) + k = torch.randn(19, 4, 16, 32) + v = torch.randn(19, 4, 16, 32) + + with torch.no_grad(): + unsliced = att(q, k, v) + with patch.object(attend_mod, "_is_dml_device", return_value=True): + sliced = att(q, k, v) + + assert torch.allclose(unsliced, sliced, atol=1e-6), "batch-sliced attention diverges" + class TestRotaryNoCatEquivalence: """The DML rotary path computes t*cos + rotate_half(t)*sin directly, From 00d7222e5b2a9e1956b86a1b7f1b5e736f1aaafd Mon Sep 17 00:00:00 2001 From: Andrew Beveridge Date: Mon, 20 Jul 2026 16:13:25 -0400 Subject: [PATCH 07/11] fix: use matmul (not einsum) for DirectML attention; add MDX23C to DML matrix MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The OOM persisted with batch slicing because torch-directml lowers einsum naively (broadcast multiply + reduce): the b×h×i×j×d intermediate is tens of GB at segment 801. torch.matmul maps to a real GEMM. Keep the batch slicing to bound the materialized similarity tensor. Also add MDX23C-8KFFT-InstVoc_HQ.ckpt (plain MDXC / TFC_TDF arch) to the DML matrix — its STFT wrapper already CPU-hops non-cuda/cpu devices, so it should pass unchanged; audibility checks only (no committed reference images), auto-downloaded on first run. Co-Authored-By: Claude Fable 5 --- .../separator/uvr_lib_v5/roformer/attend.py | 15 ++++++----- tests/integration/test_windows_directml.py | 25 ++++++++++++++++--- 2 files changed, 30 insertions(+), 10 deletions(-) diff --git a/audio_separator/separator/uvr_lib_v5/roformer/attend.py b/audio_separator/separator/uvr_lib_v5/roformer/attend.py index 02d84b4..5bf7b8f 100644 --- a/audio_separator/separator/uvr_lib_v5/roformer/attend.py +++ b/audio_separator/separator/uvr_lib_v5/roformer/attend.py @@ -109,17 +109,20 @@ def forward(self, q, k, v): return self.flash_attn(q, k, v) if _is_dml_device(device): - # Materializing the full (b h i j) similarity tensor at once - # (~1.3GB/layer at segment 801 × 62 bands) OOMs torch-directml's - # allocator. Slice over the batch dim to bound peak memory; the - # result is mathematically identical. + # Use matmul (a real GEMM) instead of einsum: torch-directml + # lowers einsum naively (broadcast multiply + reduce), whose + # b×h×i×j×d intermediate is tens of GB at segment 801 — the + # 'DML allocator out of memory' crash. Also slice the batch dim + # to bound the materialized (b h i j) similarity tensor. The + # result is mathematically identical to the einsum path below. outs = [] step = 8 for i in range(0, q.shape[0], step): - sim = einsum(f"b h i d, b h j d -> b h i j", q[i : i + step], k[i : i + step]) * scale + qs, ks, vs = q[i : i + step], k[i : i + step], v[i : i + step] + sim = torch.matmul(qs, ks.transpose(-1, -2)) * scale attn = sim.softmax(dim=-1) attn = self.attn_dropout(attn) - outs.append(einsum(f"b h i j, b h j d -> b h i d", attn, v[i : i + step])) + outs.append(torch.matmul(attn, vs)) return torch.cat(outs, dim=0) # similarity diff --git a/tests/integration/test_windows_directml.py b/tests/integration/test_windows_directml.py index a3ad976..0ccd421 100644 --- a/tests/integration/test_windows_directml.py +++ b/tests/integration/test_windows_directml.py @@ -38,7 +38,7 @@ INPUT_FILE = "tests/inputs/mardy20s.flac" REFERENCE_DIR = "tests/inputs/reference" -# (model, expected output files, is_roformer) +# (model, expected output files, is_roformer, validate_reference) DML_MODEL_PARAMS = [ ( # The exact model from the issue #292 report @@ -48,6 +48,7 @@ "mardy20s_(Vocals)_model_bs_roformer_ep_317_sdr_12.flac", ], True, + True, ), ( "mel_band_roformer_karaoke_aufr33_viperx_sdr_10.1956.ckpt", @@ -56,6 +57,7 @@ "mardy20s_(Vocals)_mel_band_roformer_karaoke_aufr33_viperx_sdr_10.flac", ], True, + True, ), ( # MDX — regression guard: already worked on DirectML before #292 @@ -65,6 +67,7 @@ "mardy20s_(Vocals)_UVR-MDX-NET-Inst_HQ_4.flac", ], False, + True, ), ( # VR — regression guard: already worked on DirectML before #292 @@ -74,6 +77,20 @@ "mardy20s_(Vocals)_2_HP-UVR.flac", ], False, + True, + ), + ( + # Plain MDXC (TFC_TDF arch) — its STFT wrapper already CPU-hops + # non-cuda/cpu devices, so it should work on DML unchanged. No + # committed reference images for this model: audibility/finiteness + # checks only. Auto-downloads on first run (not in the baked set). + "MDX23C-8KFFT-InstVoc_HQ.ckpt", + [ + "mardy20s_(Instrumental)_MDX23C-8KFFT-InstVoc_HQ.flac", + "mardy20s_(Vocals)_MDX23C-8KFFT-InstVoc_HQ.flac", + ], + False, + False, ), ] @@ -94,8 +111,8 @@ def _assert_audible_and_finite(path): assert rms > RMS_FLOOR, f"{path} is (near-)silent: rms={rms:.2e}" -@pytest.mark.parametrize("model,expected_files,is_roformer", DML_MODEL_PARAMS) -def test_dml_separation(model, expected_files, is_roformer): +@pytest.mark.parametrize("model,expected_files,is_roformer,validate_reference", DML_MODEL_PARAMS) +def test_dml_separation(model, expected_files, is_roformer, validate_reference): for f in expected_files: if os.path.exists(f): os.remove(f) @@ -126,7 +143,7 @@ def test_dml_separation(model, expected_files, is_roformer): assert os.path.getsize(output_file) > 0, f"Output file {output_file} is empty" _assert_audible_and_finite(output_file) - if os.environ.get("SKIP_AUDIO_VALIDATION") != "1": + if validate_reference and os.environ.get("SKIP_AUDIO_VALIDATION") != "1": waveform_match, spectrogram_match = validate_audio_output( output_file, REFERENCE_DIR, WAVEFORM_THRESHOLD, SPECTROGRAM_THRESHOLD ) From 3cc43aeb2bdf4d9217faf60912ae17005047d5f6 Mon Sep 17 00:00:00 2001 From: Andrew Beveridge Date: Mon, 20 Jul 2026 16:24:27 -0400 Subject: [PATCH 08/11] test: reduced MDXC segment size on DirectML (allocator can't sustain 801 fp32) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MDX23C (plain conv TFC_TDF) also OOM'd, proving the limit is torch-directml's allocator rather than our attention path. Run ckpt models at segment 256 — the documented resource knob for constrained GPUs. MDX and VR pass at defaults (0.93 waveform similarity). Co-Authored-By: Claude Fable 5 --- tests/integration/test_windows_directml.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/tests/integration/test_windows_directml.py b/tests/integration/test_windows_directml.py index 0ccd421..494720b 100644 --- a/tests/integration/test_windows_directml.py +++ b/tests/integration/test_windows_directml.py @@ -38,7 +38,13 @@ INPUT_FILE = "tests/inputs/mardy20s.flac" REFERENCE_DIR = "tests/inputs/reference" -# (model, expected output files, is_roformer, validate_reference) +# torch-directml's allocator can't sustain the default MDXC segment size +# (801 for these roformer models) at fp32 — 'DML allocator out of memory' — +# so ckpt-based models run with a reduced segment. This is the documented +# resource knob for constrained GPUs, not a DML-specific hack. +DML_SEGMENT_ARGS = ["--mdxc_override_model_segment_size", "--mdxc_segment_size", "256"] + +# (model, expected output files, is_roformer, validate_reference, extra_args) DML_MODEL_PARAMS = [ ( # The exact model from the issue #292 report @@ -49,6 +55,7 @@ ], True, True, + DML_SEGMENT_ARGS, ), ( "mel_band_roformer_karaoke_aufr33_viperx_sdr_10.1956.ckpt", @@ -58,6 +65,7 @@ ], True, True, + DML_SEGMENT_ARGS, ), ( # MDX — regression guard: already worked on DirectML before #292 @@ -68,6 +76,7 @@ ], False, True, + [], ), ( # VR — regression guard: already worked on DirectML before #292 @@ -78,6 +87,7 @@ ], False, True, + [], ), ( # Plain MDXC (TFC_TDF arch) — its STFT wrapper already CPU-hops @@ -91,6 +101,7 @@ ], False, False, + DML_SEGMENT_ARGS, ), ] @@ -111,14 +122,14 @@ def _assert_audible_and_finite(path): assert rms > RMS_FLOOR, f"{path} is (near-)silent: rms={rms:.2e}" -@pytest.mark.parametrize("model,expected_files,is_roformer,validate_reference", DML_MODEL_PARAMS) -def test_dml_separation(model, expected_files, is_roformer, validate_reference): +@pytest.mark.parametrize("model,expected_files,is_roformer,validate_reference,extra_args", DML_MODEL_PARAMS) +def test_dml_separation(model, expected_files, is_roformer, validate_reference, extra_args): for f in expected_files: if os.path.exists(f): os.remove(f) result = subprocess.run( - [resolve_cli_executable(), "--use_directml", "--log_level", "debug", "-m", model, INPUT_FILE], + [resolve_cli_executable(), "--use_directml", "--log_level", "debug", *extra_args, "-m", model, INPUT_FILE], capture_output=True, text=True, check=False, From 6dd57c74c87760aebd28e4896882b31726cfea21 Mon Sep 17 00:00:00 2001 From: Andrew Beveridge Date: Mon, 20 Jul 2026 16:39:04 -0400 Subject: [PATCH 09/11] fix: release torch-directml allocator memory between demix batches Per-chunk memory now fits (segment 256 completed 2 chunks; MDX23C got to 3/35) but torch-directml's allocator doesn't reuse freed blocks across iterations, so long chunk loops still die with 'DML allocator out of memory'. Call torch_directml.empty_cache() (hasattr-guarded; added in its 0.2.x line) plus a gc pass after each batch, gated on privateuseone. No-op on every other device. Co-Authored-By: Claude Fable 5 --- .../separator/architectures/mdxc_separator.py | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/audio_separator/separator/architectures/mdxc_separator.py b/audio_separator/separator/architectures/mdxc_separator.py index 1ddb499..a9aa76c 100644 --- a/audio_separator/separator/architectures/mdxc_separator.py +++ b/audio_separator/separator/architectures/mdxc_separator.py @@ -1,3 +1,4 @@ +import gc import os import sys @@ -13,6 +14,28 @@ # Roformer direct constructors removed; loading handled via RoformerLoader in CommonSeparator. +def _release_dml_memory_if_needed(device): + """Work around torch-directml's cross-iteration allocator leak. + + torch-directml (privateuseone) does not reliably reuse freed blocks + across inference iterations — long chunk loops die with 'DML allocator + out of memory' after a few batches even though each batch fits. Its + 0.2.x API added torch_directml.empty_cache() for exactly this; call it + (plus a gc pass to drop deferred references) after each batch on DML + only. No-op everywhere else and on older torch-directml. (Issue #292) + """ + if getattr(device, "type", None) != "privateuseone": + return + gc.collect() + try: + import torch_directml + + if hasattr(torch_directml, "empty_cache"): + torch_directml.empty_cache() + except Exception: # pragma: no cover — defensive: never break inference + pass + + class MDXCSeparator(CommonSeparator): """ MDXCSeparator is responsible for separating audio sources using MDXC models. @@ -326,6 +349,7 @@ def demix(self, mix: np.ndarray) -> dict: part = part.to(device) x = self.model_run(part.unsqueeze(0))[0] x = x.cpu() + _release_dml_memory_if_needed(device) # Perform overlap_add on CPU if i + chunk_size > mix.shape[1]: # Fixed to correctly add to the end of the tensor @@ -398,6 +422,9 @@ def demix(self, mix: np.ndarray) -> dict: accumulated_outputs[..., count * hop_size : count * hop_size + chunk_size] += individual_output_cpu count += 1 + del single_batch_result + _release_dml_memory_if_needed(self.torch_device) + self.logger.debug("Calculating inferenced outputs based on accumulated outputs and overlap") inferenced_outputs = accumulated_outputs[..., chunk_size - hop_size : -(pad_size + chunk_size - hop_size)] / self.overlap self.logger.debug("Deleting accumulated outputs to free up memory") From 530f65d706160d93017b98ddee93e8032c17148b Mon Sep 17 00:00:00 2001 From: Andrew Beveridge Date: Mon, 20 Jul 2026 16:50:31 -0400 Subject: [PATCH 10/11] fix: MDXC/RoFormer models run on CPU under DirectML (allocator limitation) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Nine real-hardware CI iterations fixed every code-level DirectML incompatibility in the RoFormer/MDXC path (state-dict loading, complex ops, SDPA, rotary embedding, einsum lowering) — inference now runs on DML — but torch-directml's allocator cannot sustain the chunked inference loops of ckpt-based models: 'DML allocator out of memory' after a few chunks regardless of segment size, batch slicing, matmul lowering, gc, or empty_cache (verified on a 16GB T4; the conv-only MDX23C hits it too, proving it's not our attention code). torch-directml is effectively unmaintained upstream (last release Sept 2024). Until upstream fixes the allocator: - MDXC-family models automatically run on CPU under DirectML, with a clear warning; MDX and VR remain GPU-accelerated (CI-verified 0.93 waveform similarity) - AUDIO_SEPARATOR_FORCE_DML_MDXC=1 attempts DML anyway (for testing on other GPUs — AMD iGPUs with large shared memory may behave better) - windows-directml CI asserts the fallback fires AND output is correct vs references; MDX23C (plain TFC_TDF MDXC) added to the matrix - README DirectML table updated to the CI-enforced truth Co-Authored-By: Claude Fable 5 --- README.md | 15 ++++--- .../separator/architectures/mdxc_separator.py | 31 +++++++++++++++ tests/integration/test_windows_directml.py | 39 ++++++++++++------- tests/unit/test_directml.py | 39 +++++++++++++++++++ 4 files changed, 103 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 1b90ba4..ab0451e 100644 --- a/README.md +++ b/README.md @@ -161,12 +161,15 @@ audio-separator path/to/audio.wav --use_directml | Architecture | Model types | Status | |---|---|---| -| MDX | `.onnx` | ✅ Confirmed working | -| MDXC (incl. the default `bs_roformer` model) | `.ckpt` / `.yaml` | ⚠️ Expected to work, community-untested | -| VR | `.pth` | ⚠️ Expected to work, community-untested | -| Demucs | — | ❓ Unverified | - -If you test any of the untested architectures, please [open an issue](https://github.com/nomadkaraoke/python-audio-separator/issues) with your `--env_info` output and logs — reports are what move these from "untested" to "confirmed". +| MDX | `.onnx` | ✅ GPU-accelerated (CI-verified on NVIDIA T4/WDDM) | +| VR | `.pth` | ✅ GPU-accelerated (CI-verified on NVIDIA T4/WDDM) | +| MDXC / RoFormer (incl. the default `bs_roformer` model) | `.ckpt` / `.yaml` | ⚠️ Runs correctly on **CPU** automatically. All code-level DirectML incompatibilities are fixed (v0.44.5), but torch-directml's allocator cannot sustain these models' inference loops (`DML allocator out of memory`, upstream limitation). Set `AUDIO_SEPARATOR_FORCE_DML_MDXC=1` to attempt GPU anyway — reports welcome, especially from AMD/Intel GPUs. | +| Demucs | — | ❌ Not supported (torch-directml lacks the fused LSTM operator) | + +These statuses are enforced by the `windows-directml` CI job, which runs real +separations on an NVIDIA T4 in WDDM mode. If `AUDIO_SEPARATOR_FORCE_DML_MDXC=1` +works on your GPU, please [open an issue](https://github.com/nomadkaraoke/python-audio-separator/issues) +with your `--env_info` output — torch-directml behaves differently across vendors. ### 🎥 FFmpeg dependency diff --git a/audio_separator/separator/architectures/mdxc_separator.py b/audio_separator/separator/architectures/mdxc_separator.py index a9aa76c..702d0a4 100644 --- a/audio_separator/separator/architectures/mdxc_separator.py +++ b/audio_separator/separator/architectures/mdxc_separator.py @@ -14,6 +14,33 @@ # Roformer direct constructors removed; loading handled via RoformerLoader in CommonSeparator. +def _mdxc_inference_device(torch_device, torch_device_cpu, logger): + """Pick the inference device for MDXC-family (incl. RoFormer) models. + + All code-level DirectML incompatibilities in these models are fixed + (state-dict loading, complex ops, SDPA, rotary embedding, einsum + lowering), but torch-directml's allocator cannot sustain their chunked + inference loops — it fails with 'DML allocator out of memory' after a + few chunks regardless of segment size (upstream limitation; verified on + a 16GB T4, torch-directml 0.2.5). Until that is fixed upstream, run + these models on CPU; MDX and VR stay GPU-accelerated via DirectML. + Set AUDIO_SEPARATOR_FORCE_DML_MDXC=1 to attempt DirectML anyway. + (Issue #292) + """ + if getattr(torch_device, "type", None) != "privateuseone": + return torch_device + if os.environ.get("AUDIO_SEPARATOR_FORCE_DML_MDXC"): + logger.warning("AUDIO_SEPARATOR_FORCE_DML_MDXC set — attempting MDXC/RoFormer on DirectML (may exhaust GPU memory).") + return torch_device + logger.warning( + "MDXC/RoFormer models currently run on CPU under DirectML: torch-directml's " + "allocator cannot sustain their chunked inference ('DML allocator out of " + "memory', upstream limitation). MDX and VR models remain GPU-accelerated. " + "Set AUDIO_SEPARATOR_FORCE_DML_MDXC=1 to attempt DirectML anyway." + ) + return torch_device_cpu if torch_device_cpu is not None else torch.device("cpu") + + def _release_dml_memory_if_needed(device): """Work around torch-directml's cross-iteration allocator leak. @@ -52,6 +79,10 @@ def __init__(self, common_config, arch_config): # The instance variable self.model_data is passed through from Separator and set in CommonSeparator self.logger.debug(f"Model data: {self.model_data}") + # DirectML: run MDXC-family models on CPU until torch-directml's + # allocator can sustain them (see _mdxc_inference_device). + self.torch_device = _mdxc_inference_device(self.torch_device, self.torch_device_cpu, self.logger) + # Arch Config is the MDXC architecture specific user configuration options, which should all be configurable by the user # either by their Separator class instantiation or by passing in a CLI parameter. # While there are similarities between architectures for some of these (e.g. batch_size), they are deliberately configured diff --git a/tests/integration/test_windows_directml.py b/tests/integration/test_windows_directml.py index 494720b..1f4e60a 100644 --- a/tests/integration/test_windows_directml.py +++ b/tests/integration/test_windows_directml.py @@ -38,13 +38,14 @@ INPUT_FILE = "tests/inputs/mardy20s.flac" REFERENCE_DIR = "tests/inputs/reference" -# torch-directml's allocator can't sustain the default MDXC segment size -# (801 for these roformer models) at fp32 — 'DML allocator out of memory' — -# so ckpt-based models run with a reduced segment. This is the documented -# resource knob for constrained GPUs, not a DML-specific hack. -DML_SEGMENT_ARGS = ["--mdxc_override_model_segment_size", "--mdxc_segment_size", "256"] - -# (model, expected output files, is_roformer, validate_reference, extra_args) +# MDXC-family (RoFormer + TFC_TDF ckpt) models automatically fall back to +# CPU under DirectML: every code-level DML incompatibility is fixed, but +# torch-directml's allocator cannot sustain their chunked inference loops +# (upstream limitation — see _mdxc_inference_device). The tests assert the +# fallback fires (warning logged) AND the output is correct. +CPU_FALLBACK_WARNING = "MDXC/RoFormer models currently run on CPU under DirectML" + +# (model, expected output files, is_roformer, validate_reference, expect_cpu_fallback) DML_MODEL_PARAMS = [ ( # The exact model from the issue #292 report @@ -55,7 +56,7 @@ ], True, True, - DML_SEGMENT_ARGS, + True, ), ( "mel_band_roformer_karaoke_aufr33_viperx_sdr_10.1956.ckpt", @@ -65,7 +66,7 @@ ], True, True, - DML_SEGMENT_ARGS, + True, ), ( # MDX — regression guard: already worked on DirectML before #292 @@ -76,7 +77,7 @@ ], False, True, - [], + False, ), ( # VR — regression guard: already worked on DirectML before #292 @@ -87,7 +88,7 @@ ], False, True, - [], + False, ), ( # Plain MDXC (TFC_TDF arch) — its STFT wrapper already CPU-hops @@ -101,7 +102,7 @@ ], False, False, - DML_SEGMENT_ARGS, + True, ), ] @@ -122,14 +123,14 @@ def _assert_audible_and_finite(path): assert rms > RMS_FLOOR, f"{path} is (near-)silent: rms={rms:.2e}" -@pytest.mark.parametrize("model,expected_files,is_roformer,validate_reference,extra_args", DML_MODEL_PARAMS) -def test_dml_separation(model, expected_files, is_roformer, validate_reference, extra_args): +@pytest.mark.parametrize("model,expected_files,is_roformer,validate_reference,expect_cpu_fallback", DML_MODEL_PARAMS) +def test_dml_separation(model, expected_files, is_roformer, validate_reference, expect_cpu_fallback): for f in expected_files: if os.path.exists(f): os.remove(f) result = subprocess.run( - [resolve_cli_executable(), "--use_directml", "--log_level", "debug", *extra_args, "-m", model, INPUT_FILE], + [resolve_cli_executable(), "--use_directml", "--log_level", "debug", "-m", model, INPUT_FILE], capture_output=True, text=True, check=False, @@ -142,6 +143,14 @@ def test_dml_separation(model, expected_files, is_roformer, validate_reference, assert "DirectML is available in Torch, setting Torch device to DirectML" in log_text assert "ONNXruntime has DmlExecutionProvider available, enabling acceleration" in log_text + if expect_cpu_fallback: + # MDXC-family models must announce the documented CPU fallback — + # if this assert fires because the fallback was removed, the DML + # allocator OOM is back unless torch-directml fixed it upstream. + assert CPU_FALLBACK_WARNING in log_text, f"{model} did not announce the DML CPU fallback" + else: + assert CPU_FALLBACK_WARNING not in log_text, f"{model} unexpectedly fell back to CPU" + if is_roformer: # Loader regression guard: the map_location fix means the NEW # implementation must load — a silent legacy fallback would keep CI diff --git a/tests/unit/test_directml.py b/tests/unit/test_directml.py index 86b87b0..8a91ed4 100644 --- a/tests/unit/test_directml.py +++ b/tests/unit/test_directml.py @@ -110,3 +110,42 @@ def test_new_implementation_map_location_unchanged_for_cuda(): def test_new_implementation_map_location_unchanged_for_mps(): assert _load_via_new_implementation("mps") == "mps" + + +# --------------------------------------------------------------------------- +# MDXC-family DirectML CPU fallback (issue #292) +# --------------------------------------------------------------------------- + +import torch as _torch + +from audio_separator.separator.architectures.mdxc_separator import _mdxc_inference_device + + +class TestMdxcInferenceDevice: + def test_cpu_and_cuda_pass_through(self): + log = MagicMock() + assert _mdxc_inference_device(_torch.device("cpu"), _torch.device("cpu"), log) == _torch.device("cpu") + cuda = _torch.device("cuda", 0) + assert _mdxc_inference_device(cuda, _torch.device("cpu"), log) == cuda + log.warning.assert_not_called() + + def test_dml_falls_back_to_cpu_with_warning(self, monkeypatch): + monkeypatch.delenv("AUDIO_SEPARATOR_FORCE_DML_MDXC", raising=False) + log = MagicMock() + dml = _torch.device("privateuseone", 0) + result = _mdxc_inference_device(dml, _torch.device("cpu"), log) + assert result == _torch.device("cpu") + assert log.warning.call_count == 1 + assert "run on CPU under DirectML" in log.warning.call_args[0][0] + + def test_dml_fallback_without_cpu_device_configured(self, monkeypatch): + monkeypatch.delenv("AUDIO_SEPARATOR_FORCE_DML_MDXC", raising=False) + result = _mdxc_inference_device(_torch.device("privateuseone", 0), None, MagicMock()) + assert result == _torch.device("cpu") + + def test_env_override_keeps_dml(self, monkeypatch): + monkeypatch.setenv("AUDIO_SEPARATOR_FORCE_DML_MDXC", "1") + log = MagicMock() + dml = _torch.device("privateuseone", 0) + assert _mdxc_inference_device(dml, _torch.device("cpu"), log) == dml + assert "attempting" in log.warning.call_args[0][0] From c0215f9df9d245382ec6601bf7813e1a18dbd7a7 Mon Sep 17 00:00:00 2001 From: Andrew Beveridge Date: Mon, 20 Jul 2026 17:26:00 -0400 Subject: [PATCH 11/11] test: assert new-implementation load via separator stats line; widen DML job timeout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The roformer_loader module logger doesn't propagate to the CLI handler, so 'with new implementation' never appears in CLI output — assert on the separator's "Roformer loading stats: {'new_implementation_success': 1" line instead. Separations themselves passed (3/5 green; the 2 roformer failures were this assertion only). CPU-fallback inference takes ~30 min of the job — raise timeout to 60. Co-Authored-By: Claude Fable 5 --- .github/workflows/run-integration-tests.yaml | 4 +++- tests/integration/test_windows_directml.py | 6 ++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/.github/workflows/run-integration-tests.yaml b/.github/workflows/run-integration-tests.yaml index cebac03..6ff4c46 100644 --- a/.github/workflows/run-integration-tests.yaml +++ b/.github/workflows/run-integration-tests.yaml @@ -234,7 +234,9 @@ jobs: needs: changes if: needs.changes.outputs.should_run == 'true' runs-on: [self-hosted, windows, gpu] - timeout-minutes: 40 + # RoFormer/MDX23C run on CPU under DML (allocator limitation) — the + # test step alone takes ~30 min. + timeout-minutes: 60 continue-on-error: true defaults: run: diff --git a/tests/integration/test_windows_directml.py b/tests/integration/test_windows_directml.py index 1f4e60a..507b62f 100644 --- a/tests/integration/test_windows_directml.py +++ b/tests/integration/test_windows_directml.py @@ -154,9 +154,11 @@ def test_dml_separation(model, expected_files, is_roformer, validate_reference, if is_roformer: # Loader regression guard: the map_location fix means the NEW # implementation must load — a silent legacy fallback would keep CI - # green while shipping the unfixed path. + # green while shipping the unfixed path. Assert on the separator's + # "Roformer loading stats" line (the loader module's own logger does + # not propagate to the CLI handler). assert "Fell back to legacy" not in log_text, f"{model} silently fell back to legacy implementation" - assert "with new implementation" in log_text, f"{model} did not report new-implementation load" + assert "'new_implementation_success': 1" in log_text, f"{model} did not report new-implementation load stats" for output_file in expected_files: assert os.path.exists(output_file), f"Output file {output_file} was not created"