cli.py

"""CLI for testing TTS and STT against the running server."""

import argparse
import subprocess
import sys
import tempfile
from pathlib import Path

import requests

DEFAULT_URL = "http://localhost:8000"


def tts(base_url: str, text: str, voice: str | None, exaggeration: float, cfg_weight: float, output: str) -> None:
    url = f"{base_url}/tts/tts"
    data = {"text": text, "exaggeration": str(exaggeration), "cfg_weight": str(cfg_weight)}
    files = {}
    if voice:
        files["voice"] = open(voice, "rb")

    print(f"Generating speech: {text!r}")
    resp = requests.post(url, data=data, files=files if files else None)
    if resp.status_code != 200:
        print(f"Error {resp.status_code}: {resp.text}", file=sys.stderr)
        return

    Path(output).write_bytes(resp.content)
    print(f"Saved to {output}")

    # Try to play it
    try:
        subprocess.run(["afplay", output], check=True)
    except FileNotFoundError:
        try:
            subprocess.run(["aplay", output], check=True)
        except FileNotFoundError:
            print("(install afplay or aplay to auto-play)")


def stt(base_url: str, audio_path: str) -> str:
    url = f"{base_url}/stt/stt"
    with open(audio_path, "rb") as f:
        resp = requests.post(url, files={"audio": f})

    if resp.status_code != 200:
        print(f"Error {resp.status_code}: {resp.text}", file=sys.stderr)
        return ""

    result = resp.json()
    print(f"Language: {result['language']} ({result['language_probability']})")
    print(f"Text: {result['text']}")
    for seg in result["segments"]:
        print(f"  [{seg['start']:.1f}s - {seg['end']:.1f}s] {seg['text']}")
    return result["text"]


def record_audio(duration: int = 5) -> str:
    """Record audio from microphone using sox/rec."""
    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    tmp.close()
    print(f"Recording {duration}s of audio... (speak now)")
    try:
        subprocess.run(
            ["rec", "-q", tmp.name, "rate", "16000", "channels", "1", "trim", "0", str(duration)],
            check=True,
        )
    except FileNotFoundError:
        print("Install sox for mic recording: brew install sox", file=sys.stderr)
        sys.exit(1)
    print("Recording done.")
    return tmp.name


def health(base_url: str) -> bool:
    try:
        resp = requests.get(f"{base_url}/health", timeout=3)
        data = resp.json()
        print(f"Server: {data}")
        return True
    except (requests.ConnectionError, requests.Timeout):
        print(f"Cannot reach server at {base_url}", file=sys.stderr)
        return False


def interactive_loop(base_url: str, voice: str | None, exaggeration: float, cfg_weight: float) -> None:
    """Interactive REPL for testing TTS/STT."""
    print("gpu-tts interactive mode")
    print("Commands:")
    print("  <text>        — synthesize speech from text (TTS)")
    print("  /stt [file]   — transcribe audio (records from mic if no file)")
    print("  /voice <file> — set voice reference WAV for cloning")
    print("  /health       — check server status")
    print("  /quit         — exit")
    print()

    current_voice = voice

    while True:
        try:
            line = input("> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break

        if not line:
            continue

        if line == "/quit":
            break
        elif line == "/health":
            health(base_url)
        elif line.startswith("/voice"):
            parts = line.split(maxsplit=1)
            if len(parts) < 2:
                if current_voice:
                    print(f"Current voice: {current_voice}")
                else:
                    print("No voice set. Usage: /voice <path-to-wav>")
            else:
                path = parts[1]
                if Path(path).exists():
                    current_voice = path
                    print(f"Voice set to: {current_voice}")
                else:
                    print(f"File not found: {path}")
        elif line.startswith("/stt"):
            parts = line.split(maxsplit=1)
            if len(parts) > 1 and Path(parts[1]).exists():
                stt(base_url, parts[1])
            else:
                audio_path = record_audio()
                stt(base_url, audio_path)
                Path(audio_path).unlink(missing_ok=True)
        else:
            tts(base_url, line, current_voice, exaggeration, cfg_weight, "output.wav")


def main() -> None:
    parser = argparse.ArgumentParser(description="GPU TTS/STT CLI")
    parser.add_argument("--url", default=DEFAULT_URL, help="Server URL")
    sub = parser.add_subparsers(dest="command")

    # tts
    tts_p = sub.add_parser("tts", help="Text to speech")
    tts_p.add_argument("text", help="Text to synthesize")
    tts_p.add_argument("--voice", help="Voice reference WAV for cloning")
    tts_p.add_argument("--exaggeration", type=float, default=0.5)
    tts_p.add_argument("--cfg-weight", type=float, default=0.5)
    tts_p.add_argument("-o", "--output", default="output.wav")

    # stt
    stt_p = sub.add_parser("stt", help="Speech to text")
    stt_p.add_argument("audio", nargs="?", help="Audio file (records from mic if omitted)")
    stt_p.add_argument("--duration", type=int, default=5, help="Recording duration in seconds")

    # interactive
    sub.add_parser("interactive", help="Interactive REPL")

    # health
    sub.add_parser("health", help="Check server health")

    args = parser.parse_args()

    if args.command == "tts":
        tts(args.url, args.text, args.voice, args.exaggeration, args.cfg_weight, args.output)
    elif args.command == "stt":
        audio = args.audio or record_audio(args.duration)
        stt(args.url, audio)
        if not args.audio:
            Path(audio).unlink(missing_ok=True)
    elif args.command == "health":
        health(args.url)
    elif args.command == "interactive":
        if health(args.url):
            interactive_loop(args.url, None, 0.5, 0.5)
    else:
        if health(args.url):
            interactive_loop(args.url, None, 0.5, 0.5)
        else:
            parser.print_help()


if __name__ == "__main__":
    main()