#!/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") ROLLING_MODEL_NAME = os.getenv("HERMES_STT_ROLLING_MODEL", "tiny") 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")) ROLLING_MIN_AUDIO_MS = max( 300, min(2_000, int(os.getenv("HERMES_STT_ROLLING_MIN_AUDIO_MS", "1800"))), ) ROLLING_INTERVAL_MS = max( 400, min(5_000, int(os.getenv("HERMES_STT_ROLLING_INTERVAL_MS", "750"))), ) ROLLING_WINDOW_MS = max( 3_000, min(30_000, int(os.getenv("HERMES_STT_ROLLING_WINDOW_MS", "12000"))), ) ROLLING_MAX_RESULT_LAG_MS = max( 500, min(5_000, int(os.getenv("HERMES_STT_ROLLING_MAX_LAG_MS", "2000"))), ) _WS_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" class InferenceGate: """Give authoritative transcripts priority over disposable partial work.""" def __init__(self) -> None: self._condition = threading.Condition() self._busy = False self._final_waiters = 0 self._last_final_wait_ms = 0.0 self._last_final_decode_ms = 0.0 self._last_background_decode_ms = 0.0 self._last_rolling_decode_ms = 0.0 self._last_speculative_decode_ms = 0.0 self._background_skips = 0 def run_final(self, callback): """Wait fairly enough that newly arriving partials cannot steal the GPU.""" wait_started = time.monotonic() with self._condition: self._final_waiters += 1 try: while self._busy: self._condition.wait() self._busy = True finally: self._final_waiters -= 1 self._last_final_wait_ms = (time.monotonic() - wait_started) * 1000 decode_started = time.monotonic() try: return callback() finally: with self._condition: self._last_final_decode_ms = (time.monotonic() - decode_started) * 1000 self._busy = False self._condition.notify_all() def try_background(self, callback, *, kind: str): """Run one useful predecode only when no final is active or waiting.""" if kind not in {"rolling", "speculative"}: raise ValueError("unknown background inference kind") with self._condition: if self._busy or self._final_waiters: self._background_skips += 1 return None self._busy = True decode_started = time.monotonic() try: return callback() finally: with self._condition: elapsed_ms = (time.monotonic() - decode_started) * 1000 self._last_background_decode_ms = elapsed_ms if kind == "rolling": self._last_rolling_decode_ms = elapsed_ms else: self._last_speculative_decode_ms = elapsed_ms self._busy = False self._condition.notify_all() def snapshot(self) -> dict: """Expose bounded timing telemetry for live latency verification.""" with self._condition: return { "busy": self._busy, "final_waiters": self._final_waiters, "last_final_wait_ms": round(self._last_final_wait_ms, 1), "last_final_decode_ms": round(self._last_final_decode_ms, 1), "last_background_decode_ms": round(self._last_background_decode_ms, 1), "last_rolling_decode_ms": round(self._last_rolling_decode_ms, 1), "last_speculative_decode_ms": round( self._last_speculative_decode_ms, 1 ), "background_skips": self._background_skips, } INFERENCE_GATE = InferenceGate() STARTUP_WARMUP_MS: float | None = None STARTUP_ROLLING_WARMUP_MS: float | None = None 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 _decode_path(model: object, path: str, language: str) -> dict: """Decode one path with the quality settings shared by partial and final work.""" 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_path(model: object, path: str, language: str) -> dict: """Run the one canonical full-utterance Whisper decode configuration.""" return INFERENCE_GATE.run_final(lambda: _decode_path(model, path, language)) def _write_pcm_wav(pcm: bytes) -> str: """Write an aligned browser PCM snapshot to a temporary Whisper input.""" 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 temp_path 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: temp_path = _write_pcm_wav(pcm) return _transcribe_path(model, temp_path, language) finally: if temp_path: try: os.unlink(temp_path) except OSError: pass def _warm_model(model: object) -> float: """Pay Whisper's lazy CUDA initialization cost before readiness.""" temp_path = _write_pcm_wav(bytes(_pcm_bytes_for_ms(1_000))) started = time.monotonic() try: _decode_path(model, temp_path, "auto") torch.cuda.synchronize() finally: try: os.unlink(temp_path) except OSError: pass return (time.monotonic() - started) * 1000 def _transcribe_pcm_rolling(model: object, pcm: bytes, language: str) -> dict | None: """Decode a rolling window only when the single GPU worker is immediately free. Partial work is disposable. It must never form an inference backlog ahead of a final utterance or another user's request, so a busy model means this speech epoch is skipped rather than retried. """ temp_path = "" try: temp_path = _write_pcm_wav(pcm) return INFERENCE_GATE.try_background( lambda: _decode_path(model, temp_path, language), kind="rolling", ) finally: if temp_path: try: os.unlink(temp_path) except OSError: pass def _transcribe_pcm_speculative( model: object, pcm: bytes, language: str ) -> dict | None: """Predecode one EOS snapshot only when the model is currently idle.""" temp_path = "" try: temp_path = _write_pcm_wav(pcm) return INFERENCE_GATE.try_background( lambda: _decode_path(model, temp_path, language), kind="speculative", ) finally: if temp_path: try: os.unlink(temp_path) except OSError: pass def _pcm_bytes_for_ms(milliseconds: int) -> int: """Convert milliseconds to an aligned byte count for the stream format.""" return ( STREAM_SAMPLE_RATE * STREAM_SAMPLE_WIDTH * milliseconds // 1000 // STREAM_SAMPLE_WIDTH * STREAM_SAMPLE_WIDTH ) def _token_identity(token: str) -> str: """Normalize punctuation/case drift before comparing consecutive decodes.""" return re.sub(r"[^\w']+", "", token.casefold(), flags=re.UNICODE) def _stable_token_prefix(previous: str, current: str) -> str: """Return only tokens repeated in the same order by consecutive decodes.""" old_tokens = previous.split() new_tokens = current.split() agreed = 0 for old, new in zip(old_tokens, new_tokens): if not _token_identity(old) or _token_identity(old) != _token_identity(new): break agreed += 1 return " ".join(new_tokens[:agreed]) 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, rolling_model: object | None = None, ) -> None: self.connection = connection self.model = model self.rolling_model = rolling_model if rolling_model is not None else 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 self._rolling_inflight = False self._rolling_last_started_at = 0.0 self._rolling_speech_bytes = 0 self._rolling_last_speech_bytes = 0 self._rolling_previous = "" self._rolling_previous_window_start = 0 self._rolling_revision = 0 self._rolling_attempted = False 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 _rolling_snapshot_locked(self) -> tuple[bytes, int]: """Return a bounded tail window and its absolute start in milliseconds.""" window_bytes = _pcm_bytes_for_ms(ROLLING_WINDOW_MS) start = max(0, len(self._pcm) - window_bytes) start -= start % STREAM_SAMPLE_WIDTH pcm = bytes(self._pcm[start:]) start_ms = start * 1000 // (STREAM_SAMPLE_RATE * STREAM_SAMPLE_WIDTH) return pcm, start_ms def _reset_rolling_history_locked(self, *, rearm: bool) -> None: """Discard provisional agreement after the client's speech epoch changes.""" self._rolling_previous = "" self._rolling_previous_window_start = 0 # Only speech in the new epoch may satisfy the rolling threshold. self._rolling_speech_bytes = 0 self._rolling_last_speech_bytes = 0 self._rolling_last_started_at = 0.0 if rearm: self._rolling_attempted = False def _begin_rolling_locked(self, now: float) -> tuple[bytes, int, int, int] | None: """Reserve one due rolling decode without ever queuing a second one.""" total_bytes = len(self._pcm) minimum_bytes = _pcm_bytes_for_ms(ROLLING_MIN_AUDIO_MS) interval_bytes = _pcm_bytes_for_ms(ROLLING_INTERVAL_MS) required_new_bytes = ( interval_bytes if self._rolling_last_speech_bytes else minimum_bytes ) if ( self._rolling_inflight or self._rolling_attempted or self._rolling_speech_bytes < minimum_bytes or self._rolling_speech_bytes - self._rolling_last_speech_bytes < required_new_bytes or ( self._rolling_last_started_at and (now - self._rolling_last_started_at) * 1000 < ROLLING_INTERVAL_MS ) ): return None pcm, window_start_ms = self._rolling_snapshot_locked() # Reserve the epoch before spawning. A busy gate, decode error or stale # result must not create a retry loop that can collide with the final. self._rolling_attempted = True self._rolling_inflight = True self._rolling_last_started_at = now self._rolling_last_speech_bytes = self._rolling_speech_bytes return pcm, window_start_ms, total_bytes, self._epoch def _start_rolling( self, pcm: bytes, window_start_ms: int, audio_bytes: int, epoch: int, ) -> None: """Run one disposable partial decode and publish only a current result.""" def worker() -> None: payload: dict | None = None try: payload = _transcribe_pcm_rolling( self.rolling_model, pcm, self.language, ) except Exception as exc: # Partial inference is best-effort. The canonical commit path # remains authoritative and reports its own actionable errors. print(f"[stt] rolling transcription skipped: {exc}", flush=True) message: dict | None = None with self._lock: self._rolling_inflight = False lag_bytes = max(0, len(self._pcm) - audio_bytes) max_lag_bytes = _pcm_bytes_for_ms(ROLLING_MAX_RESULT_LAG_MS) current = ( payload is not None and not self._closed and not self._committing and not self._at_eos and self._epoch == epoch and lag_bytes <= max_lag_bytes ) if current: transcript = str(payload.get("text") or "") stable = "" if self._rolling_previous_window_start == window_start_ms: stable = _stable_token_prefix( self._rolling_previous, transcript, ) self._rolling_previous = transcript self._rolling_previous_window_start = window_start_ms self._rolling_revision += 1 message = self._response( "partial", transcript=transcript, stable_transcript=stable, language=payload.get("language") or "", speculative=True, rolling=True, revision=self._rolling_revision, epoch=epoch, window_start_ms=window_start_ms, ) self._lock.notify_all() if message is not None: try: self.connection.send_json(message) except (BrokenPipeError, ConnectionError, OSError): self.cancel() threading.Thread( target=worker, name=f"stt-rolling-{self.turn_id}", daemon=True, ).start() def append(self, pcm: bytes) -> None: """Append PCM while overlapping rolling understanding with active speech.""" 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 rolling: tuple[bytes, int, int, int] | None = None 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._reset_rolling_history_locked(rearm=True) self._heard_speech = True self._at_eos = False self._silence_bytes = 0 self._last_speech_byte = len(self._pcm) self._rolling_speech_bytes += len(pcm) rolling = self._begin_rolling_locked(time.monotonic()) 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 self._reset_rolling_history_locked(rearm=False) auto_speculate = True if auto_speculate: self.speculate() elif rolling is not None: self._start_rolling(*rolling) 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 self._reset_rolling_history_locked(rearm=True) 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_speculative(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 if payload is None: with self._lock: if self._pending_key == key: self._pending_key = "" self._lock.notify_all() 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, "rolling_model": ROLLING_MODEL_NAME, "device": "cuda" if torch.cuda.is_available() else "cpu", "inference": INFERENCE_GATE.snapshot(), "startup": { "warmed": ( STARTUP_WARMUP_MS is not None and STARTUP_ROLLING_WARMUP_MS is not None ), "warmup_ms": round(STARTUP_WARMUP_MS or 0.0, 1), "rolling_warmup_ms": round( STARTUP_ROLLING_WARMUP_MS or 0.0, 1, ), }, "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, "rolling": { "enabled": True, "minimum_audio_ms": ROLLING_MIN_AUDIO_MS, "interval_ms": ROLLING_INTERVAL_MS, "window_ms": ROLLING_WINDOW_MS, "stable_prefix": 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] self.server.rolling_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.""" global STARTUP_ROLLING_WARMUP_MS, STARTUP_WARMUP_MS 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} for final CUDA", flush=True) model = whisper.load_model(MODEL_NAME, device="cuda", download_root=str(CACHE_DIR)) print(f"[stt] loading Whisper {ROLLING_MODEL_NAME} for rolling CUDA", flush=True) rolling_model = whisper.load_model( ROLLING_MODEL_NAME, device="cuda", download_root=str(CACHE_DIR), ) STARTUP_ROLLING_WARMUP_MS = _warm_model(rolling_model) print( "[stt] rolling CUDA warm-up completed in " f"{STARTUP_ROLLING_WARMUP_MS:.1f}ms", flush=True, ) # Warm the authoritative model last so the first real final never inherits # lazy CUDA work from either model's startup path. STARTUP_WARMUP_MS = _warm_model(model) print( f"[stt] final CUDA warm-up completed in {STARTUP_WARMUP_MS:.1f}ms", flush=True ) server = ThreadingHTTPServer((HOST, PORT), SpeechHandler) server.model = model # type: ignore[attr-defined] server.rolling_model = rolling_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()