diff --git a/abogen/domain/device.py b/abogen/domain/device.py new file mode 100644 index 0000000..9b905ad --- /dev/null +++ b/abogen/domain/device.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import platform as _platform + + +def select_device() -> str: + """Return the best available compute device (``"mps"``, ``"cuda"``, or ``"cpu"``). + + Checks ``torch`` availability at runtime so this can be called from + any context without requiring torch at import time. + """ + try: + import torch # type: ignore[import-not-found] + except Exception: + return "cpu" + + system = _platform.system() + if system == "Darwin" and _platform.processor() == "arm": + try: + if torch.backends.mps.is_available(): # type: ignore[union-attr] + return "mps" + except Exception: + pass + return "cpu" + + try: + if torch.cuda.is_available(): # type: ignore[union-attr] + return "cuda" + except Exception: + pass + return "cpu" diff --git a/abogen/pyqt/gui.py b/abogen/pyqt/gui.py index b5d7ed7..229c989 100644 --- a/abogen/pyqt/gui.py +++ b/abogen/pyqt/gui.py @@ -7,6 +7,7 @@ import base64 import re from abogen.pyqt.queue_manager_gui import QueueManager from abogen.pyqt.queued_item import QueuedItem +from abogen.domain.device import select_device as _select_device import abogen.hf_tracker as hf_tracker import hashlib # Added for cache path generation from PyQt6.QtWidgets import ( @@ -2428,10 +2429,7 @@ class abogen(QWidget): # Determine device based on GPU availability if gpu_ok: - if platform.system() == "Darwin" and platform.processor() == "arm": - device = "mps" - else: - device = "cuda" + device = _select_device() else: device = "cpu" @@ -2877,10 +2875,7 @@ class abogen(QWidget): # Determine device based on GPU availability if self.gpu_ok: - if platform.system() == "Darwin" and platform.processor() == "arm": - device = "mps" - else: - device = "cuda" + device = _select_device() else: device = "cpu" diff --git a/abogen/webui/conversion_runner.py b/abogen/webui/conversion_runner.py index 2050917..5bb9049 100644 --- a/abogen/webui/conversion_runner.py +++ b/abogen/webui/conversion_runner.py @@ -109,6 +109,7 @@ from abogen.domain.output_paths import ( resolve_output_directory as _resolve_output_directory, resolve_project_layout as _resolve_project_layout, ) +from abogen.domain.device import select_device as _select_device from abogen.domain.audio_helpers import ( build_ffmpeg_command as _build_ffmpeg_command, to_float32 as _to_float32, @@ -1188,15 +1189,6 @@ def _load_pipeline(job: Job): return create_pipeline("kokoro", lang_code=job.language, device=device) -def _select_device() -> str: - import platform - - system = platform.system() - if system == "Darwin" and platform.processor() == "arm": - return "mps" - return "cuda" - - def _prepare_output_dir(job: Job) -> Path: from platformdirs import user_desktop_dir # type: ignore[import-not-found] diff --git a/abogen/webui/routes/utils/preview.py b/abogen/webui/routes/utils/preview.py index eccd05e..75b59c0 100644 --- a/abogen/webui/routes/utils/preview.py +++ b/abogen/webui/routes/utils/preview.py @@ -6,6 +6,8 @@ import soundfile as sf from flask import current_app, send_file from flask.typing import ResponseReturnValue +from abogen.domain.device import select_device as _select_device + SPLIT_PATTERN = r"\n+" SAMPLE_RATE = 24000 @@ -25,31 +27,6 @@ def clear_preview_pipelines() -> None: _preview_pipelines.clear() -def _select_device() -> str: - import platform - - try: - import torch # type: ignore[import-not-found] - except Exception: - return "cpu" - - system = platform.system() - if system == "Darwin" and platform.processor() == "arm": - try: - if torch.backends.mps.is_available(): - return "mps" - except Exception: - pass - return "cpu" - - try: - if torch.cuda.is_available(): - return "cuda" - except Exception: - pass - return "cpu" - - def _resolve_pipeline(language: str, use_gpu: bool) -> Tuple[Any, bool]: devices: List[str] = ["cpu"] if use_gpu: diff --git a/tests/test_select_device.py b/tests/test_select_device.py new file mode 100644 index 0000000..d5f13c0 --- /dev/null +++ b/tests/test_select_device.py @@ -0,0 +1,94 @@ +import sys +from unittest.mock import patch, MagicMock + + +class TestSelectDevice: + """Tests for domain.device.select_device.""" + + def test_returns_mps_on_apple_silicon_when_available(self) -> None: + from abogen.domain.device import select_device + + mock_platform = MagicMock() + mock_platform.system.return_value = "Darwin" + mock_platform.processor.return_value = "arm" + + mock_torch = MagicMock() + mock_torch.backends.mps.is_available.return_value = True + mock_torch.cuda.is_available.return_value = False + + with patch("abogen.domain.device._platform", mock_platform), \ + patch.dict(sys.modules, {"torch": mock_torch}): + result = select_device() + assert result == "mps" + + def test_returns_cpu_on_apple_silicon_when_mps_unavailable(self) -> None: + from abogen.domain.device import select_device + + mock_platform = MagicMock() + mock_platform.system.return_value = "Darwin" + mock_platform.processor.return_value = "arm" + + mock_torch = MagicMock() + mock_torch.backends.mps.is_available.return_value = False + mock_torch.cuda.is_available.return_value = False + + with patch("abogen.domain.device._platform", mock_platform), \ + patch.dict(sys.modules, {"torch": mock_torch}): + result = select_device() + assert result == "cpu" + + def test_returns_cuda_when_available(self) -> None: + from abogen.domain.device import select_device + + mock_platform = MagicMock() + mock_platform.system.return_value = "Linux" + mock_platform.processor.return_value = "x86_64" + + mock_torch = MagicMock() + mock_torch.backends.mps.is_available.return_value = False + mock_torch.cuda.is_available.return_value = True + + with patch("abogen.domain.device._platform", mock_platform), \ + patch.dict(sys.modules, {"torch": mock_torch}): + result = select_device() + assert result == "cuda" + + def test_returns_cpu_when_cuda_unavailable(self) -> None: + from abogen.domain.device import select_device + + mock_platform = MagicMock() + mock_platform.system.return_value = "Linux" + mock_platform.processor.return_value = "x86_64" + + mock_torch = MagicMock() + mock_torch.backends.mps.is_available.return_value = False + mock_torch.cuda.is_available.return_value = False + + with patch("abogen.domain.device._platform", mock_platform), \ + patch.dict(sys.modules, {"torch": mock_torch}): + result = select_device() + assert result == "cpu" + + def test_returns_cpu_when_torch_not_installed(self) -> None: + from abogen.domain.device import select_device + + mock_platform = MagicMock() + mock_platform.system.return_value = "Linux" + mock_platform.processor.return_value = "x86_64" + + with patch("abogen.domain.device._platform", mock_platform), \ + patch.dict(sys.modules, {"torch": None}): + result = select_device() + assert result == "cpu" + + def test_handles_torch_import_error(self) -> None: + from abogen.domain.device import select_device + + mock_platform = MagicMock() + mock_platform.system.return_value = "Windows" + mock_platform.processor.return_value = "AMD64" + + with patch("abogen.domain.device._platform", mock_platform), \ + patch.dict(sys.modules, {"torch": None}): + result = select_device() + assert result == "cpu"