feat: expose --use_directml CLI flag (experimental)

README.md

@@ -27,6 +27,7 @@ The simplest (and probably most used) use case for this package is to separate a
     - [🎮 Nvidia GPU with CUDA or 🧪 Google Colab](#-nvidia-gpu-with-cuda-or--google-colab)
     - [ Apple Silicon, macOS Sonoma+ with M1 or newer CPU (CoreML acceleration)](#-apple-silicon-macos-sonoma-with-m1-or-newer-cpu-coreml-acceleration)
     - [🐢 No hardware acceleration, CPU only](#-no-hardware-acceleration-cpu-only)
+    - [🪟 Windows AMD / Intel GPU with DirectML (experimental)](#-windows-amd--intel-gpu-with-directml-experimental)
     - [🎥 FFmpeg dependency](#-ffmpeg-dependency)
   - [GPU / CUDA specific installation steps with Pip](#gpu--cuda-specific-installation-steps-with-pip)
     - [Multiple CUDA library versions may be needed](#multiple-cuda-library-versions-may-be-needed)
@@ -139,6 +140,34 @@ Docker:
 beveradb/audio-separator
 ```
 
+### 🪟 Windows AMD / Intel GPU with DirectML (experimental)
+
+> **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.
+
+Install with the `dml` extra:
+```sh
+pip install "audio-separator[dml]"
+```
+
+Then enable it explicitly with the `--use_directml` flag:
+```sh
+audio-separator path/to/audio.wav --use_directml
+```
+
+💬 If DirectML is configured correctly you should see this log message when running `audio-separator --env_info`:
+ `ONNXruntime has DmlExecutionProvider available, enabling acceleration`
+
+**Model architecture status on 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".
+
 ### 🎥 FFmpeg dependency
 
 💬 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
 ```sh
 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]
                        [--output_bitrate OUTPUT_BITRATE] [--output_dir OUTPUT_DIR] [--model_file_dir MODEL_FILE_DIR] [--download_model_only] [--invert_spect] [--normalization NORMALIZATION]
-                       [--amplification AMPLIFICATION] [--single_stem SINGLE_STEM] [--sample_rate SAMPLE_RATE] [--use_soundfile] [--use_autocast] [--custom_output_names CUSTOM_OUTPUT_NAMES]
+                       [--amplification AMPLIFICATION] [--single_stem SINGLE_STEM] [--sample_rate SAMPLE_RATE] [--use_soundfile] [--use_autocast] [--use_directml] [--custom_output_names CUSTOM_OUTPUT_NAMES]
                        [--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]
                        [--vr_window_size VR_WINDOW_SIZE] [--vr_aggression VR_AGGRESSION] [--vr_enable_tta] [--vr_high_end_process] [--vr_enable_post_process]
                        [--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:
   --sample_rate SAMPLE_RATE                              Modify the sample rate of the output audio (default: 44100). Example: --sample_rate=44100
   --use_soundfile                                        Use soundfile to write audio output (default: False). Example: --use_soundfile
   --use_autocast                                         Use PyTorch autocast for faster inference (default: False). Do not use for CPU inference. Example: --use_autocast
+  --use_directml                                         Use DirectML for hardware-accelerated inference on Windows AMD/Intel GPUs (experimental; requires the 'dml' extra). Example: --use_directml
   --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"}'
 
 MDX Architecture Parameters:
@@ -616,6 +646,7 @@ You can also rename specific stems:
 - **`sample_rate`:** (Optional) Set the sample rate of the output audio. `Default: 44100`
 - **`use_soundfile`:** (Optional) Use soundfile for output writing, can solve OOM issues, especially on longer audio.
 - **`use_autocast`:** (Optional) Flag to use PyTorch autocast for faster inference. Do not use for CPU inference. `Default: False`
+- **`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`
 - **`mdx_params`:** (Optional) MDX Architecture Specific Attributes & Defaults. `Default: {"hop_length": 1024, "segment_size": 256, "overlap": 0.25, "batch_size": 1, "enable_denoise": False}`
 - **`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}`
 - **`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

@@ -401,6 +401,14 @@ def setup_torch_device(self, system_info):
             self.torch_device = self.torch_device_cpu
             self.onnx_execution_provider = ["CPUExecutionProvider"]
 
+            # Discoverability hint: DirectML is an explicit opt-in (experimental). If the
+            # DirectML packages are installed but the feature wasn't enabled, tell the user how.
+            if not self.use_directml and (has_torch_dml_installed or self.get_package_distribution("onnxruntime-directml") is not None):
+                self.logger.info(
+                    "DirectML packages detected but DirectML is not enabled. "
+                    "Pass use_directml=True (or --use_directml on the CLI) to enable experimental DirectML acceleration."
+                )
+
     def configure_cuda(self, ort_providers):
         """
         This method configures the CUDA device for PyTorch and ONNX Runtime, if available.

audio_separator/utils/cli.py

@@ -61,6 +61,7 @@ def main():
     sample_rate_help = "Modify the sample rate of the output audio (default: %(default)s). Example: --sample_rate=44100"
     use_soundfile_help = "Use soundfile to write audio output (default: %(default)s). Example: --use_soundfile"
     use_autocast_help = "Use PyTorch autocast for faster inference (default: %(default)s). Do not use for CPU inference. Example: --use_autocast"
+    use_directml_help = "Use DirectML for hardware-accelerated inference on Windows AMD/Intel GPUs (experimental; requires the 'dml' extra). Example: --use_directml"
     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"
     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"
     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():
     common_params.add_argument("--sample_rate", type=int, default=44100, help=sample_rate_help)
     common_params.add_argument("--use_soundfile", action="store_true", help=use_soundfile_help)
     common_params.add_argument("--use_autocast", action="store_true", help=use_autocast_help)
+    common_params.add_argument("--use_directml", action="store_true", help=use_directml_help)
     common_params.add_argument("--chunk_duration", type=float, default=None, help=chunk_duration_help)
     common_params.add_argument(
         "--ensemble_algorithm",
@@ -247,6 +249,7 @@ def main():
         sample_rate=args.sample_rate,
         use_soundfile=args.use_soundfile,
         use_autocast=args.use_autocast,
+        use_directml=args.use_directml,
         chunk_duration=args.chunk_duration,
         ensemble_algorithm=args.ensemble_algorithm,
         ensemble_weights=args.ensemble_weights,

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "audio-separator"
-version = "0.44.2"
+version = "0.44.3"
 description = "Easy to use audio stem separation, using various models from UVR trained primarily by @Anjok07"
 authors = ["Andrew Beveridge <andrew@beveridge.uk>"]
 license = "MIT"

tests/unit/test_cli.py

@@ -41,6 +41,7 @@ def common_expected_args():
         "sample_rate": 44100,
         "use_soundfile": False,
         "use_autocast": False,
+        "use_directml": False,
         "chunk_duration": None,
         "ensemble_algorithm": None,
         "ensemble_weights": None,
@@ -257,6 +258,23 @@ def test_cli_use_autocast_argument(common_expected_args):
             mock_separator.assert_called_once_with(**expected_args)
 
 
+# Test using use_directml argument
+def test_cli_use_directml_argument(common_expected_args):
+    test_args = ["cli.py", "test_audio.mp3", "--use_directml"]
+    with patch("sys.argv", test_args):
+        with patch("audio_separator.separator.Separator") as mock_separator:
+            mock_separator_instance = mock_separator.return_value
+            mock_separator_instance.separate.return_value = ["output_file.mp3"]
+            main()
+
+            # Update expected args for this specific test
+            expected_args = common_expected_args.copy()
+            expected_args["use_directml"] = True
+
+            # Assertions
+            mock_separator.assert_called_once_with(**expected_args)
+
+
 # Test using custom_output_names arguments
 def test_cli_custom_output_names_argument(common_expected_args):
     custom_names = {

tests/unit/test_directml.py

@@ -0,0 +1,38 @@
+import logging
+import platform
+from unittest.mock import MagicMock, patch
+
+from audio_separator.separator import Separator
+
+HINT = "DirectML packages detected but DirectML is not enabled"
+
+
+def _run_setup(use_directml, dml_installed):
+    """Construct a Separator without auto device-setup, then drive setup_torch_device
+    with CUDA and MPS forced unavailable so the CPU-fallback path always runs."""
+    sep = Separator(info_only=True)
+    sep.use_directml = use_directml
+
+    def fake_dist(name):
+        if name in ("torch_directml", "onnxruntime-directml") and dml_installed:
+            return MagicMock()
+        return None
+
+    with patch.object(sep, "get_package_distribution", side_effect=fake_dist), \
+         patch("torch.cuda.is_available", return_value=False), \
+         patch("torch.backends.mps.is_available", return_value=False), \
+         patch("audio_separator.separator.separator.ort.get_available_providers", return_value=["CPUExecutionProvider"]):
+        sep.setup_torch_device(platform.uname())
+    return sep
+
+
+def test_directml_hint_shown_when_packages_present_but_disabled(caplog):
+    with caplog.at_level(logging.INFO):
+        _run_setup(use_directml=False, dml_installed=True)
+    assert any(HINT in r.message for r in caplog.records)
+
+
+def test_directml_hint_absent_when_no_packages(caplog):
+    with caplog.at_level(logging.INFO):
+        _run_setup(use_directml=False, dml_installed=False)
+    assert not any(HINT in r.message for r in caplog.records)