tts_engine.py

import argparse
import asyncio
import wave
from pathlib import Path
from typing import Any
from urllib.request import urlopen

import edge_tts
import numpy as np
from kokoro_onnx import Kokoro
from supertonic import TTS


DEFAULT_TEXT = "Hello. This is the default sample used to validate edge-tts output."
DEFAULT_VOICE = "en-US-AriaNeural"
DEFAULT_KOKORO_VOICE = "af_heart"
DEFAULT_KOKORO_LANG = "en-us"
KOKORO_MODEL_URL = "https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0/kokoro-v1.0.onnx"
KOKORO_VOICES_URL = "https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0/voices-v1.0.bin"
KOKORO_MODELS_DIR = Path(__file__).resolve().parent / "models" / "kokoro"
KOKORO_MODEL_FILE = KOKORO_MODELS_DIR / "kokoro-v1.0.onnx"
KOKORO_VOICES_FILE = KOKORO_MODELS_DIR / "voices-v1.0.bin"
KOKORO_VOICE_NAMES = [
    "af_alloy", "af_aoede", "af_bella", "af_heart", "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_puck", "am_santa", "bf_alice", "bf_emma", "bf_isabella", "bf_lily", "bm_daniel", "bm_fable", "bm_george",
    "bm_lewis", "ef_dora", "em_alex", "em_santa", "ff_siwis", "hf_alpha", "hf_beta", "hm_omega", "hm_psi",
    "if_sara", "im_nicola", "jf_alpha", "jf_gongitsune", "jf_nezumi", "jf_tebukuro", "jm_kumo", "pf_dora",
    "pm_alex", "pm_santa", "zf_xiaobei", "zf_xiaoni", "zf_xiaoxiao", "zf_xiaoyi", "zm_yunjian", "zm_yunxi",
    "zm_yunxia", "zm_yunyang",
]
SUPERTONIC_VOICE_NAMES = [
    "F1", "F2", "F3", "F4", "F5",
    "M1", "M2", "M3", "M4", "M5",
]

_KOKORO_ENGINES: dict[str, Kokoro] = {}
_SUPERTONIC_ENGINE: TTS | None = None


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Edge TTS test utility.")
    parser.add_argument("--text", help="Text content to synthesize.")
    parser.add_argument("--file", type=Path, help="Load synthesis text from a file.")
    parser.add_argument("--voice", default=DEFAULT_VOICE, help="Edge TTS voice name.")
    parser.add_argument("--rate", default="+0%", help="Speech rate, for example: +10%%")
    parser.add_argument("--volume", default="+0%", help="Speech volume, for example: +0%%")
    parser.add_argument("--pitch", default="+0Hz", help="Speech pitch, for example: +0Hz")
    parser.add_argument("--output", type=Path, default=Path("output.mp3"), help="Output audio file path.")
    parser.add_argument("--list-voices", action="store_true", help="List available voices.")
    parser.add_argument("--filter", help="Filter voices by keyword, for example: en-US")
    return parser.parse_args()


def load_text(args: argparse.Namespace) -> str:
    if args.text:
        return args.text.strip()
    if args.file:
        return args.file.read_text(encoding="utf-8").strip()
    return DEFAULT_TEXT


async def get_voices(keyword: str | None = None) -> list[dict[str, Any]]:
    voices = await edge_tts.list_voices()
    if keyword:
        keyword = keyword.lower()
        voices = [
            voice
            for voice in voices
            if keyword in voice["ShortName"].lower() or keyword in voice["Locale"].lower()
        ]

    return voices


def _kokoro_group_from_voice(voice_name: str) -> str:
    normalized = (voice_name or DEFAULT_KOKORO_VOICE).strip().lower()
    return normalized.split("_", 1)[0] if "_" in normalized else "af"


def _download_file(url: str, destination: Path) -> None:
    destination.parent.mkdir(parents=True, exist_ok=True)
    if destination.is_file():
        return
    with urlopen(url, timeout=90) as response:
        destination.write_bytes(response.read())


def _ensure_kokoro_engine(voice_name: str = DEFAULT_KOKORO_VOICE) -> Kokoro:
    group = _kokoro_group_from_voice(voice_name)
    cached = _KOKORO_ENGINES.get(group)
    if cached is not None:
        return cached

    _download_file(KOKORO_MODEL_URL, KOKORO_MODEL_FILE)
    _download_file(KOKORO_VOICES_URL, KOKORO_VOICES_FILE)
    engine = Kokoro(model_path=str(KOKORO_MODEL_FILE), voices_path=str(KOKORO_VOICES_FILE))
    _KOKORO_ENGINES[group] = engine
    return engine


