feat: select CUDA device consistently across inference backends

README.md

@@ -95,6 +95,25 @@ If the GPU isn't being detected, make sure your docker runtime environment is pa
 
 ### 🎮 Nvidia GPU with CUDA or 🧪 Google Colab
 
+To select a particular GPU on a multi-GPU machine, use `--cuda_device_index`:
+
+```sh
+audio-separator input.wav --cuda_device_index 1
+```
+
+The equivalent Python option is `Separator(cuda_device_index=1)`. The index is
+zero-based and relative to the devices visible through `CUDA_VISIBLE_DEVICES`.
+It configures both PyTorch (`cuda:1` in this example) and ONNX Runtime's CUDA
+provider (`device_id=1`). The provider must be installed for ONNX acceleration;
+selecting an index does not install or enable an unavailable ONNX provider.
+
+An explicitly selected index must be a non-negative integer within the available
+CUDA device range. An unavailable CUDA backend or out-of-range index raises
+`ValueError` rather than silently selecting another backend. Omit the option to
+retain automatic device selection. Existing positional Python arguments retain
+their order; this optional argument is appended to the constructor.
+
+
 **Supported CUDA Versions:** 11.8 and 12.2
 
 💬 If successfully configured, you should see this log message when running `audio-separator --env_info`:

audio_separator/separator/architectures/mdx_separator.py

@@ -124,6 +124,8 @@ def load_model(self):
             session_providers = ort_inference_session.get_providers()
 
             requested_provider = self.onnx_execution_provider[0] if self.onnx_execution_provider else None
