scripts/serve_s2s_api.py

#!/usr/bin/env python3
"""Serve the Urdu S2S API."""

from __future__ import annotations

import argparse
import os
from pathlib import Path
import sys

ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

from urdu_s2s.api import DEFAULT_BASELINE_CSV, create_app  # noqa: E402


def env(name: str, default: str) -> str:
    return os.environ.get(name, default)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--candidate-csv",
        type=Path,
        default=Path(env("S2S_CANDIDATE_CSV", str(DEFAULT_BASELINE_CSV))),
    )
    parser.add_argument("--host", default=env("S2S_HOST", "127.0.0.1"))
    parser.add_argument("--port", type=int, default=int(env("S2S_PORT", "8000")))
    parser.add_argument("--default-mode", default=env("S2S_DEFAULT_MODE", "text_only"))
    parser.add_argument("--whisper-model", default=env("S2S_WHISPER_MODEL", "large-v3"))
    parser.add_argument("--whisper-language", default=env("S2S_WHISPER_LANGUAGE", "ur"))
    parser.add_argument("--whisper-device", default=env("S2S_WHISPER_DEVICE", "cpu"))
    parser.add_argument("--whisper-compute-type", default=env("S2S_WHISPER_COMPUTE_TYPE", "int8"))
    parser.add_argument(
        "--voice-prompt-audio-path",
        type=Path,
        default=Path(
            env(
                "S2S_VOICE_PROMPT_AUDIO_PATH",
                "data/processed/voice_anchors/chatterbox_praxy_v1/bench_025.wav",
            )
        ),
    )
    parser.add_argument("--chatterbox-device", default=env("S2S_CHATTERBOX_DEVICE", "cuda"))
    parser.add_argument("--chatterbox-t3-model", default=env("S2S_CHATTERBOX_T3_MODEL", "v3"))
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        import uvicorn
    except ModuleNotFoundError as exc:
        raise SystemExit(
            "Missing API dependencies. Install them with: "
            "python3 -m pip install -r requirements-api.txt"
        ) from exc

    candidate_csv = args.candidate_csv if args.candidate_csv.is_absolute() else ROOT / args.candidate_csv
    voice_prompt_audio_path = (
        args.voice_prompt_audio_path
        if args.voice_prompt_audio_path.is_absolute()
        else ROOT / args.voice_prompt_audio_path
    )
    app = create_app(
        candidate_csv=candidate_csv,
        repo_root=ROOT,
        default_mode=args.default_mode,
        default_whisper_model=args.whisper_model,
        default_whisper_language=args.whisper_language,
        default_whisper_device=args.whisper_device,
        default_whisper_compute_type=args.whisper_compute_type,
        default_voice_prompt_audio_path=voice_prompt_audio_path,
        default_chatterbox_device=args.chatterbox_device,
        default_chatterbox_t3_model=args.chatterbox_t3_model,
    )
    uvicorn.run(app, host=args.host, port=args.port)
    return 0


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