Maven/deploy/cw2/serve.py

"""CrisperWhisper 2.0 turbo as an HTTP service, for Maven's stt.Pair.

Two endpoints and no framework.

  GET  /health      200 once the model is loaded, 503 while it is loading.
  POST /transcribe  raw 16kHz mono PCM in, {"text","confidence"} out.

The body is the PCM itself rather than JSON. A minute of 16kHz mono is under
2MB raw and about 2.6MB base64, and the format is fixed at the Maven seam, so
headers carry it more cheaply than an envelope.

Why this exists at all: whisper.cpp cannot load CW2. It derives its language
count from the vocabulary size, and CW2's 51897 tokens shift seven special
token ids. So mavsttd stays whisper.cpp on homesrv and this runs beside the
model on workpc, where it scores 10.4% WER in Russian against the floor's 27.5%
(docs/evals/2026-08-09-crisperwhisper2-russian-wer.md in the Maven repo).

Intended mode, not verbatim. The owner asked for what he meant to say, not
every stutter on the way there.
"""

import hmac
import json
import logging
import os
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import numpy as np

HOST = os.environ.get("CW2_HOST", "0.0.0.0")
PORT = int(os.environ.get("CW2_PORT", "8081"))
SIZE = os.environ.get("CW2_SIZE", "turbo")
MODE = os.environ.get("CW2_MODE", "intended")
TOKEN = os.environ.get("CW2_TOKEN", "")
# 25MB is about thirteen minutes of 16kHz mono. Longer than any utterance and
# short enough that a wrong caller cannot exhaust memory.
MAX_BODY = int(os.environ.get("CW2_MAX_BODY", str(25 * 1024 * 1024)))

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s cw2: %(message)s", stream=sys.stderr
)
log = logging.getLogger("cw2")

_model = None
# The card holds one model and transcribes one utterance at a time. The lock is
# what makes a second caller wait rather than corrupt the first.
_lock = threading.Lock()


def load_model():
    global _model
    from crisperwhisper import CrisperWhisperModel

    t0 = time.perf_counter()
    # backend is forced. With ctranslate2 importable, "auto" picks ct2, which is
    # CUDA-only and this card is AMD.
    m = CrisperWhisperModel(
        SIZE, backend="transformers", compute_type="float16", device="cuda"
    )
    _model = m
    log.info("loaded %s in %.1fs, mode=%s", SIZE, time.perf_counter() - t0, MODE)


def authorised(headers):
    if not TOKEN:
        return True
    got = headers.get("Authorization", "")
    return hmac.compare_digest(got, "Bearer " + TOKEN)


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, fmt, *args):
        log.info(fmt, *args)

    def _send(self, code, payload):
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        if self.path.rstrip("/") != "/health":
            self._send(404, {"error": "not found"})
            return
        if _model is None:
            self._send(503, {"status": "loading"})
            return
        self._send(200, {"status": "ok", "model": SIZE, "mode": MODE})

    def do_POST(self):
        if self.path.rstrip("/") != "/transcribe":
            self._send(404, {"error": "not found"})
            return
        if not authorised(self.headers):
            self._send(401, {"error": "unauthorised"})
            return
        if _model is None:
            self._send(503, {"error": "loading"})
            return

        length = int(self.headers.get("Content-Length", "0"))
        if length <= 0 or length > MAX_BODY:
            self._send(413, {"error": "bad body length"})
            return
        raw = self.rfile.read(length)

        rate = int(self.headers.get("X-Sample-Rate", "16000"))
        channels = int(self.headers.get("X-Channels", "1"))
        bits = int(self.headers.get("X-Sample-Bits", "16"))
        lang = self.headers.get("X-Language", "ru") or "ru"
        if channels != 1 or bits != 16:
            self._send(400, {"error": "want 16-bit mono pcm"})
            return

        # int16 little-endian to the float32 the encoder wants.
        wav = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
        if wav.size == 0:
            self._send(200, {"text": "", "confidence": 0.0})
            return

        t0 = time.perf_counter()
        try:
            with _lock:
                res = _model.transcribe(wav, sr=rate, language=lang, mode=MODE)
        except Exception as exc:  # noqa: BLE001 - the caller falls back to mavsttd
            log.exception("transcribe failed")
            self._send(500, {"error": str(exc)})
            return
        elapsed = time.perf_counter() - t0
        text = (res.text or "").strip()
        log.info("%.2fs audio in %.2fs: %r", wav.size / rate, elapsed, text[:60])
        # The model reports no calibrated score. 1.0 would be a claim, and the
        # Maven side reads confidence only to log it.
        self._send(200, {"text": text, "confidence": 0.0})


def main():
    if not TOKEN:
        log.warning("no CW2_TOKEN set: anything on the LAN can post audio here")
    # Bind before loading, so a restart answers 503 rather than refusing the
    # connection. Both make Maven fall back, but only one of them says why.
    srv = ThreadingHTTPServer((HOST, PORT), Handler)
    threading.Thread(target=load_model, daemon=True).start()
    log.info("listening on %s:%d", HOST, PORT)
    srv.serve_forever()


if __name__ == "__main__":
    main()