137 lines
4.9 KiB
Python
137 lines
4.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Small OpenAI-compatible Whisper service for the dedicated Jetson."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import cgi
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import threading
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import whisper
|
|
|
|
|
|
HOST = os.getenv("HERMES_STT_HOST", "0.0.0.0")
|
|
PORT = int(os.getenv("HERMES_STT_PORT", "9000"))
|
|
MODEL_NAME = os.getenv("HERMES_STT_MODEL", "small")
|
|
CACHE_DIR = Path(os.getenv("HERMES_STT_CACHE", "/cache/whisper"))
|
|
MAX_AUDIO_BYTES = 30 * 1024 * 1024
|
|
MODEL_LOCK = threading.Lock()
|
|
|
|
|
|
def _json(handler: BaseHTTPRequestHandler, status: int, payload: dict) -> None:
|
|
body = json.dumps(payload).encode("utf-8")
|
|
handler.send_response(status)
|
|
handler.send_header("Content-Type", "application/json")
|
|
handler.send_header("Content-Length", str(len(body)))
|
|
handler.send_header("Cache-Control", "no-store")
|
|
handler.end_headers()
|
|
handler.wfile.write(body)
|
|
|
|
|
|
class SpeechHandler(BaseHTTPRequestHandler):
|
|
"""Serve health and transcription without exposing a general runtime."""
|
|
|
|
server_version = "AtlasWhisper/1"
|
|
|
|
def log_message(self, message: str, *args: object) -> None:
|
|
print(f"[stt] {self.address_string()} {message % args}", flush=True)
|
|
|
|
def do_GET(self) -> None:
|
|
if self.path != "/health":
|
|
_json(self, 404, {"error": "not found"})
|
|
return
|
|
_json(
|
|
self,
|
|
200,
|
|
{
|
|
"ok": True,
|
|
"model": MODEL_NAME,
|
|
"device": "cuda" if torch.cuda.is_available() else "cpu",
|
|
},
|
|
)
|
|
|
|
def do_POST(self) -> None:
|
|
if self.path != "/v1/audio/transcriptions":
|
|
_json(self, 404, {"error": "not found"})
|
|
return
|
|
content_length = int(self.headers.get("Content-Length", "0") or "0")
|
|
if content_length <= 0 or content_length > MAX_AUDIO_BYTES:
|
|
_json(self, 413, {"error": "audio payload is missing or too large"})
|
|
return
|
|
|
|
content_type = self.headers.get("Content-Type", "")
|
|
if not content_type.lower().startswith("multipart/form-data"):
|
|
_json(self, 400, {"error": "multipart/form-data is required"})
|
|
return
|
|
|
|
form = cgi.FieldStorage(
|
|
fp=self.rfile,
|
|
headers=self.headers,
|
|
environ={
|
|
"REQUEST_METHOD": "POST",
|
|
"CONTENT_TYPE": content_type,
|
|
"CONTENT_LENGTH": str(content_length),
|
|
},
|
|
)
|
|
audio = form["file"] if "file" in form else None
|
|
if audio is None or not getattr(audio, "file", None):
|
|
_json(self, 400, {"error": "file is required"})
|
|
return
|
|
|
|
suffix = Path(getattr(audio, "filename", "audio.wav") or "audio.wav").suffix
|
|
suffix = suffix if suffix in {".wav", ".webm", ".ogg", ".mp3", ".m4a"} else ".wav"
|
|
language = str(form.getfirst("language", "auto") or "auto").strip().lower()
|
|
temp_path = ""
|
|
try:
|
|
with tempfile.NamedTemporaryFile(prefix="atlas-stt-", suffix=suffix, delete=False) as temp:
|
|
temp_path = temp.name
|
|
while True:
|
|
chunk = audio.file.read(1024 * 1024)
|
|
if not chunk:
|
|
break
|
|
temp.write(chunk)
|
|
|
|
with MODEL_LOCK:
|
|
result = self.server.model.transcribe( # type: ignore[attr-defined]
|
|
temp_path,
|
|
language=None if language in {"", "auto"} else language,
|
|
task="transcribe",
|
|
fp16=torch.cuda.is_available(),
|
|
condition_on_previous_text=False,
|
|
temperature=0,
|
|
verbose=False,
|
|
)
|
|
transcript = str(result.get("text") or "").strip()
|
|
_json(self, 200, {"text": transcript, "model": MODEL_NAME})
|
|
except Exception as exc:
|
|
print(f"[stt] transcription failed: {exc}", flush=True)
|
|
_json(self, 500, {"error": "transcription failed"})
|
|
finally:
|
|
if temp_path:
|
|
try:
|
|
os.unlink(temp_path)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def main() -> None:
|
|
"""Warm Whisper once, then serve concurrent clients through one GPU lock."""
|
|
CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("CUDA is required for the Atlas Whisper service")
|
|
print(f"[stt] loading Whisper {MODEL_NAME} into CUDA", flush=True)
|
|
model = whisper.load_model(MODEL_NAME, device="cuda", download_root=str(CACHE_DIR))
|
|
server = ThreadingHTTPServer((HOST, PORT), SpeechHandler)
|
|
server.model = model # type: ignore[attr-defined]
|
|
print(f"[stt] ready on {HOST}:{PORT}", flush=True)
|
|
server.serve_forever(poll_interval=0.25)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|