le-mystere-professeur-zacus/tools/kyutai-stt/kyutai_stt_server.py

# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "moshi_mlx==0.2.12",
#     "fastapi",
#     "uvicorn",
#     "sphn",
#     "sentencepiece",
#     "huggingface_hub",
#     "numpy",
#     "python-multipart",
# ]
# ///

import argparse
import json
import logging
import os
import tempfile
import threading

import mlx.core as mx
import mlx.nn as nn
import numpy as np
import sentencepiece
import sphn
import uvicorn
from fastapi import FastAPI, File, HTTPException, Request, UploadFile
from fastapi.responses import JSONResponse
from huggingface_hub import hf_hub_download
from moshi_mlx import models, utils

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

HF_REPO = "kyutai/stt-1b-en_fr-mlx"
DEFAULT_HOST = "0.0.0.0"
DEFAULT_PORT = 8300

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Global model state (loaded once at startup)
# ---------------------------------------------------------------------------

_lock = threading.Lock()
_ready = False

_lm_model: models.Lm = None          # type: ignore[assignment]
_lm_config = None
_audio_tokenizer: models.mimi.Mimi = None   # type: ignore[assignment]
_text_tokenizer: sentencepiece.SentencePieceProcessor = None   # type: ignore[assignment]

# ---------------------------------------------------------------------------
# FastAPI app
# ---------------------------------------------------------------------------

app = FastAPI(title="Kyutai STT Server", version="1.0.0")


@app.get("/health")
def health():
    return {"status": "ok", "model": HF_REPO, "ready": _ready}


@app.post("/transcribe")
async def transcribe(request: Request, file: UploadFile = File(default=None)):
    """Accept WAV either as raw body (Content-Type: audio/wav) or multipart field 'file'."""
    if not _ready:
        raise HTTPException(status_code=503, detail="Model not ready yet")

    content_type = request.headers.get("content-type", "")

    if file is not None:
        # multipart/form-data upload
        audio_bytes = await file.read()
    elif "audio" in content_type or "octet-stream" in content_type or "multipart" not in content_type:
        # Raw body (Content-Type: audio/wav or application/octet-stream)
        audio_bytes = await request.body()
    else:
        raise HTTPException(
            status_code=400,
            detail="Send WAV as raw body (Content-Type: audio/wav) or multipart field 'file'",
        )

    if not audio_bytes:
        raise HTTPException(status_code=400, detail="Empty audio body")

    # Write to a temp file because sphn.read requires a file path
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
        tmp.write(audio_bytes)
        tmp_path = tmp.name

    try:
        text = _run_inference(tmp_path)
    finally:
        os.unlink(tmp_path)

    return JSONResponse({"text": text})


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

def _run_inference(wav_path: str) -> str:
    """Run STT inference on a WAV file. Thread-safe via global lock."""
    global _audio_tokenizer

    audio, _ = sphn.read(wav_path, sample_rate=24000)
    # Pad with 0.5 s of silence so the model can flush its final tokens. Callers
    # (the gateway) already append ~0.8 s of trailing silence, so a short pad here
    # is enough to flush — and every padded second costs ~1.5x realtime of STT
    # latency, which dominates the conversation turn time.
    audio = np.concatenate([audio, np.zeros((1, 12000), dtype=audio.dtype)], axis=-1)
    audio_mx = mx.array(audio)

    with _lock:
        # --- Reset stateful objects per request ---
        # Mimi has reset_state() to clear its streaming KV-cache without reloading weights.
        _audio_tokenizer.reset_state()

        # The Lm holds a PERSISTENT rotating KV-cache (`transformer_cache`) shared by
        # every LmGen. LmGen.__init__ resets only its own gen_sequence/step_idx, NOT
        # this cache — so across requests it keeps accumulating stream positions while
        # each new LmGen restarts at step_idx=0. The positions drift out of alignment
        # and the cache saturates, until the model emits only padding tokens (0/3) and
        # every transcription comes back empty. Reset it per request, exactly as
        # Lm.warmup() does at startup. (depformer_cache reset defensively too.)
        for c in _lm_model.transformer_cache:
            c.reset()
        for c in _lm_model.depformer_cache:
            c.reset()

        # LmGen is cheap to construct (no weight loading, just references the model).
        gen = models.LmGen(
            model=_lm_model,
            max_steps=4096,
            text_sampler=utils.Sampler(top_k=25, temp=0),
            audio_sampler=utils.Sampler(top_k=250, temp=0.8),
            check=False,
        )

        pieces = []
        for start_idx in range(0, audio_mx.shape[-1] // 1920 * 1920, 1920):
            block = audio_mx[:, None, start_idx : start_idx + 1920]
            other_audio_tokens = _audio_tokenizer.encode_step(block).transpose(0, 2, 1)
            text_token = gen.step(other_audio_tokens[0])
            text_token = text_token[0].item()
            if text_token not in (0, 3):
                piece = _text_tokenizer.id_to_piece(text_token)  # type: ignore[arg-type]
                piece = piece.replace("▁", " ")
                pieces.append(piece)

    return "".join(pieces).strip()


# ---------------------------------------------------------------------------
# Model loading (called once before server starts)
# ---------------------------------------------------------------------------

def _load_model():
    global _lm_model, _lm_config, _audio_tokenizer, _text_tokenizer, _ready

    log.info("Downloading / locating model files for %s …", HF_REPO)

    lm_config_path = hf_hub_download(HF_REPO, "config.json")
    with open(lm_config_path) as fobj:
        lm_config_dict = json.load(fobj)

    mimi_weights = hf_hub_download(HF_REPO, lm_config_dict["mimi_name"])
    moshi_name = lm_config_dict.get("moshi_name", "model.safetensors")
    moshi_weights = hf_hub_download(HF_REPO, moshi_name)
    text_tokenizer_path = hf_hub_download(HF_REPO, lm_config_dict["tokenizer_name"])

    _lm_config = models.LmConfig.from_config_dict(lm_config_dict)

    log.info("Loading LM weights from %s …", moshi_weights)
    model = models.Lm(_lm_config)
    model.set_dtype(mx.bfloat16)
    if moshi_weights.endswith(".q4.safetensors"):
        nn.quantize(model, bits=4, group_size=32)
    elif moshi_weights.endswith(".q8.safetensors"):
        nn.quantize(model, bits=8, group_size=64)
    model.load_weights(moshi_weights, strict=True)
    _lm_model = model

    log.info("Loading text tokenizer from %s …", text_tokenizer_path)
    _text_tokenizer = sentencepiece.SentencePieceProcessor(text_tokenizer_path)  # type: ignore[call-arg]

    log.info("Loading audio tokenizer (Mimi) from %s …", mimi_weights)
    _audio_tokenizer = models.mimi.Mimi(models.mimi_202407(32))
    _audio_tokenizer.load_pytorch_weights(str(mimi_weights), strict=True)

    log.info("Warming up the model …")
    _lm_model.warmup()

    _ready = True
    log.info("=== model ready — listening for /transcribe requests ===")


# ---------------------------------------------------------------------------
# Entrypoint
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Kyutai STT FastAPI server")
    parser.add_argument("--host", default=DEFAULT_HOST)
    parser.add_argument("--port", type=int, default=DEFAULT_PORT)
    args = parser.parse_args()

    _load_model()

    uvicorn.run(app, host=args.host, port=args.port, workers=1, log_level="info")