cosyvoice3-endpoint / handler.py
Ram
Add base64/io/soundfile imports for TTS handler
6d65ae9
Raw
History Blame Contribute Delete
2.03 kB
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 # "/repository"
cosy_model_dir = str(Path(self.model_dir) / "model") # "/repository/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}]