esh-pfi-infrastructure/stacks/zonos/adapter/server.py

"""OpenAI-ish /v1/audio/speech adapter in front of the Zonos Python SDK.

Why this exists: the upstream Zyphra/Zonos repo ships only a Gradio WebUI
+ Python SDK — no REST endpoint. asset-engine routes generation over a
clean JSON POST (every other TTS in the catalog speaks /v1/audio/speech),
so this thin FastAPI layer adapts the SDK's generate path to that wire.

Wire shape (request):  POST /v1/audio/speech  application/json
Wire shape (response): JSON envelope {audio: <base64>, audio_format:
                       "audio/wav", seed: <int>}  — mirrors the
                       kokoro-captioned envelope, plus `seed` so
                       asset-engine's regenerate/fork can pin it
                       (Zonos is the first genuinely seedable TTS in the
                       fleet: catalog reproducibility.seed_field=seed).

The model is loaded once at startup. Generation is serialized behind an
asyncio lock — the SDK model and the speaker-embedding cache are not
concurrency-safe, and there is one GPU.
"""

from __future__ import annotations

import asyncio
import base64
import io
import os
from typing import Optional

import soundfile as sf
import torch
import torchaudio
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field

from zonos.conditioning import make_cond_dict, supported_language_codes
from zonos.model import Zonos
from zonos.utils import DEFAULT_DEVICE as DEVICE

MODEL_ID = os.getenv("ZONOS_MODEL", "Zyphra/Zonos-v0.1-transformer")
VOICES_DIR = os.getenv("ZONOS_VOICES_DIR", "/app/voices")

# Blessed defaults lifted verbatim from upstream gradio_interface.py — the
# model authors picked these for human-facing UX (CATALOG-CONTRACT source
# precedence: Gradio UI for blessed defaults/ranges).
SAMPLER = dict(top_p=0.0, top_k=0, min_p=0.0, linear=0.5, conf=0.40, quad=0.00)
MAX_NEW_TOKENS = 86 * 30  # ~30s ceiling, as upstream

app = FastAPI(title="zonos-adapter", version="0.1.0")
_lock = asyncio.Lock()
_model: Optional[Zonos] = None
# (voice_path -> embedding) cache; recomputing the speaker embedding per
# request is the dominant avoidable cost for repeated clones.
_spk_cache: dict[str, torch.Tensor] = {}


class SpeechRequest(BaseModel):
    input: str = Field(..., description="Text to synthesize")
    model: str = Field(MODEL_ID, description="Only the transformer variant is loaded")
    voice: Optional[str] = Field(
        None, description="Reference clip name under VOICES_DIR (cloning); omit for default voice"
    )
    language: str = Field("en-us", description="eSpeak language code")
    response_format: str = Field("wav", description="wav | flac")
    seed: Optional[int] = Field(
        None, description="Omit for a random seed; the seed actually used is returned"
    )
    # Conditioning knobs (ranges mirror the Gradio sliders).
    speaking_rate: float = Field(15.0, ge=5.0, le=30.0)
    pitch_std: float = Field(45.0, ge=0.0, le=300.0)
    fmax: float = Field(24000.0, ge=0.0, le=24000.0)
    cfg_scale: float = Field(2.0, ge=1.0, le=5.0)
    # 8-float emotion vector [happy, sad, disgust, fear, surprise, anger,
    # other, neutral]; omit to leave emotion unconditional (Gradio default).
    emotion: Optional[list[float]] = Field(None, min_length=8, max_length=8)


def _get_model() -> Zonos:
    global _model
    if _model is None:
        _model = Zonos.from_pretrained(MODEL_ID, device=DEVICE)
        _model.requires_grad_(False).eval()
    return _model


def _speaker_embedding(model: Zonos, voice: str) -> torch.Tensor:
    path = voice if os.path.isabs(voice) else os.path.join(VOICES_DIR, voice)
    if not os.path.isfile(path):
        raise HTTPException(404, f"voice not found: {voice}")
    if path not in _spk_cache:
        wav, sr = torchaudio.load(path)
        emb = model.make_speaker_embedding(wav, sr).to(DEVICE, dtype=torch.bfloat16)
        _spk_cache[path] = emb
    return _spk_cache[path]


_FORMATS = {"wav": "audio/wav", "flac": "audio/flac"}


@app.get("/health")
async def health() -> dict:
    return {"status": "ok", "model": MODEL_ID, "loaded": _model is not None}


@app.get("/v1/audio/voices")
async def voices() -> dict:
    try:
        names = sorted(f for f in os.listdir(VOICES_DIR) if f.lower().endswith((".wav", ".flac", ".mp3")))
    except FileNotFoundError:
        names = []
    return {"voices": names}


@app.post("/v1/audio/speech")
async def speech(req: SpeechRequest) -> JSONResponse:
    if req.response_format not in _FORMATS:
        raise HTTPException(400, f"unsupported response_format: {req.response_format}")
    if req.language not in supported_language_codes:
        raise HTTPException(400, f"unsupported language: {req.language}")

    async with _lock:
        model = _get_model()

        # Resolve seed: explicit pins it; otherwise draw one and report it
        # back so asset-engine can regenerate/fork deterministically.
        seed = req.seed if req.seed is not None else int(torch.randint(0, 2**32 - 1, (1,)).item())
        torch.manual_seed(seed)

        unconditional_keys: list[str] = []
        speaker = None
        if req.voice:
            speaker = _speaker_embedding(model, req.voice)
        else:
            unconditional_keys.append("speaker")
        if req.emotion is not None:
            emotion = torch.tensor([float(x) for x in req.emotion], device=DEVICE)
        else:
            emotion = None
            unconditional_keys.append("emotion")

        cond = make_cond_dict(
            text=req.input,
            language=req.language,
            speaker=speaker,
            emotion=emotion,
            fmax=float(req.fmax),
            pitch_std=float(req.pitch_std),
            speaking_rate=float(req.speaking_rate),
            device=DEVICE,
            unconditional_keys=unconditional_keys,
        )
        conditioning = model.prepare_conditioning(cond)

        with torch.inference_mode():
            codes = model.generate(
                prefix_conditioning=conditioning,
                max_new_tokens=MAX_NEW_TOKENS,
                cfg_scale=float(req.cfg_scale),
                batch_size=1,
                sampling_params=SAMPLER,
            )
            wav = model.autoencoder.decode(codes).cpu().detach()
        sr = model.autoencoder.sampling_rate
        if wav.dim() == 2 and wav.size(0) > 1:
            wav = wav[0:1, :]
        samples = wav.squeeze().to(torch.float32).numpy()

    buf = io.BytesIO()
    sf.write(buf, samples, sr, format=req.response_format.upper())
    audio_b64 = base64.b64encode(buf.getvalue()).decode("ascii")

    return JSONResponse(
        {"audio": audio_b64, "audio_format": _FORMATS[req.response_format], "seed": seed}
    )