Streaming /tts/stream route + ONNX engine (misaki G2P), torch fallback

README.md

@@ -1,11 +1,11 @@
----
-title: Kokoro TTS
-emoji: 🗣️
-colorFrom: indigo
-colorTo: purple
-sdk: gradio
-sdk_version: 6.6.0
-app_file: app.py
-pinned: false
-python_version: "3.10"
----
+---
+title: Kokoro TTS
+emoji: 🗣️
+colorFrom: indigo
+colorTo: purple
+sdk: gradio
+sdk_version: 6.6.0
+app_file: app.py
+pinned: false
+python_version: "3.10"
+---

app.py

@@ -1,69 +1,250 @@
-import gradio as gr
-from kokoro import KPipeline
-import numpy as np
-
-pipelines = {
-    "a": KPipeline(lang_code="a"),
-    "b": KPipeline(lang_code="b"),
-    "p": KPipeline(lang_code="p"),
-}
-
-def get_pipeline_for_voice(voice_id: str) -> KPipeline:
-    if voice_id.startswith(("af_", "am_")):
-        return pipelines["a"]
-    elif voice_id.startswith(("bf_", "bm_")):
-        return pipelines["b"]
-    elif voice_id.startswith(("pf_", "pm_")):
-        return pipelines["p"]
-    return pipelines["a"]
-
-def generate_tts(text: str, voice_id: str, speed: float):
-    try:
-        text = (text or "").strip()
-        if not text:
-            return None
-        MAX_CHARS = 600
-        if len(text) > MAX_CHARS:
-            text = text[:MAX_CHARS]
-        split_pattern = None if len(text) < 80 else r"(?<=[.!?…])\s+"
-        pipeline = get_pipeline_for_voice(voice_id)
-        generator = pipeline(text, voice=voice_id, speed=speed, split_pattern=split_pattern)
-        segments = []
-        for _gs, _ps, audio in generator:
-            seg = np.array(audio, dtype=np.float32)
-            if seg.ndim > 1:
-                seg = seg.mean(axis=-1)
-            if seg.size > 0:
-                segments.append(seg)
-        if not segments:
-            return None
-        return (24000, np.concatenate(segments).astype(np.float32))
-    except Exception as e:
-        print(f"[TTS] Erro: {e}")
-        return None
-
-VOICES = [
-    "af_heart", "af_alloy", "af_aoede", "af_bella", "af_jessica",
-    "af_kore", "af_nicole", "af_nova", "af_river", "af_sarah", "af_sky",
-    "am_adam", "am_echo", "am_eric", "am_fenrir", "am_liam",
-    "am_michael", "am_onyx", "am_orion", "am_puck",
-    "bf_alice", "bf_emma", "bf_isabella",
-    "bm_george", "bm_lewis",
-    "pf_dora", "pm_alex", "pm_santa",
-]
-
-demo = gr.Interface(
-    fn=generate_tts,
-    inputs=[
-        gr.Textbox(label="Texto", lines=3, placeholder="Digite o texto aqui..."),
-        gr.Dropdown(choices=VOICES, value="af_heart", label="Voz"),
-        gr.Slider(minimum=0.5, maximum=2.0, value=1.0, step=0.1, label="Velocidade"),
-    ],
-    outputs=gr.Audio(label="Audio Gerado"),
-    title="Kokoro-82M TTS",
-    description="Ingles Americano / Britanico e Portugues Brasileiro.",
-    cache_examples=False,
-)
-
-if __name__ == "__main__":
-    demo.launch()
+"""Kokoro-82M TTS Space: Gradio UI + a streaming HTTP endpoint.
+
+- Gradio UI and its API (/generate_tts) are unchanged, so existing clients
+  (gradio_client.predict) keep working.
+- POST /tts/stream returns raw PCM16 mono audio WHILE it is being generated:
+  the text is cut into pieces (the first sentence at its first clause, so the
+  first audio is short), each piece is synthesized and sent as soon as it is
+  ready. No Gradio queue, no WAV file to download, resampled to 16 kHz here.
+
+Engine (KOKORO_ENGINE): "onnx" (default) runs the same Kokoro-82M v1.0
+weights through ONNX Runtime (kokoro-onnx), with phonemes from Kokoro's own
+G2P (misaki) so pronunciation is unchanged; "torch" is the original
+PyTorch KPipeline. If ONNX can't start, it falls back to torch.
+
+The Space runs on 2 shared vCPUs: one synthesis at a time (a lock per piece,
+so concurrent requests interleave piece by piece instead of waiting for a
+whole reply).
+"""
+import os
+import re
+import threading
+import time
+import urllib.request
+
+import gradio as gr
+import numpy as np
+import torch
+import uvicorn
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+from kokoro import KPipeline
+from scipy.signal import resample_poly
+
+THREADS = max(1, os.cpu_count() or 1)
+torch.set_num_threads(THREADS)
+
+# Optional shared secret (Space secret TTS_API_KEY): when set, /tts/stream
+# requires the X-API-Key header. The Gradio UI stays public as before.
+API_KEY = os.getenv("TTS_API_KEY", "").strip()
+MAX_CHARS = 600
+FIRST_PIECE_MAX = 60  # first sentence longer than this is cut at its first clause
+SAMPLE_RATE = 24000   # Kokoro output
+
+_LANGS = {"a": ("af_", "am_"), "b": ("bf_", "bm_"), "p": ("pf_", "pm_")}
+_pipelines: dict = {}
+_pipelines_lock = threading.Lock()
+_synth_lock = threading.Lock()
+
+# ---------------------------------------------------------------- engine
+
+ONNX_RELEASE = "https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0"
+ONNX_FILES = ("kokoro-v1.0.onnx", "voices-v1.0.bin")
+MODEL_DIR = os.getenv("KOKORO_MODEL_DIR", os.path.join(os.path.dirname(os.path.abspath(__file__)), "models"))
+ENGINE = os.getenv("KOKORO_ENGINE", "onnx").strip().lower()
+_onnx = None  # kokoro_onnx.Kokoro when the ONNX engine is active
+
+
+def _start_onnx():
+    import onnxruntime as ort
+    from kokoro_onnx import Kokoro
+
+    os.makedirs(MODEL_DIR, exist_ok=True)
+    for name in ONNX_FILES:
+        path = os.path.join(MODEL_DIR, name)
+        if not os.path.exists(path) or os.path.getsize(path) == 0:
+            t0 = time.perf_counter()
+            urllib.request.urlretrieve(f"{ONNX_RELEASE}/{name}", path + ".part")
+            os.replace(path + ".part", path)
+            print(f"[TTS] downloaded {name} in {time.perf_counter() - t0:.1f}s")
+    opts = ort.SessionOptions()
+    opts.intra_op_num_threads = THREADS
+    opts.inter_op_num_threads = 1
+    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    session = ort.InferenceSession(
+        os.path.join(MODEL_DIR, ONNX_FILES[0]), sess_options=opts, providers=["CPUExecutionProvider"]
+    )
+    return Kokoro.from_session(session, os.path.join(MODEL_DIR, ONNX_FILES[1]))
+
+
+if ENGINE == "onnx":
+    try:
+        _onnx = _start_onnx()
+        print(f"[TTS] engine: onnx ({THREADS} threads)")
+    except Exception as e:
+        print(f"[TTS] onnx failed ({e}); falling back to torch")
+        ENGINE = "torch"
+if ENGINE != "onnx":
+    print(f"[TTS] engine: torch ({THREADS} threads)")
+
+
+def get_pipeline_for_voice(voice_id: str) -> KPipeline:
+    """American English is loaded at start-up; British / Portuguese on first
+    use. With ONNX the pipeline is G2P only (model=False)."""
+    lang = next((code for code, prefixes in _LANGS.items() if voice_id.startswith(prefixes)), "a")
+    with _pipelines_lock:
+        if lang not in _pipelines:
+            _pipelines[lang] = KPipeline(lang_code=lang, model=_onnx is None)
+        return _pipelines[lang]
+
+
+get_pipeline_for_voice("af_heart")
+
+
+def _synth(text: str, voice_id: str, speed: float) -> np.ndarray:
+    """Float32 mono audio at 24 kHz for one piece of text."""
+    pipeline = get_pipeline_for_voice(voice_id)
+    segments = []
+    with _synth_lock:
+        if _onnx is not None:
+            for result in pipeline(text, voice=voice_id, split_pattern=None):
+                if result.phonemes:
+                    audio, _sr = _onnx.create(result.phonemes, voice=voice_id, speed=speed, is_phonemes=True)
+                    segments.append(np.asarray(audio, dtype=np.float32))
+        else:
+            for _gs, _ps, audio in pipeline(text, voice=voice_id, speed=speed, split_pattern=None):
+                seg = np.asarray(audio, dtype=np.float32)
+                if seg.ndim > 1:
+                    seg = seg.mean(axis=-1)
+                if seg.size:
+                    segments.append(seg)
+    return np.concatenate(segments) if segments else np.zeros(0, dtype=np.float32)
+
+
+_SENTENCE = re.compile(r"(?<=[.!?…])\s+")
+_CLAUSE = re.compile(r"(?<=[,;:])\s+")
+
+
+def split_for_streaming(text: str) -> list:
+    """Sentences, with the first one cut at its first clause when long, so
+    the first audio arrives after synthesizing only a few words."""
+    sentences = [s.strip() for s in _SENTENCE.split(text) if s.strip()]
+    if not sentences:
+        return []
+    first = sentences[0]
+    if len(first) > FIRST_PIECE_MAX:
+        parts = _CLAUSE.split(first, maxsplit=1)
+        if len(parts) == 2 and len(parts[0]) >= 12 and len(parts[1]) >= 12:
+            return [parts[0], parts[1]] + sentences[1:]
+    return sentences
+
+
+def _to_pcm16(audio: np.ndarray, rate: int) -> bytes:
+    if rate == 16000:
+        audio = resample_poly(audio, 2, 3)  # 24 kHz -> 16 kHz
+    return (np.clip(audio, -1.0, 1.0) * 32767).astype(np.int16).tobytes()
+
+
+# ---------------------------------------------------------------- Gradio UI
+
+
+def generate_tts(text: str, voice_id: str, speed: float):
+    try:
+        text = (text or "").strip()[:MAX_CHARS]
+        if not text:
+            return None
+        audio = np.concatenate([_synth(p, voice_id, speed) for p in split_for_streaming(text)])
+        return (SAMPLE_RATE, audio) if audio.size else None
+    except Exception as e:
+        print(f"[TTS] Erro: {e}")
+        return None
+
+
+VOICES = [
+    "af_heart", "af_alloy", "af_aoede", "af_bella", "af_jessica",
+    "af_kore", "af_nicole", "af_nova", "af_river", "af_sarah", "af_sky",
+    "am_adam", "am_echo", "am_eric", "am_fenrir", "am_liam",
+    "am_michael", "am_onyx", "am_orion", "am_puck",
+    "bf_alice", "bf_emma", "bf_isabella",
+    "bm_george", "bm_lewis",
+    "pf_dora", "pm_alex", "pm_santa",
+]
+
+demo = gr.Interface(
+    fn=generate_tts,
+    inputs=[
+        gr.Textbox(label="Texto", lines=3, placeholder="Digite o texto aqui..."),
+        gr.Dropdown(choices=VOICES, value="af_heart", label="Voz"),
+        gr.Slider(minimum=0.5, maximum=2.0, value=1.0, step=0.1, label="Velocidade"),
+    ],
+    outputs=gr.Audio(label="Audio Gerado"),
+    title="Kokoro-82M TTS",
+    description="Ingles Americano / Britanico e Portugues Brasileiro.",
+)
+
+# ---------------------------------------------------------------- HTTP API
+
+app = FastAPI()
+
+
+@app.get("/tts/health")
+def tts_health():
+    return {"ok": True, "engine": ENGINE, "threads": THREADS, "voices": len(VOICES), "loaded": sorted(_pipelines)}
+
+
+@app.post("/tts/stream")
+async def tts_stream(request: Request):
+    """Body: {"text", "voice"="af_heart", "speed"=1.0, "sample_rate"=16000}.
+    Response: raw PCM16 mono little-endian, streamed piece by piece."""
+    if API_KEY and request.headers.get("x-api-key") != API_KEY:
+        raise HTTPException(status_code=401, detail="invalid api key")
+    try:
+        body = await request.json()
+    except Exception:
+        raise HTTPException(status_code=400, detail="json body required")
+    text = str(body.get("text") or "").strip()[:MAX_CHARS]
+    voice = str(body.get("voice") or "af_heart")
+    if voice not in VOICES:
+        raise HTTPException(status_code=400, detail="unknown voice")
+    try:
+        speed = min(2.0, max(0.5, float(body.get("speed") or 1.0)))
+        rate = int(body.get("sample_rate") or 16000)
+    except (TypeError, ValueError):
+        raise HTTPException(status_code=400, detail="invalid speed or sample_rate")
+    if rate not in (16000, 24000):
+        raise HTTPException(status_code=400, detail="sample_rate must be 16000 or 24000")
+    if not text:
+        return JSONResponse({"detail": "empty text"}, status_code=400)
+
+    pieces = split_for_streaming(text)
+
+    def generate():
+        t0 = time.perf_counter()
+        for i, piece in enumerate(pieces):
+            pcm = _to_pcm16(_synth(piece, voice, speed), rate)
+            if i == 0:
+                print(f"[TTS stream] first piece in {time.perf_counter() - t0:.2f}s ({len(piece)} chars)")
+            yield pcm
+
+    return StreamingResponse(
+        generate(),  # sync generator: Starlette runs it in a worker thread
+        media_type="application/octet-stream",
+        headers={"X-Sample-Rate": str(rate), "X-Pieces": str(len(pieces)), "Cache-Control": "no-store"},
+    )
+
+
+app = gr.mount_gradio_app(app, demo, path="/")
+
+
+def _warmup():
+    try:
+        _synth("Hello there.", "af_heart", 1.0)
+        print("[TTS] warm-up done")
+    except Exception as e:
+        print(f"[TTS] warm-up failed: {e}")
+
+
+threading.Thread(target=_warmup, daemon=True).start()
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")))

requirements.txt

@@ -1,6 +1,10 @@
-kokoro>=0.9.4
-soundfile>=0.12.1
-numpy>=1.24.0
-gradio>=4.0.0
-torch
-scipy
+kokoro>=0.9.4
+kokoro-onnx>=0.4.9
+onnxruntime>=1.20
+soundfile>=0.12.1
+numpy>=1.24.0
+gradio>=4.0.0
+torch
+scipy
+fastapi
+uvicorn