feat(gateway): Kyutai STT/TTS + voice reply loop
- 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)
This commit was merged in pull request #162.
This commit is contained in:
+1
-1
Submodule ESP32_ZACUS updated: 37db47ad7b...aa7ae277ed
@@ -0,0 +1,126 @@
|
||||
import io
|
||||
import struct
|
||||
import sys
|
||||
import wave
|
||||
from pathlib import Path
|
||||
|
||||
REPO = Path(__file__).resolve().parents[2]
|
||||
sys.path.insert(0, str(REPO / "tools" / "zacus-gateway"))
|
||||
import main
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
|
||||
def _make_wav_16k(duration_s: float = 0.1) -> bytes:
|
||||
"""Build a minimal valid WAV: 16-bit PCM, 16000 Hz, mono."""
|
||||
n = int(16000 * duration_s)
|
||||
pcm = struct.pack(f"<{n}h", *([0] * n))
|
||||
buf = io.BytesIO()
|
||||
with wave.open(buf, "wb") as w:
|
||||
w.setnchannels(1)
|
||||
w.setsampwidth(2)
|
||||
w.setframerate(16000)
|
||||
w.writeframes(pcm)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
_FAKE_WAV = _make_wav_16k()
|
||||
|
||||
|
||||
def _patch_backends(monkeypatch, *, heard="Bonjour, j'ai une question.", said="Allô, oui ?"):
|
||||
async def fake_transcribe(http, audio_wav):
|
||||
return heard
|
||||
|
||||
async def fake_chat_reply(http, persona, history):
|
||||
return said
|
||||
|
||||
async def fake_voice_tts_16k(http, text, voice):
|
||||
return _FAKE_WAV, 0.1, False
|
||||
|
||||
monkeypatch.setattr(main, "_transcribe_kyutai", fake_transcribe)
|
||||
monkeypatch.setattr(main, "_chat_reply", fake_chat_reply)
|
||||
monkeypatch.setattr(main, "_voice_tts_16k", fake_voice_tts_16k)
|
||||
|
||||
|
||||
def test_voice_reply_200(monkeypatch):
|
||||
_patch_backends(monkeypatch, heard="C'est une urgence.", said="J'arrive tout de suite.")
|
||||
with TestClient(main.app) as client:
|
||||
resp = client.post(
|
||||
"/v1/voice/reply",
|
||||
data={"session_id": "r1", "number": "17"},
|
||||
files={"audio": ("rec.wav", _FAKE_WAV, "audio/wav")},
|
||||
headers={"Authorization": f"Bearer {main.settings.token}"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert resp.headers["content-type"] == "audio/wav"
|
||||
assert resp.headers["x-zacus-heard"] == "C'est une urgence."
|
||||
assert resp.headers["x-zacus-said"] == "J'arrive tout de suite."
|
||||
assert resp.content[:4] == b"RIFF"
|
||||
|
||||
|
||||
def test_voice_reply_appends_transcription_to_session(monkeypatch):
|
||||
_patch_backends(monkeypatch, heard="Le code est 1234.", said="Bien noté.")
|
||||
sid = "r-session"
|
||||
main.VOICE_SESSIONS.end(sid) # start clean
|
||||
with TestClient(main.app) as client:
|
||||
resp = client.post(
|
||||
"/v1/voice/reply",
|
||||
data={"session_id": sid, "number": "17"},
|
||||
files={"audio": ("rec.wav", _FAKE_WAV, "audio/wav")},
|
||||
headers={"Authorization": f"Bearer {main.settings.token}"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
history = main.VOICE_SESSIONS.history(sid)
|
||||
roles = [(m["role"], m["content"]) for m in history]
|
||||
assert ("user", "Le code est 1234.") in roles
|
||||
assert ("assistant", "Bien noté.") in roles
|
||||
main.VOICE_SESSIONS.end(sid)
|
||||
|
||||
|
||||
def test_voice_reply_unknown_number_404(monkeypatch):
|
||||
_patch_backends(monkeypatch)
|
||||
with TestClient(main.app) as client:
|
||||
resp = client.post(
|
||||
"/v1/voice/reply",
|
||||
data={"session_id": "r2", "number": "99"},
|
||||
files={"audio": ("rec.wav", _FAKE_WAV, "audio/wav")},
|
||||
headers={"Authorization": f"Bearer {main.settings.token}"},
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
def test_voice_reply_stt_down_returns_502(monkeypatch):
|
||||
async def failing_transcribe(http, audio_wav):
|
||||
raise RuntimeError("connection refused")
|
||||
|
||||
async def fake_chat_reply(http, persona, history):
|
||||
return "..."
|
||||
|
||||
async def fake_voice_tts_16k(http, text, voice):
|
||||
return _FAKE_WAV, 0.1, False
|
||||
|
||||
monkeypatch.setattr(main, "_transcribe_kyutai", failing_transcribe)
|
||||
monkeypatch.setattr(main, "_chat_reply", fake_chat_reply)
|
||||
monkeypatch.setattr(main, "_voice_tts_16k", fake_voice_tts_16k)
|
||||
|
||||
with TestClient(main.app) as client:
|
||||
resp = client.post(
|
||||
"/v1/voice/reply",
|
||||
data={"session_id": "r502", "number": "17"},
|
||||
files={"audio": ("rec.wav", _FAKE_WAV, "audio/wav")},
|
||||
headers={"Authorization": f"Bearer {main.settings.token}"},
|
||||
)
|
||||
assert resp.status_code == 502
|
||||
detail = resp.json().get("detail", "")
|
||||
assert "voice backend unreachable" in detail
|
||||
assert "Traceback" not in detail
|
||||
|
||||
|
||||
def test_voice_reply_requires_token():
|
||||
with TestClient(main.app) as client:
|
||||
resp = client.post(
|
||||
"/v1/voice/reply",
|
||||
data={"session_id": "r3", "number": "17"},
|
||||
files={"audio": ("rec.wav", _FAKE_WAV, "audio/wav")},
|
||||
)
|
||||
assert resp.status_code in (401, 403)
|
||||
@@ -0,0 +1,201 @@
|
||||
# /// 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 2 s of silence so the model can flush its final tokens
|
||||
audio = np.concatenate([audio, np.zeros((1, 48000), 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()
|
||||
|
||||
# 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")
|
||||
@@ -0,0 +1,326 @@
|
||||
# /// 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")
|
||||
@@ -0,0 +1,100 @@
|
||||
# /// script
|
||||
# requires-python = ">=3.12"
|
||||
# dependencies = [
|
||||
# "huggingface_hub",
|
||||
# "moshi_mlx==0.2.12",
|
||||
# "numpy",
|
||||
# "sentencepiece",
|
||||
# "sounddevice",
|
||||
# "sphn",
|
||||
# ]
|
||||
# ///
|
||||
|
||||
import argparse
|
||||
import json
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
import sentencepiece
|
||||
import sphn
|
||||
from huggingface_hub import hf_hub_download
|
||||
from moshi_mlx import models, utils
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("in_file", help="The file to transcribe.")
|
||||
parser.add_argument("--max-steps", default=4096)
|
||||
parser.add_argument("--hf-repo")
|
||||
parser.add_argument(
|
||||
"--vad", action="store_true", help="Enable VAD (Voice Activity Detection)."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
audio, _ = sphn.read(args.in_file, sample_rate=24000)
|
||||
if args.hf_repo is None:
|
||||
if args.vad:
|
||||
args.hf_repo = "kyutai/stt-1b-en_fr-candle"
|
||||
else:
|
||||
args.hf_repo = "kyutai/stt-1b-en_fr-mlx"
|
||||
lm_config = hf_hub_download(args.hf_repo, "config.json")
|
||||
with open(lm_config, "r") as fobj:
|
||||
lm_config = json.load(fobj)
|
||||
mimi_weights = hf_hub_download(args.hf_repo, lm_config["mimi_name"])
|
||||
moshi_name = lm_config.get("moshi_name", "model.safetensors")
|
||||
moshi_weights = hf_hub_download(args.hf_repo, moshi_name)
|
||||
text_tokenizer = hf_hub_download(args.hf_repo, lm_config["tokenizer_name"])
|
||||
|
||||
lm_config = models.LmConfig.from_config_dict(lm_config)
|
||||
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)
|
||||
|
||||
print(f"loading model weights from {moshi_weights}")
|
||||
if args.hf_repo.endswith("-candle"):
|
||||
model.load_pytorch_weights(moshi_weights, lm_config, strict=True)
|
||||
else:
|
||||
model.load_weights(moshi_weights, strict=True)
|
||||
|
||||
print(f"loading the text tokenizer from {text_tokenizer}")
|
||||
text_tokenizer = sentencepiece.SentencePieceProcessor(text_tokenizer) # type: ignore
|
||||
|
||||
print(f"loading the audio tokenizer {mimi_weights}")
|
||||
audio_tokenizer = models.mimi.Mimi(models.mimi_202407(32))
|
||||
audio_tokenizer.load_pytorch_weights(str(mimi_weights), strict=True)
|
||||
print("warming up the model")
|
||||
model.warmup()
|
||||
gen = models.LmGen(
|
||||
model=model,
|
||||
max_steps=args.max_steps,
|
||||
text_sampler=utils.Sampler(top_k=25, temp=0),
|
||||
audio_sampler=utils.Sampler(top_k=250, temp=0.8),
|
||||
check=False,
|
||||
)
|
||||
|
||||
print(f"starting inference {audio.shape}")
|
||||
audio = mx.concat([mx.array(audio), mx.zeros((1, 48000))], axis=-1)
|
||||
last_print_was_vad = False
|
||||
for start_idx in range(0, audio.shape[-1] // 1920 * 1920, 1920):
|
||||
block = audio[:, None, start_idx : start_idx + 1920]
|
||||
other_audio_tokens = audio_tokenizer.encode_step(block).transpose(0, 2, 1)
|
||||
if args.vad:
|
||||
text_token, vad_heads = gen.step_with_extra_heads(other_audio_tokens[0])
|
||||
if vad_heads:
|
||||
pr_vad = vad_heads[2][0, 0, 0].item()
|
||||
if pr_vad > 0.5 and not last_print_was_vad:
|
||||
print(" [end of turn detected]")
|
||||
last_print_was_vad = True
|
||||
else:
|
||||
text_token = gen.step(other_audio_tokens[0])
|
||||
text_token = text_token[0].item()
|
||||
audio_tokens = gen.last_audio_tokens()
|
||||
_text = None
|
||||
if text_token not in (0, 3):
|
||||
_text = text_tokenizer.id_to_piece(text_token) # type: ignore
|
||||
_text = _text.replace("▁", " ")
|
||||
print(_text, end="", flush=True)
|
||||
last_print_was_vad = False
|
||||
print()
|
||||
@@ -0,0 +1,210 @@
|
||||
# /// script
|
||||
# requires-python = ">=3.12"
|
||||
# dependencies = [
|
||||
# "huggingface_hub",
|
||||
# "moshi_mlx==0.2.12",
|
||||
# "numpy",
|
||||
# "sounddevice",
|
||||
# ]
|
||||
# ///
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import queue
|
||||
import sys
|
||||
import time
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
import numpy as np
|
||||
import sentencepiece
|
||||
import sounddevice as sd
|
||||
import sphn
|
||||
from moshi_mlx import models
|
||||
from moshi_mlx.client_utils import make_log
|
||||
from moshi_mlx.models.tts import (
|
||||
DEFAULT_DSM_TTS_REPO,
|
||||
DEFAULT_DSM_TTS_VOICE_REPO,
|
||||
TTSModel,
|
||||
)
|
||||
from moshi_mlx.utils.loaders import hf_get
|
||||
|
||||
|
||||
def log(level: str, msg: str):
|
||||
print(make_log(level, msg))
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Run Kyutai TTS using the MLX implementation"
|
||||
)
|
||||
parser.add_argument("inp", type=str, help="Input file, use - for stdin")
|
||||
parser.add_argument(
|
||||
"out", type=str, help="Output file to generate, use - for playing the audio"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hf-repo",
|
||||
type=str,
|
||||
default=DEFAULT_DSM_TTS_REPO,
|
||||
help="HF repo in which to look for the pretrained models.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--voice-repo",
|
||||
default=DEFAULT_DSM_TTS_VOICE_REPO,
|
||||
help="HF repo in which to look for pre-computed voice embeddings.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--voice", default="expresso/ex03-ex01_happy_001_channel1_334s.wav"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantize",
|
||||
type=int,
|
||||
help="The quantization to be applied, e.g. 8 for 8 bits.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
mx.random.seed(299792458)
|
||||
|
||||
log("info", "retrieving checkpoints")
|
||||
|
||||
raw_config = hf_get("config.json", args.hf_repo)
|
||||
with open(hf_get(raw_config), "r") as fobj:
|
||||
raw_config = json.load(fobj)
|
||||
|
||||
mimi_weights = hf_get(raw_config["mimi_name"], args.hf_repo)
|
||||
moshi_name = raw_config.get("moshi_name", "model.safetensors")
|
||||
moshi_weights = hf_get(moshi_name, args.hf_repo)
|
||||
tokenizer = hf_get(raw_config["tokenizer_name"], args.hf_repo)
|
||||
lm_config = models.LmConfig.from_config_dict(raw_config)
|
||||
# There is a bug in moshi_mlx <= 0.3.0 handling of the ring kv cache.
|
||||
# The following line gets around it for now.
|
||||
lm_config.transformer.max_seq_len = lm_config.transformer.context
|
||||
model = models.Lm(lm_config)
|
||||
model.set_dtype(mx.bfloat16)
|
||||
|
||||
log("info", f"loading model weights from {moshi_weights}")
|
||||
model.load_pytorch_weights(str(moshi_weights), lm_config, strict=True)
|
||||
|
||||
if args.quantize is not None:
|
||||
log("info", f"quantizing model to {args.quantize} bits")
|
||||
nn.quantize(model.depformer, bits=args.quantize)
|
||||
for layer in model.transformer.layers:
|
||||
nn.quantize(layer.self_attn, bits=args.quantize)
|
||||
nn.quantize(layer.gating, bits=args.quantize)
|
||||
|
||||
log("info", f"loading the text tokenizer from {tokenizer}")
|
||||
text_tokenizer = sentencepiece.SentencePieceProcessor(str(tokenizer)) # type: ignore
|
||||
|
||||
log("info", f"loading the audio tokenizer {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)
|
||||
|
||||
cfg_coef_conditioning = None
|
||||
tts_model = TTSModel(
|
||||
model,
|
||||
audio_tokenizer,
|
||||
text_tokenizer,
|
||||
voice_repo=args.voice_repo,
|
||||
temp=0.6,
|
||||
cfg_coef=1,
|
||||
max_padding=8,
|
||||
initial_padding=2,
|
||||
final_padding=2,
|
||||
padding_bonus=0,
|
||||
raw_config=raw_config,
|
||||
)
|
||||
if tts_model.valid_cfg_conditionings:
|
||||
# Model was trained with CFG distillation.
|
||||
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
|
||||
mimi = tts_model.mimi
|
||||
|
||||
log("info", f"reading input from {args.inp}")
|
||||
if args.inp == "-":
|
||||
if sys.stdin.isatty(): # Interactive
|
||||
print("Enter text to synthesize (Ctrl+D to end input):")
|
||||
text_to_tts = sys.stdin.read().strip()
|
||||
else:
|
||||
with open(args.inp, "r", encoding="utf-8") as fobj:
|
||||
text_to_tts = fobj.read().strip()
|
||||
|
||||
all_entries = [tts_model.prepare_script([text_to_tts])]
|
||||
if tts_model.multi_speaker:
|
||||
voices = [tts_model.get_voice_path(args.voice)]
|
||||
else:
|
||||
voices = []
|
||||
all_attributes = [
|
||||
tts_model.make_condition_attributes(voices, cfg_coef_conditioning)
|
||||
]
|
||||
|
||||
wav_frames = queue.Queue()
|
||||
_frames_cnt = 0
|
||||
|
||||
def _on_frame(frame):
|
||||
nonlocal _frames_cnt
|
||||
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)
|
||||
_frames_cnt += 1
|
||||
print(f"generated {_frames_cnt / 12.5:.2f}s", end="\r", flush=True)
|
||||
|
||||
def run():
|
||||
log("info", "starting the inference loop")
|
||||
begin = time.time()
|
||||
result = 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,
|
||||
)
|
||||
frames = mx.concat(result.frames, axis=-1)
|
||||
total_duration = frames.shape[0] * frames.shape[-1] / mimi.frame_rate
|
||||
time_taken = time.time() - begin
|
||||
total_speed = total_duration / time_taken
|
||||
log("info", f"[LM] took {time_taken:.2f}s, total speed {total_speed:.2f}x")
|
||||
return result
|
||||
|
||||
if args.out == "-":
|
||||
|
||||
def audio_callback(outdata, _a, _b, _c):
|
||||
try:
|
||||
pcm_data = wav_frames.get(block=False)
|
||||
outdata[:, 0] = pcm_data
|
||||
except queue.Empty:
|
||||
outdata[:] = 0
|
||||
|
||||
with sd.OutputStream(
|
||||
samplerate=mimi.sample_rate,
|
||||
blocksize=1920,
|
||||
channels=1,
|
||||
callback=audio_callback,
|
||||
):
|
||||
run()
|
||||
time.sleep(3)
|
||||
while True:
|
||||
if wav_frames.qsize() == 0:
|
||||
break
|
||||
time.sleep(1)
|
||||
else:
|
||||
run()
|
||||
frames = []
|
||||
while True:
|
||||
try:
|
||||
frames.append(wav_frames.get_nowait())
|
||||
except queue.Empty:
|
||||
break
|
||||
wav = np.concat(frames, -1)
|
||||
sphn.write_wav(args.out, wav, mimi.sample_rate)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+256
-23
@@ -8,6 +8,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import secrets
|
||||
@@ -21,7 +22,7 @@ from typing import AsyncIterator, Literal
|
||||
|
||||
import httpx
|
||||
import yaml
|
||||
from fastapi import Depends, FastAPI, HTTPException, Request, UploadFile, File, WebSocket, WebSocketDisconnect, status
|
||||
from fastapi import Depends, FastAPI, Form, HTTPException, Request, UploadFile, File, WebSocket, WebSocketDisconnect, status
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import JSONResponse, PlainTextResponse, Response
|
||||
from pydantic import BaseModel
|
||||
@@ -118,6 +119,8 @@ class Settings(BaseSettings):
|
||||
ailiance_tts_voice: str = "nova"
|
||||
ailiance_chat_model: str = "gpt-4o-mini"
|
||||
kokoro_tts_url: str = "http://macm1:8520" # vllm-mlx native /v1/audio/speech (stock voices)
|
||||
kyutai_stt_url: str = "http://127.0.0.1:8300" # local Kyutai stt-1b-en_fr MLX server (/transcribe)
|
||||
kyutai_tts_url: str = "http://127.0.0.1:8302" # local Kyutai tts-1.6b-en_fr MLX server (/tts, native FR)
|
||||
request_timeout: float = 10.0
|
||||
probe_interval: float = 2.0
|
||||
# Browser clients (atelier DeployPanel). Comma-separated origins; "*" for
|
||||
@@ -192,12 +195,24 @@ class VoiceSessions:
|
||||
VOICE_SESSIONS = VoiceSessions()
|
||||
|
||||
|
||||
# Spoken-telephone style: the reply is sent straight to TTS and played in an
|
||||
# earpiece, so it must be short, plain spoken French — no markdown, no stage
|
||||
# directions, no lists, no headings.
|
||||
_PHONE_STYLE = (
|
||||
"\n\nCONTRAINTE DE FORME : tu réponds au téléphone. Réponds en UNE seule "
|
||||
"phrase courte (15 mots maximum), en français parlé naturel. JAMAIS de "
|
||||
"markdown (pas de **, *, #, ---), JAMAIS de didascalies entre crochets, "
|
||||
"JAMAIS de listes ni de titres. Uniquement les mots que le personnage "
|
||||
"dirait à voix haute, et fais court."
|
||||
)
|
||||
|
||||
|
||||
async def _chat_reply(http: httpx.AsyncClient, persona: str, history: list[dict]) -> str:
|
||||
"""Call ailiance chat completions with the NPC persona and return the reply text."""
|
||||
messages = [{"role": "system", "content": persona}] + history
|
||||
messages = [{"role": "system", "content": persona + _PHONE_STYLE}] + history
|
||||
resp = await http.post(
|
||||
f"{settings.ailiance_tts_url.rstrip('/')}/v1/chat/completions",
|
||||
json={"model": settings.ailiance_chat_model, "messages": messages, "max_tokens": 200},
|
||||
json={"model": settings.ailiance_chat_model, "messages": messages, "max_tokens": 48},
|
||||
timeout=60.0,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
@@ -215,9 +230,26 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
||||
app.state.health_cache = HealthCache(app.state.http)
|
||||
app.state.audio_jobs = {} # job_id -> AudioJob (in-process async TTS queue)
|
||||
STAGED_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
async def _prewarm_greetings() -> None:
|
||||
"""Synthesise every directory greeting into the TTS cache in the
|
||||
background so the first phone call to each number plays instantly."""
|
||||
for number, entry in PHONE_DIRECTORY.items():
|
||||
greeting = entry.get("greeting")
|
||||
if not greeting:
|
||||
continue
|
||||
try:
|
||||
await _voice_tts_16k(app.state.http, greeting,
|
||||
entry.get("voice", settings.ailiance_tts_voice))
|
||||
logging.info("TTS prewarm: cached greeting for %s", number)
|
||||
except Exception as exc: # noqa: BLE001 — prewarm is best-effort
|
||||
logging.warning("TTS prewarm failed for %s: %s", number, str(exc)[:100])
|
||||
|
||||
prewarm_task = asyncio.create_task(_prewarm_greetings())
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
prewarm_task.cancel()
|
||||
await app.state.http.aclose()
|
||||
|
||||
|
||||
@@ -809,6 +841,77 @@ async def _synthesise_wav(http: httpx.AsyncClient, body: AudioGenRequest) -> byt
|
||||
return resp.content
|
||||
|
||||
|
||||
try: # high-quality anti-aliased resampling (avoids the harsh aliasing of naive linear interp)
|
||||
import numpy as _np
|
||||
import soxr as _soxr
|
||||
except Exception: # pragma: no cover
|
||||
_np = None
|
||||
_soxr = None
|
||||
|
||||
|
||||
def _softknee_compressor(x, sr: int, threshold: float, ratio: float = 3.0,
|
||||
knee_db: float = 6.0, attack_ms: float = 5.0,
|
||||
release_ms: float = 120.0):
|
||||
"""Feed-forward soft-knee compressor (float samples in/out).
|
||||
|
||||
Evens out the voice dynamics (quiet parts up, loud parts down) so the
|
||||
perceived level is uniform before the limiter. `threshold` is a linear
|
||||
amplitude in the same units as x. Smooth (soft) knee around the threshold.
|
||||
"""
|
||||
n = len(x)
|
||||
if n == 0:
|
||||
return x
|
||||
eps = 1e-9
|
||||
thr_db = 20.0 * _np.log10(threshold + eps)
|
||||
half_knee = knee_db / 2.0
|
||||
aa = _np.exp(-1.0 / (sr * attack_ms / 1000.0))
|
||||
ar = _np.exp(-1.0 / (sr * release_ms / 1000.0))
|
||||
absx = _np.abs(x)
|
||||
# Peak envelope follower (one-pole, separate attack/release) — Python loop.
|
||||
env = _np.empty(n, dtype=_np.float64)
|
||||
e = 0.0
|
||||
for i in range(n):
|
||||
a = float(absx[i])
|
||||
coef = aa if a > e else ar
|
||||
e = a + (e - a) * coef
|
||||
env[i] = e
|
||||
lvl_db = 20.0 * _np.log10(env + eps)
|
||||
over = lvl_db - thr_db
|
||||
gain_db = _np.zeros(n, dtype=_np.float64)
|
||||
above = over > half_knee
|
||||
in_knee = _np.abs(over) <= half_knee
|
||||
gain_db[above] = (1.0 / ratio - 1.0) * over[above]
|
||||
k = over[in_knee] + half_knee
|
||||
gain_db[in_knee] = (1.0 / ratio - 1.0) * (k * k) / (2.0 * knee_db)
|
||||
return x * (10.0 ** (gain_db / 20.0))
|
||||
|
||||
|
||||
def _lookahead_limiter(x, ceiling: float, sr: int,
|
||||
lookahead_ms: float = 5.0, release_ms: float = 60.0):
|
||||
"""Brick-wall look-ahead limiter (float in/out). Guarantees |out| <= ceiling
|
||||
without hard clipping: gain ducks *before* a peak (look-ahead = instant
|
||||
attack) and recovers smoothly (release one-pole)."""
|
||||
n = len(x)
|
||||
if n == 0:
|
||||
return x
|
||||
la = max(1, int(sr * lookahead_ms / 1000.0))
|
||||
absx = _np.abs(x)
|
||||
desired = _np.ones(n, dtype=_np.float64)
|
||||
over = absx > ceiling
|
||||
desired[over] = ceiling / absx[over]
|
||||
env = desired.copy()
|
||||
for s in range(1, la + 1):
|
||||
env[: n - s] = _np.minimum(env[: n - s], desired[s:])
|
||||
rel = _np.exp(-1.0 / (sr * release_ms / 1000.0))
|
||||
g = _np.empty(n, dtype=_np.float64)
|
||||
prev = 1.0
|
||||
for i in range(n):
|
||||
t = float(env[i])
|
||||
prev = t if t < prev else t + (prev - t) * rel
|
||||
g[i] = prev
|
||||
return x * g
|
||||
|
||||
|
||||
def _wav_to_16k_mono(wav_bytes: bytes, max_seconds: float | None = None) -> tuple[bytes, float, bool]:
|
||||
"""Pure-Python WAV resampler: any-rate mono/stereo → 16000 Hz mono 16-bit PCM.
|
||||
|
||||
@@ -859,10 +962,16 @@ def _wav_to_16k_mono(wav_bytes: bytes, max_seconds: float | None = None) -> tupl
|
||||
for i in range(len(interleaved) // src_channels)
|
||||
]
|
||||
|
||||
# Resample to 16000 Hz via linear interpolation
|
||||
# Resample to 16000 Hz. Prefer soxr (anti-aliased, VHQ) — naive linear interp
|
||||
# downsampling folds >8 kHz content back into the band → harsh "saturated"
|
||||
# artefacts. Fall back to linear interp only if soxr/numpy are unavailable.
|
||||
dst_rate = 16000
|
||||
if src_rate == dst_rate:
|
||||
out_samples = mono
|
||||
elif _soxr is not None and _np is not None:
|
||||
arr = _np.asarray(mono, dtype=_np.float32)
|
||||
res = _soxr.resample(arr, src_rate, dst_rate, quality="VHQ")
|
||||
out_samples = res.astype(_np.int32).tolist()
|
||||
else:
|
||||
ratio = src_rate / dst_rate
|
||||
src_len = len(mono)
|
||||
@@ -886,8 +995,29 @@ def _wav_to_16k_mono(wav_bytes: bytes, max_seconds: float | None = None) -> tupl
|
||||
|
||||
duration_s = len(out_samples) / dst_rate
|
||||
|
||||
# Clamp to 16-bit range
|
||||
out_samples = [max(-32768, min(32767, s)) for s in out_samples]
|
||||
# ── Dynamics chain: normalize → soft-knee compressor → look-ahead limiter ──
|
||||
# Evens out the voice (uniform perceived level) and guarantees no clipping.
|
||||
if _np is not None and out_samples:
|
||||
FS = 32767.0
|
||||
x = _np.asarray(out_samples, dtype=_np.float64)
|
||||
peak = float(_np.max(_np.abs(x)))
|
||||
if peak > 0:
|
||||
x *= min(0.95 * FS / peak, 12.0) # normalize to a known scale
|
||||
# Soft-knee compressor: tame peaks above ~0.30 FS, 3:1, then make up gain.
|
||||
x = _softknee_compressor(x, dst_rate, threshold=0.30 * FS, ratio=3.0, knee_db=6.0)
|
||||
peak2 = float(_np.max(_np.abs(x)))
|
||||
if peak2 > 0:
|
||||
x *= min(0.95 * FS / peak2, 12.0) # makeup → back to 0.95 FS
|
||||
# Brick-wall look-ahead limiter at 0.95 FS (no hard clip, catches overshoot).
|
||||
x = _lookahead_limiter(x, ceiling=0.95 * FS, sr=dst_rate)
|
||||
out_samples = _np.clip(_np.round(x), -32768, 32767).astype(_np.int32).tolist()
|
||||
else:
|
||||
# Fallback (no numpy): simple peak normalize + hard clamp.
|
||||
peak = max((abs(s) for s in out_samples), default=0)
|
||||
if peak > 0:
|
||||
gain = min(0.70 * 32767 / peak, 12.0)
|
||||
out_samples = [int(s * gain) for s in out_samples]
|
||||
out_samples = [max(-32768, min(32767, s)) for s in out_samples]
|
||||
|
||||
# Pack back to WAV
|
||||
pcm = struct.pack(f"<{len(out_samples)}h", *out_samples)
|
||||
@@ -1467,15 +1597,9 @@ async def voice_say(body: VoiceSayRequest, request: Request, _: None = Depends(r
|
||||
ip = BOARDS[body.board]["ip"]
|
||||
|
||||
# 1. Synthesise WAV via selected backend
|
||||
# Build a minimal AudioGenRequest-compatible object for _synthesise_wav
|
||||
class _SynthReq:
|
||||
def __init__(self, text: str, backend: str, voice: str | None) -> None:
|
||||
self.text = text
|
||||
self.backend = backend
|
||||
self.voice = voice or ""
|
||||
|
||||
try:
|
||||
raw_wav = await _synthesise_wav(request.app.state.http, _SynthReq(body.text, body.backend, body.voice)) # type: ignore[arg-type]
|
||||
synth_req = types.SimpleNamespace(text=body.text, backend=body.backend, voice=body.voice or "")
|
||||
raw_wav = await _synthesise_wav(request.app.state.http, synth_req) # type: ignore[arg-type]
|
||||
except Exception as exc:
|
||||
raise HTTPException(502, f"TTS failed ({body.backend}): {str(exc)[:200]}") from exc
|
||||
|
||||
@@ -1621,11 +1745,79 @@ def _validate_yaml(text: str) -> ValidationResult:
|
||||
|
||||
# ---------- P5: PLIP telephone — voice turn ----------
|
||||
|
||||
async def _tts_kyutai(http: httpx.AsyncClient, text: str, voice: str | None = None) -> bytes:
|
||||
"""Synthesise text via the local Kyutai tts-1.6b-en_fr MLX server (native FR).
|
||||
|
||||
Returns raw WAV bytes (24 kHz mono). `voice` is an optional Kyutai voice path
|
||||
(e.g. 'cml-tts/fr/...'); omit to use the server's default French voice.
|
||||
The OpenAI-style voice names in the directory (nova/alloy) are NOT Kyutai
|
||||
voices, so they are ignored here — the server picks its default FR voice.
|
||||
"""
|
||||
payload: dict = {"text": text}
|
||||
if voice and "/" in voice: # only forward genuine Kyutai voice paths
|
||||
payload["voice"] = voice
|
||||
resp = await http.post(
|
||||
f"{settings.kyutai_tts_url.rstrip('/')}/tts",
|
||||
json=payload,
|
||||
timeout=60.0,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return resp.content
|
||||
|
||||
|
||||
# TTS is the latency bottleneck (Kyutai MLX ~0.3x realtime). Greetings are fixed
|
||||
# strings, so cache the finished 16 kHz WAV by (voice, text): first synth is slow,
|
||||
# every repeat is instant. Replies are dynamic and won't hit the cache.
|
||||
_TTS_CACHE: dict[str, tuple[bytes, float, bool]] = {}
|
||||
_TTS_CACHE_MAX = 64
|
||||
|
||||
|
||||
async def _voice_tts_16k(http: httpx.AsyncClient, text: str, voice: str) -> tuple[bytes, float, bool]:
|
||||
"""Synthesise text via ailiance TTS and resample to 16 kHz mono WAV."""
|
||||
fake_body = types.SimpleNamespace(backend="ailiance", text=text, voice=voice)
|
||||
raw = await _synthesise_wav(http, fake_body) # type: ignore[arg-type]
|
||||
return _wav_to_16k_mono(raw, max_seconds=7.5)
|
||||
"""Synthesise text and resample to 16 kHz mono WAV for the PLIP.
|
||||
|
||||
Cached by (voice, text). Primary backend: local Kyutai TTS (native French);
|
||||
fallback: ailiance tts-1 if Kyutai is unreachable (anglophone, but not silent).
|
||||
"""
|
||||
key = f"{voice}|{text}"
|
||||
cached = _TTS_CACHE.get(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
try:
|
||||
raw = await _tts_kyutai(http, text, voice)
|
||||
except Exception as exc:
|
||||
logging.warning("Kyutai TTS unreachable (%s) — falling back to ailiance", str(exc)[:120])
|
||||
fake_body = types.SimpleNamespace(backend="ailiance", text=text, voice=voice)
|
||||
raw = await _synthesise_wav(http, fake_body) # type: ignore[arg-type]
|
||||
result = _wav_to_16k_mono(raw, max_seconds=7.5)
|
||||
|
||||
if len(_TTS_CACHE) >= _TTS_CACHE_MAX:
|
||||
_TTS_CACHE.clear() # crude bound — greetings are few, this rarely trips
|
||||
_TTS_CACHE[key] = result
|
||||
return result
|
||||
|
||||
|
||||
async def _transcribe_kyutai(http: httpx.AsyncClient, audio_wav: bytes) -> str:
|
||||
"""Transcribe a WAV via the local Kyutai stt-1b-en_fr MLX server (/transcribe).
|
||||
|
||||
The server resamples internally (Mimi @ 24 kHz), so any sample rate is fine.
|
||||
Returns the recognised text (possibly empty on silence/inaudible input).
|
||||
"""
|
||||
resp = await http.post(
|
||||
f"{settings.kyutai_stt_url.rstrip('/')}/transcribe",
|
||||
content=audio_wav,
|
||||
headers={"Content-Type": "audio/wav"},
|
||||
timeout=30.0,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return (resp.json().get("text") or "").strip()
|
||||
|
||||
|
||||
def _ascii_header(s: str) -> str:
|
||||
"""Make a string safe as an HTTP header value: ASCII-only AND no control
|
||||
characters (newlines/CR/tab break the HTTP framing)."""
|
||||
ascii_s = s.encode("ascii", errors="replace").decode("ascii")
|
||||
return "".join(ch if 0x20 <= ord(ch) < 0x7F else " " for ch in ascii_s).strip()
|
||||
|
||||
|
||||
class VoiceTurnRequest(BaseModel):
|
||||
@@ -1653,8 +1845,8 @@ async def voice_turn(body: VoiceTurnRequest, request: Request, _: None = Depends
|
||||
if body.kind == "greeting":
|
||||
said = entry.get("greeting") or await _chat_reply(http, persona, [])
|
||||
else:
|
||||
# Stage 3: transcription-driven reply; greeting path covers Stage 2.
|
||||
# Append the user turn first so the LLM sees the conversation correctly.
|
||||
# Text-only reply (no audio input — heard is always empty here).
|
||||
# The transcription-driven path lives in POST /v1/voice/reply.
|
||||
VOICE_SESSIONS.append(sid, "user", heard or "(inaudible)")
|
||||
history = VOICE_SESSIONS.history(sid)
|
||||
said = await _chat_reply(http, persona, history)
|
||||
@@ -1666,9 +1858,50 @@ async def voice_turn(body: VoiceTurnRequest, request: Request, _: None = Depends
|
||||
except Exception as exc:
|
||||
raise HTTPException(502, f"voice backend unreachable: {str(exc)[:120]}") from exc
|
||||
|
||||
# HTTP headers must be ASCII; encode non-ASCII chars to avoid codec errors.
|
||||
def _ascii_header(s: str) -> str:
|
||||
return s.encode("ascii", errors="replace").decode("ascii")
|
||||
return Response(
|
||||
content=wav16,
|
||||
media_type="audio/wav",
|
||||
headers={
|
||||
"X-Zacus-Heard": _ascii_header(heard[:200]),
|
||||
"X-Zacus-Said": _ascii_header(said[:200]),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@app.post("/v1/voice/reply")
|
||||
async def voice_reply(
|
||||
request: Request,
|
||||
session_id: str = Form(...),
|
||||
number: str = Form(...),
|
||||
audio: UploadFile = File(...),
|
||||
_: None = Depends(require_token),
|
||||
) -> Response:
|
||||
"""Stage 3 conversational turn: the player's recorded speech drives the reply.
|
||||
|
||||
Multipart form: session_id, number, audio (WAV). The audio is transcribed
|
||||
by the local Kyutai STT, appended to the session as the user turn, the NPC
|
||||
persona LLM produces a reply, and the reply is synthesised to a 16 kHz WAV.
|
||||
X-Zacus-Heard carries the transcription so the firmware/dashboard can log it.
|
||||
"""
|
||||
entry = PHONE_DIRECTORY.get(number)
|
||||
if entry is None:
|
||||
raise HTTPException(404, f"unknown phone number '{number}'")
|
||||
|
||||
http: httpx.AsyncClient = request.app.state.http
|
||||
persona = entry.get("persona", "")
|
||||
|
||||
try:
|
||||
audio_bytes = await audio.read()
|
||||
heard = await _transcribe_kyutai(http, audio_bytes)
|
||||
VOICE_SESSIONS.append(session_id, "user", heard or "(inaudible)")
|
||||
history = VOICE_SESSIONS.history(session_id)
|
||||
said = await _chat_reply(http, persona, history)
|
||||
VOICE_SESSIONS.append(session_id, "assistant", said)
|
||||
wav16, _, _ = await _voice_tts_16k(http, said, entry.get("voice", settings.ailiance_tts_voice))
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise HTTPException(502, f"voice backend unreachable: {str(exc)[:120]}") from exc
|
||||
|
||||
return Response(
|
||||
content=wav16,
|
||||
|
||||
Reference in New Issue
Block a user