VoxCPM2

"""
Demonstrate VoxCPM2 voice caching over /v1/audio/speech.

The first request sends ``voice=<name>`` plus ``ref_audio``. The server stores
that reference clip as an uploaded voice, then generates via the named-voice
path and warms VoxCPM2's ref_audio_feat cache. Later requests send only the
same ``voice`` name and should hit the cache.

Usage:
    pip install httpx

    python speech_voice_cache_client.py \
        --ref-audio /path/to/reference.wav \
        --voice-name demo_voice \
        --text "This first HTTP request should warm the voice cache." \
        --text "This second HTTP request should reuse the cached voice." \
        --stream
"""

from __future__ import annotations

import argparse
import base64
import mimetypes
import os
import time
from pathlib import Path

import httpx

DEFAULT_API_BASE = "https://model-qrj02553.api.baseten.co/deployment/wp5d088/sync"


def encode_audio_to_data_url(ref_audio_path: str) -> str:
    audio_path = Path(ref_audio_path)
    if not audio_path.exists():
        raise FileNotFoundError(f"Reference audio not found: {audio_path}")

    mime_type = mimetypes.guess_type(audio_path.name)[0] or "audio/wav"
    with audio_path.open("rb") as audio_file:
        encoded = base64.b64encode(audio_file.read()).decode("utf-8")
    return f"data:{mime_type};base64,{encoded}"


def post_speech(
    *,
    client: httpx.Client,
    api_base: str,
    api_key: str,
    model: str,
    text: str,
    voice_name: str,
    output_path: Path,
    response_format: str,
    stream: bool,
    ref_audio: str | None = None,
    ref_text: str | None = None,
) -> float:
    payload: dict[str, object] = {
        "model": model,
        "input": text,
        "voice": voice_name,
        "response_format": response_format,
        "stream": stream,
    }
    if ref_audio is not None:
        payload["ref_audio"] = ref_audio
    if ref_text is not None:
        payload["ref_text"] = ref_text

    url = f"{api_base.rstrip('/')}/v1/audio/speech"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    start = time.perf_counter()
    if stream:
        bytes_written = 0
        with client.stream(
            "POST",
            url,
            json=payload,
            headers={"Authorization": f"Bearer {api_key}"},
        ) as response:
            if response.status_code != 200:
                error_body = response.read().decode("utf-8", errors="replace")
                raise RuntimeError(f"Speech request failed ({response.status_code}): {error_body[:500]}")

            with output_path.open("wb") as output_file:
                for chunk in response.iter_bytes():
                    if not chunk:
                        continue
                    output_file.write(chunk)
                    bytes_written += len(chunk)

        elapsed = time.perf_counter() - start
        print(f"  streamed {output_path} ({bytes_written:,} bytes) in {elapsed:.2f}s")
        return elapsed

    response = client.post(
        url,
        json=payload,
        headers={"Authorization": f"Bearer {api_key}"},
    )
    elapsed = time.perf_counter() - start

    if response.status_code != 200:
        raise RuntimeError(f"Speech request failed ({response.status_code}): {response.text[:500]}")

    output_path.write_bytes(response.content)
    print(f"  saved {output_path} ({len(response.content):,} bytes) in {elapsed:.2f}s")
    return elapsed


def main() -> None:
    parser = argparse.ArgumentParser(description="VoxCPM2 /v1/audio/speech voice-cache demo client")
    parser.add_argument("--api-base", default=DEFAULT_API_BASE)
    parser.add_argument("--api-key", default=os.getenv('BASETEN_API_KEY_INTERNAL'))
    parser.add_argument("--model", default="voxcpm2")
    parser.add_argument("--voice-name", default="voxcpm_http_cache_demo")
    parser.add_argument("--ref-audio", help="Local reference audio sent only on the first request")
    parser.add_argument("--ref-text", default=None, help="Optional reference transcript for first request")
    parser.add_argument(
        "--skip-upload",
        action="store_true",
        help="Assume --voice-name already exists; send no ref_audio on any request",
    )
    parser.add_argument(
        "--text",
        action="append",
        default=[],
        help="Text to synthesize. Repeat to make multiple cached requests.",
    )
    parser.add_argument("--output-dir", default="voxcpm2_speech_cache_output")
    parser.add_argument("--response-format", default="wav", choices=["wav", "pcm", "flac", "mp3", "aac", "opus"])
    parser.add_argument("--stream", action="store_true", help="Send stream=True and write the streamed response to disk")
    args = parser.parse_args()

    if not args.skip_upload and not args.ref_audio:
        raise ValueError("Pass --ref-audio for the first request, or --skip-upload for an existing voice.")
    if not args.text:
        args.text = [
            "This first HTTP speech request should upload the voice and warm the cache.",
            "This second HTTP speech request uses only the voice name and should hit the cache.",
        ]

    ref_audio_data_url = None if args.skip_upload else encode_audio_to_data_url(args.ref_audio)
    output_dir = Path(args.output_dir)

    print("HTTP /v1/audio/speech VoxCPM2 cache demo")
    if ref_audio_data_url is None:
        print(f"Using existing voice {args.voice_name!r}; no ref_audio will be sent.")
    else:
        print(f"Request 1 will send voice={args.voice_name!r} plus ref_audio.")
    if args.stream:
        print("Streaming enabled; requests will send stream=True and write chunks as they arrive.")
    print("Later requests send only the voice name. Watch server logs for VoxCPM2 cache STORE/HIT.")

    with httpx.Client(timeout=300) as client:
        for idx, text in enumerate(args.text, start=1):
            print(f"\nRequest {idx}/{len(args.text)}")
            print("  expected cache behavior:", "miss + store" if idx == 1 else "hit")
            post_speech(
                client=client,
                api_base=args.api_base,
                api_key=args.api_key,
                model=args.model,
                text=text,
                voice_name=args.voice_name,
                output_path=output_dir / f"request_{idx:02d}.{args.response_format}",
                response_format=args.response_format,
                stream=args.stream,
                ref_audio=ref_audio_data_url if idx == 1 else None,
                ref_text=args.ref_text if idx == 1 else None,
            )


if __name__ == "__main__":
    main()