6e20f2cd79
The STT/TTS worker _ensure() treated a spawned-but-not-yet-handshaked subprocess as ready, so a voice turn arriving during warmup read the same stdout StreamReader concurrently with the warmup handshake and crashed with "readuntil() called while another coroutine is already waiting for incoming data". Add a _start_lock + _ready flag so (re)start and the ready handshake run atomically and callers wait for real readiness before reading stdout. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
231 lines
10 KiB
Python
231 lines
10 KiB
Python
"""Real STT via faster-whisper, run as a persistent out-of-venv worker.
|
|
|
|
faster-whisper needs its own interpreter (whisper312) and loading the model
|
|
costs several seconds, so we keep one worker process alive and stream
|
|
transcription requests to it (see whisper_worker.py for the protocol) — the
|
|
same warm-worker shape as MeloTTS on the TTS side.
|
|
|
|
`transcribe(wav)` is the core engine: hand it a wav path, get the recognised
|
|
text back. It is what closes the voice round trip (TTS wav -> STT text) and what
|
|
the Discord voice path will call once per detected utterance.
|
|
|
|
`utterances()` turns this into a SpeechToText source: it pulls finished-utterance
|
|
wav paths from an injected `audio_source` and yields a transcribed Utterance for
|
|
each. Until the Discord voice receiver is wired, `audio_source` is None and the
|
|
loop simply idles (like the eyes-free perception loop), while `transcribe()`
|
|
stays usable directly.
|
|
|
|
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 | 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)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import collections
|
|
import json
|
|
import logging
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
from typing import AsyncIterator
|
|
|
|
from ..interfaces import Utterance
|
|
|
|
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,
|
|
*,
|
|
python: str | None = None,
|
|
model: str | None = None,
|
|
device: str | None = None,
|
|
language: str | None = None,
|
|
audio_source: AsyncIterator[str] | None = None,
|
|
) -> 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", "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)
|
|
self.audio_source = audio_source
|
|
self._proc: asyncio.subprocess.Process | None = None
|
|
self._lock = asyncio.Lock()
|
|
# Serialises worker (re)start + the ready handshake so a caller that
|
|
# arrives mid-warmup waits for readiness instead of reading the same
|
|
# stdout StreamReader concurrently (asyncio forbids overlapping reads).
|
|
self._start_lock = asyncio.Lock()
|
|
self._ready = False # True only after the ready handshake completes
|
|
self.load_ms: int | None = None
|
|
self.resolved_device: str | None = None # "cuda" | "cpu", known after start
|
|
# Keep the worker's most recent stderr so a crash reports its real cause
|
|
# instead of a bare JSONDecodeError. Bounded so it can't grow unbounded.
|
|
self._stderr_tail: collections.deque[str] = collections.deque(maxlen=40)
|
|
self._stderr_task: asyncio.Task | None = None
|
|
|
|
async def _drain_stderr(self, stream: asyncio.StreamReader) -> None:
|
|
# The worker redirects fd1 -> fd2, so model/library chatter lands on
|
|
# stderr. If we PIPE but never read it, the pipe buffer fills and the
|
|
# worker blocks. So we drain continuously, keeping only the last lines.
|
|
try:
|
|
while True:
|
|
line = await stream.readline()
|
|
if not line:
|
|
return
|
|
self._stderr_tail.append(line.decode(errors="replace").rstrip())
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception: # draining must never crash the caller
|
|
return
|
|
|
|
def _stderr_hint(self) -> str:
|
|
tail = "\n".join(self._stderr_tail)
|
|
return f" worker stderr tail:\n{tail}" if tail else " (worker produced no stderr)"
|
|
|
|
async def warmup(self) -> None:
|
|
"""Load the model now so the first real utterance is transcribed warm."""
|
|
await self._ensure()
|
|
|
|
async def _ensure(self) -> None:
|
|
# Fast path: only skip when the worker is not just spawned but fully
|
|
# handshaked. Checking `_proc` alone would let a caller sail past while
|
|
# another coroutine (e.g. warmup) is still awaiting the ready line on
|
|
# this same stdout, causing overlapping StreamReader reads.
|
|
if self._proc is not None and self._proc.returncode is None and self._ready:
|
|
return
|
|
async with self._start_lock:
|
|
# Re-check under the lock: another coroutine may have finished the
|
|
# (re)start + handshake while we waited.
|
|
if self._proc is not None and self._proc.returncode is None and self._ready:
|
|
return
|
|
self._ready = False
|
|
env = {
|
|
**os.environ,
|
|
"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])
|
|
self._proc = await asyncio.create_subprocess_exec(
|
|
self.python, "-m", "wsai.backends.whisper_worker",
|
|
cwd=repo_root, env=env,
|
|
stdin=asyncio.subprocess.PIPE,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
self._stderr_tail.clear()
|
|
assert self._proc.stderr is not None
|
|
self._stderr_task = asyncio.create_task(self._drain_stderr(self._proc.stderr))
|
|
ready = await self._proc.stdout.readline()
|
|
if not ready: # worker died before signalling ready
|
|
await self._proc.wait()
|
|
raise RuntimeError(
|
|
f"whisper worker exited before ready (code {self._proc.returncode})."
|
|
f"{self._stderr_hint()}"
|
|
)
|
|
try:
|
|
info = json.loads(ready.decode())
|
|
except json.JSONDecodeError as exc:
|
|
raise RuntimeError(
|
|
f"whisper worker sent invalid ready line {ready!r}: {exc}."
|
|
f"{self._stderr_hint()}"
|
|
) from exc
|
|
if not info.get("ready"):
|
|
raise RuntimeError(
|
|
f"whisper worker failed to start: {info}.{self._stderr_hint()}"
|
|
)
|
|
self.load_ms = info.get("ms")
|
|
self.resolved_device = info.get("device")
|
|
self._ready = True
|
|
log.info(
|
|
"whisper worker ready in %s ms on %s (model %s)",
|
|
self.load_ms, info.get("device"), info.get("model"),
|
|
)
|
|
|
|
async def transcribe(self, wav_path: str, *, language: str | None = None) -> str:
|
|
"""Transcribe one wav file to text using the warm worker."""
|
|
await self._ensure()
|
|
req: dict[str, object] = {"wav": wav_path}
|
|
lang = language if language is not None else self.language
|
|
if lang:
|
|
req["language"] = lang
|
|
s = time.monotonic()
|
|
async with self._lock:
|
|
assert self._proc and self._proc.stdin and self._proc.stdout
|
|
self._proc.stdin.write((json.dumps(req) + "\n").encode())
|
|
await self._proc.stdin.drain()
|
|
resp = await self._proc.stdout.readline()
|
|
if not resp:
|
|
raise RuntimeError(f"whisper worker closed unexpectedly.{self._stderr_hint()}")
|
|
res = json.loads(resp.decode())
|
|
if not res.get("ok"):
|
|
raise RuntimeError(f"whisper transcribe failed: {res.get('error')}")
|
|
log.debug(
|
|
"transcribe %d ms (worker %s ms): %s",
|
|
int((time.monotonic() - s) * 1000), res.get("ms"), res.get("text", "")[:60],
|
|
)
|
|
return res.get("text", "")
|
|
|
|
async def utterances(self) -> AsyncIterator[Utterance]:
|
|
"""Yield an Utterance per finished-utterance wav from `audio_source`.
|
|
|
|
With no audio source wired yet (Discord voice receiver pending) this
|
|
idles and returns, leaving the voice loop dormant but valid."""
|
|
if self.audio_source is None:
|
|
return
|
|
async for wav_path in self.audio_source:
|
|
text = await self.transcribe(wav_path)
|
|
if text:
|
|
yield Utterance(text=text, ts=time.monotonic(), source="voice")
|
|
|
|
async def aclose(self) -> None:
|
|
if self._proc is not None and self._proc.returncode is None:
|
|
try:
|
|
self._proc.terminate()
|
|
await asyncio.wait_for(self._proc.wait(), timeout=5)
|
|
except (ProcessLookupError, asyncio.TimeoutError):
|
|
pass
|
|
if self._stderr_task is not None:
|
|
self._stderr_task.cancel()
|
|
try:
|
|
await self._stderr_task
|
|
except (asyncio.CancelledError, Exception):
|
|
pass
|
|
self._stderr_task = None
|
|
self._proc = None
|
|
self._ready = False
|