| from pathlib import Path |
| from typing import Any, Dict, List |
| import sys |
| import numpy as np |
|
|
| import base64 |
| import io |
| import soundfile as sf |
|
|
|
|
| REPO_ROOT = Path(__file__).resolve().parent |
| if str(REPO_ROOT) not in sys.path: |
| sys.path.insert(0, str(REPO_ROOT)) |
|
|
| from cosyvoice.cli.cosyvoice import CosyVoice |
|
|
| class EndpointHandler: |
| def __init__(self, model_dir: str): |
| self.model_dir = model_dir |
| cosy_model_dir = str(Path(self.model_dir) / "model") |
|
|
| self.cosyvoice = CosyVoice( |
| cosy_model_dir, |
| load_jit=False, |
| load_trt=False, |
| fp16=False, |
| ) |
|
|
|
|
| def _b64_to_bytes(self, s: str) -> bytes: |
| return base64.b64decode(s) |
|
|
| def _bytes_to_b64_wav(self, audio: np.ndarray, sr: int) -> str: |
| buf = io.BytesIO() |
| sf.write(buf, audio.T, sr, format="WAV") |
| return base64.b64encode(buf.getvalue()).decode("utf-8") |
|
|
| def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
| """ |
| Expected payload: |
| { |
| "inputs": { |
| "text": "...", |
| "ref_text": "...", |
| "ref_audio_b64": "<base64 wav>" |
| } |
| } |
| """ |
| inputs = data["inputs"] |
| text = inputs["text"] |
| ref_text = inputs.get("ref_text", "") |
| ref_audio_b64 = inputs["ref_audio_b64"] |
|
|
| ref_bytes = self._b64_to_bytes(ref_audio_b64) |
| ref_path = "/tmp/ref.wav" |
| with open(ref_path, "wb") as f: |
| f.write(ref_bytes) |
|
|
| chunks = [] |
| for chunk in self.cosyvoice.inference_zero_shot( |
| text, |
| ref_text, |
| ref_path, |
| stream=False, |
| ): |
| audio = chunk.get("tts_speech", chunk.get("audio")) |
| chunks.append(audio) |
|
|
| audio = np.concatenate(chunks, axis=-1) |
| sr = self.cosyvoice.sample_rate |
| audio_b64 = self._bytes_to_b64_wav(audio, sr) |
|
|
| return [{"audio_wav_base64": audio_b64, "sample_rate": sr}] |
|
|
|
|
