serving/openai_embeddings_server.py

#!/usr/bin/env python3
"""OpenAI-compatible /v1/embeddings server for VibeVoice-Embed.

Reference implementation, deliberately small: one model, one lock, no batching tricks.
``input`` carries AUDIO rather than text — a data: URI, bare base64, or a server-side
file path (or a list of them). Vectors come back unnormalised, matching the model's
contract; pass ``"normalize": true`` to get unit vectors.

    pip install torch torchaudio transformers fastapi uvicorn
    python openai_embeddings_server.py --model lemuriandezapada/VibeVoice-Embed --port 8080

    curl -s localhost:8080/v1/embeddings \
      -H 'Content-Type: application/json' \
      -d '{"input": "'"$(base64 -w0 clip.wav)"'", "normalize": true}'

Why no dynamic batching: the encoder runs two orders of magnitude faster than realtime
on a modest GPU, so request transport dominates end-to-end latency long before the
forward pass does. Batched clips of unequal length would also need mask-aware pooling
(the model supports it via ``padding_mask``) — correct, but pointless complexity at
this model size. Scale out with replicas instead.
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import io
import logging
import os
import time
from typing import Optional, Union

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("vibevoice-embed")

MAX_AUDIO_BYTES = 32 * 1024 * 1024
MAX_BATCH = 64


def load_model(model_id: str, dtype: str, device: str):
    import torch
    from transformers import AutoModel

    torch_dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16}[dtype]
    model = AutoModel.from_pretrained(model_id, torch_dtype=torch_dtype, trust_remote_code=True)
    model.to(device)
    model.eval()
    params = sum(p.numel() for p in model.parameters())
    logger.info("model ready: %.1fM params, %s, %s, dim=%d",
                params / 1e6, torch_dtype, device, model.config.vae_dim)
    return model


def decode_audio(value: str, sampling_rate: int):
    """data: URI, bare base64, or path -> (mono float32 waveform (1, T), seconds)."""
    text = (value or "").strip()
    if not text:
        raise ValueError("empty audio input")

    import torch
    import torchaudio

    if text.startswith("data:"):
        raw = base64.b64decode(text.split(",", 1)[-1], validate=False)
        source = io.BytesIO(raw)
    elif os.path.exists(text):
        if os.path.getsize(text) > MAX_AUDIO_BYTES:
            raise ValueError("audio input is too large")
        source = text
    else:
        raw = base64.b64decode(text, validate=False)
        if len(raw) > MAX_AUDIO_BYTES:
            raise ValueError("audio input is too large")
        source = io.BytesIO(raw)

    wav, sr = torchaudio.load(source)
    if wav.numel() == 0:
        raise ValueError("audio input decoded to zero samples")
    if wav.shape[0] > 1:
        wav = wav.mean(0, keepdim=True)
    if sr != sampling_rate:
        wav = torchaudio.functional.resample(wav, sr, sampling_rate)
    return wav.to(torch.float32), wav.shape[-1] / sampling_rate


def build_app(model, served_name: str, device: str):
    import torch
    from fastapi import FastAPI, HTTPException
    from pydantic import BaseModel, Field

    sampling_rate = int(model.config.sampling_rate)
    frame_rate = sampling_rate / model.config.hop_length

    class EmbeddingRequest(BaseModel):
        model: Optional[str] = None
        input: Union[str, list[str]]
        normalize: bool = False
        encoding_format: str = Field(default="float")

    app = FastAPI(title="VibeVoice-Embed")
    # One request at a time on the one model; the GPU work never touches the event loop.
    lock = asyncio.Lock()

    @app.get("/health")
    async def health():
        return {"status": "ok"}

    @app.get("/v1/models")
    async def models():
        return {"object": "list", "data": [{
            "id": served_name,
            "object": "model",
            "owned_by": "vibevoice",
            "dimensions": int(model.config.vae_dim),
            "sample_rate": sampling_rate,
            "frame_rate": frame_rate,
        }]}

    def embed_all(waveforms):
        with torch.no_grad():
            return [
                model(wav.unsqueeze(0).to(device)).pooler_output.squeeze(0)
                for wav in waveforms
            ]

    @app.post("/v1/embeddings")
    async def embeddings(body: EmbeddingRequest):
        items = [body.input] if isinstance(body.input, str) else list(body.input)
        if not items:
            raise HTTPException(400, "input is required")
        if len(items) > MAX_BATCH:
            raise HTTPException(400, f"at most {MAX_BATCH} clips per request, got {len(items)}")

        try:
            decoded = [decode_audio(item, sampling_rate) for item in items]
        except ValueError as exc:
            raise HTTPException(400, str(exc)) from exc
        except Exception as exc:
            raise HTTPException(400, f"could not decode audio: {exc}") from exc

        t0 = time.time()
        async with lock:
            vectors = await asyncio.to_thread(embed_all, [wav for wav, _ in decoded])
        if body.normalize:
            vectors = [torch.nn.functional.normalize(v, dim=0) for v in vectors]

        seconds = sum(secs for _, secs in decoded)
        logger.info("%d clip(s), %.2fs audio in %.2fs", len(items), seconds, time.time() - t0)
        return {
            "object": "list",
            "data": [
                {"object": "embedding", "index": i, "embedding": v.cpu().tolist()}
                for i, v in enumerate(vectors)
            ],
            "model": served_name,
            # Audio has no tokens; encoder frames are the honest cost analogue.
            "usage": {
                "prompt_tokens": int(round(seconds * frame_rate)),
                "total_tokens": int(round(seconds * frame_rate)),
                "audio_seconds": round(seconds, 3),
            },
        }

    return app


def main() -> None:
    parser = argparse.ArgumentParser(description="Serve VibeVoice-Embed voice embeddings")
    parser.add_argument("--model", default="lemuriandezapada/VibeVoice-Embed")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8080)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"],
                        help="float32 by default: the encoder is small and the pooled "
                             "vector is compared at tolerances bfloat16 cannot hold")
    args = parser.parse_args()

    model = load_model(args.model, args.dtype, args.device)

    import uvicorn

    uvicorn.run(build_app(model, args.model, args.device),
                host=args.host, port=args.port, log_level="info")


if __name__ == "__main__":
    main()