#!/usr/bin/env python3 """Small OpenAI-compatible Whisper service for the dedicated Jetson.""" from __future__ import annotations import cgi import base64 import hashlib import json import os import re import socket import struct import tempfile import threading import time import wave 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 STREAM_SAMPLE_RATE = 16_000 STREAM_SAMPLE_WIDTH = 2 STREAM_CHANNELS = 1 MAX_STREAM_SECONDS = int(os.getenv("HERMES_STT_STREAM_MAX_SECONDS", "90")) MAX_STREAM_AUDIO_BYTES = min( MAX_AUDIO_BYTES, STREAM_SAMPLE_RATE * STREAM_SAMPLE_WIDTH * MAX_STREAM_SECONDS, ) MAX_WS_MESSAGE_BYTES = 256 * 1024 STREAM_IDLE_SECONDS = int(os.getenv("HERMES_STT_STREAM_IDLE_SECONDS", "120")) VAD_END_SILENCE_MS = int(os.getenv("HERMES_STT_VAD_END_SILENCE_MS", "650")) VAD_TAIL_MS = int(os.getenv("HERMES_STT_VAD_TAIL_MS", "220")) MODEL_LOCK = threading.Lock() _WS_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" def _repetitive_token(token: str) -> bool: """Identify long periodic Whisper hallucinations caused by steady noise.""" letters = "".join(re.findall(r"[a-z]+", token.lower())) if len(letters) < 10: return False for period in range(1, 5): pattern = letters[:period] matches = sum( character == pattern[index % period] for index, character in enumerate(letters) ) if matches / len(letters) >= 0.86: return True return ( max(letters.count(character) for character in set(letters)) / len(letters) >= 0.78 ) def _clean_transcript(result: dict) -> str: """Drop noise-only segments and repetitive tokens while retaining speech.""" segments = result.get("segments") if not isinstance(segments, list): segments = [{"text": result.get("text") or ""}] kept: list[str] = [] for segment in segments: if not isinstance(segment, dict): continue text = str(segment.get("text") or "").strip() if not text: continue no_speech = float(segment.get("no_speech_prob") or 0.0) average_logprob = float(segment.get("avg_logprob") or 0.0) if no_speech >= 0.55 and average_logprob <= -0.55: continue words = [word for word in text.split() if not _repetitive_token(word)] if words: kept.append(" ".join(words)) return " ".join(kept).strip() def _detected_language(result: object) -> str: """Return the bare ISO-639 code Whisper decoded with, or nothing at all. ``whisper.transcribe`` reports the language it auto-detected (or the one it was told to use) as a plain lowercase token such as ``en``/``ru``/``yue``. Anything that is not that exact shape is dropped rather than guessed at, so a surprising model result can never become a downstream voice selector. """ if not isinstance(result, dict): return "" value = result.get("language") if not isinstance(value, str): return "" code = value.strip().lower() if not 2 <= len(code) <= 3 or not code.isascii() or not code.isalpha(): return "" return code def _transcription_payload(result: dict) -> dict: """Build the transcription contract: text plus the model's own language.""" return { "text": _clean_transcript(result), "model": MODEL_NAME, "language": _detected_language(result), } def _transcribe_path(model: object, path: str, language: str) -> dict: """Run the one canonical full-utterance Whisper decode configuration.""" with MODEL_LOCK: result = model.transcribe( path, language=None if language in {"", "auto"} else language, task="transcribe", fp16=torch.cuda.is_available(), condition_on_previous_text=False, temperature=0, compression_ratio_threshold=2.0, logprob_threshold=-0.8, no_speech_threshold=0.5, verbose=False, ) return _transcription_payload(result) def _transcribe_pcm(model: object, pcm: bytes, language: str) -> dict: """Decode one complete 16-kHz mono PCM snapshot through the canonical path.""" temp_path = "" try: with tempfile.NamedTemporaryFile( prefix="atlas-stt-stream-", suffix=".wav", delete=False ) as temp: temp_path = temp.name with wave.open(temp_path, "wb") as wav_file: wav_file.setnchannels(STREAM_CHANNELS) wav_file.setsampwidth(STREAM_SAMPLE_WIDTH) wav_file.setframerate(STREAM_SAMPLE_RATE) wav_file.writeframes(pcm) return _transcribe_path(model, temp_path, language) finally: if temp_path: try: os.unlink(temp_path) except OSError: pass class WebSocketError(Exception): """Protocol error carrying an RFC 6455 close code safe to expose.""" def __init__(self, message: str, close_code: int = 1002) -> None: super().__init__(message) self.close_code = close_code class WebSocketConnection: """Minimal server-side RFC 6455 framing for the private streaming endpoint.""" def __init__(self, handler: BaseHTTPRequestHandler) -> None: self.handler = handler self._send_lock = threading.Lock() self._closed = False @staticmethod def accept_value(key: str) -> str: """Return the RFC 6455 handshake response for a client nonce.""" digest = hashlib.sha1((key + _WS_GUID).encode("ascii")).digest() return base64.b64encode(digest).decode("ascii") def _read_exact(self, size: int) -> bytes: data = bytearray() while len(data) < size: chunk = self.handler.rfile.read(size - len(data)) if not chunk: raise EOFError("websocket peer disconnected") data.extend(chunk) return bytes(data) def receive(self) -> tuple[int, bytes]: """Read and reassemble one client message, answering ping inline.""" opcode = 0 payload = bytearray() while True: header = self._read_exact(2) first, second = header final = bool(first & 0x80) reserved = first & 0x70 frame_opcode = first & 0x0F masked = bool(second & 0x80) length = second & 0x7F if reserved or not masked: raise WebSocketError("invalid websocket frame") if length == 126: length = struct.unpack("!H", self._read_exact(2))[0] elif length == 127: length = struct.unpack("!Q", self._read_exact(8))[0] if frame_opcode >= 0x8 and (not final or length > 125): raise WebSocketError("invalid websocket control frame") if ( length > MAX_WS_MESSAGE_BYTES or len(payload) + length > MAX_WS_MESSAGE_BYTES ): raise WebSocketError("websocket message is too large", 1009) mask = self._read_exact(4) frame = bytearray(self._read_exact(length)) for index in range(length): frame[index] ^= mask[index % 4] if frame_opcode == 0x8: raise EOFError("websocket peer closed") if frame_opcode == 0x9: self.send(bytes(frame), opcode=0xA) continue if frame_opcode == 0xA: continue if frame_opcode in {0x1, 0x2}: if opcode: raise WebSocketError("nested websocket message") opcode = frame_opcode elif frame_opcode != 0x0 or not opcode: raise WebSocketError("unexpected websocket continuation") payload.extend(frame) if final: return opcode, bytes(payload) def send(self, payload: bytes, opcode: int) -> None: """Send one unmasked server frame while serializing background writers.""" with self._send_lock: if self._closed: return length = len(payload) header = bytearray([0x80 | opcode]) if length < 126: header.append(length) elif length <= 0xFFFF: header.extend((126,)) header.extend(struct.pack("!H", length)) else: header.extend((127,)) header.extend(struct.pack("!Q", length)) self.handler.wfile.write(bytes(header) + payload) self.handler.wfile.flush() def send_json(self, payload: dict) -> None: """Send a compact UTF-8 JSON message.""" self.send(json.dumps(payload, separators=(",", ":")).encode("utf-8"), 0x1) def close(self, code: int = 1000, reason: str = "") -> None: """Close once; callers may safely race disconnect and worker completion.""" reason_bytes = reason.encode("utf-8")[:123] with self._send_lock: if self._closed: return try: payload = struct.pack("!H", code) + reason_bytes self.handler.wfile.write(bytes([0x88, len(payload)]) + payload) self.handler.wfile.flush() except (BrokenPipeError, ConnectionError, OSError): pass self._closed = True def _pcm_rms(pcm: bytes) -> float: """Return normalized RMS for aligned signed-16-bit little-endian PCM.""" sample_count = len(pcm) // 2 if not sample_count: return 0.0 samples = struct.unpack(f"<{sample_count}h", pcm[: sample_count * 2]) square_mean = sum(sample * sample for sample in samples) / sample_count return (square_mean**0.5) / 32768.0 class StreamingTranscription: """Own one turn's bounded PCM, VAD state and speculative Whisper work.""" def __init__(self, connection: WebSocketConnection, model: object) -> None: self.connection = connection self.model = model self.turn_id = "" self.language = "auto" self._pcm = bytearray() self._lock = threading.Condition() self._started = False self._closed = False self._committing = False self._heard_speech = False self._at_eos = False self._silence_bytes = 0 self._last_speech_byte = 0 self._noise_floor = 0.004 self._client_active = False self._epoch = 0 self._pending_key = "" self._cached_key = "" self._cached_payload: dict | None = None def _response(self, message_type: str, **values: object) -> dict: return {"type": message_type, "turn_id": self.turn_id, **values} def error(self, message: str, *, fatal: bool = False) -> None: """Report a turn-scoped error without leaking implementation details.""" self.connection.send_json(self._response("error", error=message, fatal=fatal)) def start(self, message: dict) -> None: """Validate and initialize the only accepted browser audio contract.""" turn_id = str(message.get("turn_id") or "") language = str(message.get("language") or "auto").strip().lower() if self._started: raise WebSocketError("stream is already started", 1008) if not re.fullmatch(r"[A-Za-z0-9_.:-]{1,128}", turn_id): raise WebSocketError("invalid turn_id", 1008) if message.get("format") != "pcm_s16le": raise WebSocketError("pcm_s16le audio is required", 1003) try: sample_rate = int(message.get("sample_rate") or 0) except (TypeError, ValueError) as exc: raise WebSocketError("invalid sample rate", 1003) from exc if sample_rate != STREAM_SAMPLE_RATE: raise WebSocketError("16000 Hz audio is required", 1003) if language != "auto" and not re.fullmatch(r"[a-z]{2,3}", language): raise WebSocketError("invalid language", 1008) self.turn_id = turn_id self.language = language self._started = True def _snapshot_locked(self) -> tuple[bytes, str]: """Keep the full utterance and discard only bounded post-speech silence. Browser ``resume`` is authoritative: once sent, every following PCM frame is preserved until the next speculative EOS marker, even when a quiet syllable falls below the server's conservative energy threshold. """ cutoff = len(self._pcm) if self._heard_speech and self._last_speech_byte: tail_bytes = STREAM_SAMPLE_RATE * STREAM_SAMPLE_WIDTH * VAD_TAIL_MS // 1000 cutoff = min(cutoff, self._last_speech_byte + tail_bytes) pcm = bytes(self._pcm[:cutoff]) key = f"{self._epoch}:{hashlib.sha256(pcm).hexdigest()}" return pcm, key def append(self, pcm: bytes) -> None: """Append PCM, track server VAD and begin speculation at probable EOS.""" if not self._started: raise WebSocketError("start must precede audio", 1008) if not pcm or len(pcm) % STREAM_SAMPLE_WIDTH: raise WebSocketError("unaligned PCM audio", 1003) auto_speculate = False with self._lock: if self._closed or self._committing: raise WebSocketError("turn is no longer accepting audio", 1008) if len(self._pcm) + len(pcm) > MAX_STREAM_AUDIO_BYTES: raise WebSocketError("stream audio is too long", 1009) self._pcm.extend(pcm) rms = _pcm_rms(pcm) threshold = max(0.012, self._noise_floor * 2.5 + 0.004) speech = self._client_active or rms >= threshold if not self._heard_speech and not speech: self._noise_floor = self._noise_floor * 0.96 + rms * 0.04 if speech: if self._at_eos: self._epoch += 1 self._cached_key = "" self._cached_payload = None self._heard_speech = True self._at_eos = False self._silence_bytes = 0 self._last_speech_byte = len(self._pcm) elif self._heard_speech: self._silence_bytes += len(pcm) required = ( STREAM_SAMPLE_RATE * STREAM_SAMPLE_WIDTH * VAD_END_SILENCE_MS // 1000 ) if self._silence_bytes >= required and not self._at_eos: self._at_eos = True auto_speculate = True if auto_speculate: self.speculate() def resume(self) -> None: """Invalidate an EOS snapshot when client-side VAD hears resumed speech.""" with self._lock: if not self._started or self._closed: return self._at_eos = False self._silence_bytes = 0 self._epoch += 1 self._client_active = True self._cached_key = "" self._cached_payload = None def speculate(self) -> None: """Decode a stable snapshot once; only an unchanged turn may consume it.""" with self._lock: if not self._started or self._closed or self._committing or not self._pcm: return self._client_active = False self._at_eos = True pcm, key = self._snapshot_locked() if not pcm or key in {self._pending_key, self._cached_key}: return self._pending_key = key def worker() -> None: try: payload = _transcribe_pcm(self.model, pcm, self.language) except Exception as exc: print(f"[stt] speculative transcription failed: {exc}", flush=True) with self._lock: if self._pending_key == key: self._pending_key = "" self._lock.notify_all() if not self._closed and not self._committing: self.error("speculative transcription failed") return should_send = False with self._lock: if self._pending_key == key: self._pending_key = "" current_pcm, current_key = self._snapshot_locked() del current_pcm if not self._closed and current_key == key: self._cached_key = key self._cached_payload = payload should_send = not self._committing self._lock.notify_all() if should_send: try: self.connection.send_json( self._response( "partial", transcript=payload["text"], language=payload["language"], speculative=True, ) ) except (BrokenPipeError, ConnectionError, OSError): self.cancel() threading.Thread( target=worker, name=f"stt-speculate-{self.turn_id}", daemon=True, ).start() def commit(self) -> None: """Emit one final result from the complete stable utterance snapshot.""" with self._lock: if not self._started or self._closed or self._committing: raise WebSocketError("turn cannot be committed", 1008) if not self._pcm: raise WebSocketError("turn has no audio", 1008) self._committing = True pcm, key = self._snapshot_locked() deadline = time.monotonic() + STREAM_IDLE_SECONDS while self._pending_key == key and not self._cached_payload: remaining = deadline - time.monotonic() if remaining <= 0: break self._lock.wait(timeout=remaining) payload = self._cached_payload if self._cached_key == key else None try: if payload is None: payload = _transcribe_pcm(self.model, pcm, self.language) self.connection.send_json( self._response( "final", transcript=payload["text"], language=payload["language"], model=payload["model"], ) ) except Exception as exc: print(f"[stt] streaming transcription failed: {exc}", flush=True) self.error("transcription failed", fatal=True) finally: with self._lock: self._closed = True self._lock.notify_all() def cancel(self) -> None: """Invalidate queued work and prevent background writers after disconnect.""" with self._lock: self._closed = True self._cached_payload = None self._cached_key = "" self._lock.notify_all() 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: path = self.path.split("?", 1)[0] if path == "/v1/audio/transcriptions/stream": self._serve_stream() return if 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", "streaming": { "enabled": True, "path": "/v1/audio/transcriptions/stream", "format": "pcm_s16le", "sample_rate": STREAM_SAMPLE_RATE, "max_seconds": MAX_STREAM_SECONDS, "server_vad": True, "speculative": True, }, }, ) def _serve_stream(self) -> None: """Upgrade one request and process a single, isolated voice turn.""" upgrade = self.headers.get("Upgrade", "").strip().lower() connection_tokens = { token.strip().lower() for token in self.headers.get("Connection", "").split(",") } key = self.headers.get("Sec-WebSocket-Key", "").strip() version = self.headers.get("Sec-WebSocket-Version", "").strip() if upgrade != "websocket" or "upgrade" not in connection_tokens or not key: _json(self, 426, {"error": "websocket upgrade is required"}) return if version != "13": self.send_response(426) self.send_header("Sec-WebSocket-Version", "13") self.send_header("Content-Length", "0") self.end_headers() return try: decoded_key = base64.b64decode(key, validate=True) except (ValueError, TypeError): decoded_key = b"" if len(decoded_key) != 16: _json(self, 400, {"error": "invalid websocket key"}) return self.send_response(101, "Switching Protocols") self.send_header("Upgrade", "websocket") self.send_header("Connection", "Upgrade") self.send_header("Sec-WebSocket-Accept", WebSocketConnection.accept_value(key)) self.end_headers() self.wfile.flush() self.close_connection = True self.connection.settimeout(STREAM_IDLE_SECONDS) websocket = WebSocketConnection(self) session = StreamingTranscription( websocket, self.server.model, # type: ignore[attr-defined] ) close_code = 1000 close_reason = "" try: while True: opcode, payload = websocket.receive() if opcode == 0x2: session.append(payload) continue try: message = json.loads(payload.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError): raise WebSocketError("control message must be JSON", 1007) if not isinstance(message, dict): raise WebSocketError("control message must be an object", 1007) message_type = str(message.get("type") or "") if message_type == "start": session.start(message) continue if message.get("turn_id") != session.turn_id: raise WebSocketError("turn_id does not match", 1008) if message_type == "speculate": session.speculate() elif message_type == "resume": session.resume() elif message_type == "commit": session.commit() break elif message_type == "cancel": break else: raise WebSocketError("unknown stream control", 1008) except WebSocketError as exc: close_code = exc.close_code close_reason = str(exc) try: session.error(str(exc), fatal=True) except (BrokenPipeError, ConnectionError, OSError): pass except (EOFError, BrokenPipeError, ConnectionError, socket.timeout, OSError): close_reason = "peer disconnected" finally: session.cancel() websocket.close(close_code, close_reason) 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) payload = _transcribe_path( self.server.model, # type: ignore[attr-defined] temp_path, language, ) _json(self, 200, payload) 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()