generate_mnn.py

#!/usr/bin/env python3
"""Generate speech with Pocket TTS MNN."""

import argparse
import time
from pathlib import Path

from pocket_tts_mnn import PocketTTSMnn


def main() -> int:
    parser = argparse.ArgumentParser(description="Pocket TTS MNN generate")
    parser.add_argument("text")
    parser.add_argument("voice", help="Reference wav path or built-in voice name")
    parser.add_argument("output", help="Output wav path")
    parser.add_argument("--models-dir", default=".")
    parser.add_argument("--lsd-steps", type=int, default=1)
    parser.add_argument("--temperature", type=float, default=0.7)
    args = parser.parse_args()

    print("Loading MNN models...")
    t0 = time.time()
    tts = PocketTTSMnn(
        models_dir=args.models_dir,
        mnn_dir=str(Path(args.models_dir) / "mnn"),
        temperature=args.temperature,
        lsd_steps=args.lsd_steps,
    )
    print(f"Loaded in {time.time() - t0:.2f}s")

    t0 = time.time()
    audio = tts.generate(args.text, voice=args.voice)
    gen = time.time() - t0
    dur = len(audio) / tts.sample_rate
    print(f"Generated {dur:.2f}s audio in {gen:.2f}s (RTFx: {dur / gen:.2f}x)")
    tts.save_audio(audio, args.output)
    print(f"Saved {args.output}")
    return 0


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