vignette/scripts/local-whisper-stt-server.py
2026-08-09 18:22:03 +09:00

657 lines
22 KiB
Python

# -*- coding: utf-8 -*-
"""로컬 상주 faster-whisper 스트리밍 STT 사이드카.
Higgs TTS 서버와 같은 형태다. 모델을 한 번만 올려 두고 loopback WebSocket으로
linear16 PCM을 받아 interim/final 전사를 돌려준다. 오디오는 디스크에 쓰지 않고
발화 단위 메모리 버퍼만 사용하며, 발화가 끝나면 즉시 버린다.
프로토콜(클라이언트 → 서버):
binary frame raw linear16 PCM
{"type": "Finalize"} 현재 발화를 즉시 확정한다
{"type": "CloseStream"} 남은 발화를 확정하고 닫는다
프로토콜(서버 → 클라이언트):
{"type": "ready", ...}
{"type": "transcript", "text": ..., "is_final": bool, "speech_final": bool,
"confidence": float|null, "words": [...], "duration": float}
{"type": "error", "detail": "..."}
`--enable` 없이는 뜨지 않고, 기본 바인딩은 loopback뿐이다.
"""
from __future__ import annotations
import argparse
import array
import asyncio
import json
import math
import sys
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Iterable, Protocol
from urllib.parse import parse_qs, urlsplit
DEFAULT_HOST = "127.0.0.1"
DEFAULT_PORT = 9882
DEFAULT_MODEL = "small"
DEFAULT_LANGUAGE = "ko"
DEFAULT_SAMPLE_RATE = 16_000
DEFAULT_CHANNELS = 1
SAMPLE_WIDTH = 2
# 발화 경계 판정
DEFAULT_ENDPOINTING_MS = 300
DEFAULT_UTTERANCE_END_MS = 1_200
DEFAULT_INTERIM_INTERVAL_MS = 700
DEFAULT_SILENCE_RMS = 320
MIN_UTTERANCE_MS = 200
MAX_UTTERANCE_SECONDS = 60.0
MAX_FRAME_BYTES = 1 << 20
ALLOWED_MODELS = (
"large-v3",
"large-v3-turbo",
"medium",
"small",
"base",
)
class TranscriptionError(RuntimeError):
"""전사 실패. 원문 오디오는 로그에 남기지 않는다."""
@dataclass(frozen=True, slots=True)
class Word:
word: str
start: float
end: float
confidence: float | None = None
def as_payload(self) -> dict[str, Any]:
payload: dict[str, Any] = {
"word": self.word,
"start": round(self.start, 3),
"end": round(self.end, 3),
}
if self.confidence is not None:
payload["confidence"] = round(self.confidence, 4)
return payload
@dataclass(frozen=True, slots=True)
class Transcription:
text: str
confidence: float | None = None
words: tuple[Word, ...] = ()
class Transcriber(Protocol):
def transcribe(
self, pcm: bytes, *, sample_rate: int, final: bool
) -> Transcription: ...
def frame_rms(pcm: bytes) -> int:
"""linear16 프레임의 RMS. 빈 입력은 무음으로 본다.
`audioop` 은 Python 3.13에서 제거되므로 표준 정수 연산으로 직접 구한다.
"""
if not pcm or len(pcm) < SAMPLE_WIDTH:
return 0
usable = len(pcm) - (len(pcm) % SAMPLE_WIDTH)
samples = array.array("h")
samples.frombytes(pcm[:usable])
if sys.byteorder != "little": # pragma: no cover - little-endian 개발 환경
samples.byteswap()
if not samples:
return 0
total = 0
for sample in samples:
total += sample * sample
return int(math.sqrt(total / len(samples)))
@dataclass
class UtteranceBuffer:
"""한 발화 분량의 PCM과 무음 누적을 추적한다.
이 클래스는 모델을 전혀 모른다. 덕분에 GPU 없이도 경계 판정을 그대로
테스트할 수 있다.
"""
sample_rate: int = DEFAULT_SAMPLE_RATE
channels: int = DEFAULT_CHANNELS
endpointing_ms: int = DEFAULT_ENDPOINTING_MS
utterance_end_ms: int = DEFAULT_UTTERANCE_END_MS
interim_interval_ms: int = DEFAULT_INTERIM_INTERVAL_MS
silence_rms: int = DEFAULT_SILENCE_RMS
_pcm: bytearray = field(default_factory=bytearray, init=False)
_silence_ms: float = field(default=0.0, init=False)
_speech_ms: float = field(default=0.0, init=False)
_since_interim_ms: float = field(default=0.0, init=False)
_saw_speech: bool = field(default=False, init=False)
@property
def bytes_per_ms(self) -> float:
return self.sample_rate * self.channels * SAMPLE_WIDTH / 1000.0
def duration_ms(self, pcm_length: int | None = None) -> float:
length = len(self._pcm) if pcm_length is None else pcm_length
return length / self.bytes_per_ms if self.bytes_per_ms else 0.0
@property
def pcm(self) -> bytes:
return bytes(self._pcm)
@property
def saw_speech(self) -> bool:
return self._saw_speech
def append(self, chunk: bytes) -> None:
if not chunk:
return
self._pcm.extend(chunk)
span_ms = self.duration_ms(len(chunk))
self._since_interim_ms += span_ms
if frame_rms(chunk) >= self.silence_rms:
self._saw_speech = True
self._speech_ms += span_ms
self._silence_ms = 0.0
else:
self._silence_ms += span_ms
def should_emit_interim(self) -> bool:
if not self._saw_speech:
return False
if self.duration_ms() < MIN_UTTERANCE_MS:
return False
return self._since_interim_ms >= self.interim_interval_ms
def mark_interim_emitted(self) -> None:
self._since_interim_ms = 0.0
def should_finalize(self) -> bool:
"""무음이 endpointing 기준을 넘거나 발화가 최대 길이에 닿으면 확정한다."""
if not self._saw_speech:
return False
if self.duration_ms() >= MAX_UTTERANCE_SECONDS * 1000:
return True
return self._silence_ms >= self.endpointing_ms
def is_stalled(self) -> bool:
"""말이 없는 채로 utterance_end_ms 를 넘긴 상태."""
return not self._saw_speech and self._silence_ms >= self.utterance_end_ms
def reset(self) -> None:
self._pcm.clear()
self._silence_ms = 0.0
self._speech_ms = 0.0
self._since_interim_ms = 0.0
self._saw_speech = False
def transcript_payload(
result: Transcription,
*,
is_final: bool,
speech_final: bool,
duration_seconds: float,
) -> dict[str, Any]:
payload: dict[str, Any] = {
"type": "transcript",
"text": result.text,
"is_final": bool(is_final),
"speech_final": bool(speech_final),
"duration": round(max(0.0, duration_seconds), 3),
"words": [word.as_payload() for word in result.words] if is_final else [],
}
payload["confidence"] = (
round(result.confidence, 4) if result.confidence is not None else None
)
return payload
def parse_stream_options(path: str) -> dict[str, Any]:
"""접속 쿼리에서 세션 옵션을 읽는다. 알 수 없는 값은 기본값으로 되돌린다."""
query = parse_qs(urlsplit(path).query)
def one(name: str) -> str:
values = query.get(name) or []
return values[0].strip() if values else ""
def integer(name: str, default: int, low: int, high: int) -> int:
raw = one(name)
if not raw:
return default
try:
return max(low, min(high, int(raw)))
except ValueError:
return default
model = one("model") or DEFAULT_MODEL
if model not in ALLOWED_MODELS:
raise ValueError("model_not_allowlisted")
language = one("language") or DEFAULT_LANGUAGE
if not language.isalpha() or len(language) > 8:
raise ValueError("language_invalid")
return {
"model": model,
"language": language,
"sample_rate": integer("sample_rate", DEFAULT_SAMPLE_RATE, 8_000, 48_000),
"channels": integer("channels", DEFAULT_CHANNELS, 1, 2),
"endpointing_ms": integer("endpointing", DEFAULT_ENDPOINTING_MS, 10, 5_000),
"utterance_end_ms": integer(
"utterance_end_ms", DEFAULT_UTTERANCE_END_MS, 1_000, 10_000
),
"interim_interval_ms": integer(
"interim_interval_ms", DEFAULT_INTERIM_INTERVAL_MS, 200, 5_000
),
"silence_rms": integer("silence_rms", DEFAULT_SILENCE_RMS, 1, 10_000),
}
class StreamSession:
"""한 학습자 발화 스트림. 모델 호출은 주입된 transcriber 로만 한다."""
def __init__(
self,
*,
options: dict[str, Any],
transcriber: Transcriber,
send: Callable[[dict[str, Any]], Any],
run_blocking: Callable[..., Any] | None = None,
) -> None:
self.options = options
self._transcriber = transcriber
self._send = send
self._run_blocking = run_blocking
self.buffer = UtteranceBuffer(
sample_rate=options["sample_rate"],
channels=options["channels"],
endpointing_ms=options["endpointing_ms"],
utterance_end_ms=options["utterance_end_ms"],
interim_interval_ms=options["interim_interval_ms"],
silence_rms=options["silence_rms"],
)
self.final_segments: list[str] = []
async def _transcribe(self, pcm: bytes, *, final: bool) -> Transcription:
def call() -> Transcription:
return self._transcriber.transcribe(
pcm, sample_rate=self.options["sample_rate"], final=final
)
if self._run_blocking is None:
return call()
return await self._run_blocking(call)
async def feed(self, chunk: bytes) -> None:
if len(chunk) > MAX_FRAME_BYTES:
raise TranscriptionError("audio_frame_too_large")
self.buffer.append(chunk)
if self.buffer.should_finalize():
await self.finalize(speech_final=True)
return
if self.buffer.should_emit_interim():
self.buffer.mark_interim_emitted()
pcm = self.buffer.pcm
result = await self._transcribe(pcm, final=False)
if result.text.strip():
await self._send(
transcript_payload(
result,
is_final=False,
speech_final=False,
duration_seconds=self.buffer.duration_ms(len(pcm)) / 1000.0,
)
)
async def finalize(self, *, speech_final: bool) -> None:
pcm = self.buffer.pcm
duration_seconds = self.buffer.duration_ms(len(pcm)) / 1000.0
saw_speech = self.buffer.saw_speech
self.buffer.reset()
if not saw_speech or self.buffer.duration_ms(len(pcm)) < MIN_UTTERANCE_MS:
return
result = await self._transcribe(pcm, final=True)
text = result.text.strip()
if not text:
return
self.final_segments.append(text)
await self._send(
transcript_payload(
result,
is_final=True,
speech_final=speech_final,
duration_seconds=duration_seconds,
)
)
def warmup_pcm(sample_rate: int, *, milliseconds: int = 500) -> bytes:
"""warmup 용 무음 linear16 버퍼."""
samples = max(1, int(sample_rate * milliseconds / 1000))
return b"\x00\x00" * samples
class FasterWhisperTranscriber:
"""상주 faster-whisper 모델 하나. 오디오는 메모리에서만 다룬다."""
def __init__(self, model: str, *, device: str, compute_type: str) -> None:
try:
import numpy as np
from faster_whisper import WhisperModel
except ImportError as exc: # pragma: no cover - 런타임 환경 의존
raise TranscriptionError("faster_whisper_unavailable") from exc
self._np = np
self.model_name = model
self.device = device
self.compute_type = compute_type
self._model = WhisperModel(model, device=device, compute_type=compute_type)
self._language = DEFAULT_LANGUAGE
def configure_language(self, language: str) -> None:
self._language = language
def warmup(self) -> None:
"""짧은 무음 한 번을 실제로 돌려 디바이스가 살아 있는지 확인한다.
`get_cuda_device_count() > 0` 만으로는 부족하다. cuDNN 이 없으면 모델은
올라가지만 첫 추론에서 죽어 스트림이 조용히 끊긴다.
"""
self.transcribe(
warmup_pcm(DEFAULT_SAMPLE_RATE),
sample_rate=DEFAULT_SAMPLE_RATE,
final=False,
)
def transcribe(
self, pcm: bytes, *, sample_rate: int, final: bool
) -> Transcription:
if not pcm:
return Transcription(text="")
audio = (
self._np.frombuffer(pcm, dtype=self._np.int16).astype(self._np.float32)
/ 32768.0
)
segments, _info = self._model.transcribe(
audio,
language=self._language,
beam_size=5 if final else 1,
word_timestamps=final,
vad_filter=False,
condition_on_previous_text=False,
)
texts: list[str] = []
words: list[Word] = []
confidences: list[float] = []
for segment in segments:
piece = (segment.text or "").strip()
if piece:
texts.append(piece)
probability = getattr(segment, "avg_logprob", None)
if probability is not None:
confidences.append(min(1.0, max(0.0, float(2 ** probability))))
for word in getattr(segment, "words", None) or ():
token = (getattr(word, "word", "") or "").strip()
start = getattr(word, "start", None)
end = getattr(word, "end", None)
if not token or start is None or end is None or end <= start:
continue
probability = getattr(word, "probability", None)
words.append(
Word(
word=token,
start=float(start),
end=float(end),
confidence=float(probability)
if probability is not None
else None,
)
)
confidence = sum(confidences) / len(confidences) if confidences else None
return Transcription(
text=" ".join(texts).strip(),
confidence=confidence,
words=tuple(words),
)
def resolve_device(requested: str, *, cuda_available: bool | None = None) -> tuple[str, str]:
"""요청한 디바이스를 실제 가용 자원으로 해석한다."""
if requested not in {"auto", "cuda", "cpu"}:
raise ValueError("device_invalid")
if cuda_available is None:
try: # pragma: no cover - 런타임 환경 의존
import ctranslate2
cuda_available = ctranslate2.get_cuda_device_count() > 0
except Exception:
cuda_available = False
if requested == "cuda":
if not cuda_available:
raise ValueError("cuda_unavailable")
return "cuda", "float16"
if requested == "cpu":
return "cpu", "int8"
return ("cuda", "float16") if cuda_available else ("cpu", "int8")
def probe_device(model: str, device: str, *, python: str | None = None) -> bool:
"""별도 프로세스에서 디바이스를 실제로 한 번 돌려 본다.
cuDNN 이 없으면 ctranslate2 는 Python 예외가 아니라 **네이티브 크래시**로
죽는다. 같은 프로세스의 try/except 로는 절대 잡을 수 없으므로, 살릴 수 없는
실패는 버릴 수 있는 자식 프로세스에서 먼저 확인한다.
"""
import subprocess
argv = [
python or sys.executable,
"-X",
"utf8",
str(Path(__file__).resolve()),
"--self-check",
"--model",
model,
"--device",
device,
]
try:
completed = subprocess.run(
argv,
check=False,
capture_output=True,
timeout=600,
shell=False,
)
except (OSError, subprocess.TimeoutExpired):
return False
return completed.returncode == 0
def build_transcriber(
model: str,
requested_device: str,
*,
factory: Callable[[str, str, str], Any],
cuda_available: bool | None = None,
prober: Callable[[str, str], bool] | None = None,
) -> Any:
"""디바이스를 해석하고 실제로 추론이 되는 transcriber 만 돌려준다.
`auto` 는 CUDA 가 보이면 먼저 자식 프로세스로 확인하고, 거기서 실패하면
CPU 로 내려간다. 명시적 `--device cuda` 는 조용히 강등하지 않는다.
"""
device, compute_type = resolve_device(requested_device, cuda_available=cuda_available)
check = prober if prober is not None else (lambda m, d: probe_device(m, d))
if device == "cuda" and not check(model, "cuda"):
if requested_device != "auto":
raise TranscriptionError("device_unusable:cuda")
device, compute_type = "cpu", "int8"
try:
transcriber = factory(model, device, compute_type)
transcriber.warmup()
return transcriber
except Exception as exc:
raise TranscriptionError(f"device_unusable:{device}") from exc
async def serve(args: argparse.Namespace) -> int: # pragma: no cover - I/O 진입점
import websockets
transcriber = build_transcriber(
args.model,
args.device,
factory=lambda model, device, compute_type: FasterWhisperTranscriber(
model, device=device, compute_type=compute_type
),
)
loop = asyncio.get_running_loop()
async def run_blocking(call: Callable[[], Transcription]) -> Transcription:
return await loop.run_in_executor(None, call)
async def handler(socket: Any) -> None:
try:
options = parse_stream_options(getattr(socket, "path", "") or "")
except ValueError as exc:
await socket.send(json.dumps({"type": "error", "detail": str(exc)}))
await socket.close()
return
transcriber.configure_language(options["language"])
async def send(payload: dict[str, Any]) -> None:
await socket.send(json.dumps(payload, ensure_ascii=False))
session = StreamSession(
options=options,
transcriber=transcriber,
send=send,
run_blocking=run_blocking,
)
await send(
{
"type": "ready",
"provider": "local_whisper",
"model": transcriber.model_name,
"device": transcriber.device,
"compute_type": transcriber.compute_type,
"language": options["language"],
"sample_rate": options["sample_rate"],
}
)
try:
async for message in socket:
if isinstance(message, bytes):
await session.feed(message)
continue
try:
control = json.loads(message)
except (TypeError, ValueError):
continue
kind = str((control or {}).get("type") or "")
if kind == "Finalize":
await session.finalize(speech_final=True)
elif kind == "CloseStream":
await session.finalize(speech_final=True)
break
except TranscriptionError as exc:
await send({"type": "error", "detail": exc.args[0] if exc.args else "error"})
finally:
with_close = getattr(socket, "close", None)
if callable(with_close):
await socket.close()
print(
json.dumps(
{
"ready": True,
"provider": "local_whisper",
"host": args.host,
"port": args.port,
"model": transcriber.model_name,
"device": transcriber.device,
"compute_type": transcriber.compute_type,
},
separators=(",", ":"),
),
flush=True,
)
async with websockets.serve(
handler,
args.host,
args.port,
max_size=MAX_FRAME_BYTES,
max_queue=16,
ping_interval=20,
ping_timeout=20,
):
await asyncio.Future()
return 0
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--host", default=DEFAULT_HOST, choices=[DEFAULT_HOST, "localhost"])
parser.add_argument("--port", type=int, default=DEFAULT_PORT)
parser.add_argument("--model", default=DEFAULT_MODEL, choices=list(ALLOWED_MODELS))
parser.add_argument("--device", default="auto", choices=["auto", "cuda", "cpu"])
parser.add_argument("--enable", action="store_true")
parser.add_argument(
"--self-check",
action="store_true",
help="load the model on --device once and exit; used as a crash-safe probe",
)
return parser
def run_self_check(args: argparse.Namespace) -> int: # pragma: no cover - 서브프로세스
device, compute_type = resolve_device(
"cpu" if args.device == "cpu" else args.device
)
transcriber = FasterWhisperTranscriber(
args.model, device=device, compute_type=compute_type
)
transcriber.warmup()
print(json.dumps({"ok": True, "device": device}, separators=(",", ":")))
return 0
def main(argv: Iterable[str] | None = None) -> int: # pragma: no cover - CLI
args = build_parser().parse_args(list(argv) if argv is not None else None)
if args.self_check:
try:
return run_self_check(args)
except Exception as exc:
print(json.dumps({"ok": False, "error": type(exc).__name__}), file=sys.stderr)
return 3
if not args.enable:
print("local whisper STT server is disabled; pass --enable", file=sys.stderr)
return 2
started = time.monotonic()
try:
return asyncio.run(serve(args))
except KeyboardInterrupt:
print(
json.dumps({"stopped": True, "uptime_seconds": round(time.monotonic() - started, 1)}),
file=sys.stderr,
)
return 0
if __name__ == "__main__": # pragma: no cover - CLI
raise SystemExit(main())