ci: windows-directml integration job on self-hosted Windows GPU runners (#292)

.github/workflows/run-integration-tests.yaml

@@ -218,10 +218,71 @@ 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]
+    # 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:
+        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
+          $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: 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))"
+          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 +297,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" ]] || \

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
 

audio_separator/separator/architectures/mdxc_separator.py

@@ -1,3 +1,4 @@
+import gc
 import os
 import sys
 
@@ -13,6 +14,55 @@
 # 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.
+
+    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.
@@ -29,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
@@ -326,6 +380,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 +453,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")

audio_separator/separator/separator.py

@@ -984,15 +984,15 @@ 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}")
                 try:
                     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
 

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,9 +103,28 @@ 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)
 
+        if _is_dml_device(device):
+            # 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):
+                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(torch.matmul(attn, vs))
+            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

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)
 

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)
 

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)}")

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
 

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 <andrew@beveridge.uk>"]
 license = "MIT"