fix(stt): respect device and compute_type from config.yaml

The local STT transcription function hardcoded device="auto" and
compute_type="auto" when instantiating WhisperModel, ignoring the
user's stt.local.device and stt.local.compute_type config values.

Closes #8319
This commit is contained in:
Tranquil-Flow 2026-04-14 03:33:37 +10:00 committed by Teknium
parent fda771498e
commit 06fc6e0c29
2 changed files with 80 additions and 5 deletions

View file

@ -611,6 +611,67 @@ class TestTranscribeLocalExtended:
assert result["success"] is False
assert "CUDA out of memory" in result["error"]
def test_config_device_and_compute_type_passed_to_whisper(self, tmp_path):
"""User-configured device and compute_type should be forwarded to WhisperModel.
Regression test for #8319: these values were hardcoded to "auto".
"""
audio = tmp_path / "test.ogg"
audio.write_bytes(b"fake")
mock_segment = MagicMock()
mock_segment.text = "hi"
mock_info = MagicMock()
mock_info.language = "en"
mock_info.duration = 1.0
mock_model = MagicMock()
mock_model.transcribe.return_value = ([mock_segment], mock_info)
mock_whisper_cls = MagicMock(return_value=mock_model)
fake_config = {
"local": {
"device": "cpu",
"compute_type": "float32",
}
}
with patch("tools.transcription_tools._HAS_FASTER_WHISPER", True), \
patch("faster_whisper.WhisperModel", mock_whisper_cls), \
patch("tools.transcription_tools._local_model", None), \
patch("tools.transcription_tools._local_model_name", None), \
patch("tools.transcription_tools._load_stt_config", return_value=fake_config):
from tools.transcription_tools import _transcribe_local
result = _transcribe_local(str(audio), "base")
assert result["success"] is True
mock_whisper_cls.assert_called_once_with("base", device="cpu", compute_type="float32")
def test_config_defaults_to_auto_when_not_set(self, tmp_path):
"""Without config, device and compute_type should default to "auto"."""
audio = tmp_path / "test.ogg"
audio.write_bytes(b"fake")
mock_segment = MagicMock()
mock_segment.text = "hi"
mock_info = MagicMock()
mock_info.language = "en"
mock_info.duration = 1.0
mock_model = MagicMock()
mock_model.transcribe.return_value = ([mock_segment], mock_info)
mock_whisper_cls = MagicMock(return_value=mock_model)
with patch("tools.transcription_tools._HAS_FASTER_WHISPER", True), \
patch("faster_whisper.WhisperModel", mock_whisper_cls), \
patch("tools.transcription_tools._local_model", None), \
patch("tools.transcription_tools._local_model_name", None), \
patch("tools.transcription_tools._load_stt_config", return_value={}):
from tools.transcription_tools import _transcribe_local
_transcribe_local(str(audio), "base")
mock_whisper_cls.assert_called_once_with("base", device="auto", compute_type="auto")
def test_multiple_segments_joined(self, tmp_path):
audio = tmp_path / "test.ogg"
audio.write_bytes(b"fake")

View file

@ -1160,7 +1160,7 @@ def _looks_like_cuda_lib_error(exc: BaseException) -> bool:
return any(marker in msg for marker in _CUDA_LIB_ERROR_MARKERS)
def _load_local_whisper_model(model_name: str):
def _load_local_whisper_model(model_name: str, device: str = "auto", compute_type: str = "auto"):
"""Load faster-whisper with graceful CUDA → CPU fallback.
faster-whisper's ``device="auto"`` picks CUDA when the ctranslate2 wheel
@ -1170,12 +1170,16 @@ def _load_local_whisper_model(model_name: str):
On those hosts the load itself sometimes succeeds and the dlopen failure
only surfaces at first ``transcribe()`` call.
We try ``auto`` first (fast CUDA path when it works), and on any CUDA
library load failure fall back to CPU + int8.
``device`` / ``compute_type`` default to ``"auto"`` so the historical
behaviour is unchanged; pass explicit values from ``stt.local.device`` /
``stt.local.compute_type`` to pin a configuration (#9088).
We try the requested config first (fast CUDA path when it works), and on
any CUDA library load failure fall back to CPU + int8.
"""
from faster_whisper import WhisperModel
try:
return WhisperModel(model_name, device="auto", compute_type="auto")
return WhisperModel(model_name, device=device, compute_type=compute_type)
except Exception as exc:
if not _looks_like_cuda_lib_error(exc):
raise
@ -1196,10 +1200,20 @@ def _transcribe_local(file_path: str, model_name: str) -> Dict[str, Any]:
return {"success": False, "transcript": "", "error": "faster-whisper not installed"}
try:
local_cfg = _load_stt_config().get("local", {})
# Lazy-load the model (downloads on first use, ~150 MB for 'base')
if _local_model is None or _local_model_name != model_name:
logger.info("Loading faster-whisper model '%s' (first load downloads the model)...", model_name)
_local_model = _load_local_whisper_model(model_name)
# Honour stt.local.device / stt.local.compute_type from config so
# users on hosts where ``auto`` mis-detects (NVIDIA libs present but
# not usable, etc.) can pin a working configuration (#9088).
# _load_local_whisper_model retains the CUDA→CPU fallback for the
# auto/CUDA paths.
_local_model = _load_local_whisper_model(
model_name,
device=local_cfg.get("device", "auto"),
compute_type=local_cfg.get("compute_type", "auto"),
)
_local_model_name = model_name
# Language: stt.local.language > stt.language > env var > auto-detect.