From 51811ad25179e649417ec9daf2c25ed3a2dd5a9b Mon Sep 17 00:00:00 2001 From: EJClaw Date: Tue, 18 Aug 2026 21:02:52 +0900 Subject: [PATCH] perf(voice): run STT+TTS on the GPU by default with CPU fallback Both voice backends defaulted to CPU. Fix the "CUDA unavailable" gaps so everything that benefits from the RTX 5050 uses it: - MeloTTS venv had CPU-only torch (2.12.0+cpu) -> installed Blackwell-capable torch/torchaudio 2.11.0+cu128 (sm_120 verified with a real GPU matmul). - faster-whisper CUDA loaded but transcribe() died with "libcublas.so.12 not found": installed nvidia-cublas-cu12 + nvidia-cudnn-cu12 into the whisper venv and inject those nvidia/*/lib dirs into the worker's LD_LIBRARY_PATH at spawn (the loader only honours it at exec). - WSAI_WHISPER_DEVICE / WSAI_MELO_DEVICE now default to "auto": pick CUDA when present, else CPU, and each worker falls back to CPU if a CUDA load fails so the voice loop never dies on a GPU-less host. Verified end-to-end through the real backend classes: both workers report "ready on cuda"; steady-state STT ~170ms (was ~1350ms CPU), TTS ~4s first call vs ~23s CPU. All 12 tests pass. Co-Authored-By: Claude Opus 4.7 --- wsai/backends/melo.py | 5 +++-- wsai/backends/melo_worker.py | 28 ++++++++++++++++++++--- wsai/backends/whisper.py | 28 +++++++++++++++++++++-- wsai/backends/whisper_worker.py | 39 +++++++++++++++++++++++++++------ 4 files changed, 86 insertions(+), 14 deletions(-) diff --git a/wsai/backends/melo.py b/wsai/backends/melo.py index 690d65a..9a3f782 100644 --- a/wsai/backends/melo.py +++ b/wsai/backends/melo.py @@ -9,7 +9,8 @@ integration later swaps the sink for "play this wav into the call". Env: WSAI_MELO_PYTHON interpreter with melo installed (default: /home/claude/jarvis-tts/melo311/bin/python) - WSAI_MELO_DEVICE cpu | cuda | auto (default cpu; cuda needs GPU approval) + WSAI_MELO_DEVICE cpu | cuda | auto (default auto: GPU if torch sees one, + else CPU; the worker falls back to CPU if CUDA fails) WSAI_TTS_OUT_DIR where wavs are written (default ~/.cache/wsai/tts) WSAI_TTS_SPEED synthesis speed multiplier (default 1.3) """ @@ -50,7 +51,7 @@ class MeloTTS: sink: Sink | None = None, ) -> None: self.python = python or os.environ.get("WSAI_MELO_PYTHON", _DEFAULT_PYTHON) - self.device = device or os.environ.get("WSAI_MELO_DEVICE", "cpu") + self.device = device or os.environ.get("WSAI_MELO_DEVICE", "auto") self.out_dir = Path(out_dir or os.environ.get("WSAI_TTS_OUT_DIR") or (Path.home() / ".cache/wsai/tts")) self.speed = float(speed if speed is not None diff --git a/wsai/backends/melo_worker.py b/wsai/backends/melo_worker.py index b0c0d64..b1ab879 100644 --- a/wsai/backends/melo_worker.py +++ b/wsai/backends/melo_worker.py @@ -38,11 +38,33 @@ def _log(*a): def main() -> None: lang = "KR" - device = os.environ.get("WSAI_MELO_DEVICE", "cpu") # "cpu" | "cuda" | "auto" - t0 = time.monotonic() + requested = os.environ.get("WSAI_MELO_DEVICE", "auto") # cpu | cuda | auto from melo.api import TTS # heavy import; only in the melo venv - tts = TTS(language=lang, device=device) + def _has_cuda() -> bool: + try: + import torch + + return torch.cuda.is_available() + except Exception: + return False + + device = requested + if requested == "auto": + device = "cuda" if _has_cuda() else "cpu" + + t0 = time.monotonic() + try: + tts = TTS(language=lang, device=device) + except Exception as exc: + # CUDA picked but unusable (CPU-only torch, missing libs, OOM): fall back + # to CPU rather than leaving the whole voice loop dead. + if device == "cuda": + _log(f"[melo_worker] CUDA load failed ({exc}); falling back to CPU") + device = "cpu" + tts = TTS(language=lang, device=device) + else: + raise speaker_id = tts.hps.data.spk2id[lang] load_ms = int((time.monotonic() - t0) * 1000) _emit({"ready": True, "ms": load_ms, "device": device}) diff --git a/wsai/backends/whisper.py b/wsai/backends/whisper.py index 19ebf81..49069ca 100644 --- a/wsai/backends/whisper.py +++ b/wsai/backends/whisper.py @@ -19,7 +19,8 @@ Env: WSAI_WHISPER_PYTHON interpreter with faster-whisper installed (default: /home/claude/jarvis-stt/whisper312/bin/python) WSAI_WHISPER_MODEL model size/name (default: small) - WSAI_WHISPER_DEVICE cpu | cuda (default cpu; cuda needs GPU approval) + WSAI_WHISPER_DEVICE cpu | cuda | auto (default auto: GPU if present, + else CPU; the worker falls back to CPU if CUDA fails) WSAI_WHISPER_LANGUAGE forced language, e.g. ko (default ko; "" = autodetect) """ @@ -41,6 +42,20 @@ log = logging.getLogger("wsai.stt.whisper") _DEFAULT_PYTHON = "/home/claude/jarvis-stt/whisper312/bin/python" +def _cuda_lib_dirs(python_exe: str) -> list[str]: + """nvidia/*/lib dirs of the worker venv (cublas, cudnn, ...), for + LD_LIBRARY_PATH so ctranslate2 can dlopen the CUDA runtime. Empty if the + interpreter has no such packages (CPU-only install).""" + import glob + + # Use the literal path, NOT .resolve(): the venv's bin/python is a symlink + # into the uv-managed interpreter, and resolving it would jump out of the + # venv and miss its site-packages/nvidia libs. + venv = Path(python_exe).parent.parent # .../bin/python -> venv root + dirs = glob.glob(str(venv / "lib" / "python*" / "site-packages" / "nvidia" / "*" / "lib")) + return sorted(set(dirs)) + + class WhisperSTT: def __init__( self, @@ -53,7 +68,7 @@ class WhisperSTT: ) -> None: self.python = python or os.environ.get("WSAI_WHISPER_PYTHON", _DEFAULT_PYTHON) self.model = model or os.environ.get("WSAI_WHISPER_MODEL", "small") - self.device = device or os.environ.get("WSAI_WHISPER_DEVICE", "cpu") + self.device = device or os.environ.get("WSAI_WHISPER_DEVICE", "auto") # "" means autodetect; a real code like "ko" forces the language. env_lang = os.environ.get("WSAI_WHISPER_LANGUAGE", "ko") self.language = language if language is not None else (env_lang or None) @@ -97,6 +112,15 @@ class WhisperSTT: "WSAI_WHISPER_MODEL": self.model, "WSAI_WHISPER_DEVICE": self.device, } + # ctranslate2 dlopens libcublas/libcudnn from the whisper venv's nvidia + # pip packages; the dynamic loader only honours LD_LIBRARY_PATH captured + # at exec, so inject those lib dirs into the child env here (harmless on + # CPU). Without this the CUDA model loads but transcribe() dies with + # "Library libcublas.so.12 is not found". + lib_dirs = _cuda_lib_dirs(self.python) + if lib_dirs: + prev = env.get("LD_LIBRARY_PATH", "") + env["LD_LIBRARY_PATH"] = ":".join(lib_dirs + ([prev] if prev else [])) if self.language: env["WSAI_WHISPER_LANGUAGE"] = self.language repo_root = str(Path(__file__).resolve().parents[2]) diff --git a/wsai/backends/whisper_worker.py b/wsai/backends/whisper_worker.py index 3fa70f4..df85c18 100644 --- a/wsai/backends/whisper_worker.py +++ b/wsai/backends/whisper_worker.py @@ -39,17 +39,42 @@ def _log(*a): def main() -> None: model_name = os.environ.get("WSAI_WHISPER_MODEL", "small") - device = os.environ.get("WSAI_WHISPER_DEVICE", "cpu") # "cpu" | "cuda" - # int8 on CPU keeps a small model fast; float16 is the usual CUDA choice. - compute = os.environ.get( - "WSAI_WHISPER_COMPUTE", "int8" if device == "cpu" else "float16" - ) + requested = os.environ.get("WSAI_WHISPER_DEVICE", "auto") # cpu | cuda | auto default_lang = os.environ.get("WSAI_WHISPER_LANGUAGE", "ko") or None - t0 = time.monotonic() from faster_whisper import WhisperModel # heavy import; only in whisper venv - model = WhisperModel(model_name, device=device, compute_type=compute) + def _has_cuda() -> bool: + try: + import ctranslate2 + + return ctranslate2.get_cuda_device_count() > 0 + except Exception: + return False + + device = requested + if requested == "auto": + device = "cuda" if _has_cuda() else "cpu" + + def _compute_for(dev: str) -> str: + # int8 on CPU keeps a small model fast; float16 is the usual CUDA choice. + return os.environ.get( + "WSAI_WHISPER_COMPUTE", "int8" if dev == "cpu" else "float16" + ) + + t0 = time.monotonic() + try: + model = WhisperModel(model_name, device=device, compute_type=_compute_for(device)) + except Exception as exc: + # CUDA picked but unusable (missing libs, OOM): fall back to CPU rather + # than leaving the whole voice loop dead. + if device == "cuda": + _log(f"[whisper_worker] CUDA load failed ({exc}); falling back to CPU") + device = "cpu" + model = WhisperModel(model_name, device=device, compute_type=_compute_for(device)) + else: + raise + compute = _compute_for(device) load_ms = int((time.monotonic() - t0) * 1000) _emit({"ready": True, "ms": load_ms, "device": device, "model": model_name}) _log(f"[whisper_worker] {model_name} ready in {load_ms} ms on {device}/{compute}")