esh-pfi-infrastructure/stacks/omnivoice/app.py

"""Thin FastAPI wrapper exposing OmniVoice (k2-fsa/OmniVoice) as an OpenAI-style
TTS for the asset-engine.

Upstream ships only a Gradio demo; we own this wrapper (same pattern as
stacks/index-tts/app.py). Voices are reference WAVs staged in
${OMNIVOICE_VOICES_DIR} (the reused chatterbox /refs/*.wav). A voice-clone prompt
is precomputed once per voice at startup (the loaded Whisper ASR auto-transcribes
each reference) and cached, so per-request latency is just generation.

Exposes OmniVoice's full generation surface:
  - clone (voice=<staged ref>) and/or voice-DESIGN (instruct=<free-text style>)
  - language (Auto + 600+), speed, duration
  - diffusion controls: num_step, guidance_scale, denoise, preprocess_prompt,
    postprocess_output, plus a generation_overrides passthrough for expert knobs
    (t_shift, layer_penalty_factor, position_temperature, class_temperature, ...).

Endpoints:
  GET  /healthz            -> readiness (200 once model + >=1 voice are loaded)
  GET  /v1/audio/voices    -> {"voices": [<name>, ...]}
  GET  /v1/audio/languages -> {"languages": ["Auto", <display name>, ...]}
  POST /v1/audio/speech    -> audio/wav
"""
import glob
import io
import logging
import os
from pathlib import Path
from typing import Any, Dict, Optional

import numpy as np
import soundfile as sf
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from pydantic import BaseModel

from omnivoice import OmniVoice, OmniVoiceGenerationConfig

try:
    from omnivoice.utils.lang_map import LANG_NAMES, lang_display_name
    LANGUAGES = ["Auto"] + sorted(lang_display_name(n) for n in LANG_NAMES)
except Exception:                                     # noqa: BLE001
    LANGUAGES = ["Auto"]

try:
    # Voice-DESIGN `instruct` is a CONTROLLED vocabulary (gender / age / pitch /
    # accent / whisper tags), NOT free prose — surfaced so callers can discover it.
    from omnivoice.utils.voice_design import _INSTRUCT_VALID_EN
    INSTRUCT_ITEMS = sorted(_INSTRUCT_VALID_EN)
except Exception:                                     # noqa: BLE001
    INSTRUCT_ITEMS = []

logging.basicConfig(level=os.environ.get("OMNIVOICE_LOG_LEVEL", "INFO"))
log = logging.getLogger("omnivoice-api")

CKPT = os.environ.get("OMNIVOICE_CKPT", "k2-fsa/OmniVoice")
VOICES_DIR = os.environ.get("OMNIVOICE_VOICES_DIR", "/app/voices")
ASR_MODEL = os.environ.get("OMNIVOICE_ASR_MODEL", "openai/whisper-large-v3-turbo")

app = FastAPI(title="OmniVoice TTS (asset-engine wrapper)")

MODEL: Optional[OmniVoice] = None
PROMPTS: dict = {}        # voice name -> VoiceClonePrompt
SR: int = 24000


class SpeechRequest(BaseModel):
    input: str
    # Voice source — at least one of voice (clone) / instruct (design) is required.
    voice: Optional[str] = None              # staged reference clip -> clone timbre
    instruct: Optional[str] = None           # free-text voice DESIGN / style
    # Generation controls (defaults mirror the upstream demo).
    language: Optional[str] = "Auto"         # "Auto" -> auto-detect
    speed: Optional[float] = None            # 0.5–1.5; ignored if duration set
    duration: Optional[float] = None         # fixed seconds; overrides speed
    num_step: int = 32                       # 4–64 diffusion steps
    guidance_scale: float = 2.0              # 0.0–4.0 CFG
    denoise: bool = True
    preprocess_prompt: bool = True
    postprocess_output: bool = True
    # Expert passthrough into OmniVoiceGenerationConfig (t_shift,
    # layer_penalty_factor, position_temperature, class_temperature,
    # audio_chunk_duration, audio_chunk_threshold). Unknown keys are dropped.
    generation_overrides: Optional[Dict[str, Any]] = None
    response_format: str = "wav"
    model: Optional[str] = None              # ignored (single model); OpenAI-compat


