| 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)} |
|
|
