server.py

"""Local Kokoro TTS HTTP server with chunked, streaming synthesis.

Long selections are split into sentence-bounded chunks. The first chunk
generates and starts playing while later chunks generate in the background,
so the time-to-first-word is roughly one short sentence regardless of total
selection length.

POST   /speak    {"text": "...", "voice": "af_heart", "speed": 1.0}
POST   /stop
GET    /health
GET    /voices
GET    /config
POST   /config   {"voice": "...", "speed": 1.0}   (partial ok)
"""
from __future__ import annotations

import json
import logging
import os
import queue
import re
import subprocess
import tempfile
import threading
import uuid
from pathlib import Path

import soundfile as sf
from flask import Flask, jsonify, request
from kokoro_onnx import Kokoro

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("tts")

ROOT = Path.home() / ".local/tts"
MODEL_DIR = ROOT / "models"
MODEL_PATH = MODEL_DIR / "kokoro-v1.0.onnx"
VOICES_PATH = MODEL_DIR / "voices-v1.0.bin"
CONFIG_PATH = ROOT / "config.json"

DEFAULTS = {"voice": "af_heart", "speed": 1.0, "lang": "en-us"}

# Chunk size targets. Smaller first chunk minimizes time-to-first-word;
# larger later chunks keep total synth calls down.
FIRST_CHUNK_MAX = 140
CHUNK_TARGET_CHARS = 380
QUEUE_LOOKAHEAD = 3   # max chunks pre-generated and waiting to play

if not MODEL_PATH.exists() or not VOICES_PATH.exists():
    raise SystemExit(
        f"Missing model files. Expected:\n  {MODEL_PATH}\n  {VOICES_PATH}\n"
        "Run download_models.sh first."
    )


def load_config() -> dict:
    cfg = dict(DEFAULTS)
    if CONFIG_PATH.exists():
        try:
            cfg.update(json.loads(CONFIG_PATH.read_text()))
        except Exception as e:
            log.warning("config load failed: %s", e)
    return cfg


def save_config(cfg: dict) -> None:
    CONFIG_PATH.write_text(json.dumps(cfg, indent=2))


def split_text(
    text: str,
    first_max: int = FIRST_CHUNK_MAX,
    max_chars: int = CHUNK_TARGET_CHARS,
) -> list[str]:
    """Split into chunks ending on sentence boundaries when possible.

    First chunk is intentionally smaller so playback starts faster;
    later chunks are larger to reduce total synthesis calls.
    """
    text = (text or "").strip()
    if not text:
        return []
    parts = re.split(r"(?<=[.!?…])\s+|\n{2,}", text)
    chunks: list[str] = []
    cur = ""
    target = first_max
    for p in parts:
        p = p.strip()
        if not p:
            continue
        if not cur:
            cur = p
            continue
        candidate = cur + " " + p
        if len(candidate) <= target:
            cur = candidate
        else:
            chunks.append(cur)
            cur = p
            target = max_chars
    if cur:
        chunks.append(cur)
    # hard-split any sentence-less walls of text
    final: list[str] = []
    for c in chunks:
        while len(c) > max_chars * 2:
            split_at = c.rfind(" ", 0, max_chars * 2)
            if split_at <= 0:
                split_at = max_chars * 2
            final.append(c[:split_at].strip())
            c = c[split_at:].lstrip()
        if c.strip():
            final.append(c.strip())
    return final


log.info("Loading Kokoro model...")
kokoro = Kokoro(str(MODEL_PATH), str(VOICES_PATH))
log.info("Warming up...")
try:
    kokoro.create("ready.", voice=DEFAULTS["voice"], speed=1.0, lang=DEFAULTS["lang"])
except Exception as e:
    log.warning("warmup failed: %s", e)
log.info("Kokoro ready.")

config = load_config()
log.info("Config: %s", config)

app = Flask(__name__)

# ─── Streaming pipeline ──────────────────────────────────────────────────

_state_lock = threading.Lock()
_session_id = 0
_proc: subprocess.Popen | None = None
_play_queue: queue.Queue = queue.Queue(maxsize=QUEUE_LOOKAHEAD)


def _drain_queue() -> None:
    while True:
        try:
            _, path = _play_queue.get_nowait()
        except queue.Empty:
            return
        if path:
            try:
                os.unlink(path)
            except OSError:
                pass


