#!/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()