146 lines
4.9 KiB
Python
146 lines
4.9 KiB
Python
#!/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()
|