116 خطوط
3.5 KiB
Python
116 خطوط
3.5 KiB
Python
"""Minimal FastAPI server exposing Chatterbox Multilingual TTS. https://huggingface.co/ResembleAI/chatterbox"""
|
|
import io
|
|
import os
|
|
import tempfile
|
|
|
|
import torch
|
|
import torchaudio as ta
|
|
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
|
|
from fastapi.responses import StreamingResponse
|
|
from pydantic import BaseModel
|
|
from chatterbox.mtl_tts import ChatterboxMultilingualTTS
|
|
|
|
_forced_device = os.environ.get("CHATTERBOX_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"
|
|
else:
|
|
DEVICE = "cpu"
|
|
|
|
EXAGGERATION = float(os.environ.get("CHATTERBOX_EXAGGERATION", "0.5"))
|
|
CFG_WEIGHT = float(os.environ.get("CHATTERBOX_CFG_WEIGHT", "0.5"))
|
|
|
|
# Full language name -> ISO code fallback, for callers that don't send language_code.
|
|
_NAME_TO_CODE = {
|
|
"arabic": "ar", "danish": "da", "german": "de", "greek": "el", "english": "en",
|
|
"spanish": "es", "finnish": "fi", "french": "fr", "hebrew": "he", "hindi": "hi",
|
|
"italian": "it", "japanese": "ja", "korean": "ko", "malay": "ms", "dutch": "nl",
|
|
"norwegian": "no", "polish": "pl", "portuguese": "pt", "russian": "ru",
|
|
"swedish": "sv", "swahili": "sw", "turkish": "tr", "chinese": "zh",
|
|
}
|
|
|
|
|
|
def resolve_language_id(language: str, language_code: str) -> str:
|
|
code = (language_code or "").strip().lower()
|
|
if code:
|
|
return code
|
|
return _NAME_TO_CODE.get((language or "").strip().lower(), "en")
|
|
|
|
|
|
app = FastAPI(title="Chatterbox server")
|
|
|
|
model: ChatterboxMultilingualTTS | None = None
|
|
|
|
|
|
@app.on_event("startup")
|
|
def load_model():
|
|
global model
|
|
print(f"Loading Chatterbox Multilingual TTS on {DEVICE}...")
|
|
model = ChatterboxMultilingualTTS.from_pretrained(device=DEVICE)
|
|
print("Chatterbox model loaded.")
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok", "device": DEVICE}
|
|
|
|
|
|
class SpeechRequest(BaseModel):
|
|
text: str
|
|
language: str = "English"
|
|
language_code: str = ""
|
|
|
|
|
|
@app.post("/speech")
|
|
def speech(req: SpeechRequest):
|
|
if model is None:
|
|
raise HTTPException(status_code=503, detail="Model not loaded yet")
|
|
|
|
wav = model.generate(
|
|
text=req.text,
|
|
language_id=resolve_language_id(req.language, req.language_code),
|
|
exaggeration=EXAGGERATION,
|
|
cfg_weight=CFG_WEIGHT,
|
|
)
|
|
|
|
buf = io.BytesIO()
|
|
ta.save(buf, wav, model.sr, format="wav")
|
|
buf.seek(0)
|
|
return StreamingResponse(buf, media_type="audio/wav")
|
|
|
|
|
|
@app.post("/speech/clone")
|
|
def speech_clone(
|
|
text: str = Form(...),
|
|
language: str = Form("English"),
|
|
language_code: str = Form(""),
|
|
ref_audio: UploadFile = File(...),
|
|
):
|
|
"""Clone a voice from ref_audio 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"
|
|
with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
|
|
tmp.write(ref_bytes)
|
|
tmp_path = tmp.name
|
|
|
|
try:
|
|
wav = model.generate(
|
|
text=text,
|
|
language_id=resolve_language_id(language, language_code),
|
|
audio_prompt_path=tmp_path,
|
|
exaggeration=EXAGGERATION,
|
|
cfg_weight=CFG_WEIGHT,
|
|
)
|
|
finally:
|
|
os.remove(tmp_path)
|
|
|
|
buf = io.BytesIO()
|
|
ta.save(buf, wav, model.sr, format="wav")
|
|
buf.seek(0)
|
|
return StreamingResponse(buf, media_type="audio/wav")
|