backend/batch.py

"""
CLI batch tool. Imports inference.py directly. No uvicorn needed.

Usage (emotion shorthand — maps to EMOTION_TEMPLATES for speaker/language):
  python batch.py --file script.txt --speaker Yash --language Gujarati \
    --emotion happy --output ./out/ [--seed 42] [--resume]

Usage (raw description — full control):
  python batch.py --file script.txt --speaker Yash --language Gujarati \
    --description "Yash speaks very fast with a very high-pitched, extremely joyful voice." --output ./out/

--emotion and --description are mutually exclusive.
"""
import argparse
import os
import sys
import json
from pathlib import Path


VALID_EMOTIONS = ["calm", "happy", "narrative", "angry", "command", "excited", "sad", "fearful"]
# Gujarati is not in the official Rasmalai emotion benchmark — all output is cross-language transfer.
# "medium" emotions have higher variability and may sound neutral.
MEDIUM_RELIABILITY: set[str] = {"angry", "command", "excited", "sad", "fearful"}


def resolve_description(args) -> str:
    if args.description:
        return args.description

    from app.constants import EMOTION_TEMPLATES, EMOTION_METADATA
    templates = EMOTION_TEMPLATES.get(args.language, {}).get(args.speaker, {})
    template = templates.get(args.emotion)
    if not template:
        sys.exit(
            f"ERROR: No template for emotion '{args.emotion}' with speaker '{args.speaker}' "
            f"in language '{args.language}'. Use --description for a custom description."
        )
    if args.emotion in MEDIUM_RELIABILITY:
        print(
            f"WARNING: '{args.emotion}' has medium reliability for Gujarati — "
            f"emotion output is not officially benchmarked and may sound neutral.",
            file=sys.stderr,
        )
    return template


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--file", required=True)
    parser.add_argument("--speaker", default="Yash")
    parser.add_argument("--language", default="Gujarati")
    parser.add_argument("--output", default="./out/")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--resume", action="store_true")

    desc_group = parser.add_mutually_exclusive_group(required=True)
    desc_group.add_argument("--description", help="Freeform description of how to speak")
    desc_group.add_argument(
        "--emotion",
        choices=VALID_EMOTIONS,
        help="Emotion shorthand — maps to EMOTION_TEMPLATES for the given speaker/language",
    )

    args = parser.parse_args()

    from app.constants import SPEAKERS_BY_LANGUAGE
    valid_speakers = SPEAKERS_BY_LANGUAGE.get(args.language, [])
    if not valid_speakers:
        sys.exit(f"ERROR: Language '{args.language}' not supported. Valid: {list(SPEAKERS_BY_LANGUAGE)}")
    if args.speaker not in valid_speakers:
        sys.exit(f"ERROR: Speaker '{args.speaker}' not valid for language '{args.language}'. Valid: {valid_speakers}")

    description = resolve_description(args)

    out_dir = Path(args.output)
    out_dir.mkdir(parents=True, exist_ok=True)

    lines = Path(args.file).read_text().strip().splitlines()

    from app.service import tts_service
    from app.config import settings
    import torch
    device = "mps" if torch.backends.mps.is_available() else "cpu"
    print(f"Loading model on {device}...")
    tts_service.load(settings.model_name, device)
    print(f"Ready. Processing {len(lines)} lines...")

    manifest = []
    for i, line in enumerate(lines):
        out_path = out_dir / f"{i+1:04d}.wav"
        if args.resume and out_path.exists():
            print(f"  [{i+1}/{len(lines)}] skip (exists): {out_path}")
            manifest.append({"index": i+1, "text": line, "file": str(out_path), "skipped": True})
            continue

        print(f"  [{i+1}/{len(lines)}] {line[:60]}...")
        audio_bytes = tts_service.run_inference(line, description)

        if len(audio_bytes) < 1000:
            print(f"    WARNING: output suspiciously small ({len(audio_bytes)} bytes), skipping")
            manifest.append({"index": i+1, "text": line, "file": None, "error": "too_short"})
            continue

        out_path.write_bytes(audio_bytes)
        manifest.append({"index": i+1, "text": line, "file": str(out_path)})

    manifest_path = out_dir / "manifest.json"
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2))
    print(f"Done. Manifest: {manifest_path}")


if __name__ == "__main__":
    main()