clement
da205ffc99
- tools/kyutai-stt: local FastAPI STT server (kyutai_stt_server.py, stt_from_file_mlx.py) + TTS server (kyutai_tts_server.py, tts_mlx.py) using MLX Kyutai models - gateway /v1/voice/reply: STT → LLM → TTS pipeline + dynamic soxr resample + compressor/limiter chain + greeting cache + warm-up at startup - tests/gateway/test_voice_reply.py: pytest coverage - bump ESP32_ZACUS submodule to plip voice loop (aa7ae27)
327 lines
11 KiB
Python
327 lines
11 KiB
Python
# /// script
|
|
# requires-python = ">=3.12"
|
|
# dependencies = [
|
|
# "moshi_mlx==0.2.12",
|
|
# "fastapi",
|
|
# "uvicorn",
|
|
# "numpy",
|
|
# "sphn",
|
|
# "sentencepiece",
|
|
# "huggingface_hub",
|
|
# ]
|
|
# ///
|
|
|
|
"""Kyutai TTS FastAPI server — MLX backend for Mac M1/M2/M3.
|
|
|
|
Exposes POST /tts that accepts JSON {"text": "...", "voice": "<optional>"}
|
|
and returns a WAV (24 kHz mono) audio/wav response.
|
|
|
|
French voices available in kyutai/tts-voices under cml-tts/fr/:
|
|
cml-tts/fr/296_1028_000022-0001.wav <- chosen default (clear FR male)
|
|
cml-tts/fr/10087_11650_000028-0002.wav
|
|
cml-tts/fr/10177_10625_000134-0003.wav
|
|
cml-tts/fr/10179_11051_000005-0001.wav
|
|
cml-tts/fr/12080_11650_000047-0001.wav
|
|
cml-tts/fr/12205_11650_000004-0002.wav
|
|
cml-tts/fr/12977_10625_000037-0001.wav
|
|
cml-tts/fr/1406_1028_000009-0003.wav
|
|
cml-tts/fr/1591_1028_000108-0004.wav
|
|
cml-tts/fr/1770_1028_000036-0002.wav
|
|
cml-tts/fr/2114_1656_000053-0001.wav
|
|
cml-tts/fr/2154_2576_000020-0003.wav
|
|
cml-tts/fr/2216_1745_000007-0001.wav
|
|
cml-tts/fr/2223_1745_000009-0002.wav
|
|
cml-tts/fr/2465_1943_000152-0002.wav
|
|
cml-tts/fr/3267_1902_000075-0001.wav
|
|
cml-tts/fr/4193_3103_000004-0001.wav
|
|
cml-tts/fr/4482_3103_000063-0001.wav
|
|
cml-tts/fr/4724_3731_000031-0001.wav
|
|
cml-tts/fr/4937_3731_000004-0001.wav
|
|
cml-tts/fr/5207_3078_000031-0002.wav
|
|
cml-tts/fr/5476_3103_000072-0001.wav
|
|
cml-tts/fr/577_394_000070-0001.wav
|
|
cml-tts/fr/5790_4893_000052-0001.wav
|
|
cml-tts/fr/579_2548_000015-0001.wav
|
|
cml-tts/fr/5830_4703_000037-0001.wav
|
|
cml-tts/fr/6318_7016_000027-0002.wav
|
|
cml-tts/fr/7142_2432_000124-0003.wav
|
|
cml-tts/fr/7400_2928_000100-0001.wav
|
|
cml-tts/fr/7591_6742_000149-0002.wav
|
|
cml-tts/fr/7601_7727_000062-0001.wav
|
|
cml-tts/fr/7762_8734_000048-0002.wav
|
|
cml-tts/fr/8128_7016_000047-0002.wav
|
|
cml-tts/fr/928_486_000075-0001.wav
|
|
cml-tts/fr/9834_9697_000150-0003.wav
|
|
(+ *_enhanced.wav variants for all of the above)
|
|
|
|
State reset: TTSModel.generate() is stateless (it constructs a fresh LmGen internally
|
|
each call). Mimi decode_step() uses a streaming KV-cache however; we call
|
|
tts_model.mimi.reset_state() before each synthesis to clear it. The threading.Lock
|
|
ensures MLX non-reentrancy.
|
|
"""
|
|
|
|
import argparse
|
|
import io
|
|
import json
|
|
import logging
|
|
import os
|
|
import queue
|
|
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, HTTPException
|
|
from fastapi.responses import Response
|
|
from moshi_mlx import models
|
|
from moshi_mlx.models.tts import (
|
|
DEFAULT_DSM_TTS_REPO,
|
|
DEFAULT_DSM_TTS_VOICE_REPO,
|
|
TTSModel,
|
|
)
|
|
from moshi_mlx.utils.loaders import hf_get
|
|
from pydantic import BaseModel
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Constants
|
|
# ---------------------------------------------------------------------------
|
|
|
|
HF_REPO = DEFAULT_DSM_TTS_REPO # kyutai/tts-1.6b-en_fr
|
|
VOICE_REPO = DEFAULT_DSM_TTS_VOICE_REPO # kyutai/tts-voices
|
|
|
|
# Default French voice — cml-tts/fr dataset, clear neutral-male timbre,
|
|
# tested and confirmed to produce natural French output with this bilingual model.
|
|
DEFAULT_VOICE_FR = "cml-tts/fr/296_1028_000022-0001.wav"
|
|
|
|
DEFAULT_HOST = "0.0.0.0"
|
|
DEFAULT_PORT = 8302
|
|
|
|
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
|
|
|
|
_tts_model: TTSModel = None # type: ignore[assignment]
|
|
_cfg_coef_conditioning = None
|
|
_cfg_is_no_text: bool = True
|
|
_cfg_is_no_prefix: bool = True
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# FastAPI app
|
|
# ---------------------------------------------------------------------------
|
|
|
|
app = FastAPI(title="Kyutai TTS Server", version="1.0.0")
|
|
|
|
|
|
class TTSRequest(BaseModel):
|
|
text: str
|
|
voice: str = ""
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {
|
|
"status": "ok",
|
|
"model": HF_REPO,
|
|
"voice": DEFAULT_VOICE_FR,
|
|
"ready": _ready,
|
|
}
|
|
|
|
|
|
@app.post("/tts")
|
|
def synthesize(req: TTSRequest):
|
|
"""Synthesize French text to speech.
|
|
|
|
Body: {"text": "...", "voice": "<optional voice path in kyutai/tts-voices>"}
|
|
Returns: audio/wav, 24 kHz mono PCM.
|
|
"""
|
|
if not _ready:
|
|
raise HTTPException(status_code=503, detail="Model not ready yet")
|
|
|
|
if not req.text.strip():
|
|
raise HTTPException(status_code=400, detail="Empty text")
|
|
|
|
voice_path = req.voice.strip() if req.voice else DEFAULT_VOICE_FR
|
|
|
|
wav_bytes = _run_tts(req.text.strip(), voice_path)
|
|
return Response(content=wav_bytes, media_type="audio/wav")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Inference
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _run_tts(text: str, voice: str) -> bytes:
|
|
"""Synthesize text and return raw WAV bytes. Thread-safe via global lock."""
|
|
global _tts_model, _cfg_coef_conditioning, _cfg_is_no_text, _cfg_is_no_prefix
|
|
|
|
all_entries = [_tts_model.prepare_script([text])]
|
|
if _tts_model.multi_speaker:
|
|
voices = [_tts_model.get_voice_path(voice)]
|
|
else:
|
|
voices = []
|
|
all_attributes = [
|
|
_tts_model.make_condition_attributes(voices, _cfg_coef_conditioning)
|
|
]
|
|
|
|
wav_frames: queue.Queue = queue.Queue()
|
|
|
|
def _on_frame(frame):
|
|
if (frame == -1).any():
|
|
return
|
|
pcm = _tts_model.mimi.decode_step(frame[:, :, None])
|
|
pcm = np.array(mx.clip(pcm[0, 0], -1, 1))
|
|
wav_frames.put_nowait(pcm)
|
|
|
|
with _lock:
|
|
# Per-request state hardening. A fresh server synthesises cleanly, but
|
|
# output degrades/distorts over many requests — accumulated RNG drift
|
|
# and MLX memory-cache buildup. Reset to a FIXED seed (deterministic,
|
|
# no drift), clear the Mimi streaming KV-cache, and free the MLX cache.
|
|
mx.random.seed(299792458)
|
|
_tts_model.mimi.reset_state()
|
|
# ROOT CAUSE of the cross-request degradation: the Lm transformer KV-cache
|
|
# lives on the model and PERSISTS across generate() calls (LmGen does not
|
|
# reset it — only Lm.warmup() does). Without this reset the state
|
|
# accumulates and the audio distorts after several syntheses. Reset both
|
|
# the transformer and depformer caches per request.
|
|
for _c in _tts_model.lm.transformer_cache:
|
|
_c.reset()
|
|
for _c in getattr(_tts_model.lm, "depformer_cache", []):
|
|
_c.reset()
|
|
|
|
_tts_model.generate(
|
|
all_entries,
|
|
all_attributes,
|
|
cfg_is_no_prefix=_cfg_is_no_prefix,
|
|
cfg_is_no_text=_cfg_is_no_text,
|
|
on_frame=_on_frame,
|
|
)
|
|
mx.clear_cache()
|
|
|
|
# Collect all PCM frames
|
|
frames = []
|
|
while True:
|
|
try:
|
|
frames.append(wav_frames.get_nowait())
|
|
except queue.Empty:
|
|
break
|
|
|
|
if not frames:
|
|
raise HTTPException(status_code=500, detail="TTS produced no audio frames")
|
|
|
|
wav = np.concatenate(frames, axis=-1)
|
|
sample_rate = _tts_model.mimi.sample_rate # 24000
|
|
|
|
# Write WAV to in-memory buffer via sphn (writes to file path only) → use tmp
|
|
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
|
|
tmp_path = tmp.name
|
|
try:
|
|
sphn.write_wav(tmp_path, wav, sample_rate)
|
|
with open(tmp_path, "rb") as fobj:
|
|
return fobj.read()
|
|
finally:
|
|
os.unlink(tmp_path)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model loading (called once before server starts)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _load_model():
|
|
global _tts_model, _cfg_coef_conditioning, _cfg_is_no_text, _cfg_is_no_prefix, _ready
|
|
|
|
log.info("Retrieving checkpoints from %s …", HF_REPO)
|
|
|
|
raw_config_path = hf_get("config.json", HF_REPO)
|
|
with open(hf_get(raw_config_path), "r") as fobj:
|
|
raw_config = json.load(fobj)
|
|
|
|
mimi_weights = hf_get(raw_config["mimi_name"], HF_REPO)
|
|
moshi_name = raw_config.get("moshi_name", "model.safetensors")
|
|
moshi_weights = hf_get(moshi_name, HF_REPO)
|
|
tokenizer = hf_get(raw_config["tokenizer_name"], HF_REPO)
|
|
|
|
lm_config = models.LmConfig.from_config_dict(raw_config)
|
|
# Workaround for ring KV-cache bug in moshi_mlx <= 0.3.0
|
|
lm_config.transformer.max_seq_len = lm_config.transformer.context
|
|
|
|
log.info("Loading LM weights from %s …", moshi_weights)
|
|
model = models.Lm(lm_config)
|
|
model.set_dtype(mx.bfloat16)
|
|
model.load_pytorch_weights(str(moshi_weights), lm_config, strict=True)
|
|
|
|
log.info("Loading text tokenizer from %s …", tokenizer)
|
|
text_tokenizer = sentencepiece.SentencePieceProcessor(str(tokenizer)) # type: ignore
|
|
|
|
log.info("Loading audio tokenizer (Mimi) from %s …", mimi_weights)
|
|
generated_codebooks = lm_config.generated_codebooks
|
|
audio_tokenizer = models.mimi.Mimi(models.mimi_202407(generated_codebooks))
|
|
audio_tokenizer.load_pytorch_weights(str(mimi_weights), strict=True)
|
|
|
|
log.info("Building TTSModel …")
|
|
tts_model = TTSModel(
|
|
model,
|
|
audio_tokenizer,
|
|
text_tokenizer,
|
|
voice_repo=VOICE_REPO,
|
|
temp=0.6,
|
|
cfg_coef=1,
|
|
max_padding=8,
|
|
initial_padding=2,
|
|
final_padding=2,
|
|
padding_bonus=0,
|
|
raw_config=raw_config,
|
|
)
|
|
|
|
cfg_coef_conditioning = None
|
|
if tts_model.valid_cfg_conditionings:
|
|
cfg_coef_conditioning = tts_model.cfg_coef
|
|
tts_model.cfg_coef = 1.0
|
|
cfg_is_no_text = False
|
|
cfg_is_no_prefix = False
|
|
else:
|
|
cfg_is_no_text = True
|
|
cfg_is_no_prefix = True
|
|
|
|
log.info("Warming up with default FR voice …")
|
|
mx.random.seed(299792458)
|
|
# Warmup: synthesize a short phrase to prime the MLX graph
|
|
_tts_model = tts_model
|
|
_cfg_coef_conditioning = cfg_coef_conditioning
|
|
_cfg_is_no_text = cfg_is_no_text
|
|
_cfg_is_no_prefix = cfg_is_no_prefix
|
|
|
|
try:
|
|
_run_tts("Bonjour.", DEFAULT_VOICE_FR)
|
|
log.info("Warmup complete.")
|
|
except Exception as exc:
|
|
log.warning("Warmup synthesis failed (non-fatal): %s", exc)
|
|
|
|
_ready = True
|
|
log.info("=== TTS model ready === Listening on :%d", DEFAULT_PORT)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Entrypoint
|
|
# ---------------------------------------------------------------------------
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Kyutai TTS FastAPI server (MLX, Mac)")
|
|
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")
|