def _cancel_current() -> bool:
    """Bump session, kill any playing afplay, drain queued chunks."""
    global _proc, _session_id
    stopped = False
    with _state_lock:
        _session_id += 1
        if _proc is not None and _proc.poll() is None:
            _proc.terminate()
            try:
                _proc.wait(timeout=1.0)
            except subprocess.TimeoutExpired:
                _proc.kill()
            stopped = True
        _proc = None
    _drain_queue()
    return stopped


def _start_session() -> int:
    """Cancel anything in flight and return a fresh session id."""
    _cancel_current()
    with _state_lock:
        return _session_id


def _producer(session: int, chunks: list[str], voice: str, speed: float, lang: str) -> None:
    for chunk in chunks:
        with _state_lock:
            if session != _session_id:
                return
        try:
            samples, sr = kokoro.create(chunk, voice=voice, speed=speed, lang=lang)
        except Exception:
            log.exception("synthesis failed for chunk")
            return

        path = str(Path(tempfile.gettempdir()) / f"tts-{uuid.uuid4().hex}.wav")
        sf.write(path, samples, sr)

        with _state_lock:
            stale = session != _session_id
        if stale:
            try:
                os.unlink(path)
            except OSError:
                pass
            return

        # Bounded queue — blocks until consumer drains a slot.
        # If session is cancelled while we're blocked, the drain will
        # free the slot, and the next iteration's session check stops us.
        while True:
            try:
                _play_queue.put((session, path), timeout=0.5)
                break
            except queue.Full:
                with _state_lock:
                    if session != _session_id:
                        try:
                            os.unlink(path)
                        except OSError:
                            pass
                        return


def _consumer_loop() -> None:
    global _proc
    while True:
        session, path = _play_queue.get()
        with _state_lock:
            current = _session_id
        if session != current or not path or not os.path.exists(path):
            if path:
                try:
                    os.unlink(path)
                except OSError:
                    pass
            continue

        proc = subprocess.Popen(["afplay", path])
        with _state_lock:
            _proc = proc
        proc.wait()
        try:
            os.unlink(path)
        except OSError:
            pass


threading.Thread(target=_consumer_loop, daemon=True).start()

# ─── HTTP routes ─────────────────────────────────────────────────────────

_cfg_lock = threading.Lock()


@app.post("/speak")
def speak():
    data = request.get_json(force=True, silent=True) or {}
    text = (data.get("text") or "").strip()
    if not text:
        return jsonify({"error": "no text"}), 400

    with _cfg_lock:
        voice = data.get("voice") or config["voice"]
        speed = float(data.get("speed") or config["speed"])
        lang = data.get("lang") or config["lang"]

    chunks = split_text(text)
    if not chunks:
        return jsonify({"error": "empty after splitting"}), 400

    session = _start_session()
    threading.Thread(
        target=_producer,
        args=(session, chunks, voice, speed, lang),
        daemon=True,
    ).start()

    log.info(
        "speak voice=%s speed=%s len=%d chunks=%d session=%d",
        voice, speed, len(text), len(chunks), session,
    )
    return jsonify({"ok": True, "chunks": len(chunks), "session": session})


@app.post("/stop")
def stop():
    return jsonify({"ok": True, "stopped": _cancel_current()})


@app.get("/health")
def health():
    return jsonify({"ok": True})


@app.get("/voices")
def voices():
    return jsonify({"voices": sorted(kokoro.get_voices())})


@app.get("/config")
def get_config():
    with _cfg_lock:
        return jsonify(config)


@app.post("/config")
def set_config():
    data = request.get_json(force=True, silent=True) or {}
    allowed = {"voice", "speed", "lang"}
    updates = {k: v for k, v in data.items() if k in allowed}
    if "speed" in updates:
        updates["speed"] = float(updates["speed"])
    if "voice" in updates and updates["voice"] not in kokoro.get_voices():
        return jsonify({"error": f"unknown voice: {updates['voice']}"}), 400

    with _cfg_lock:
        config.update(updates)
        save_config(config)
        snapshot = dict(config)
    log.info("config updated: %s", updates)
    return jsonify(snapshot)


if __name__ == "__main__":
    app.run(host="127.0.0.1", port=8765, threaded=True)