scripts/verify_kokoro_voices.py

"""Discover Kokoro voice IDs available for the configured Kokoro repo.

Run after installing the TTS dependency group:

    uv sync --extra dev --extra tts
    uv run python scripts/verify_kokoro_voices.py
"""

from __future__ import annotations

import importlib
import json
import sys
from collections.abc import Iterable
from pathlib import Path
from typing import Any

KOKORO_REPO_ID = "hexgrad/Kokoro-82M"
SELECTED_VOICE_IDS = {
    "am_michael",
    "am_onyx",
    "af_bella",
    "am_puck",
    "am_fenrir",
    "bm_george",
    "hf_alpha",
    "hm_omega",
}


def main() -> None:
    report = discover_kokoro_voices()
    print(json.dumps(report, indent=2, sort_keys=True))
    voices = report.get("voices") or []
    if not voices:
        raise SystemExit("No Kokoro voices discovered. Inspect the report above.")
    missing = sorted(SELECTED_VOICE_IDS.difference(voices))
    if missing:
        raise SystemExit(f"Selected VoiceID values missing from Kokoro repo: {missing}")


def discover_kokoro_voices() -> dict[str, Any]:
    report: dict[str, Any] = {
        "repo_id": KOKORO_REPO_ID,
        "package": None,
        "voices": [],
        "selected": sorted(SELECTED_VOICE_IDS),
        "sources": [],
        "errors": [],
    }
    try:
        kokoro = importlib.import_module("kokoro")
    except Exception as exc:  # pragma: no cover - depends on optional install
        report["errors"].append(f"import kokoro failed: {exc!r}")
        return report

    report["package"] = {
        "module": getattr(kokoro, "__name__", "kokoro"),
        "file": getattr(kokoro, "__file__", None),
        "version": getattr(kokoro, "__version__", None),
    }

    voices: set[str] = set()
    voices.update(_voices_from_module_attrs(kokoro, report))
    voices.update(_voices_from_hub(report))
    voices.update(_voices_from_package_files(kokoro, report))

    report["voices"] = sorted(voices)
    return report


def _voices_from_module_attrs(module: Any, report: dict[str, Any]) -> set[str]:
    voices: set[str] = set()
    for attr_name in ("VOICES", "voices", "VOICE_NAMES", "VOICE_IDS"):
        value = getattr(module, attr_name, None)
        extracted = _extract_voice_names(value)
        if extracted:
            report["sources"].append(f"kokoro.{attr_name}")
            voices.update(extracted)
    return voices


def _voices_from_hub(report: dict[str, Any]) -> set[str]:
    try:
        from huggingface_hub import list_repo_files
    except Exception as exc:  # pragma: no cover - optional dependency guard
        report["errors"].append(f"import huggingface_hub.list_repo_files failed: {exc!r}")
        return set()
    try:
        files = list_repo_files(KOKORO_REPO_ID)
    except Exception as exc:  # pragma: no cover - network/cache dependent
        report["errors"].append(f"list_repo_files({KOKORO_REPO_ID!r}) failed: {exc!r}")
        return set()
    voices = {
        Path(file_name).stem
        for file_name in files
        if file_name.startswith("voices/") and file_name.endswith(".pt")
    }
    if voices:
        report["sources"].append(f"{KOKORO_REPO_ID}:voices/*.pt")
    return voices


def _voices_from_package_files(module: Any, report: dict[str, Any]) -> set[str]:
    module_file = getattr(module, "__file__", None)
    if not module_file:
        return set()
    root = Path(module_file).resolve().parent
    patterns = ("*.pt", "*.pth", "*.bin", "*.onnx", "*.npz")
    voices: set[str] = set()
    for pattern in patterns:
        for path in root.rglob(pattern):
            stem = path.stem
            if stem and not stem.startswith(("model", "kokoro")):
                voices.add(stem)
    if voices:
        report["sources"].append(f"package files under {root}")
    return voices


def _extract_voice_names(value: Any) -> set[str]:
    if value is None:
        return set()
    if isinstance(value, str):
        return {value}
    if isinstance(value, dict):
        return {str(key) for key in value.keys()}
    if isinstance(value, Iterable):
        return {str(item) for item in value if isinstance(item, str)}
    return set()


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(130)