diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 00000000..f48f5b5a --- /dev/null +++ b/.dockerignore @@ -0,0 +1,16 @@ +.git +.github +.venv +venv +__pycache__ +*.pyc +*.pyo +.pytest_cache +.ruff_cache +dist +build +*.egg-info +models +*.flac +*.wav +*.mp3 diff --git a/Dockerfile.rocm b/Dockerfile.rocm new file mode 100644 index 00000000..7a63cf02 --- /dev/null +++ b/Dockerfile.rocm @@ -0,0 +1,51 @@ +# Use ROCm base image with Python +FROM rocm/dev-ubuntu-22.04:7.2-complete + +# Set the working directory in the container +WORKDIR /workdir + +# Install necessary packages +RUN apt-get update && apt-get install -y --no-install-recommends \ + ffmpeg \ + python3 \ + python3-pip \ + && rm -rf /var/lib/apt/lists/* + +RUN python3 -m pip install --upgrade pip + +# Install PyTorch with ROCm support +RUN --mount=type=cache,target=/root/.cache \ + pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm7.2 + +# Install audio-separator from this local checkout so the ROCm support in the +# repository is used, not a released PyPI version that may predate it. +# The [rocm] extra installs onnxruntime-gpu. Its CUDAExecutionProvider is CUDA-linked +# and cannot activate on a ROCm host, so MDX-Net ONNX models run through the PyTorch +# HIP path (see MDXSeparator._prefer_pytorch_inference_on_rocm); VR, Demucs, and +# MDXC/RoFormer already run on the ROCm GPU through PyTorch. +COPY . /src +RUN --mount=type=cache,target=/root/.cache \ + cd /src && pip3 install ".[rocm]" + +# GPU compatibility overrides are opt-in and unset by default. PYTORCH_ROCM_ARCH is a +# kernel build target that is not needed with prebuilt ROCm PyTorch wheels, and +# HSA_OVERRIDE_GFX_VERSION is only needed for consumer GPUs whose gfx target the +# installed ROCm build does not natively support (for example RX 6600, gfx1032). +# Supply it at run time with `-e`: +# docker run -e HSA_OVERRIDE_GFX_VERSION=10.3.0 ... + +# Run as a non-root user. ROCm device access still comes from the caller's +# --device /dev/kfd --device /dev/dri and --group-add=video --group-add=render, which +# add the host GPU groups to this user. APP_UID/APP_GID default to 1000; override at build time +# (--build-arg APP_UID=$(id -u) APP_GID=$(id -g)) or pass --user at run time so the +# host-owned /workdir bind mount stays writable. +ARG APP_UID=1000 +ARG APP_GID=1000 +RUN groupadd --gid "${APP_GID}" appuser \ + && useradd --uid "${APP_UID}" --gid "${APP_GID}" --create-home --shell /bin/bash appuser \ + && mkdir -p /workdir \ + && chown -R appuser:appuser /workdir +USER appuser + +# Run audio-separator when the container launches +ENTRYPOINT ["audio-separator"] diff --git a/README.md b/README.md index 8b6c73d3..b892ee2a 100644 --- a/README.md +++ b/README.md @@ -28,10 +28,12 @@ The simplest (and probably most used) use case for this package is to separate a - [ Apple Silicon, macOS Sonoma+ with M1 or newer (CoreML and MPS acceleration)](#-apple-silicon-macos-sonoma-with-m1-or-newer-coreml-and-mps-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) + - [🖥️ AMD GPU with ROCm (Linux)](#amd-gpu-rocm) - [Inference precision and compilation](#inference-precision-and-compilation) - [🎥 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) + - [ROCm specific troubleshooting](#rocm-specific-troubleshooting) - [Usage 🚀](#usage-) - [Command Line Interface (CLI)](#command-line-interface-cli) - [Listing and Filtering Available Models](#listing-and-filtering-available-models) @@ -70,6 +72,7 @@ The simplest (and probably most used) use case for this package is to separate a - Ability to inference using a pre-trained model in PTH or ONNX format. - CLI support for easy use in scripts and batch processing. - Python API for integration into other projects. +- GPU acceleration on NVIDIA CUDA, AMD ROCm (Linux), Apple Silicon, Windows DirectML, or CPU. ## Installation 🛠️ @@ -198,6 +201,88 @@ 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. + + +### 🖥️ AMD GPU with ROCm (Linux) + +**Supported ROCm versions:** 6.0+ (tested on ROCm 7.2) + +💬 If successfully configured, you should see one of these messages when running `audio-separator --env_info`: + +- **PyTorch models (VR, Demucs, MDXC/RoFormer):** `ROCm (AMD GPU) is available in Torch, setting Torch device to CUDA (ROCm presents as CUDA)` +- **MDX-Net with native AMD ONNX acceleration (`onnxruntime-migraphx` installed):** `ONNXruntime has MIGraphXExecutionProvider available, enabling acceleration` +- **MDX-Net without a native AMD ONNX provider (the usual case):** `ROCm detected and no AMD ONNX GPU provider is selected; running MDX-Net through PyTorch (HIP) instead of the ONNX Runtime CUDA/CPU fallback.` + +Pip (complete installation): +```sh +# Step 1: Install PyTorch with ROCm support (change the ROCm version as needed) +pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm7.2 + +# Step 2: Install audio-separator with the rocm extra +pip install "audio-separator[rocm]" +``` + +> ⚠️ **Important:** PyTorch with ROCm support must be installed **before** audio-separator. +> The `[rocm]` extra installs `onnxruntime-gpu`; its `CUDAExecutionProvider` is CUDA-linked +> and cannot load on a ROCm host, so MDX-Net ONNX models are run through the PyTorch HIP path. + +**Required ROCm packages:** +- ROCm system libraries (see https://rocm.docs.amd.com for distro-specific instructions) +- PyTorch: `torch`, `torchvision`, `torchaudio` built for ROCm +- ONNX Runtime: `onnxruntime-gpu` (installed automatically by the `[rocm]` extra) + +**Basic ROCm setup:** + +Set `HSA_OVERRIDE_GFX_VERSION` **before** running Python or `audio-separator` when your GPU needs the override. The table also lists the `PYTORCH_ROCM_ARCH` kernel build target for reference. + +| GPU | `HSA_OVERRIDE_GFX_VERSION` | `PYTORCH_ROCM_ARCH` | +|-----|---------------------------|---------------------| +| RX 6600 / 6600 XT (gfx1032) | `10.3.0` | `gfx1030` | +| RX 7900 XT/XTX (gfx1100) | `11.0.0` | `gfx1100` | +| MI250 (gfx90a) | not needed | `gfx90a` | +| MI300X (gfx942) | not needed | `gfx942` | + +```sh +export HSA_OVERRIDE_GFX_VERSION=10.3.0 +``` + +**How GPU acceleration works on AMD:** + +audio-separator automatically detects ROCm PyTorch at startup (via `torch.version.hip`), and PyTorch exposes the AMD GPU through its `cuda` device: + +1. PyTorch models (VR, Demucs, MDXC/RoFormer) run on the AMD GPU directly through HIP. +2. For MDX-Net (`.onnx`) models, `MIGraphXExecutionProvider` is used when `onnxruntime-migraphx` is installed. Otherwise the ONNX graph is converted to PyTorch with `onnx2torch` and runs on the AMD GPU. Official `onnxruntime-gpu` wheels contain only a CUDA-linked `CUDAExecutionProvider`, which cannot activate on a ROCm host. +3. Set `AUDIO_SEPARATOR_ROCM_ONNX` to `ort`, `cpu`, or `0` (case-insensitive) to keep the ONNX Runtime path and skip the PyTorch conversion. The provider still comes from the runtime's provider list, which on a ROCm host is the CUDA-linked `CUDAExecutionProvider`; ONNX Runtime then falls back to `CPUExecutionProvider` if that provider cannot load. + +**Docker (build from source):** +```sh +docker build -f Dockerfile.rocm -t audio-separator:rocm . +# Run as your host user so the bind-mounted /workdir stays writable, and grant the +# GIDs that own the GPU device nodes. +docker run -it \ + --device=/dev/kfd --device=/dev/dri \ + --group-add "$(stat -c '%g' /dev/kfd)" \ + --group-add "$(stat -c '%g' /dev/dri/renderD128)" \ + --user "$(id -u):$(id -g)" \ + -v "$(pwd):/workdir" audio-separator:rocm input.wav +# Add -e HSA_OVERRIDE_GFX_VERSION=10.3.0 only if your GPU needs the override. + +# Rootless podman: keep the host supplementary groups instead of mapping them, +# alongside --userns=keep-id so the bind mount maps to your host user. +podman run -it --userns=keep-id --group-add=keep-groups \ + --device=/dev/kfd --device=/dev/dri \ + --user "$(id -u):$(id -g)" \ + -v "$(pwd):/workdir" audio-separator:rocm input.wav +``` + +The image creates a non-root `appuser` (uid/gid 1000 by default). Rebuild with +`--build-arg APP_UID=$(id -u) --build-arg APP_GID=$(id -g)` to bake in a different +uid/gid, or pass `--user` as above. The rootless podman command keeps the host +supplementary groups with `--group-add=keep-groups`, which cannot be combined with +named `--group-add` options. + +**Note:** `onnxruntime-rocm` (which provided `ROCMExecutionProvider`) was removed from ONNX Runtime since 1.23 and is not recommended. For native ONNX GPU acceleration on AMD, install `onnxruntime-migraphx` (it requires the MIGraphX libraries from ROCm); otherwise MDX-Net uses the PyTorch HIP path automatically. + ### Inference precision and compilation Three opt-in flags control PyTorch inference: @@ -262,7 +347,7 @@ brew update; brew install ffmpeg ## GPU / CUDA specific installation steps with Pip -In theory, all you should need to do to get `audio-separator` working with a GPU is install it with the `[gpu]` extra as above. +In theory, all you should need to do to get `audio-separator` working with a GPU is install it with the `[gpu]` extra (CUDA/NVIDIA) or the `[rocm]` extra (AMD/ROCm) as above, after installing the matching PyTorch build. However, sometimes getting both PyTorch and ONNX Runtime working with CUDA support can be a bit tricky so it may not work that easily. @@ -302,6 +387,47 @@ python -m pip install ort-nightly-gpu --index-url=https://aiinfra.pkgs.visualstu > Note: if anyone knows how to make this cleaner so we can support both different platform-specific dependencies for hardware acceleration without a separate installation process for each, please let me know or raise a PR! +### ROCm specific troubleshooting + +**Prerequisites:** +1. ROCm system libraries installed (minimum 6.0+, recommended 7.2+) + - Check: `rocminfo | head -20` should show your GPU + - Install from: https://rocm.docs.amd.com +2. PyTorch with ROCm support (installed from PyTorch's ROCm index URL) + - Check: `python -c "import torch; print(torch.__version__)"` should contain `+rocm` +3. `onnxruntime-gpu` installed (automatic via `pip install "audio-separator[rocm]"`) + +**If GPU acceleration is not working:** + +Run `audio-separator --env_info` and check for these clues: + +| Symptom | Likely Cause | Fix | +|---------|-------------|-----| +| MDX-Net runs on the CPU | `AUDIO_SEPARATOR_ROCM_ONNX` is set to `ort`, `cpu`, or `0`, forcing the ONNX Runtime path | Remove the override. Without a usable AMD ONNX provider, MDX-Net converts to PyTorch and runs through HIP by default, so CPU execution is not expected with `onnxruntime-gpu` | +| MDX-Net still runs on the CPU with no override set | Conversion did not run, or a conversion error occurred | Check the log for `running MDX-Net through PyTorch (HIP)`; confirm `torch.version.hip` is set and the resolved device is `cuda` | +| `MIGraphXExecutionProvider` not in providers | MIGraphX libraries or wheel not installed | `pip uninstall -y onnxruntime onnxruntime-gpu && pip install onnxruntime-migraphx` and `sudo apt install migraphx` for native ONNX acceleration | +| No `CUDAExecutionProvider` in providers | Normal on ROCm | `CUDAExecutionProvider` is CUDA-only and is not expected to activate on a ROCm host | +| `torch.cuda.is_available()` is `False` | Wrong PyTorch build | Reinstall PyTorch: `pip install torch --index-url https://download.pytorch.org/whl/rocm7.2` | +| `librocm_smi64.so` errors | Missing ROCm runtime libs | `sudo apt install rocm-libs` or distro equivalent | +| `HSA_OVERRIDE_GFX_VERSION not set` warning | Consumer GPU needs env variables | See the GPU table above in the AMD section | +| STFT assertions / crashes | ROCm 5.x bug (fixed in 7.x) | Upgrade to ROCm 7.2+ | + +**Clean reinstall:** +```sh +pip uninstall torch torchvision torchaudio onnxruntime onnxruntime-gpu +pip cache purge +pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm7.2 +pip install "audio-separator[rocm]" +``` + +**Docker issues:** +- Use the provided `Dockerfile.rocm` and ensure proper device mounting: + ```sh + docker run -it --device=/dev/kfd --device=/dev/dri --group-add "$(stat -c '%g' /dev/kfd)" --group-add "$(stat -c '%g' /dev/dri/renderD128)" --user "$(id -u):$(id -g)" ... + ``` +- The Dockerfile uses ROCm 7.2 and `onnxruntime-gpu` (not `onnxruntime-rocm`) for compatibility. +- GPU compatibility env vars are not set in the image; pass `-e HSA_OVERRIDE_GFX_VERSION=...` at run time only if your GPU needs it. + ## Usage 🚀 ### Command Line Interface (CLI) diff --git a/audio_separator/separator/architectures/mdx_separator.py b/audio_separator/separator/architectures/mdx_separator.py index af5e311e..acffdc00 100644 --- a/audio_separator/separator/architectures/mdx_separator.py +++ b/audio_separator/separator/architectures/mdx_separator.py @@ -20,6 +20,7 @@ class MDXSeparator(CommonSeparator): """ def __init__(self, common_config, arch_config): + """Initialize the MDX separator with shared and MDX-specific parameters.""" # Any configuration values which can be shared between architectures should be set already in CommonSeparator, # e.g. user-specified functionality choices (self.output_single_stem) or common model parameters (self.primary_stem_name) super().__init__(config=common_config) @@ -112,8 +113,12 @@ def load_model(self): """ self.logger.debug("Loading ONNX model for inference...") - if self.segment_size == self.dim_t: - self.uses_pytorch_inference = False + use_onnx_runtime = self.segment_size == self.dim_t + if use_onnx_runtime and self._prefer_pytorch_inference_on_rocm(): + use_onnx_runtime = False + + ort_inference_session = None + if use_onnx_runtime: ort_session_options = ort.SessionOptions() if self.log_level > 10: ort_session_options.log_severity_level = 3 @@ -122,28 +127,109 @@ def load_model(self): ort_inference_session = ort.InferenceSession(self.model_path, providers=self.onnx_execution_provider, sess_options=ort_session_options) session_providers = ort_inference_session.get_providers() - requested_provider = self.onnx_execution_provider[0] if self.onnx_execution_provider else None + if requested_provider and requested_provider not in session_providers: self.logger.warning( f"ONNX Runtime could not activate requested provider {requested_provider}; " - f"session is using {session_providers}. This usually means required CUDA/cuDNN " - f"runtime libraries are not visible to the dynamic loader." + f"session is using {session_providers}. This usually means required runtime " + f"libraries are not visible to the dynamic loader." ) + if self._should_convert_inactive_amd_provider_on_rocm(requested_provider): + use_onnx_runtime = False + ort_inference_session = None else: self.logger.debug(f"ONNX Runtime session providers: {session_providers}") + if use_onnx_runtime: self.model_run = lambda spek: ort_inference_session.run(None, {"input": spek.cpu().numpy()})[0] self.logger.debug("Model loaded successfully using ONNXruntime inferencing session.") else: - if platform.system() == 'Windows': + if platform.system() == "Windows": onnx_model = onnx.load(self.model_path) self.model_run = onnx2torch.convert(onnx_model) else: self.model_run = onnx2torch.convert(self.model_path) - + self.model_run.to(self.torch_device).eval() - self.logger.warning("Model converted from onnx to pytorch due to segment size not matching dim_t, processing may be slower.") + self.logger.warning( + "Model converted from ONNX to PyTorch (segment size does not match dim_t, " + "or ROCm has no usable ONNX GPU provider); processing may be slower." + ) + + self.uses_pytorch_inference = not use_onnx_runtime + + def _rocm_gpu_selected(self): + """ + Return True when ROCm PyTorch is active and the resolved device is the GPU. + + torch.version.hip is set for any ROCm build of PyTorch, including a CPU-only run + on such a build. Requiring the selected device to be "cuda" keeps the ONNX + Runtime path on CPU, so a ROCm host without a usable GPU is unaffected. + """ + if getattr(torch.version, "hip", None) is None: + return False + return getattr(self.torch_device, "type", None) == "cuda" + + def _prefer_pytorch_inference_on_rocm(self): + """ + Return True when an ROCm host should run the converted PyTorch graph. + + Official onnxruntime-gpu wheels link against CUDA runtime libraries, which are + not present on an ROCm host, so requesting CUDAExecutionProvider silently falls + back to the CPU provider. When no real AMD ONNX GPU provider is selected, + converting the graph and running it on the HIP device keeps MDX-Net on the GPU. + Non-ROCm hosts, and ROCm hosts whose selected device is not the GPU, return + False, so CPU/CUDA/MPS/DirectML are unchanged. + """ + if not self._rocm_gpu_selected(): + return False + + if self._onnx_runtime_override_requested(): + return False + + provider = (self.onnx_execution_provider or [None])[0] + if provider in {"MIGraphXExecutionProvider", "ROCMExecutionProvider"}: + return False + + self.logger.info( + "ROCm detected and no AMD ONNX GPU provider is selected; running MDX-Net " + "through PyTorch (HIP) instead of the ONNX Runtime CUDA/CPU fallback." + ) + return True + + def _should_convert_inactive_amd_provider_on_rocm(self, requested_provider): + """ + Return True when a requested but inactive AMD ONNX provider should run on HIP. + + A requested MIGraphXExecutionProvider or ROCMExecutionProvider can still fail to + activate (missing native libraries or an unsupported gfx target) and silently + leave the session on CPUExecutionProvider. On ROCm, converting the graph and + running it on the HIP device keeps MDX-Net on the GPU. The explicit + AUDIO_SEPARATOR_ROCM_ONNX override keeps the ONNX Runtime path. + """ + if requested_provider not in {"MIGraphXExecutionProvider", "ROCMExecutionProvider"}: + return False + if not self._rocm_gpu_selected(): + return False + if self._onnx_runtime_override_requested(): + return False + + self.logger.info( + "ROCm detected but the requested AMD ONNX provider did not activate; running " + "MDX-Net through PyTorch (HIP) instead of the ONNX Runtime CPU fallback." + ) + return True + + def _onnx_runtime_override_requested(self): + """ + Return True when AUDIO_SEPARATOR_ROCM_ONNX asks to keep the ONNX Runtime path. + + Accepted values (case-insensitive): "ort", "cpu", and "0". They only disable the + PyTorch conversion; the provider still comes from onnx_execution_provider, which + can be a CUDA-named provider that ONNX Runtime falls back to CPU. + """ + return os.environ.get("AUDIO_SEPARATOR_ROCM_ONNX", "").strip().lower() in {"cpu", "ort", "0"} def separate(self, audio_file_path, custom_output_names=None): """ diff --git a/audio_separator/separator/separator.py b/audio_separator/separator/separator.py index ece2f10a..98b2f391 100644 --- a/audio_separator/separator/separator.py +++ b/audio_separator/separator/separator.py @@ -55,6 +55,7 @@ def _iter_directory_audio_files(directory): + """Yield audio file paths found under the given directory tree.""" for root, _dirs, files in os.walk(directory): for filename in files: if filename.lower().endswith(SUPPORTED_AUDIO_EXTENSIONS): @@ -420,6 +421,7 @@ def log_onnxruntime_packages(self): onnxruntime_silicon_package = self.get_package_distribution("onnxruntime-silicon") onnxruntime_cpu_package = self.get_package_distribution("onnxruntime") onnxruntime_dml_package = self.get_package_distribution("onnxruntime-directml") + onnxruntime_rocm_package = self.get_package_distribution("onnxruntime-rocm") if onnxruntime_gpu_package is not None: self.logger.info(f"ONNX Runtime GPU package installed with version: {onnxruntime_gpu_package.version}") @@ -429,6 +431,8 @@ def log_onnxruntime_packages(self): self.logger.info(f"ONNX Runtime CPU package installed with version: {onnxruntime_cpu_package.version}") if onnxruntime_dml_package is not None: self.logger.info(f"ONNX Runtime DirectML package installed with version: {onnxruntime_dml_package.version}") + if onnxruntime_rocm_package is not None: + self.logger.info(f"ONNX Runtime ROCm package installed with version: {onnxruntime_rocm_package.version}") def setup_torch_device(self, system_info): """ @@ -441,7 +445,13 @@ def setup_torch_device(self, system_info): self.torch_device_cpu = torch.device("cpu") if torch.cuda.is_available(): - self.configure_cuda(ort_providers) + # ROCm PyTorch also exposes AMD GPUs through torch.cuda, so it is detected via + # torch.version.hip and routed to the ROCm-specific configuration. Non-ROCm + # CUDA users keep the original configure_cuda path unchanged. + if getattr(torch.version, "hip", None) is not None: + self.configure_rocm(ort_providers) + else: + self.configure_cuda(ort_providers) hardware_acceleration_enabled = True elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available() and system_info.processor == "arm": self.configure_mps(ort_providers) @@ -477,6 +487,35 @@ def configure_cuda(self, ort_providers): else: self.logger.warning("CUDAExecutionProvider not available in ONNXruntime, so acceleration will NOT be enabled") + def configure_rocm(self, ort_providers): + """ + This method configures the ROCm device for PyTorch and ONNX Runtime, if available. + + ROCm PyTorch exposes AMD GPUs through torch.cuda, so self.torch_device is still "cuda". + ONNX Runtime removed the ROCm execution provider in 1.23; MIGraphXExecutionProvider is + the current AMD ONNX GPU provider. Official onnxruntime-gpu wheels only expose a + CUDA-linked CUDAExecutionProvider, which cannot activate on a ROCm host. + """ + self.logger.info("ROCm (AMD GPU) is available in Torch, setting Torch device to CUDA (ROCm presents as CUDA)") + self.torch_device = torch.device("cuda") + for provider in ("MIGraphXExecutionProvider", "ROCMExecutionProvider", "CUDAExecutionProvider"): + if provider in ort_providers: + self.onnx_execution_provider = [provider] + if provider == "CUDAExecutionProvider": + self.logger.info( + "ONNXruntime only exposes CUDAExecutionProvider on this ROCm host; that " + "CUDA-linked provider cannot activate here, so MDX-Net will use the PyTorch " + "HIP path and other architectures are unaffected." + ) + else: + self.logger.info(f"ONNXruntime has {provider} available, enabling acceleration") + break + else: + self.logger.warning( + "No ONNX Runtime GPU execution provider is available for ROCm; MDX-Net will " + "use the PyTorch HIP path and other PyTorch architectures are unaffected." + ) + def configure_mps(self, ort_providers): """ This method configures the Apple Silicon MPS/CoreML device for PyTorch and ONNX Runtime, if available. @@ -1417,6 +1456,7 @@ def get_simplified_model_list(self, filter_sort_by: Optional[str] = None): # Sort by SDR score if available, putting None values last def sort_key(item): + """Sort key that puts missing SDR scores last.""" sdr = item[1]["SDR"][sort_by_lower] return (0 if sdr is None else 1, sdr if sdr is not None else float("-inf")) diff --git a/audio_separator/separator/uvr_lib_v5/demucs/hdemucs_rocm.py b/audio_separator/separator/uvr_lib_v5/demucs/hdemucs_rocm.py new file mode 100644 index 00000000..d1c18fdc --- /dev/null +++ b/audio_separator/separator/uvr_lib_v5/demucs/hdemucs_rocm.py @@ -0,0 +1,13 @@ +"""ROCm entry point for the Demucs spectrogram and hybrid model. + +Upstream ``HDemucs`` already keeps complex and spectral work on the GPU for cuda +devices, and PyTorch reports AMD GPUs under ROCm as cuda. There is no ROCm-specific +behaviour to fork, so ``HDemucsROCm`` is a thin subclass that preserves the name used +by the ROCm integration and inherits the upstream constructor and forward pass. +""" + +from .hdemucs import HDemucs + + +class HDemucsROCm(HDemucs): + """ROCm-named alias of the upstream HDemucs model (no behaviour changes).""" diff --git a/poetry.lock b/poetry.lock index fc2ed8d3..631a2622 100644 --- a/poetry.lock +++ b/poetry.lock @@ -388,7 +388,7 @@ description = "Colored terminal output for Python's logging module" optional = true python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*" groups = ["main"] -markers = "extra == \"cpu\" or extra == \"gpu\" or extra == \"dml\"" +markers = "extra == \"cpu\" or extra == \"gpu\" or extra == \"rocm\" or extra == \"dml\"" files = [ {file = "coloredlogs-15.0.1-py2.py3-none-any.whl", hash = "sha256:612ee75c546f53e92e70049c9dbfcc18c935a2b9a53b66085ce9ef6a6e5c0934"}, {file = "coloredlogs-15.0.1.tar.gz", hash = "sha256:7c991aa71a4577af2f82600d8f8f3a89f936baeaf9b50a9c197da014e5bf16b0"}, @@ -1109,7 +1109,7 @@ description = "The FlatBuffers serialization format for Python" optional = true python-versions = "*" groups = ["main"] -markers = "extra == \"cpu\" or extra == \"gpu\" or extra == \"dml\"" +markers = "extra == \"cpu\" or extra == \"gpu\" or extra == \"rocm\" or extra == \"dml\"" files = [ {file = "flatbuffers-25.9.23-py2.py3-none-any.whl", hash = "sha256:255538574d6cb6d0a79a17ec8bc0d30985913b87513a01cce8bcdb6b4c44d0e2"}, {file = "flatbuffers-25.9.23.tar.gz", hash = "sha256:676f9fa62750bb50cf531b42a0a2a118ad8f7f797a511eda12881c016f093b12"}, @@ -1552,7 +1552,7 @@ description = "Human friendly output for text interfaces using Python" optional = true python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*" groups = ["main"] -markers = "extra == \"cpu\" or extra == \"gpu\" or extra == \"dml\"" +markers = "extra == \"cpu\" or extra == \"gpu\" or extra == \"rocm\" or extra == \"dml\"" files = [ {file = "humanfriendly-10.0-py2.py3-none-any.whl", hash = "sha256:1697e1a8a8f550fd43c2865cd84542fc175a61dcb779b6fee18cf6b6ccba1477"}, {file = "humanfriendly-10.0.tar.gz", hash = "sha256:6b0b831ce8f15f7300721aa49829fc4e83921a9a301cc7f606be6686a2288ddc"}, @@ -2474,7 +2474,7 @@ description = "CUBLAS native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "(python_version >= \"3.14\" or platform_machine != \"aarch64\" and platform_machine != \"x86_64\") and platform_system == \"Linux\" and (platform_machine == \"arm64\" or python_version >= \"3.14\" or sys_platform == \"linux\") and (platform_machine != \"aarch64\" and platform_machine != \"x86_64\" or sys_platform != \"linux\") and (platform_machine != \"x86_64\" or sys_platform != \"linux\" and sys_platform != \"win32\") and (python_version >= \"3.14\" or sys_platform == \"darwin\" or sys_platform == \"linux\")" +markers = "(platform_machine == \"arm64\" or sys_platform == \"linux\" or python_version >= \"3.14\") and (platform_machine != \"aarch64\" and platform_machine != \"x86_64\" or python_version >= \"3.14\") and platform_system == \"Linux\" and (platform_machine != \"x86_64\" or sys_platform != \"linux\" and sys_platform != \"win32\") and (platform_machine != \"aarch64\" and platform_machine != \"x86_64\" or sys_platform != \"linux\") and (python_version >= \"3.14\" or sys_platform == \"darwin\" or sys_platform == \"linux\")" files = [ {file = "nvidia_cublas-13.6.0.2-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:946f6a252b1cc72d8de912c75975fd6d8ba44f67d4e5044fe764ddb909f4a688"}, {file = "nvidia_cublas-13.6.0.2-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b82c80c886cea6da6e149a5c3bdba274f12b7e4ec4b00a050b916b0446fb4153"}, @@ -2547,7 +2547,7 @@ description = "NVRTC native runtime libraries" optional = false python-versions = ">=3" groups = ["main"] -markers = "(python_version >= \"3.14\" or platform_machine != \"aarch64\" and platform_machine != \"x86_64\") and platform_system == \"Linux\" and (platform_machine == \"arm64\" or python_version >= \"3.14\" or sys_platform == \"linux\") and (platform_machine != \"aarch64\" and platform_machine != \"x86_64\" or sys_platform != \"linux\") and (platform_machine != \"x86_64\" or sys_platform != \"linux\" and sys_platform != \"win32\") and (python_version >= \"3.14\" or sys_platform == \"darwin\" or sys_platform == \"linux\")" +markers = "(platform_machine == \"arm64\" or sys_platform == \"linux\" or python_version >= \"3.14\") and (platform_machine != \"aarch64\" and platform_machine != \"x86_64\" or python_version >= \"3.14\") and platform_system == \"Linux\" and (platform_machine != \"x86_64\" or sys_platform != \"linux\" and sys_platform != \"win32\") and (platform_machine != \"aarch64\" and platform_machine != \"x86_64\" or sys_platform != \"linux\") and (python_version >= \"3.14\" or sys_platform == \"darwin\" or sys_platform == \"linux\")" files = [ {file = "nvidia_cuda_nvrtc-13.3.33-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:82530788b8c6164a54d3fd9ae8bcca8893d397c4aeb998861982a03bbe41e204"}, {file = "nvidia_cuda_nvrtc-13.3.33-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7b05ecda494c6dabc44231a608b060a71008a730d9dfda932cc508e6d29159e0"}, @@ -3065,7 +3065,7 @@ description = "ONNX Runtime is a runtime accelerator for Machine Learning models optional = true python-versions = ">=3.10" groups = ["main"] -markers = "extra == \"gpu\"" +markers = "extra == \"gpu\" or extra == \"rocm\"" files = [ {file = "onnxruntime_gpu-1.23.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:756c27f922924d8f8bab0fe75f0046b84c3eb65574d74ac00900864cc40998c7"}, {file = "onnxruntime_gpu-1.23.0-cp310-cp310-win_amd64.whl", hash = "sha256:84fea9f666b898c8450c2b918c5d6c85f4c8c505840b948e6ddfa71531bf6a2f"}, @@ -3420,7 +3420,7 @@ description = "A python implementation of GNU readline." optional = true python-versions = ">=3.8" groups = ["main"] -markers = "sys_platform == \"win32\" and (extra == \"cpu\" or extra == \"gpu\" or extra == \"dml\")" +markers = "sys_platform == \"win32\" and (extra == \"cpu\" or extra == \"gpu\" or extra == \"rocm\" or extra == \"dml\")" files = [ {file = "pyreadline3-3.5.4-py3-none-any.whl", hash = "sha256:eaf8e6cc3c49bcccf145fc6067ba8643d1df34d604a1ec0eccbf7a18e6d3fae6"}, {file = "pyreadline3-3.5.4.tar.gz", hash = "sha256:8d57d53039a1c75adba8e50dd3d992b28143480816187ea5efbd5c78e6c885b7"}, @@ -4513,8 +4513,9 @@ zstd = ["zstandard (>=0.18.0)"] cpu = ["onnxruntime"] dml = ["onnxruntime-directml", "torch_directml"] gpu = ["onnxruntime-gpu"] +rocm = ["onnxruntime-gpu"] [metadata] lock-version = "2.1" python-versions = ">=3.10,!=3.14.1" -content-hash = "191745b98ea3a5a70729fa29797229a680d1c9dbb1d878029bbc7080a66ad7ae" +content-hash = "a1560edc009041ddcfa234d2fded74eade87612b60c43af21cc1fba5564dae07" diff --git a/pyproject.toml b/pyproject.toml index 3cee1c09..840cb5ed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,10 @@ dynamic = [ cpu = ["onnxruntime>=1.17"] gpu = ["onnxruntime-gpu>=1.17"] dml = ["onnxruntime-directml>=1.17", "torch_directml"] +# ROCm/AMD GPUs use onnxruntime-gpu's CUDAExecutionProvider through ROCm's HIP +# compatibility layer; onnxruntime-rocm (ROCMExecutionProvider) is only needed +# for MI-series GPUs and is not published for newer Python versions. +rocm = ["onnxruntime-gpu>=1.17"] [tool.poetry] requires-poetry = ">=2.0.0" diff --git a/tests/unit/test_mdx_rocm_inference.py b/tests/unit/test_mdx_rocm_inference.py new file mode 100644 index 00000000..316b1c7e --- /dev/null +++ b/tests/unit/test_mdx_rocm_inference.py @@ -0,0 +1,114 @@ +"""Unit tests for the ROCm MDX inference-path guard (no GPU or torch runtime needed).""" + +import logging +import os +from unittest.mock import patch + +import torch + +from audio_separator.separator.architectures.mdx_separator import MDXSeparator + +_NO_OVERRIDE = {"AUDIO_SEPARATOR_ROCM_ONNX": ""} + + +def _separator(provider, device="cuda"): + """Build an MDXSeparator shell with the attributes the guard reads.""" + separator = MDXSeparator.__new__(MDXSeparator) + separator.logger = logging.getLogger("test_mdx_rocm_inference") + separator.onnx_execution_provider = [provider] if provider else None + separator.torch_device = torch.device(device) + return separator + + +def test_non_rocm_host_never_converts(): + """A host without ROCm PyTorch keeps the ONNX Runtime path.""" + separator = _separator("CUDAExecutionProvider") + with patch.object(torch.version, "hip", None, create=True): + assert separator._prefer_pytorch_inference_on_rocm() is False + + +def test_rocm_with_only_cuda_provider_converts(): + """On a ROCm GPU, a CUDA-named provider cannot activate, so convert to PyTorch.""" + separator = _separator("CUDAExecutionProvider") + with patch.object(torch.version, "hip", "7.2", create=True), patch.dict(os.environ, _NO_OVERRIDE): + assert separator._prefer_pytorch_inference_on_rocm() is True + + +def test_rocm_with_migraphx_keeps_onnx_runtime(): + """A real AMD ONNX GPU provider keeps the ONNX Runtime path.""" + separator = _separator("MIGraphXExecutionProvider") + with patch.object(torch.version, "hip", "7.2", create=True), patch.dict(os.environ, _NO_OVERRIDE): + assert separator._prefer_pytorch_inference_on_rocm() is False + + +def test_rocm_with_legacy_rocm_provider_keeps_onnx_runtime(): + """The legacy ROCm execution provider also keeps the ONNX Runtime path.""" + separator = _separator("ROCMExecutionProvider") + with patch.object(torch.version, "hip", "7.2", create=True), patch.dict(os.environ, _NO_OVERRIDE): + assert separator._prefer_pytorch_inference_on_rocm() is False + + +def test_rocm_conversion_can_be_opted_out(): + """AUDIO_SEPARATOR_ROCM_ONNX restores the ONNX Runtime path.""" + separator = _separator("CUDAExecutionProvider") + with patch.object(torch.version, "hip", "7.2", create=True), patch.dict(os.environ, {"AUDIO_SEPARATOR_ROCM_ONNX": "ort"}): + assert separator._prefer_pytorch_inference_on_rocm() is False + + +def test_missing_provider_defaults_to_conversion_on_rocm(): + """An unset provider list on a ROCm GPU is treated as no AMD ONNX provider.""" + separator = _separator(None) + with patch.object(torch.version, "hip", "7.2", create=True), patch.dict(os.environ, _NO_OVERRIDE): + assert separator._prefer_pytorch_inference_on_rocm() is True + + +def test_rocm_cpu_only_does_not_convert(): + """A ROCm PyTorch build running on CPU keeps the ONNX Runtime path.""" + separator = _separator("CUDAExecutionProvider", device="cpu") + with patch.object(torch.version, "hip", "7.2", create=True), patch.dict(os.environ, _NO_OVERRIDE): + assert separator._prefer_pytorch_inference_on_rocm() is False + assert separator._should_convert_inactive_amd_provider_on_rocm("MIGraphXExecutionProvider") is False + + +def test_inactive_amd_provider_converts_on_rocm(): + """A requested AMD provider that did not activate falls back to PyTorch on ROCm.""" + separator = _separator("MIGraphXExecutionProvider") + with patch.object(torch.version, "hip", "7.2", create=True), patch.dict(os.environ, _NO_OVERRIDE): + assert separator._should_convert_inactive_amd_provider_on_rocm("MIGraphXExecutionProvider") is True + assert separator._should_convert_inactive_amd_provider_on_rocm("ROCMExecutionProvider") is True + + +def test_inactive_amd_provider_not_converted_on_non_rocm(): + """Non-ROCm hosts keep the ONNX Runtime session even if the provider was inactive.""" + separator = _separator("MIGraphXExecutionProvider") + with patch.object(torch.version, "hip", None, create=True): + assert separator._should_convert_inactive_amd_provider_on_rocm("MIGraphXExecutionProvider") is False + + +def test_inactive_amd_provider_conversion_respects_override(): + """The explicit override keeps the ONNX Runtime path even for an inactive AMD provider.""" + separator = _separator("MIGraphXExecutionProvider") + with patch.object(torch.version, "hip", "7.2", create=True), patch.dict(os.environ, {"AUDIO_SEPARATOR_ROCM_ONNX": "0"}): + assert separator._should_convert_inactive_amd_provider_on_rocm("MIGraphXExecutionProvider") is False + + +def test_inactive_cuda_provider_does_not_trigger_amd_conversion(): + """The inactive-AMD-provider rule only applies to AMD provider names.""" + separator = _separator("CUDAExecutionProvider") + with patch.object(torch.version, "hip", "7.2", create=True): + assert separator._should_convert_inactive_amd_provider_on_rocm("CUDAExecutionProvider") is False + + +if __name__ == "__main__": + test_non_rocm_host_never_converts() + test_rocm_with_only_cuda_provider_converts() + test_rocm_with_migraphx_keeps_onnx_runtime() + test_rocm_with_legacy_rocm_provider_keeps_onnx_runtime() + test_rocm_conversion_can_be_opted_out() + test_missing_provider_defaults_to_conversion_on_rocm() + test_rocm_cpu_only_does_not_convert() + test_inactive_amd_provider_converts_on_rocm() + test_inactive_amd_provider_not_converted_on_non_rocm() + test_inactive_amd_provider_conversion_respects_override() + test_inactive_cuda_provider_does_not_trigger_amd_conversion() + print("all MDX ROCm guard checks passed") diff --git a/tests/unit/test_rocm_device_setup.py b/tests/unit/test_rocm_device_setup.py new file mode 100644 index 00000000..7a528570 --- /dev/null +++ b/tests/unit/test_rocm_device_setup.py @@ -0,0 +1,88 @@ +""" +Unit tests for ROCm (AMD GPU) device routing in Separator.setup_torch_device. + +ROCm PyTorch reports AMD GPUs through torch.cuda and advertises torch.version.hip. +These tests confirm that ROCm input routes to configure_rocm without changing the +CUDA/CPU paths, and that MIGraphXExecutionProvider is preferred over the legacy +ROCMExecutionProvider and CUDAExecutionProvider when available. +""" + +import logging +from unittest.mock import patch + +import torch + +import audio_separator.separator.separator as separator_module +from audio_separator.separator.separator import Separator + + +def _new_separator() -> Separator: + """Build a Separator without running the heavy __init__ side effects.""" + separator = Separator.__new__(Separator) + separator.logger = logging.getLogger("test_rocm_device_setup") + separator.use_directml = False + separator.torch_device_cpu = torch.device("cpu") + separator.torch_device = torch.device("cpu") + separator.onnx_execution_provider = None + return separator + + +def test_rocm_torch_routes_to_configure_rocm_with_cuda_provider(): + """ROCm PyTorch is detected via torch.version.hip and stays on the cuda device.""" + separator = _new_separator() + with patch("torch.cuda.is_available", return_value=True), patch.object( + torch.version, "hip", "7.2", create=True + ), patch.object( + separator_module.ort, "get_available_providers", return_value=["CUDAExecutionProvider", "CPUExecutionProvider"] + ): + separator.setup_torch_device(object()) + + assert separator.torch_device == torch.device("cuda") + assert separator.onnx_execution_provider == ["CUDAExecutionProvider"] + + +def test_rocm_prefers_native_rocm_execution_provider_when_available(): + """A native AMD ONNX provider is selected ahead of CUDAExecutionProvider.""" + separator = _new_separator() + with patch("torch.cuda.is_available", return_value=True), patch.object( + torch.version, "hip", "7.2", create=True + ), patch.object( + separator_module.ort, "get_available_providers", return_value=["ROCMExecutionProvider", "CPUExecutionProvider"] + ): + separator.setup_torch_device(object()) + + assert separator.onnx_execution_provider == ["ROCMExecutionProvider"] + + +def test_non_rocm_cuda_still_uses_configure_cuda(): + """Without torch.version.hip, CUDA hosts keep the original configure_cuda path.""" + separator = _new_separator() + with patch("torch.cuda.is_available", return_value=True), patch.object( + torch.version, "hip", None, create=True + ), patch.object( + separator_module.ort, "get_available_providers", return_value=["CUDAExecutionProvider", "CPUExecutionProvider"] + ), patch.object(separator, "configure_cuda", wraps=separator.configure_cuda) as mock_cuda: + separator.setup_torch_device(object()) + + mock_cuda.assert_called_once() + assert separator.onnx_execution_provider == ["CUDAExecutionProvider"] + + +def test_cpu_only_has_no_gpu_provider(): + """A host with no GPU accelerator falls back to the CPU execution provider.""" + separator = _new_separator() + with patch("torch.cuda.is_available", return_value=False), patch.object( + separator_module.ort, "get_available_providers", return_value=["CPUExecutionProvider"] + ): + separator.setup_torch_device(type("SystemInfo", (), {"processor": "x86_64"})()) + + assert separator.torch_device == torch.device("cpu") + assert separator.onnx_execution_provider == ["CPUExecutionProvider"] + + +if __name__ == "__main__": + test_rocm_torch_routes_to_configure_rocm_with_cuda_provider() + test_rocm_prefers_native_rocm_execution_provider_when_available() + test_non_rocm_cuda_still_uses_configure_cuda() + test_cpu_only_has_no_gpu_provider() + print("all ROCm device-setup checks passed")