File size: 2,310 Bytes
cb81b58 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | import torch
import torchaudio
import io
import base64
from typing import Dict, Any
from chatterbox.mtl_tts import ChatterboxMultilingualTTS
class EndpointHandler:
def __init__(self, path=""):
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"🚀 Loading Chatterbox-Multilingual model on {device}...")
self.model = ChatterboxMultilingualTTS.from_pretrained(device=device)
self.device = device
print("✅ Model loaded successfully")
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
inputs = data.pop("inputs", data)
text = inputs.get("text")
if not text:
return {"error": "Missing required parameter: text"}
reference_audio = inputs.get("reference_audio")
if not reference_audio:
return {"error": "Missing required parameter: reference_audio"}
cfg_weight = inputs.get("cfg_weight", 1.0)
exaggeration = inputs.get("exaggeration", 1.0)
language_id = inputs.get("language", "en")
if isinstance(reference_audio, str) and reference_audio.startswith("data:"):
header, encoded = reference_audio.split(",", 1)
reference_audio_bytes = base64.b64decode(encoded)
elif isinstance(reference_audio, str):
reference_audio_bytes = base64.b64decode(reference_audio)
else:
reference_audio_bytes = reference_audio
with open("/tmp/reference_audio.tmp", "wb") as f:
f.write(reference_audio_bytes)
try:
wav = self.model.generate(
text,
audio_prompt_path="/tmp/reference_audio.tmp",
language_id=language_id,
exaggeration=exaggeration,
cfg_weight=cfg_weight
)
buffer = io.BytesIO()
torchaudio.save(buffer, wav, self.model.sr, format="wav")
buffer.seek(0)
audio_base64 = base64.b64encode(buffer.read()).decode("utf-8")
return {
"audio": audio_base64,
"sample_rate": self.model.sr,
"format": "wav"
}
except Exception as e:
return {"error": str(e)}
|