92 строки
2.8 KiB
Python
92 строки
2.8 KiB
Python
"""Minimal FastAPI server exposing OmniVoice zero-shot voice cloning (/speech/clone only). https://github.com/k2-fsa/OmniVoice"""
|
|
import io
|
|
import os
|
|
import tempfile
|
|
|
|
import soundfile as sf
|
|
import torch
|
|
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
|
from fastapi.responses import StreamingResponse
|
|
from omnivoice import OmniVoice
|
|
|
|
MODEL_NAME = os.environ.get("OMNIVOICE_MODEL", "k2-fsa/OmniVoice")
|
|
|
|
_forced_device = os.environ.get("OMNIVOICE_DEVICE", "").strip()
|
|
if _forced_device:
|
|
DEVICE = _forced_device
|
|
elif torch.cuda.is_available(): # ROCm builds expose the GPU via the CUDA API too
|
|
DEVICE = "cuda:0"
|
|
elif getattr(torch, "xpu", None) is not None and torch.xpu.is_available():
|
|
DEVICE = "xpu"
|
|
else:
|
|
DEVICE = "cpu"
|
|
# float16 is unsupported on CPU and broken on some XPU/oneDNN SDPA builds.
|
|
DTYPE = torch.float16 if DEVICE.startswith("cuda") else torch.float32
|
|
|
|
# OmniVoice hardcodes fp16 for its internal ASR sub-model on cuda/xpu, which breaks there; keep it on CPU by default.
|
|
ASR_DEVICE = os.environ.get("OMNIVOICE_ASR_DEVICE", "cpu")
|
|
|
|
app = FastAPI(title="OmniVoice server")
|
|
|
|
model: OmniVoice | None = None
|
|
|
|
|
|
@app.on_event("startup")
|
|
def load_model():
|
|
global model
|
|
print(f"Loading OmniVoice model {MODEL_NAME} on {DEVICE}...")
|
|
model = OmniVoice.from_pretrained(
|
|
MODEL_NAME,
|
|
device_map=DEVICE,
|
|
dtype=DTYPE,
|
|
asr_device=ASR_DEVICE,
|
|
)
|
|
print(f"OmniVoice model loaded. ASR (ref_text auto-transcribe) on {ASR_DEVICE}.")
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {
|
|
"status": "ok",
|
|
"model": MODEL_NAME,
|
|
"device": DEVICE,
|
|
}
|
|
|
|
|
|
@app.post("/speech/clone")
|
|
def speech_clone(
|
|
text: str = Form(...),
|
|
language: str = Form(""),
|
|
ref_text: str = Form(""),
|
|
instruct: str = Form(""),
|
|
speed: float = Form(1.0),
|
|
ref_audio: UploadFile = File(...),
|
|
):
|
|
"""Clone a voice from ref_audio (ref_text auto-transcribed via Whisper if omitted) and synthesize text with it."""
|
|
if model is None:
|
|
raise HTTPException(status_code=503, detail="Model not loaded yet")
|
|
|
|
ref_bytes = ref_audio.file.read()
|
|
|
|
suffix = os.path.splitext(ref_audio.filename or "")[1] or ".wav" # generate() needs a file path, not bytes
|
|
with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
|
|
tmp.write(ref_bytes)
|
|
tmp_path = tmp.name
|
|
|
|
try:
|
|
audio = model.generate(
|
|
text=text,
|
|
language=language or None,
|
|
ref_audio=tmp_path,
|
|
ref_text=ref_text or None,
|
|
instruct=instruct or None,
|
|
speed=speed,
|
|
)
|
|
finally:
|
|
os.remove(tmp_path)
|
|
|
|
buf = io.BytesIO()
|
|
sf.write(buf, audio[0], model.sampling_rate, format="WAV")
|
|
buf.seek(0)
|
|
return StreamingResponse(buf, media_type="audio/wav")
|