def get_kokoro_voices(keyword: str | None = None) -> list[dict[str, Any]]:
    names = KOKORO_VOICE_NAMES
    voices = [
        {
            "ShortName": name,
            "Locale": "multi",
            "Gender": "Unknown",
            "FriendlyName": f"Kokoro {name}",
        }
        for name in names
    ]
    if not keyword:
        return voices
    needle = keyword.lower()
    return [v for v in voices if needle in v["ShortName"].lower()]


def _write_float_audio_to_wav(output: Path, samples: np.ndarray, sample_rate: int) -> None:
    pcm = np.clip(samples, -1.0, 1.0)
    pcm_int16 = (pcm * 32767.0).astype(np.int16)
    with wave.open(str(output), "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(sample_rate)
        wav_file.writeframes(pcm_int16.tobytes())


async def synthesize_kokoro_to_file(
    *,
    text: str,
    output: Path,
    voice: str = DEFAULT_KOKORO_VOICE,
    speed: float = 1.0,
    lang: str = DEFAULT_KOKORO_LANG,
) -> Path:
    if not text:
        raise ValueError("Input text cannot be empty.")

    output.parent.mkdir(parents=True, exist_ok=True)
    engine = _ensure_kokoro_engine(voice)
    audio, sample_rate = engine.create(text=text, voice=voice, speed=speed, lang=lang)
    _write_float_audio_to_wav(output=output, samples=audio, sample_rate=sample_rate)
    return output.resolve()


def prewarm_kokoro() -> None:
    _ensure_kokoro_engine(DEFAULT_KOKORO_VOICE)


def _ensure_supertonic_engine() -> TTS:
    global _SUPERTONIC_ENGINE
    if _SUPERTONIC_ENGINE is None:
        _SUPERTONIC_ENGINE = TTS(auto_download=True)
    return _SUPERTONIC_ENGINE


def get_supertonic_voices(keyword: str | None = None) -> list[dict[str, Any]]:
    names = SUPERTONIC_VOICE_NAMES
    voices = [
        {
            "ShortName": name,
            "Locale": "en",
            "Gender": "Female" if name.startswith("F") else "Male",
            "FriendlyName": f"Supertonic {name}",
        }
        for name in names
    ]
    if not keyword:
        return voices
    needle = keyword.lower()
    return [v for v in voices if needle in v["ShortName"].lower()]


async def synthesize_supertonic_to_file(
    *,
    text: str,
    output: Path,
    voice: str = "M1",
    lang: str = "en",
) -> Path:
    if not text:
        raise ValueError("Input text cannot be empty.")

    output.parent.mkdir(parents=True, exist_ok=True)
    engine = _ensure_supertonic_engine()
    style = engine.get_voice_style(voice_name=voice)
    wav, duration = engine.synthesize(text, voice_style=style, lang=lang)
    engine.save_audio(wav, str(output))
    return output.resolve()


async def list_voices(keyword: str | None) -> None:
    voices = await get_voices(keyword)

    if not voices:
        print("No matching voice was found.")
        return

    for voice in voices:
        print(f'{voice["ShortName"]} | {voice["Locale"]} | {voice["Gender"]}')


async def synthesize_to_file(
    *,
    text: str,
    output: Path,
    voice: str = DEFAULT_VOICE,
    rate: str = "+0%",
    volume: str = "+0%",
    pitch: str = "+0Hz",
) -> Path:
    if not text:
        raise ValueError("Input text cannot be empty.")

    output.parent.mkdir(parents=True, exist_ok=True)

    communicate = edge_tts.Communicate(
        text=text,
        voice=voice,
        rate=rate,
        volume=volume,
        pitch=pitch,
    )
    await communicate.save(str(output))
    return output.resolve()


async def generate_tts(args: argparse.Namespace) -> None:
    output_path = await synthesize_to_file(
        text=load_text(args),
        output=args.output,
        voice=args.voice,
        rate=args.rate,
        volume=args.volume,
        pitch=args.pitch,
    )
    print(f"Audio file generated: {output_path}")
    print(f"Voice: {args.voice}")


async def async_main() -> None:
    args = parse_args()
    if args.list_voices:
        await list_voices(args.filter)
        return
    await generate_tts(args)


if __name__ == "__main__":
    asyncio.run(async_main())