@app.on_event("startup")
def _load() -> None:
    global MODEL, SR
    device = "cuda" if torch.cuda.is_available() else "cpu"
    log.info("loading OmniVoice %s on %s (asr=%s)", CKPT, device, ASR_MODEL)
    MODEL = OmniVoice.from_pretrained(
        CKPT, device_map=device, load_asr=True, asr_model_name=ASR_MODEL
    )
    SR = int(getattr(MODEL, "sampling_rate", 24000))

    for wav in sorted(glob.glob(os.path.join(VOICES_DIR, "*.wav"))):
        name = Path(wav).stem
        if name.startswith("_"):
            continue          # skip _*.wav (chatterbox test/deploy artifacts)
        try:
            PROMPTS[name] = MODEL.create_voice_clone_prompt(ref_audio=wav)
            log.info("voice ready: %s", name)
        except Exception as exc:                          # noqa: BLE001
            log.warning("voice %s failed to load: %s", name, exc)
    log.info("%d voices loaded; %d languages", len(PROMPTS), len(LANGUAGES))


@app.get("/healthz")
def healthz():
    if MODEL is None or not PROMPTS:
        raise HTTPException(status_code=503, detail="not ready")
    return {"status": "ok", "voices": len(PROMPTS), "languages": len(LANGUAGES),
            "sampling_rate": SR}


@app.get("/v1/audio/voices")
def voices():
    return {"voices": sorted(PROMPTS.keys())}


@app.get("/v1/audio/languages")
def languages():
    return {"languages": LANGUAGES}


@app.get("/v1/audio/instruct-items")
def instruct_items():
    # Valid comma-separable voice-DESIGN attribute tags (English).
    return {"instruct_items": INSTRUCT_ITEMS}


@app.post("/v1/audio/speech")
def speech(req: SpeechRequest):
    if MODEL is None:
        raise HTTPException(status_code=503, detail="model not loaded")
    if not req.input or not req.input.strip():
        raise HTTPException(status_code=400, detail="input is empty")
    if req.response_format not in ("wav", "", None):
        raise HTTPException(status_code=400, detail="only response_format=wav is supported")

    has_instruct = bool(req.instruct and req.instruct.strip())
    if not req.voice and not has_instruct:
        raise HTTPException(
            status_code=400,
            detail="provide a 'voice' (clone a staged reference) and/or 'instruct' (design a voice)",
        )

    cfg = {
        "num_step": req.num_step,
        "guidance_scale": req.guidance_scale,
        "denoise": req.denoise,
        "preprocess_prompt": req.preprocess_prompt,
        "postprocess_output": req.postprocess_output,
        **(req.generation_overrides or {}),
    }
    gen = OmniVoiceGenerationConfig.from_dict(cfg)

    kw: Dict[str, Any] = {"text": req.input.strip(), "generation_config": gen}

    if req.voice:
        prompt = PROMPTS.get(req.voice)
        if prompt is None:
            raise HTTPException(
                status_code=404,
                detail=f"unknown voice '{req.voice}'; have {sorted(PROMPTS)}",
            )
        kw["voice_clone_prompt"] = prompt
    if has_instruct:
        kw["instruct"] = req.instruct.strip()
    if req.language and req.language != "Auto":
        kw["language"] = req.language
    if req.speed is not None:
        kw["speed"] = req.speed
    if req.duration is not None:
        kw["duration"] = req.duration

    try:
        out = MODEL.generate(**kw)
    except Exception as exc:                              # noqa: BLE001
        raise HTTPException(status_code=400, detail=f"{type(exc).__name__}: {exc}")

    audio = out[0] if isinstance(out, (list, tuple)) else out
    audio = np.asarray(audio, dtype=np.float32)

    buf = io.BytesIO()
    sf.write(buf, audio, SR, format="WAV", subtype="PCM_16")
    return Response(content=buf.getvalue(), media_type="audio/wav")