#!/usr/bin/env python3 # coding: utf‑8 """ CosyVoice gRPC back‑end – updated to mirror the FastAPI logic * loads CosyVoice2 with TRT / FP16 first (falls back to CosyVoice) * inference_zero_shot ➜ adds stream=False + speed * inference_instruct ➜ keeps original “speaker‑ID” path * inference_instruct2 ➜ new: prompt‑audio + speed (no speaker‑ID) """ import io, tempfile, requests, soundfile as sf, torchaudio import os import sys from concurrent import futures import argparse import logging import grpc import numpy as np import torch import cosyvoice_pb2 import cosyvoice_pb2_grpc # ──────────────────────────────────────────────────────────────────────────────── # set‑up # ──────────────────────────────────────────────────────────────────────────────── logging.getLogger("matplotlib").setLevel(logging.WARNING) logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) sys.path.extend([ f"{ROOT_DIR}/../../..", f"{ROOT_DIR}/../../../third_party/Matcha-TTS", ]) from cosyvoice.cli.cosyvoice import CosyVoice2 # noqa: E402 # ──────────────────────────────────────────────────────────────────────────────── # helpers # ──────────────────────────────────────────────────────────────────────────────── def _bytes_to_tensor(wav_bytes: bytes) -> torch.Tensor: """ Convert int16 little‑endian PCM bytes → torch.FloatTensor in range [‑1,1] """ speech = torch.from_numpy( np.frombuffer(wav_bytes, dtype=np.int16) ).unsqueeze(0).float() / (2 ** 15) return speech # [1, T] def _yield_audio(model_output): """ Generator that converts CosyVoice output → protobuf Response messages. """ for seg in model_output: pcm16 = (seg["tts_speech"].numpy() * (2 ** 15)).astype(np.int16) resp = cosyvoice_pb2.Response(tts_audio=pcm16.tobytes()) yield resp # ──────────────────────────────────────────────────────────────────────────────── # gRPC service # ──────────────────────────────────────────────────────────────────────────────── class CosyVoiceServiceImpl(cosyvoice_pb2_grpc.CosyVoiceServicer): def __init__(self, args): # try CosyVoice2 first (preferred runtime: TRT / FP16) try: self.cosyvoice = CosyVoice2(args.model_dir, load_jit=False, load_trt=True, fp16=True) logging.info("Loaded CosyVoice2 (TRT / FP16).") except Exception: raise TypeError("No valid CosyVoice model found!") # --------------------------------------------------------------------- # single bi‑di streaming RPC # --------------------------------------------------------------------- def Inference(self, request, context): """Route to the correct model call based on the oneof field present.""" # 1. Supervised fine‑tuning if request.HasField("sft_request"): logging.info("Received SFT inference request") mo = self.cosyvoice.inference_sft( request.sft_request.tts_text, request.sft_request.spk_id ) yield from _yield_audio(mo) return # 2. Zero‑shot speaker cloning (bytes OR S3 URL) if request.HasField("zero_shot_request"): logging.info("Received zero‑shot inference request") zr = request.zero_shot_request tmp_path = None # initialise so we can delete later try: # ───── determine payload type ────────────────────────────────────── if zr.prompt_audio.startswith(b'http'): # —— remote URL —— --------------------------------------------- url = zr.prompt_audio.decode('utf‑8') logging.info("Downloading prompt audio from %s", url) resp = requests.get(url, timeout=10) resp.raise_for_status() # save to a temp file with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f: f.write(resp.content) tmp_path = f.name # load, mono‑ise, resample → tensor [1, T] wav, sr = sf.read(tmp_path, dtype="float32") if wav.ndim > 1: wav = wav.mean(axis=1) if sr != 16_000: wav = torchaudio.functional.resample( torch.from_numpy(wav).unsqueeze(0), sr, 16_000 )[0].numpy() prompt = torch.from_numpy(wav).unsqueeze(0) else: # —— legacy raw PCM bytes —— ----------------------------------- prompt = _bytes_to_tensor(zr.prompt_audio) # ───── call the model ────────────────────────────────────────────── speed = getattr(zr, "speed", 1.0) mo = self.cosyvoice.inference_zero_shot( zr.tts_text, zr.prompt_text, prompt, stream=False, speed=speed, ) finally: # clean up any temporary file we created if tmp_path and os.path.exists(tmp_path): try: os.remove(tmp_path) except Exception as e: logging.warning("Could not remove temp file %s: %s", tmp_path, e) yield from _yield_audio(mo) return # 3. Cross‑lingual if request.HasField("cross_lingual_request"): logging.info("Received cross‑lingual inference request") cr = request.cross_lingual_request tmp_path = None try: if cr.prompt_audio.startswith(b'http'): # S3 URL case url = cr.prompt_audio.decode('utf‑8') logging.info("Downloading cross‑lingual prompt from %s", url) resp = requests.get(url, timeout=10) resp.raise_for_status() with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f: f.write(resp.content) tmp_path = f.name wav, sr = sf.read(tmp_path, dtype='float32') if wav.ndim > 1: wav = wav.mean(axis=1) if sr != 16_000: wav = torchaudio.functional.resample( torch.from_numpy(wav).unsqueeze(0), sr, 16_000 )[0].numpy() prompt = torch.from_numpy(wav).unsqueeze(0) else: # legacy raw bytes prompt = _bytes_to_tensor(cr.prompt_audio) mo = self.cosyvoice.inference_cross_lingual( cr.tts_text, prompt ) finally: if tmp_path and os.path.exists(tmp_path): try: os.remove(tmp_path) except Exception as e: logging.warning("Could not remove temp file %s: %s", tmp_path, e) yield from _yield_audio(mo) return # 4. Instruction‑TTS (two flavours) if request.HasField("instruct_request"): ir = request.instruct_request # ────────────────────────────────────────────────────────────────── # 4‑a) instruct‑2 (has prompt_audio → bytes OR S3 URL) # ────────────────────────────────────────────────────────────────── if ir.HasField("prompt_audio"): logging.info("Received instruct‑2 inference request") tmp_path = None try: if ir.prompt_audio.startswith(b'http'): # treat as URL, download then load url = ir.prompt_audio.decode('utf‑8') logging.info("Downloading prompt audio from %s", url) resp = requests.get(url, timeout=10) resp.raise_for_status() with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f: f.write(resp.content) tmp_path = f.name wav, sr = sf.read(tmp_path, dtype='float32') if wav.ndim > 1: wav = wav.mean(axis=1) if sr != 16_000: wav = torchaudio.functional.resample( torch.from_numpy(wav).unsqueeze(0), sr, 16_000 )[0].numpy() prompt = torch.from_numpy(wav).unsqueeze(0) else: # legacy raw‑bytes payload prompt = _bytes_to_tensor(ir.prompt_audio) speed = getattr(ir, "speed", 1.0) mo = self.cosyvoice.inference_instruct2( ir.tts_text, ir.instruct_text, prompt, stream=False, speed=speed ) finally: if tmp_path and os.path.exists(tmp_path): try: os.remove(tmp_path) except Exception as e: logging.warning("Could not remove temp file %s: %s", tmp_path, e) # ────────────────────────────────────────────────────────────────── # 4‑b) classic instruct (speaker‑ID, no prompt audio) # ────────────────────────────────────────────────────────────────── else: logging.info("Received instruct inference request") mo = self.cosyvoice.inference_instruct( ir.tts_text, ir.spk_id, ir.instruct_text ) yield from _yield_audio(mo) return # unknown request type context.abort(grpc.StatusCode.INVALID_ARGUMENT, "Unsupported request type in oneof field.") # ──────────────────────────────────────────────────────────────────────────────── # entry‑point # ──────────────────────────────────────────────────────────────────────────────── def serve(args): server = grpc.server( futures.ThreadPoolExecutor(max_workers=args.max_conc), maximum_concurrent_rpcs=args.max_conc ) cosyvoice_pb2_grpc.add_CosyVoiceServicer_to_server( CosyVoiceServiceImpl(args), server ) server.add_insecure_port(f"0.0.0.0:{args.port}") server.start() logging.info("CosyVoice gRPC server listening on 0.0.0.0:%d", args.port) server.wait_for_termination() if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--port", type=int, default=8000) parser.add_argument("--max_conc", type=int, default=4, help="maximum concurrent requests / threads") parser.add_argument("--model_dir", type=str, default="pretrained_models/CosyVoice2-0.5B", help="local path or ModelScope repo id") serve(parser.parse_args())