kokoro_cli.py

"""Minimal Kokoro v1.1-zh CLI wrapper for local_tts_server.

Usage:
    python kokoro_cli.py <text_file> <out_file> <voice> <speed>

Reads text from <text_file>, synthesizes with kokoro, writes WAV to <out_file>.
"""

from __future__ import annotations

import argparse
import os
import sys
import wave
from pathlib import Path

import numpy as np

from local_tts_profiles import (
    DEFAULT_KOKORO_REPO_ID,
    DEFAULT_KOKORO_VOICE,
    available_kokoro_voices,
    find_kokoro_model_file,
    resolve_kokoro_model_dir,
    resolve_kokoro_voice_file,
)

SAMPLE_RATE = 24000


def _audio_from_result(result):
    if hasattr(result, "audio"):
        return result.audio
    if isinstance(result, tuple) and result:
        return result[-1]
    return None


def _infer_lang_code(voice: str) -> str:
    # Kokoro uses single-letter lang codes: z=zh, a=en-us, b=en-gb.
    if voice.startswith(("a", "af", "am")):
        return "a"
    if voice.startswith(("b", "bf", "bm")):
        return "b"
    return "z"


def _speed_callable(base_speed: float):
    """Mitigate rushed long Chinese phoneme sequences in v1.1-zh."""

    base = base_speed if base_speed > 0 else 1.0

    def speed_by_len(len_ps: int) -> float:
        speed = 1.0
        if len_ps > 83 and len_ps < 183:
            speed = 1.0 - (len_ps - 83) / 500.0
        elif len_ps >= 183:
            speed = 0.8
        return max(0.5, speed * base)

    return speed_by_len


def synthesize(text_path: str, out_path: str, voice: str, speed: float) -> int:
    try:
        import torch
        from kokoro import KModel, KPipeline
    except ImportError:
        print(
            'kokoro v1.1-zh deps missing. Run: uv pip install "kokoro>=0.8.2" "misaki[zh]>=0.8.2"',
            file=sys.stderr,
        )
        return 1

    text = Path(text_path).read_text(encoding="utf-8").strip()
    if not text:
        print("Empty text file", file=sys.stderr)
        return 1

    model_dir = resolve_kokoro_model_dir()
    repo_id = os.getenv("LOCAL_TTS_KOKORO_REPO_ID", DEFAULT_KOKORO_REPO_ID).strip() or DEFAULT_KOKORO_REPO_ID
    voice = (voice or "").strip() or os.getenv("LOCAL_TTS_KOKORO_DEFAULT_VOICE", DEFAULT_KOKORO_VOICE)
    available_voices = set(available_kokoro_voices(model_dir))
    if available_voices and voice not in available_voices:
        fallback_voice = os.getenv("LOCAL_TTS_KOKORO_DEFAULT_VOICE", DEFAULT_KOKORO_VOICE).strip() or DEFAULT_KOKORO_VOICE
        if fallback_voice not in available_voices:
            fallback_voice = sorted(available_voices)[0]
        print(
            f"Kokoro voice '{voice}' not found in local model dir; falling back to '{fallback_voice}'.",
            file=sys.stderr,
        )
        voice = fallback_voice
    pipeline_voice = resolve_kokoro_voice_file(voice, model_dir)
    lang = _infer_lang_code(voice)
    device = os.getenv("LOCAL_TTS_KOKORO_DEVICE", "").strip()
    if not device:
        device = "cuda" if torch.cuda.is_available() else "cpu"

    if model_dir:
        config_path = model_dir / "config.json"
        model_path = find_kokoro_model_file(model_dir)
        if not config_path.is_file() or model_path is None:
            print(
                f"Invalid LOCAL_TTS_KOKORO_MODEL_DIR: {model_dir} "
                "(expected config.json and a .pth model file)",
                file=sys.stderr,
            )
            return 1
        model = KModel(repo_id=repo_id, config=str(config_path), model=str(model_path)).to(device).eval()
    else:
        model = KModel(repo_id=repo_id).to(device).eval()

    en_pipeline = None
    en_callable = None
    if lang == "z":
        en_pipeline = KPipeline(lang_code="a", repo_id=repo_id, model=False)

        def en_callable(text_part: str):
            if text_part == "Kokoro":
                return "kˈOkəɹO"
            if text_part == "Sol":
                return "sˈOl"
            return next(en_pipeline(text_part)).phonemes

    pipeline = KPipeline(
        lang_code=lang,
        repo_id=repo_id,
        model=model,
        en_callable=en_callable,
    )
    effective_speed = _speed_callable(speed) if lang == "z" else speed
    generator = pipeline(text, voice=pipeline_voice, speed=effective_speed)

    chunks: list[np.ndarray] = []
    for result in generator:
        audio = _audio_from_result(result)
        if audio is not None:
            chunks.append(np.asarray(audio, dtype=np.float32))

    if not chunks:
        print("No audio generated", file=sys.stderr)
        return 1

    pcm = np.concatenate(chunks)
    pcm = np.clip(pcm, -1.0, 1.0)
    pcm_int16 = (pcm * 32767.0).astype(np.int16)

    with wave.open(out_path, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(SAMPLE_RATE)
        wf.writeframes(pcm_int16.tobytes())

    print(
        f"Wrote {out_path}: {len(pcm_int16)} samples @ {SAMPLE_RATE} Hz "
        f"repo={repo_id} model_dir={model_dir or '<hf-cache>'} voice={voice} device={device}"
    )
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description="Kokoro CLI wrapper for local_tts")
    parser.add_argument("text_file")
    parser.add_argument("out_file")
    parser.add_argument("voice")
    parser.add_argument("speed", type=float)
    args = parser.parse_args()
    return synthesize(args.text_file, args.out_file, args.voice, args.speed)


if __name__ == "__main__":
    raise SystemExit(main())