""" OmniVoice TTS — fully local FastAPI server for HuggingFace Spaces (T4 / dedicated GPU). Loads `k2-fsa/OmniVoice` once at startup and serves four endpoints: GET / → service descriptor (JSON) GET /health → liveness probe POST /tts → one-shot WAV (audio/wav, 24 kHz mono int16) POST /tts/stream → chunked WAV (header + raw PCM tail) — same payload WS /ws/tts → real-time PCM frames + JSON status messages Voice-design (no reference audio) and voice-clone (with reference audio) modes are both supported. Attributes (gender / age / pitch / accent / dialect / style) are forwarded to the OmniVoice model via the `instruct` parameter, mirroring the upstream Gradio demo's behaviour. This file is self-contained: it does NOT call out to the public Gradio Space. The model runs in-process on the local GPU. """ from __future__ import annotations import asyncio import io import json import logging import os import struct import tempfile import threading from collections import OrderedDict from typing import Any, AsyncIterator, Optional import numpy as np import torch from fastapi import FastAPI, HTTPException, UploadFile, File, Form, WebSocket, WebSocketDisconnect from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse, Response, StreamingResponse from pydantic import BaseModel, Field from omnivoice import OmniVoice, OmniVoiceGenerationConfig logging.basicConfig( level=os.environ.get("LOG_LEVEL", "INFO"), format="%(asctime)s %(levelname)s %(name)s: %(message)s", ) logger = logging.getLogger("omnivoice-tts") # ─── Model loading ──────────────────────────────────────────────────────────── CHECKPOINT = os.environ.get("OMNIVOICE_MODEL", "k2-fsa/OmniVoice") DEVICE = os.environ.get("OMNIVOICE_DEVICE", "cuda" if torch.cuda.is_available() else "cpu") DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32 # Whisper ASR is only needed to auto-transcribe a reference clip when ref_text # is omitted. Our workflow always supplies ref_text, so default to NOT loading it # (~1 GB less to download/load → faster cold starts). Set OMNIVOICE_LOAD_ASR=1 to # re-enable auto-transcription. LOAD_ASR = os.environ.get("OMNIVOICE_LOAD_ASR", "0") == "1" logger.info("Loading OmniVoice from %s onto %s (dtype=%s, asr=%s)", CHECKPOINT, DEVICE, DTYPE, LOAD_ASR) model = OmniVoice.from_pretrained( CHECKPOINT, device_map=DEVICE, dtype=DTYPE, load_asr=LOAD_ASR, ) SAMPLE_RATE: int = int(model.sampling_rate) logger.info("Model loaded — sample_rate=%d Hz", SAMPLE_RATE) # ─── Voice registry (clone once, reuse many times) ──────────────────────────── # A client registers a reference voice once via POST /voices under a chosen name; # we compute the OmniVoice voice-clone prompt a single time and cache it here. # Subsequent /tts (and /tts/stream, /ws/tts) requests pass `voice_id` to reuse it, # skipping all reference re-encoding (and the Whisper ASR pass when ref_text was # supplied). Prompts are persisted to VOICES_DIR so they survive restarts. MAX_VOICES = int(os.environ.get("OMNIVOICE_MAX_VOICES", "50")) VOICES_DIR = os.environ.get("OMNIVOICE_VOICES_DIR", "/data/voices") _VOICE_CACHE: "OrderedDict[str, Any]" = OrderedDict() _VOICE_LOCK = threading.Lock() os.makedirs(VOICES_DIR, exist_ok=True) def _voice_path(name: str) -> str: """Return the .pt path for a given voice name (safe filename).""" safe = name.replace("/", "_").replace("\\", "_") return os.path.join(VOICES_DIR, f"{safe}.pt") def _persist_voice(name: str, prompt: Any) -> None: """Save a voice-clone prompt to disk using torch.save.""" try: torch.save(prompt, _voice_path(name)) except Exception: logger.warning("Could not persist voice %r to disk", name, exc_info=True) def _delete_persisted_voice(name: str) -> None: """Remove the on-disk .pt file for a voice (best-effort).""" try: os.unlink(_voice_path(name)) except Exception: pass def _load_persisted_voices() -> None: """At startup, reload all saved voice prompts from VOICES_DIR into the cache.""" try: files = [f for f in os.listdir(VOICES_DIR) if f.endswith(".pt")] except Exception: return loaded = 0 for fname in files: name = fname[:-3] # strip .pt try: prompt = torch.load(os.path.join(VOICES_DIR, fname), map_location=DEVICE, weights_only=False) _VOICE_CACHE[name] = prompt loaded += 1 except Exception: logger.warning("Failed to reload voice %r from disk", name, exc_info=True) if loaded: logger.info("Reloaded %d persisted voice(s) from %s", loaded, VOICES_DIR) _load_persisted_voices() def _register_voice(name: str, ref_audio_path: str, ref_text: Optional[str]) -> int: """Build and cache a voice-clone prompt under `name` (LRU-evicting). Returns cache size.""" prompt = model.create_voice_clone_prompt(ref_audio=ref_audio_path, ref_text=ref_text) _persist_voice(name, prompt) with _VOICE_LOCK: _VOICE_CACHE[name] = prompt # overwrite if name already exists _VOICE_CACHE.move_to_end(name) # mark most-recently-used while len(_VOICE_CACHE) > MAX_VOICES: evicted, _ = _VOICE_CACHE.popitem(last=False) # drop least-recently-used _delete_persisted_voice(evicted) logger.info("Evicted LRU voice %r (cache full at %d)", evicted, MAX_VOICES) return len(_VOICE_CACHE) def _get_voice(name: str) -> Optional[Any]: """Return the cached prompt for `name`, marking it most-recently-used, or None.""" with _VOICE_LOCK: prompt = _VOICE_CACHE.get(name) if prompt is not None: _VOICE_CACHE.move_to_end(name) return prompt # ─── Voice-design attribute vocabulary (matches upstream Gradio demo) ───────── # Each entry is "English / 中文". For accents we pass the English half; for # dialects we pass the Chinese half — mirrors omnivoice/cli/demo.py logic. GENDERS = {"Auto", "Male / 男", "Female / 女"} AGES = { "Auto", "Child / 儿童", "Teenager / 少年", "Young Adult / 青年", "Middle-aged / 中年", "Elderly / 老年", } PITCHES = { "Auto", "Very Low Pitch / 极低音调", "Low Pitch / 低音调", "Moderate Pitch / 中音调", "High Pitch / 高音调", "Very High Pitch / 极高音调", } STYLES = {"Auto", "Whisper / 耳语"} ACCENTS = { "Auto", "American Accent / 美式口音", "Australian Accent / 澳大利亚口音", "British Accent / 英国口音", "Chinese Accent / 中国口音", "Canadian Accent / 加拿大口音", "Indian Accent / 印度口音", "Korean Accent / 韩国口音", "Portuguese Accent / 葡萄牙口音", "Russian Accent / 俄罗斯口音", "Japanese Accent / 日本口音", } DIALECTS = { "Auto", "Henan Dialect / 河南话", "Shaanxi Dialect / 陕西话", "Sichuan Dialect / 四川话", "Guizhou Dialect / 贵州话", "Yunnan Dialect / 云南话", "Guilin Dialect / 桂林话", "Jinan Dialect / 济南话", "Shijiazhuang Dialect / 石家庄话", "Gansu Dialect / 甘肃话", "Ningxia Dialect / 宁夏话", "Qingdao Dialect / 青岛话", "Northeast Dialect / 东北话", } LANGUAGES_ADVERTISED = ["Auto", "English", "Hindi", "Panjabi", "Urdu", "Western Panjabi"] def _attr_part(value: Optional[str], is_dialect: bool = False) -> Optional[str]: """Pick the right half of a bilingual attribute label, or None for Auto.""" if not value or value == "Auto": return None if " / " in value: en, zh = value.split(" / ", 1) return zh.strip() if is_dialect else en.strip() return value.strip() def build_instruct( *, gender: Optional[str] = None, age: Optional[str] = None, pitch: Optional[str] = None, style: Optional[str] = None, accent: Optional[str] = None, dialect: Optional[str] = None, extra: Optional[str] = None, ) -> Optional[str]: """Compose the OmniVoice `instruct` string from voice-design attributes.""" parts = [] for v, is_dialect in ( (gender, False), (age, False), (pitch, False), (style, False), (accent, False), (dialect, True), ): p = _attr_part(v, is_dialect=is_dialect) if p: parts.append(p) if extra and extra.strip(): parts.append(extra.strip()) return ", ".join(parts) if parts else None # ─── Synthesis core ─────────────────────────────────────────────────────────── class TTSRequest(BaseModel): text: str = Field(..., min_length=1, description="Text to synthesize.") language: str = Field("Auto", description="Language name (English label, e.g. 'Urdu', 'Panjabi') or 'Auto'.") # Voice-design attributes — accept either the canonical bilingual label # ("Female / 女") or the friendly English half ("Female"). Auto = no constraint. gender: str = "Auto" age: str = "Auto" pitch: str = "Auto" style: str = "Auto" accent: str = "Auto" dialect: str = "Auto" instruct: Optional[str] = None # raw override; takes precedence if set voice_id: Optional[str] = None # reuse a voice registered via POST /voices # Generation config speed: float = 1.0 duration: Optional[float] = None nfe_steps: int = Field(32, ge=4, le=64) guidance: float = Field(2.0, ge=0.0, le=4.0) # Sampling temperatures (model defaults: class=0.0 greedy, position=5.0). # class_temperature 0.0 = deterministic/greedy → can sound flat/robotic; # raise toward ~0.3–0.7 for more natural variation. class_temperature: float = Field(0.3, ge=0.0, le=2.0) position_temperature: float = Field(5.0, ge=0.0, le=10.0) denoise: bool = True preprocess_prompt: bool = True postprocess_output: bool = True def _normalize_label(value: str, valid: set) -> str: """Accept friendly aliases ('female', 'young', 'low') in addition to bilingual labels.""" if not value or value == "Auto": return "Auto" if value in valid: return value lc = value.strip().lower() aliases = { "male": "Male / 男", "m": "Male / 男", "female": "Female / 女", "f": "Female / 女", "child": "Child / 儿童", "teen": "Teenager / 少年", "teenager": "Teenager / 少年", "young": "Young Adult / 青年", "young_adult": "Young Adult / 青年", "middle": "Middle-aged / 中年", "middle_aged": "Middle-aged / 中年", "elderly": "Elderly / 老年", "old": "Elderly / 老年", "very_low": "Very Low Pitch / 极低音调", "very low": "Very Low Pitch / 极低音调", "low": "Low Pitch / 低音调", "moderate": "Moderate Pitch / 中音调", "medium": "Moderate Pitch / 中音调", "high": "High Pitch / 高音调", "very_high": "Very High Pitch / 极高音调", "very high": "Very High Pitch / 极高音调", "whisper": "Whisper / 耳语", } return aliases.get(lc, value if value in valid else "Auto") def _synthesize( req: TTSRequest, *, ref_audio_path: Optional[str] = None, ref_text: Optional[str] = None, voice_clone_prompt: Optional[Any] = None, ) -> np.ndarray: """Run OmniVoice inference and return int16 PCM samples (1-D ndarray).""" # Normalize friendly aliases to canonical bilingual labels gender = _normalize_label(req.gender, GENDERS) age = _normalize_label(req.age, AGES) pitch = _normalize_label(req.pitch, PITCHES) style = _normalize_label(req.style, STYLES) accent = req.accent if req.accent in ACCENTS else "Auto" dialect = req.dialect if req.dialect in DIALECTS else "Auto" instruct = req.instruct or build_instruct( gender=gender, age=age, pitch=pitch, style=style, accent=accent, dialect=dialect, ) lang = req.language if req.language and req.language != "Auto" else None gen_config = OmniVoiceGenerationConfig( num_step=req.nfe_steps, guidance_scale=req.guidance, class_temperature=req.class_temperature, position_temperature=req.position_temperature, denoise=req.denoise, preprocess_prompt=req.preprocess_prompt, postprocess_output=req.postprocess_output, ) kw: dict[str, Any] = dict( text=req.text.strip(), language=lang, generation_config=gen_config, ) if req.speed and req.speed != 1.0: kw["speed"] = float(req.speed) if req.duration and req.duration > 0: kw["duration"] = float(req.duration) if instruct: kw["instruct"] = instruct if voice_clone_prompt is not None: kw["voice_clone_prompt"] = voice_clone_prompt elif ref_audio_path: kw["voice_clone_prompt"] = model.create_voice_clone_prompt( ref_audio=ref_audio_path, ref_text=ref_text, ) cloning = voice_clone_prompt is not None or bool(ref_audio_path) logger.info("Synthesize: lang=%s len=%d instruct=%r clone=%s voice_id=%s", lang, len(req.text), instruct, cloning, req.voice_id) with torch.inference_mode(): audio = model.generate(**kw) # audio[0] is float32 in [-1, 1]; convert to int16 PCM. pcm = np.clip(audio[0], -1.0, 1.0) return (pcm * 32767).astype(np.int16) # ─── WAV helpers ────────────────────────────────────────────────────────────── def _wav_header(num_samples: int, sample_rate: int = SAMPLE_RATE, channels: int = 1) -> bytes: """Build a 44-byte WAV header for int16 mono PCM.""" sample_width = 2 data_size = num_samples * channels * sample_width return ( b"RIFF" + struct.pack(" bytes: return _wav_header(len(pcm_int16)) + pcm_int16.tobytes() # ─── FastAPI app ────────────────────────────────────────────────────────────── app = FastAPI( title="OmniVoice TTS", description="In-process OmniVoice TTS server (no upstream proxy).", version="1.0.0", ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) _DESCRIPTOR = { "service": "OmniVoice TTS", "model": CHECKPOINT, "device": DEVICE, "sample_rate": SAMPLE_RATE, "endpoints": { "GET /health": "Health probe", "POST /tts": "One-shot WAV (audio/wav)", "POST /tts/stream": "Chunked WAV (header + raw PCM tail)", "WS /ws/tts": "Real-time binary PCM chunks + JSON status", "POST /tts/clone": "Ad-hoc voice cloning (multipart form: ref_audio + text)", "POST /voices": "Register a voice once (multipart: name + ref_audio + ref_text?)", "GET /voices": "List cached voice ids", "DELETE /voices/{name}": "Free a registered voice", }, "languages": LANGUAGES_ADVERTISED, "genders": sorted(GENDERS), "ages": sorted(AGES), "pitches": sorted(PITCHES), "styles": sorted(STYLES), "accents": sorted(ACCENTS), "dialects": sorted(DIALECTS), } @app.get("/") async def root(): return JSONResponse(_DESCRIPTOR) @app.get("/health") async def health(): return {"status": "ok", "model_loaded": True, "device": DEVICE} def _resolve_voice(voice_id: Optional[str]) -> Optional[Any]: """Look up a registered voice prompt, or raise 404. Returns None when no voice_id given.""" if not voice_id: return None prompt = _get_voice(voice_id) if prompt is None: raise HTTPException( status_code=404, detail=f"voice_id {voice_id!r} not registered; POST /voices first", ) return prompt @app.post("/tts") async def tts_one_shot(req: TTSRequest): """Return the full WAV in one response.""" voice_prompt = _resolve_voice(req.voice_id) try: # Run inference in a thread so the event loop stays responsive. pcm = await asyncio.to_thread(_synthesize, req, voice_clone_prompt=voice_prompt) except Exception as e: logger.exception("synth failed") raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}") if pcm.size == 0: raise HTTPException(status_code=500, detail="empty audio") return Response(content=_wrap_wav(pcm), media_type="audio/wav") @app.post("/tts/stream") async def tts_stream(req: TTSRequest): """ Chunked HTTP: emit a WAV header up-front (with a placeholder data size) followed by the PCM bytes. Most decoders that read sequentially handle this fine. Lower TTFB than /tts because the header flushes immediately. """ voice_prompt = _resolve_voice(req.voice_id) try: pcm = await asyncio.to_thread(_synthesize, req, voice_clone_prompt=voice_prompt) except Exception as e: logger.exception("stream synth failed") raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}") async def gen() -> AsyncIterator[bytes]: yield _wav_header(len(pcm)) # Yield ~0.2 s frames so clients see audio progressively. frame = SAMPLE_RATE // 5 raw = pcm.tobytes() bytes_per_frame = frame * 2 # int16 for i in range(0, len(raw), bytes_per_frame): yield raw[i: i + bytes_per_frame] return StreamingResponse(gen(), media_type="audio/wav") @app.post("/tts/clone") async def tts_clone( text: str = Form(...), ref_audio: UploadFile = File(...), language: str = Form("Auto"), ref_text: Optional[str] = Form(None), speed: float = Form(1.0), nfe_steps: int = Form(32), guidance: float = Form(2.0), class_temperature: float = Form(0.3), position_temperature: float = Form(5.0), denoise: bool = Form(True), ): """Voice cloning with an uploaded reference audio (ad-hoc, one-off).""" suffix = os.path.splitext(ref_audio.filename or "ref.wav")[1] or ".wav" with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: tmp.write(await ref_audio.read()) tmp_path = tmp.name try: req = TTSRequest( text=text, language=language, speed=speed, nfe_steps=nfe_steps, guidance=guidance, denoise=denoise, class_temperature=class_temperature, position_temperature=position_temperature, ) pcm = await asyncio.to_thread( _synthesize, req, ref_audio_path=tmp_path, ref_text=ref_text, ) finally: try: os.unlink(tmp_path) except Exception: pass return Response(content=_wrap_wav(pcm), media_type="audio/wav") # ─── Voice registry endpoints (clone once, reuse via voice_id) ──────────────── @app.post("/voices") async def register_voice( name: str = Form(..., min_length=1), ref_audio: UploadFile = File(...), ref_text: Optional[str] = Form(None), ): """ Register a reference voice ONCE under `name`. The voice-clone prompt is computed here and cached in memory; reuse it on /tts, /tts/stream and /ws/tts by passing "voice_id": "". Re-registering the same name overwrites it. Omitting ref_text triggers Whisper ASR (needs OMNIVOICE_LOAD_ASR=1). """ suffix = os.path.splitext(ref_audio.filename or "ref.wav")[1] or ".wav" with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: tmp.write(await ref_audio.read()) tmp_path = tmp.name try: cached = await asyncio.to_thread(_register_voice, name, tmp_path, ref_text) except Exception as e: logger.exception("voice registration failed") raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}") finally: try: os.unlink(tmp_path) except Exception: pass return {"voice_id": name, "voices_cached": cached, "max_voices": MAX_VOICES} @app.get("/voices") async def list_voices(): """List currently cached voice ids (most-recently-used last).""" with _VOICE_LOCK: return {"voices": list(_VOICE_CACHE.keys()), "max_voices": MAX_VOICES} @app.get("/debug/voices-dir") async def debug_voices_dir(): """Show what .pt files are actually on disk in VOICES_DIR.""" result: dict = {"voices_dir": VOICES_DIR, "data_writable": os.access("/data", os.W_OK)} try: files = os.listdir(VOICES_DIR) result["files"] = { f: os.path.getsize(os.path.join(VOICES_DIR, f)) for f in sorted(files) } except Exception as e: result["error"] = str(e) return result @app.get("/debug/gpu") async def debug_gpu(): """Report live VRAM / RAM usage so we can right-size the GPU.""" out: dict = {"device": DEVICE, "dtype": str(DTYPE), "load_asr": LOAD_ASR} try: if torch.cuda.is_available(): free, total = torch.cuda.mem_get_info() out["gpu"] = { "name": torch.cuda.get_device_name(0), "vram_allocated_gb": round(torch.cuda.memory_allocated() / 1e9, 2), "vram_reserved_gb": round(torch.cuda.memory_reserved() / 1e9, 2), "vram_used_total_gb": round((total - free) / 1e9, 2), # incl. other procs/context "vram_total_gb": round(total / 1e9, 2), } except Exception as e: out["gpu_error"] = str(e) try: import psutil p = psutil.Process(os.getpid()) vm = psutil.virtual_memory() out["ram_rss_gb"] = round(p.memory_info().rss / 1e9, 2) out["ram_used_total_gb"] = round((vm.total - vm.available) / 1e9, 2) out["ram_total_gb"] = round(vm.total / 1e9, 2) except Exception as e: out["ram_error"] = str(e) return out @app.delete("/voices/{name}") async def delete_voice(name: str): """Free one registered voice. 404 if it isn't cached.""" with _VOICE_LOCK: if _VOICE_CACHE.pop(name, None) is None: raise HTTPException(status_code=404, detail=f"voice_id {name!r} not registered") remaining = len(_VOICE_CACHE) _delete_persisted_voice(name) return {"deleted": name, "voices_cached": remaining} @app.websocket("/ws/tts") async def ws_tts(ws: WebSocket): """ Bidirectional WebSocket: Client sends JSON request (TTSRequest schema). Server replies: {"type":"started","sample_rate":24000} ... {"type":"complete","duration_s":1.23} On error: {"type":"error","message":"..."} Connection closes after one synthesis (matches the Parler pattern). """ await ws.accept() try: raw = await ws.receive_text() try: payload = json.loads(raw) except json.JSONDecodeError as e: await ws.send_json({"type": "error", "message": f"bad json: {e}"}) return try: req = TTSRequest(**payload) except Exception as e: await ws.send_json({"type": "error", "message": f"bad request: {e}"}) return voice_prompt = _get_voice(req.voice_id) if req.voice_id else None if req.voice_id and voice_prompt is None: await ws.send_json({ "type": "error", "message": f"voice_id {req.voice_id!r} not registered; POST /voices first", }) return await ws.send_json({"type": "started", "sample_rate": SAMPLE_RATE}) try: pcm = await asyncio.to_thread(_synthesize, req, voice_clone_prompt=voice_prompt) except Exception as e: logger.exception("ws synth failed") await ws.send_json({"type": "error", "message": f"{type(e).__name__}: {e}"}) return # Stream as ~0.5 s binary chunks so the client can begin playback early. frame_bytes = SAMPLE_RATE * 2 // 2 # 0.5 s of int16 mono raw_pcm = pcm.tobytes() for i in range(0, len(raw_pcm), frame_bytes): await ws.send_bytes(raw_pcm[i: i + frame_bytes]) await ws.send_json({ "type": "complete", "duration_s": round(len(pcm) / SAMPLE_RATE, 3), }) except WebSocketDisconnect: logger.info("ws client disconnected") except Exception as e: logger.exception("ws unexpected error") try: await ws.send_json({"type": "error", "message": str(e)}) except Exception: pass finally: try: await ws.close() except Exception: pass if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))