Skip to content

Commit ee1fcee

Browse files
authored
feat: expose --use_directml CLI flag (experimental) (#290)
* docs: design spec for exposing DirectML via CLI (experimental tier) Contributor (Vageesha Gupta) reported DirectML support exists (PR #211) but is unreachable from the CLI. Design: expose --use_directml as explicit opt-in, add conditional discoverability hint, document as experimental with an honest per-architecture status, ship a patch release, and reply inviting validation of untested architectures. * docs: implementation plan for DirectML CLI exposure + spec correction Correct spec to drop remote-CLI scope: DirectML is local-only and the remote separation servers run on CUDA/CPU, so plumbing use_directml through audio_separator/remote/* would be dead config. Add 5-task TDD plan (CLI flag, discoverability hint, README, version bump, reply draft). * feat: expose --use_directml CLI flag (experimental) Wire the existing Separator(use_directml=...) support (PR #211) to the CLI, mirroring --use_autocast. Explicit opt-in only. Adds the fixture key + a wiring test in test_cli.py. * feat: hint that --use_directml is available when DML packages are installed In the CPU-fallback path of setup_torch_device, if torch_directml or onnxruntime-directml is installed but DirectML wasn't enabled, log a one-line INFO hint pointing to --use_directml. Fires during 'audio-separator --env_info' on DirectML-capable machines. * docs: add experimental DirectML (Windows AMD/Intel GPU) install section New install section with the dml extra, --use_directml usage, env_info confirmation line, and an honest per-architecture status matrix. Also adds --use_directml to the ToC, the full CLI-options reference, and the Separator class parameters list. * chore: bump version to 0.44.3 * chore: drop internal planning docs from PR
1 parent eecd260 commit ee1fcee

6 files changed

Lines changed: 100 additions & 2 deletions

File tree

README.md

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ The simplest (and probably most used) use case for this package is to separate a
2727
- [🎮 Nvidia GPU with CUDA or 🧪 Google Colab](#-nvidia-gpu-with-cuda-or--google-colab)
2828
- [ Apple Silicon, macOS Sonoma+ with M1 or newer CPU (CoreML acceleration)](#-apple-silicon-macos-sonoma-with-m1-or-newer-cpu-coreml-acceleration)
2929
- [🐢 No hardware acceleration, CPU only](#-no-hardware-acceleration-cpu-only)
30+
- [🪟 Windows AMD / Intel GPU with DirectML (experimental)](#-windows-amd--intel-gpu-with-directml-experimental)
3031
- [🎥 FFmpeg dependency](#-ffmpeg-dependency)
3132
- [GPU / CUDA specific installation steps with Pip](#gpu--cuda-specific-installation-steps-with-pip)
3233
- [Multiple CUDA library versions may be needed](#multiple-cuda-library-versions-may-be-needed)
@@ -139,6 +140,34 @@ Docker:
139140
beveradb/audio-separator
140141
```
141142

143+
### 🪟 Windows AMD / Intel GPU with DirectML (experimental)
144+
145+
> **Experimental / community-supported.** DirectML acceleration was contributed by the community and is not tested in CI or by the maintainer (it requires Windows plus an AMD or Intel GPU). It is opt-in and will never affect CUDA, Apple Silicon, or CPU users.
146+
147+
Install with the `dml` extra:
148+
```sh
149+
pip install "audio-separator[dml]"
150+
```
151+
152+
Then enable it explicitly with the `--use_directml` flag:
153+
```sh
154+
audio-separator path/to/audio.wav --use_directml
155+
```
156+
157+
💬 If DirectML is configured correctly you should see this log message when running `audio-separator --env_info`:
158+
`ONNXruntime has DmlExecutionProvider available, enabling acceleration`
159+
160+
**Model architecture status on DirectML:**
161+
162+
| Architecture | Model types | Status |
163+
|---|---|---|
164+
| MDX | `.onnx` | ✅ Confirmed working |
165+
| MDXC (incl. the default `bs_roformer` model) | `.ckpt` / `.yaml` | ⚠️ Expected to work, community-untested |
166+
| VR | `.pth` | ⚠️ Expected to work, community-untested |
167+
| Demucs || ❓ Unverified |
168+
169+
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".
170+
142171
### 🎥 FFmpeg dependency
143172

144173
💬 To test if `audio-separator` has been successfully configured to use FFmpeg, run `audio-separator --env_info`. The log will show `FFmpeg installed`.
@@ -418,7 +447,7 @@ Presets are defined in `audio_separator/ensemble_presets.json` — contributions
418447
```sh
419448
usage: audio-separator [-h] [-v] [-d] [-e] [-l] [--log_level LOG_LEVEL] [--list_filter LIST_FILTER] [--list_limit LIST_LIMIT] [--list_format {pretty,json}] [-m MODEL_FILENAME] [--output_format OUTPUT_FORMAT]
420449
[--output_bitrate OUTPUT_BITRATE] [--output_dir OUTPUT_DIR] [--model_file_dir MODEL_FILE_DIR] [--download_model_only] [--invert_spect] [--normalization NORMALIZATION]
421-
[--amplification AMPLIFICATION] [--single_stem SINGLE_STEM] [--sample_rate SAMPLE_RATE] [--use_soundfile] [--use_autocast] [--custom_output_names CUSTOM_OUTPUT_NAMES]
450+
[--amplification AMPLIFICATION] [--single_stem SINGLE_STEM] [--sample_rate SAMPLE_RATE] [--use_soundfile] [--use_autocast] [--use_directml] [--custom_output_names CUSTOM_OUTPUT_NAMES]
422451
[--mdx_segment_size MDX_SEGMENT_SIZE] [--mdx_overlap MDX_OVERLAP] [--mdx_batch_size MDX_BATCH_SIZE] [--mdx_hop_length MDX_HOP_LENGTH] [--mdx_enable_denoise] [--vr_batch_size VR_BATCH_SIZE]
423452
[--vr_window_size VR_WINDOW_SIZE] [--vr_aggression VR_AGGRESSION] [--vr_enable_tta] [--vr_high_end_process] [--vr_enable_post_process]
424453
[--vr_post_process_threshold VR_POST_PROCESS_THRESHOLD] [--demucs_segment_size DEMUCS_SEGMENT_SIZE] [--demucs_shifts DEMUCS_SHIFTS] [--demucs_overlap DEMUCS_OVERLAP]
@@ -460,6 +489,7 @@ Common Separation Parameters:
460489
--sample_rate SAMPLE_RATE Modify the sample rate of the output audio (default: 44100). Example: --sample_rate=44100
461490
--use_soundfile Use soundfile to write audio output (default: False). Example: --use_soundfile
462491
--use_autocast Use PyTorch autocast for faster inference (default: False). Do not use for CPU inference. Example: --use_autocast
492+
--use_directml Use DirectML for hardware-accelerated inference on Windows AMD/Intel GPUs (experimental; requires the 'dml' extra). Example: --use_directml
463493
--custom_output_names CUSTOM_OUTPUT_NAMES Custom names for all output files in JSON format (default: None). Example: --custom_output_names='{"Vocals": "vocals_output", "Drums": "drums_output"}'
464494
465495
MDX Architecture Parameters:
@@ -616,6 +646,7 @@ You can also rename specific stems:
616646
- **`sample_rate`:** (Optional) Set the sample rate of the output audio. `Default: 44100`
617647
- **`use_soundfile`:** (Optional) Use soundfile for output writing, can solve OOM issues, especially on longer audio.
618648
- **`use_autocast`:** (Optional) Flag to use PyTorch autocast for faster inference. Do not use for CPU inference. `Default: False`
649+
- **`use_directml`:** (Optional) Flag to use DirectML for hardware-accelerated inference on Windows AMD/Intel GPUs (experimental; requires the `dml` extra and only takes effect when CUDA and Apple Silicon MPS are unavailable). `Default: False`
619650
- **`mdx_params`:** (Optional) MDX Architecture Specific Attributes & Defaults. `Default: {"hop_length": 1024, "segment_size": 256, "overlap": 0.25, "batch_size": 1, "enable_denoise": False}`
620651
- **`vr_params`:** (Optional) VR Architecture Specific Attributes & Defaults. `Default: {"batch_size": 1, "window_size": 512, "aggression": 5, "enable_tta": False, "enable_post_process": False, "post_process_threshold": 0.2, "high_end_process": False}`
621652
- **`demucs_params`:** (Optional) Demucs Architecture Specific Attributes & Defaults. `Default: {"segment_size": "Default", "shifts": 2, "overlap": 0.25, "segments_enabled": True}` _(Note: `segment_size` "Default" uses the model's internal default, typically 40 for older Demucs models and 10 for Demucs v4/htdemucs)_

audio_separator/separator/separator.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -401,6 +401,14 @@ def setup_torch_device(self, system_info):
401401
self.torch_device = self.torch_device_cpu
402402
self.onnx_execution_provider = ["CPUExecutionProvider"]
403403

404+
# Discoverability hint: DirectML is an explicit opt-in (experimental). If the
405+
# DirectML packages are installed but the feature wasn't enabled, tell the user how.
406+
if not self.use_directml and (has_torch_dml_installed or self.get_package_distribution("onnxruntime-directml") is not None):
407+
self.logger.info(
408+
"DirectML packages detected but DirectML is not enabled. "
409+
"Pass use_directml=True (or --use_directml on the CLI) to enable experimental DirectML acceleration."
410+
)
411+
404412
def configure_cuda(self, ort_providers):
405413
"""
406414
This method configures the CUDA device for PyTorch and ONNX Runtime, if available.

audio_separator/utils/cli.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ def main():
6161
sample_rate_help = "Modify the sample rate of the output audio (default: %(default)s). Example: --sample_rate=44100"
6262
use_soundfile_help = "Use soundfile to write audio output (default: %(default)s). Example: --use_soundfile"
6363
use_autocast_help = "Use PyTorch autocast for faster inference (default: %(default)s). Do not use for CPU inference. Example: --use_autocast"
64+
use_directml_help = "Use DirectML for hardware-accelerated inference on Windows AMD/Intel GPUs (experimental; requires the 'dml' extra). Example: --use_directml"
6465
chunk_duration_help = "Split audio into chunks of this duration in seconds (default: %(default)s = no chunking). Useful for processing very long audio files on systems with limited memory. Recommended: 600 (10 minutes) for files >1 hour. Chunks are concatenated without overlap/crossfade. Example: --chunk_duration=600"
6566
ensemble_algorithm_help = "Algorithm to use for ensembling multiple models (default: avg_wave). Choices: avg_wave, median_wave, min_wave, max_wave, avg_fft, median_fft, min_fft, max_fft, uvr_max_spec, uvr_min_spec, ensemble_wav. Example: --ensemble_algorithm=uvr_max_spec"
6667
ensemble_weights_help = "Weights for ensembling multiple models (default: equal). Number of weights must match number of models. Example: --ensemble_weights 1.0 0.5"
@@ -76,6 +77,7 @@ def main():
7677
common_params.add_argument("--sample_rate", type=int, default=44100, help=sample_rate_help)
7778
common_params.add_argument("--use_soundfile", action="store_true", help=use_soundfile_help)
7879
common_params.add_argument("--use_autocast", action="store_true", help=use_autocast_help)
80+
common_params.add_argument("--use_directml", action="store_true", help=use_directml_help)
7981
common_params.add_argument("--chunk_duration", type=float, default=None, help=chunk_duration_help)
8082
common_params.add_argument(
8183
"--ensemble_algorithm",
@@ -247,6 +249,7 @@ def main():
247249
sample_rate=args.sample_rate,
248250
use_soundfile=args.use_soundfile,
249251
use_autocast=args.use_autocast,
252+
use_directml=args.use_directml,
250253
chunk_duration=args.chunk_duration,
251254
ensemble_algorithm=args.ensemble_algorithm,
252255
ensemble_weights=args.ensemble_weights,

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
44

55
[tool.poetry]
66
name = "audio-separator"
7-
version = "0.44.2"
7+
version = "0.44.3"
88
description = "Easy to use audio stem separation, using various models from UVR trained primarily by @Anjok07"
99
authors = ["Andrew Beveridge <andrew@beveridge.uk>"]
1010
license = "MIT"

tests/unit/test_cli.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ def common_expected_args():
4141
"sample_rate": 44100,
4242
"use_soundfile": False,
4343
"use_autocast": False,
44+
"use_directml": False,
4445
"chunk_duration": None,
4546
"ensemble_algorithm": None,
4647
"ensemble_weights": None,
@@ -257,6 +258,23 @@ def test_cli_use_autocast_argument(common_expected_args):
257258
mock_separator.assert_called_once_with(**expected_args)
258259

259260

261+
# Test using use_directml argument
262+
def test_cli_use_directml_argument(common_expected_args):
263+
test_args = ["cli.py", "test_audio.mp3", "--use_directml"]
264+
with patch("sys.argv", test_args):
265+
with patch("audio_separator.separator.Separator") as mock_separator:
266+
mock_separator_instance = mock_separator.return_value
267+
mock_separator_instance.separate.return_value = ["output_file.mp3"]
268+
main()
269+
270+
# Update expected args for this specific test
271+
expected_args = common_expected_args.copy()
272+
expected_args["use_directml"] = True
273+
274+
# Assertions
275+
mock_separator.assert_called_once_with(**expected_args)
276+
277+
260278
# Test using custom_output_names arguments
261279
def test_cli_custom_output_names_argument(common_expected_args):
262280
custom_names = {

tests/unit/test_directml.py

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
import logging
2+
import platform
3+
from unittest.mock import MagicMock, patch
4+
5+
from audio_separator.separator import Separator
6+
7+
HINT = "DirectML packages detected but DirectML is not enabled"
8+
9+
10+
def _run_setup(use_directml, dml_installed):
11+
"""Construct a Separator without auto device-setup, then drive setup_torch_device
12+
with CUDA and MPS forced unavailable so the CPU-fallback path always runs."""
13+
sep = Separator(info_only=True)
14+
sep.use_directml = use_directml
15+
16+
def fake_dist(name):
17+
if name in ("torch_directml", "onnxruntime-directml") and dml_installed:
18+
return MagicMock()
19+
return None
20+
21+
with patch.object(sep, "get_package_distribution", side_effect=fake_dist), \
22+
patch("torch.cuda.is_available", return_value=False), \
23+
patch("torch.backends.mps.is_available", return_value=False), \
24+
patch("audio_separator.separator.separator.ort.get_available_providers", return_value=["CPUExecutionProvider"]):
25+
sep.setup_torch_device(platform.uname())
26+
return sep
27+
28+
29+
def test_directml_hint_shown_when_packages_present_but_disabled(caplog):
30+
with caplog.at_level(logging.INFO):
31+
_run_setup(use_directml=False, dml_installed=True)
32+
assert any(HINT in r.message for r in caplog.records)
33+
34+
35+
def test_directml_hint_absent_when_no_packages(caplog):
36+
with caplog.at_level(logging.INFO):
37+
_run_setup(use_directml=False, dml_installed=False)
38+
assert not any(HINT in r.message for r in caplog.records)

0 commit comments

Comments
 (0)