diff --git a/README.md b/README.md index e51e32d..1b90ba4 100644 --- a/README.md +++ b/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)_ diff --git a/audio_separator/separator/separator.py b/audio_separator/separator/separator.py index 8302488..6b1259d 100644 --- a/audio_separator/separator/separator.py +++ b/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. diff --git a/audio_separator/utils/cli.py b/audio_separator/utils/cli.py index c45fcaf..6c33fae 100755 --- a/audio_separator/utils/cli.py +++ b/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, diff --git a/pyproject.toml b/pyproject.toml index a9de5ca..cde5b4f 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.2" +version = "0.44.3" 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_cli.py b/tests/unit/test_cli.py index cae4518..ce49eef 100644 --- a/tests/unit/test_cli.py +++ b/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 = { diff --git a/tests/unit/test_directml.py b/tests/unit/test_directml.py new file mode 100644 index 0000000..2bbf862 --- /dev/null +++ b/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)