+            if isinstance(requested_provider, tuple):
+                requested_provider = requested_provider[0]
             if requested_provider and requested_provider not in session_providers:
                 self.logger.warning(
                     f"ONNX Runtime could not activate requested provider {requested_provider}; "

audio_separator/separator/common_separator.py

@@ -489,9 +489,10 @@ def clear_gpu_cache(self):
         if self.torch_device == torch.device("mps"):
             self.logger.debug("Clearing MPS cache...")
             torch.mps.empty_cache()
-        if self.torch_device == torch.device("cuda"):
+        if self.torch_device.type == "cuda":
             self.logger.debug("Clearing CUDA cache...")
-            torch.cuda.empty_cache()
+            with torch.cuda.device(self.torch_device):
+                torch.cuda.empty_cache()
 
     def clear_file_specific_paths(self):
         """

audio_separator/separator/separator.py

@@ -54,6 +54,11 @@
 SUPPORTED_AUDIO_EXTENSIONS = (".wav", ".flac", ".mp3", ".ogg", ".opus", ".m4a", ".aiff", ".ac3")
 
 
+def _validate_cuda_device_index(index):
+    if index is not None and (isinstance(index, bool) or not isinstance(index, int) or index < 0):
+        raise ValueError("cuda_device_index must be a non-negative integer or None")
+
+
 def _iter_directory_audio_files(directory):
     for root, _dirs, files in os.walk(directory):
         for filename in files:
@@ -84,6 +89,7 @@ class Separator:
         output_single_stem (str): Option to output a single stem.
         invert_using_spec (bool): Flag to invert using spectrogram.
         sample_rate (int): The sample rate of the audio.
+        cuda_device_index (int or None): Optional visible CUDA index used by PyTorch and ONNX Runtime.
         use_soundfile (bool): Use soundfile for audio writing, can solve OOM issues.
         use_autocast (bool): Use PyTorch autocast when the loaded model and device support it.
         use_torch_compile (bool): Compile verified repeated model blocks when supported.
@@ -146,8 +152,11 @@ def __init__(
         info_only=False,
         use_torch_compile=False,
         use_native_fp16=False,
+        cuda_device_index=None,
     ):
         """Initialize the separator."""
+        _validate_cuda_device_index(cuda_device_index)
+        self.cuda_device_index = cuda_device_index
         if use_autocast and use_native_fp16:
             raise ValueError("use_autocast and use_native_fp16 are mutually exclusive precision modes.")
 
@@ -434,6 +443,8 @@ def setup_torch_device(self, system_info):
         """
         This method sets up the PyTorch and/or ONNX Runtime inferencing device, using GPU hardware acceleration if available.
         """
+        if self.cuda_device_index is not None and not torch.cuda.is_available():
+            raise ValueError("cuda_device_index was specified, but CUDA is not available")
         hardware_acceleration_enabled = False
         ort_providers = ort.get_available_providers()
         has_torch_dml_installed = self.get_package_distribution("torch_directml")
@@ -465,15 +476,19 @@ def setup_torch_device(self, system_info):
                     "Pass use_directml=True (or --use_directml on the CLI) to enable experimental DirectML acceleration."
                 )
 
-    def configure_cuda(self, ort_providers):
+    def configure_cuda(self, ort_providers, device_index=None):
         """
         This method configures the CUDA device for PyTorch and ONNX Runtime, if available.
         """
-        self.logger.info("CUDA is available in Torch, setting Torch device to CUDA")
-        self.torch_device = torch.device("cuda")
+        index = self.cuda_device_index if device_index is None else device_index
+        _validate_cuda_device_index(index)
+        if index is not None and index >= torch.cuda.device_count():
+            raise ValueError(f"cuda_device_index {index} is outside the available CUDA device range")
+        self.torch_device = torch.device("cuda" if index is None else f"cuda:{index}")
+        self.logger.info(f"CUDA is available in Torch, setting Torch device to {self.torch_device}")
         if "CUDAExecutionProvider" in ort_providers:
             self.logger.info("ONNXruntime has CUDAExecutionProvider available, enabling acceleration")
-            self.onnx_execution_provider = ["CUDAExecutionProvider"]
+            self.onnx_execution_provider = ["CUDAExecutionProvider"] if index is None else [("CUDAExecutionProvider", {"device_id": index})]
         else:
             self.logger.warning("CUDAExecutionProvider not available in ONNXruntime, so acceleration will NOT be enabled")
 

audio_separator/utils/cli.py

@@ -91,6 +91,7 @@ def main():
     precision_params.add_argument("--use_autocast", action="store_true", help=use_autocast_help)
     precision_params.add_argument("--use_native_fp16", action="store_true", help=use_native_fp16_help)
     common_params.add_argument("--use_torch_compile", action="store_true", help=use_torch_compile_help)
+    common_params.add_argument("--cuda_device_index", type=int, default=None, help="CUDA device index for both PyTorch and ONNX Runtime (relative to CUDA_VISIBLE_DEVICES).")
     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(
@@ -270,6 +271,7 @@ def main():
         use_native_fp16=args.use_native_fp16,
         use_torch_compile=args.use_torch_compile,
         use_directml=args.use_directml,
+        cuda_device_index=args.cuda_device_index,
         chunk_duration=args.chunk_duration,
         ensemble_algorithm=args.ensemble_algorithm,
         ensemble_weights=args.ensemble_weights,

tests/unit/test_cli.py

@@ -46,6 +46,7 @@ def common_expected_args():
         "use_native_fp16": False,
         "use_torch_compile": False,
         "use_directml": False,
+        "cuda_device_index": None,
         "chunk_duration": None,
         "ensemble_algorithm": None,
         "ensemble_weights": None,
@@ -476,3 +477,9 @@ def test_cli_list_presets(capsys):
     captured = capsys.readouterr()
     assert "vocal_balanced" in captured.out
     assert "karaoke" in captured.out
+
+
+def test_cli_passes_cuda_device_index():
+    with patch("sys.argv", ["audio-separator", "test.wav", "--cuda_device_index", "2"]), patch("audio_separator.separator.Separator") as separator_class:
+        main()
+    assert separator_class.call_args.kwargs["cuda_device_index"] == 2

tests/unit/test_cuda_device_selection.py

@@ -0,0 +1,91 @@
+"""Keep explicit CUDA selection consistent across inference backends."""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from audio_separator.separator import Separator
+
+
+@pytest.mark.parametrize("index", [0, 1, 3])
+def test_selected_device_reaches_torch_and_onnx(index):
+    separator = Separator(info_only=True)
+    with patch("torch.cuda.device_count", return_value=4):
+        separator.configure_cuda(["CUDAExecutionProvider"], device_index=index)
+    assert str(separator.torch_device) == f"cuda:{index}"
+    assert separator.onnx_execution_provider == [("CUDAExecutionProvider", {"device_id": index})]
+
+
+@pytest.mark.parametrize("index", [-1, True, 1.5, "1"])
+def test_invalid_index_rejected_before_initialization(index):
+    with pytest.raises(ValueError, match="cuda_device_index"):
+        Separator(info_only=True, cuda_device_index=index)
+
+
+def test_constructor_selection_is_used_by_device_setup():
+    separator = Separator(info_only=True, cuda_device_index=2)
+    with (
+        patch("torch.cuda.is_available", return_value=True),
+        patch("torch.cuda.device_count", return_value=3),
+        patch("audio_separator.separator.separator.ort.get_available_providers", return_value=["CUDAExecutionProvider"]),
+    ):
+        separator.setup_torch_device(SimpleNamespace(processor="test"))
+    assert str(separator.torch_device) == "cuda:2"
+    assert separator.onnx_execution_provider == [("CUDAExecutionProvider", {"device_id": 2})]
+
+
+def test_explicit_index_cannot_silently_fall_back_to_cpu():
+    separator = Separator(info_only=True, cuda_device_index=0)
+    with patch("torch.cuda.is_available", return_value=False), pytest.raises(ValueError, match="CUDA"):
+        separator.setup_torch_device(SimpleNamespace(processor="test"))
+
+
+def test_out_of_range_selection_does_not_mutate_devices():
+    separator = Separator(info_only=True)
+    separator.torch_device = "unchanged"
+    with patch("torch.cuda.device_count", return_value=2), pytest.raises(ValueError, match="cuda_device_index"):
+        separator.configure_cuda(["CUDAExecutionProvider"], device_index=2)
+    assert separator.torch_device == "unchanged"
+
+
+def test_default_selection_keeps_existing_provider_configuration():
+    separator = Separator(info_only=True)
+    separator.configure_cuda(["CUDAExecutionProvider"])
+    assert str(separator.torch_device) == "cuda"
+    assert separator.onnx_execution_provider == ["CUDAExecutionProvider"]
+
+
+@pytest.mark.parametrize("active_providers, warns", [(["CUDAExecutionProvider"], False), (["CPUExecutionProvider"], True)])
+def test_mdx_checks_provider_name_and_preserves_device_options(active_providers, warns):
+    from unittest.mock import Mock
+    from audio_separator.separator.architectures.mdx_separator import MDXSeparator
+
+    separator = MDXSeparator.__new__(MDXSeparator)
+    separator.logger = Mock()
+    separator.segment_size = separator.dim_t = 256
+    separator.log_level = 20
+    separator.model_path = "model.onnx"
+    separator.onnx_execution_provider = [("CUDAExecutionProvider", {"device_id": 2})]
+    with patch("audio_separator.separator.architectures.mdx_separator.ort.InferenceSession") as session:
+        session.return_value.get_providers.return_value = active_providers
+        separator.load_model()
+    assert session.call_args.kwargs["providers"] == [("CUDAExecutionProvider", {"device_id": 2})]
+    assert separator.logger.warning.called is warns
+
+
+@pytest.mark.parametrize("device_name", ["cuda", "cuda:0", "cuda:2"])
+def test_cache_cleanup_uses_selected_cuda_context(device_name):
+    from unittest.mock import Mock
+    import torch
+    from audio_separator.separator.common_separator import CommonSeparator
+
+    separator = CommonSeparator.__new__(CommonSeparator)
+    separator.logger = Mock()
+    separator.torch_device = torch.device(device_name)
+    with patch("torch.cuda.device") as context, patch("torch.cuda.empty_cache") as empty_cache:
+        separator.clear_gpu_cache()
+    context.assert_called_once_with(separator.torch_device)
+    context.return_value.__enter__.assert_called_once()
+    empty_cache.assert_called_once()
+    context.return_value.__exit__.assert_called_once()

tests/unit/test_separator_api_compatibility.py

@@ -33,6 +33,7 @@ def test_execution_options_are_appended_to_constructor_signature():
         "info_only",
         "use_torch_compile",
         "use_native_fp16",
+        "cuda_device_index",
     ]
     assert inspect.signature(Separator.__init__).parameters["use_torch_compile"].default is False
     assert inspect.signature(Separator.__init__).parameters["use_native_fp16"].default is False