feat(zonos): OpenAI-ish REST adapter for asset-engine routing
Upstream Zonos ships only Gradio + Python SDK — no REST surface — so
asset-engine (which routes a clean JSON POST to /v1/audio/speech) can't
target it directly. Add a thin FastAPI adapter (stacks/zonos/adapter/):
POST /v1/audio/speech in front of the Zonos SDK, built FROM local/zonos
to reuse torch/CUDA/SDK. Returns a JSON envelope {audio, audio_format,
seed} — the seed rides back so asset-engine regenerate/fork can pin it
(Zonos is the fleet's first genuinely seedable TTS). compose gains a
zonos-api service on 8201; .env.example gains the port + voices dir.
This commit is contained in:
@@ -0,0 +1,174 @@
|
||||
"""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}
|
||||
)
|
||||
Reference in New Issue
Block a user