diff --git a/README.md b/README.md index 8b6c73d..56af29d 100644 --- a/README.md +++ b/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`: diff --git a/audio_separator/separator/architectures/mdx_separator.py b/audio_separator/separator/architectures/mdx_separator.py index af5e311..17920e4 100644 --- a/audio_separator/separator/architectures/mdx_separator.py +++ b/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}; " diff --git a/audio_separator/separator/common_separator.py b/audio_separator/separator/common_separator.py index be20426..fe86dbf 100644 --- a/audio_separator/separator/common_separator.py +++ b/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): """ diff --git a/audio_separator/separator/separator.py b/audio_separator/separator/separator.py index ece2f10..cde8422 100644 --- a/audio_separator/separator/separator.py +++ b/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") diff --git a/audio_separator/utils/cli.py b/audio_separator/utils/cli.py index 927b7ad..bec733b 100755 --- a/audio_separator/utils/cli.py +++ b/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, diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index ebd419f..ba8272e 100644 --- a/tests/unit/test_cli.py +++ b/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 diff --git a/tests/unit/test_cuda_device_selection.py b/tests/unit/test_cuda_device_selection.py new file mode 100644 index 0000000..2c6a936 --- /dev/null +++ b/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() diff --git a/tests/unit/test_separator_api_compatibility.py b/tests/unit/test_separator_api_compatibility.py index 2875957..774f500 100644 --- a/tests/unit/test_separator_api_compatibility.py +++ b/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