First Version

2026-07-10 15:02:09 +01:00
commit 38494a85a9
562 changed files with 196313 additions and 0 deletions
@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+"""Kokoro-ONNX TTS inference server.
+
+Started by the plugin-tts-kokoro Rust plugin. Prints "PORT:<n>" to stdout
+once the HTTP server is bound so the plugin knows which port to connect to.
+
+Model files (kokoro-v1.0.onnx + voices-v1.0.bin) are downloaded from GitHub
+releases on first run and cached in --model-dir.
+
+Endpoints
+---------
+POST /synthesize
+    Body: {"text": "...", "voice": "if_sara", "lang": "it", "speed": 1.0}
+    Returns: audio/wav bytes
+
+GET /health
+    Returns: {"status": "ok"}
+"""
+
+import argparse
+import io
+import os
+import socket
+import sys
+import threading
+
+import numpy as np
+import soundfile as sf
+import urllib.request
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import Response
+from pydantic import BaseModel
+from kokoro_onnx import Kokoro
+import uvicorn
+
+# ── Model URLs ────────────────────────────────────────────────────────────────
+
+MODEL_BASE_URL = "https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0"
+MODEL_ONNX_FILE   = "kokoro-v1.0.onnx"
+MODEL_VOICES_FILE = "voices-v1.0.bin"
+
+VALID_VOICES = {
+    "af_heart", "af_bella", "af_nicole", "af_sarah", "af_sky",
+    "am_adam", "am_michael",
+    "bf_emma", "bf_isabella", "bm_george", "bm_lewis",
+    "if_sara", "im_nicola",
+    "jf_alpha", "jf_gongitsune", "jm_kumo",
+    "zf_xiaobei", "zm_yunxi",
+}
+
+# ── Globals set at startup ────────────────────────────────────────────────────
+
+kokoro_model: Kokoro | None = None
+default_voice = "if_sara"
+default_lang  = "it"
+default_speed = 1.0
+
+# ── Model loading ─────────────────────────────────────────────────────────────
+
+def download_if_missing(model_dir: str, filename: str) -> str:
+    path = os.path.join(model_dir, filename)
+    if not os.path.exists(path):
+        url = f"{MODEL_BASE_URL}/{filename}"
+        print(f"[kokoro] downloading {filename} from {url}", flush=True)
+        urllib.request.urlretrieve(url, path)
+        print(f"[kokoro] downloaded {filename}", flush=True)
+    return path
+
+
+def load_model(model_dir: str) -> None:
+    global kokoro_model
+    os.makedirs(model_dir, exist_ok=True)
+    onnx_path   = download_if_missing(model_dir, MODEL_ONNX_FILE)
+    voices_path = download_if_missing(model_dir, MODEL_VOICES_FILE)
+    print(f"[kokoro] loading model from {model_dir}", flush=True)
+    kokoro_model = Kokoro(onnx_path, voices_path)
+    print("[kokoro] model ready", flush=True)
+
+# ── FastAPI server ────────────────────────────────────────────────────────────
+
+app = FastAPI()
+
+
+class SynthesizeRequest(BaseModel):
+    text:  str
+    voice: str | None = None
+    lang:  str | None = None
+    speed: float | None = None
+
+
+@app.get("/health")
+def health():
+    return {"status": "ok"}
+
+
+@app.post("/synthesize")
+def synthesize(req: SynthesizeRequest):
+    if kokoro_model is None:
+        raise HTTPException(status_code=503, detail="model not loaded")
+
+    voice = req.voice if req.voice in VALID_VOICES else default_voice
+    lang  = req.lang  or default_lang
+    speed = req.speed or default_speed
+
+    try:
+        samples, sample_rate = kokoro_model.create(req.text, voice=voice, speed=speed, lang=lang)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    buf = io.BytesIO()
+    sf.write(buf, samples, sample_rate, format="WAV")
+    return Response(content=buf.getvalue(), media_type="audio/wav")
+
+# ── Entry point ───────────────────────────────────────────────────────────────
+
+def find_free_port() -> int:
+    with socket.socket() as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-dir",     default="models/kokoro")
+    parser.add_argument("--default-voice", default="if_sara")
+    parser.add_argument("--default-lang",  default="it")
+    parser.add_argument("--default-speed", type=float, default=1.0)
+    args = parser.parse_args()
+
+    global default_voice, default_lang, default_speed
+    default_voice = args.default_voice
+    default_lang  = args.default_lang
+    default_speed = args.default_speed
+
+    load_model(args.model_dir)
+
+    port = find_free_port()
+    # Signal to the Rust parent that we are ready.
+    print(f"PORT:{port}", flush=True)
+
+    uvicorn.run(app, host="127.0.0.1", port=port, log_level="warning")
+
+
+if __name__ == "__main__":
+    main()