feat(device): auto-detect cuda/mps/cpu with env override

server/device.py

@@ -0,0 +1,32 @@
+"""Device auto-detection for Chatterbox.
+
+Order: env override → cuda → mps → cpu.
+"""
+from __future__ import annotations
+
+import os
+
+import torch
+
+
+_VALID = {"cuda", "mps", "cpu"}
+
+
+def _cuda_available() -> bool:
+    return torch.cuda.is_available()
+
+
+def _mps_available() -> bool:
+    backend = getattr(torch.backends, "mps", None)
+    return bool(backend and backend.is_available())
+
+
+def select_device() -> str:
+    forced = (os.getenv("CHATTERBOX_DEVICE") or "").strip().lower()
+    if forced in _VALID:
+        return forced
+    if _cuda_available():
+        return "cuda"
+    if _mps_available():
+        return "mps"
+    return "cpu"

tests/test_device.py

@@ -0,0 +1,46 @@
+import os
+from unittest.mock import patch
+
+from server.device import select_device
+
+
+def test_env_override_cuda():
+    with patch.dict(os.environ, {"CHATTERBOX_DEVICE": "cuda"}):
+        assert select_device() == "cuda"
+
+
+def test_env_override_mps():
+    with patch.dict(os.environ, {"CHATTERBOX_DEVICE": "MPS"}):
+        assert select_device() == "mps"
+
+
+def test_env_override_cpu():
+    with patch.dict(os.environ, {"CHATTERBOX_DEVICE": "cpu"}):
+        assert select_device() == "cpu"
+
+
+def test_invalid_env_falls_through_to_autodetect():
+    with patch.dict(os.environ, {"CHATTERBOX_DEVICE": "tpu"}, clear=False):
+        with patch("server.device._cuda_available", return_value=True):
+            assert select_device() == "cuda"
+
+
+def test_autodetect_prefers_cuda_over_mps():
+    with patch.dict(os.environ, {}, clear=True):
+        with patch("server.device._cuda_available", return_value=True), \
+             patch("server.device._mps_available", return_value=True):
+            assert select_device() == "cuda"
+
+
+def test_autodetect_uses_mps_when_no_cuda():
+    with patch.dict(os.environ, {}, clear=True):
+        with patch("server.device._cuda_available", return_value=False), \
+             patch("server.device._mps_available", return_value=True):
+            assert select_device() == "mps"
+
+
+def test_autodetect_falls_back_to_cpu():
+    with patch.dict(os.environ, {}, clear=True):
+        with patch("server.device._cuda_available", return_value=False), \
+             patch("server.device._mps_available", return_value=False):
+            assert select_device() == "cpu"