77d7cd8b56
The first CUDA inference pays a large lazy cost (kernel autotune/cudnn) — ~10s for a cold TTS synth — which would blow the voice loop's ~1s budget on the very first reply. Each worker now runs one dummy inference (TTS: a short phrase; STT: 1s of silence) after model load and before emitting "ready", so "ready" means "hot". Warmup failures are logged and never block startup. Verified: first real call after startup is now TTS ~238ms / STT ~189ms (was ~11s cold for TTS). 12 tests pass. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
123 lines
4.6 KiB
Python
123 lines
4.6 KiB
Python
"""Persistent faster-whisper STT worker.
|
|
|
|
faster-whisper (ctranslate2) lives in its own Python (whisper312); loading the
|
|
model takes seconds, so we load it ONCE here and then serve transcription
|
|
requests over stdin/stdout. This process is launched with the whisper312
|
|
interpreter by wsai.backends.whisper.WhisperSTT.
|
|
|
|
Like the MeloTTS worker, model/backend chatter could corrupt the JSON protocol,
|
|
so on startup we split the streams: a private duplicate of the original stdout
|
|
carries the protocol, and fd 1 is redirected to fd 2 so any library print lands
|
|
on stderr instead (where the parent drains it for diagnostics).
|
|
|
|
Protocol (one JSON object per line, on the protocol channel):
|
|
<- {"wav": "/abs/path.wav", "language": "ko"}
|
|
-> {"ok": true, "text": "...", "language": "ko", "ms": 123}
|
|
-> {"ok": false, "error": "..."}
|
|
On startup, once the model is ready, it emits exactly one line:
|
|
-> {"ready": true, "ms": <load-ms>, "device": "cpu", "model": "small"}
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
# Split protocol from library noise BEFORE importing anything heavy.
|
|
_proto = os.fdopen(os.dup(1), "w", buffering=1) # private copy of real stdout
|
|
os.dup2(2, 1) # fd1 -> stderr, so stray library prints don't hit the protocol
|
|
|
|
|
|
def _emit(obj: dict) -> None:
|
|
_proto.write(json.dumps(obj, ensure_ascii=False) + "\n")
|
|
_proto.flush()
|
|
|
|
|
|
def _log(*a):
|
|
print(*a, file=sys.stderr, flush=True)
|
|
|
|
|
|
def main() -> None:
|
|
model_name = os.environ.get("WSAI_WHISPER_MODEL", "small")
|
|
requested = os.environ.get("WSAI_WHISPER_DEVICE", "auto") # cpu | cuda | auto
|
|
default_lang = os.environ.get("WSAI_WHISPER_LANGUAGE", "ko") or None
|
|
|
|
from faster_whisper import WhisperModel # heavy import; only in whisper venv
|
|
|
|
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)
|
|
|
|
# Warm up before signalling ready: the first CUDA transcribe pays a large
|
|
# lazy cost (kernel autotune), which would slow the first real utterance.
|
|
# Run a dummy transcribe on 1s of silence here so "ready" means "hot".
|
|
warmup_ms = None
|
|
try:
|
|
import numpy as np
|
|
|
|
w = time.monotonic()
|
|
segs, _ = model.transcribe(np.zeros(16000, dtype=np.float32), language=default_lang)
|
|
for _ in segs: # segments are lazy; drain to force the actual compute
|
|
pass
|
|
warmup_ms = int((time.monotonic() - w) * 1000)
|
|
except Exception as exc:
|
|
_log(f"[whisper_worker] warmup skipped: {exc}")
|
|
|
|
_emit({"ready": True, "ms": load_ms, "device": device, "model": model_name, "warmup_ms": warmup_ms})
|
|
_log(f"[whisper_worker] {model_name} ready in {load_ms} ms on {device}/{compute} (warmup {warmup_ms} ms)")
|
|
|
|
for line in sys.stdin:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
req = json.loads(line)
|
|
wav = req["wav"]
|
|
language = req.get("language", default_lang)
|
|
s = time.monotonic()
|
|
segments, info = model.transcribe(
|
|
wav,
|
|
language=language,
|
|
beam_size=int(req.get("beam_size", 5)),
|
|
vad_filter=bool(req.get("vad_filter", True)),
|
|
)
|
|
text = "".join(seg.text for seg in segments).strip()
|
|
ms = int((time.monotonic() - s) * 1000)
|
|
_emit({"ok": True, "text": text, "language": info.language, "ms": ms})
|
|
except Exception as exc: # keep the worker alive across bad requests
|
|
_emit({"ok": False, "error": f"{type(exc).__name__}: {exc}"})
|
|
_log(f"[whisper_worker] error: {exc}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|