cosyvoice_worker.py
Rama main · réplica de lectura del repositorio original.
#!/usr/bin/env python3
"""CosyVoice 3 sidecar — runs in the isolated `cosyvoice` conda env and exposes
a small HTTP API on 127.0.0.1:8101 so PolyTTS's CosyvoiceEngine can drive it
without importing CosyVoice's conflicting deps (CosyVoice pins torch 2.3.1; the
main PolyTTS venv runs torch 2.12 for qwen/voxcpm).
Run: conda run -n cosyvoice python cosyvoice_worker.py
PolyTTS orchestrates VRAM: POST /load loads the model, /unload frees it.
"""
import os
import sys
import io
import gc
import threading
COSYVOICE_REPO = os.environ.get("COSYVOICE_REPO", "/home/ubuntu/CosyVoice")
MODEL_DIR = os.environ.get("COSYVOICE_MODEL_DIR",
f"{COSYVOICE_REPO}/pretrained_models/Fun-CosyVoice3-0.5B")
PORT = int(os.environ.get("COSYVOICE_PORT", "8101"))
# CosyVoice package + its Matcha-TTS submodule must be importable
for p in (COSYVOICE_REPO, f"{COSYVOICE_REPO}/third_party/Matcha-TTS"):
if p not in sys.path:
sys.path.insert(0, p)
import numpy as np
import soundfile as sf
import torch
import uvicorn
from fastapi import FastAPI
from fastapi.responses import Response, JSONResponse
from pydantic import BaseModel
from cosyvoice.cli.cosyvoice import AutoModel
SYS_PREFIX = "You are a helpful assistant"
app = FastAPI(title="CosyVoice sidecar")
_model = None
_lock = threading.Lock()
def _ensure_loaded():
global _model
if _model is None:
with _lock:
if _model is None:
print(f"[cosyvoice] loading model from {MODEL_DIR} ...", flush=True)
_model = AutoModel(model_dir=MODEL_DIR)
print(f"[cosyvoice] loaded. sr={_model.sample_rate}", flush=True)
return _model
class TTSReq(BaseModel):
text: str
voice_wav_path: str
ref_text: str = ""
instruct: str | None = None # emotion/style instruction, e.g. "请用温暖的语气"
speed: float | None = None # 0.5..2.0 (passed as part of instruct if set)
# Group N sentences per generation so prosody flows within a chunk, then join
# chunks with a controlled pause. 1 = per-sentence (choppy); large = whole-text
# (CosyVoice decides boundaries, can cut mid-sentence / under-pause).
sentences_per_chunk: int | None = None
interchunk_silence: float | None = None
# Insert a CosyVoice control token (e.g. "[breath]" / "[quick_breath]") at
# every within-chunk sentence boundary so the model breathes there naturally
# (fixes under-paused sentence ends inside multi-sentence chunks).
breath_token: str | None = None
# Generate sentence-by-sentence and join with a guaranteed inter-sentence
# silence. CosyVoice's whole-text mode groups multiple sentences per chunk and
# hard-concat ran some sentence ends into the next with insufficient pause.
INTERSENTENCE_SILENCE = float(os.environ.get("COSYVOICE_INTERSENTENCE_SILENCE", "0.35"))
_SENT_ENDERS = "。！？!?\n"
def _split_sentences(text: str):
out, buf = [], ""
for ch in text:
buf += ch
if ch in _SENT_ENDERS:
s = buf.strip()
if s:
out.append(s)
buf = ""
if buf.strip():
out.append(buf.strip())
return out or [text]
def _gen_one(m, text, inst, ref_text, voice_wav_path):
if inst:
prompt = f"{SYS_PREFIX}. {inst}<|endofprompt|>"
return m.inference_instruct2(text, prompt, voice_wav_path, stream=False)
prompt = f"{SYS_PREFIX}<|endofprompt|>{ref_text}"
return m.inference_zero_shot(text, prompt, voice_wav_path, stream=False)
def _synth(req: TTSReq):
m = _ensure_loaded()
inst = req.instruct or ""
if req.speed:
inst = (inst + " " if inst else "") + f"语速设为{req.speed}"
sents = _split_sentences(req.text)
spc = req.sentences_per_chunk or int(os.environ.get("COSYVOICE_SENTENCES_PER_CHUNK", "2"))
gap_s = (req.interchunk_silence if req.interchunk_silence is not None
else float(os.environ.get("COSYVOICE_INTERCHUNK_SILENCE", "0.5")))
breath = req.breath_token or os.environ.get("COSYVOICE_BREATH_TOKEN", "")
# Pack sentences into sentence-aligned chunks (no mid-sentence cuts); insert
# a breath token between sentences within a chunk so every boundary breathes.
joiner = f"{breath}" if breath else ""
chunks_text = [joiner.join(sents[i:i + spc]) for i in range(0, len(sents), spc)]
gap = torch.zeros(int(m.sample_rate * gap_s)) if len(chunks_text) > 1 else None
with _lock: # serialize GPU access across the whole multi-chunk job
pieces = []
for ct in chunks_text:
seg = []
for j in _gen_one(m, ct, inst, req.ref_text, req.voice_wav_path):
t = j["tts_speech"]
if t.dim() > 1:
t = t.squeeze(0)
seg.append(t.cpu().float())
if seg:
pieces.append(torch.cat(seg))
print(f"[cosyvoice] {len(sents)} sentences -> {len(chunks_text)} chunks "
f"(<= {spc}/chunk), gap={gap_s}s, breath={breath or 'none'}", flush=True)
if not pieces:
return np.zeros(0, dtype=np.float32), m.sample_rate
audio = pieces[0]
for p in pieces[1:]:
audio = torch.cat([audio, gap, p])
return np.asarray(audio.numpy(), dtype=np.float32), m.sample_rate
@app.post("/load")
def load():
_ensure_loaded()
return {"loaded": True, "sample_rate": _model.sample_rate}
@app.post("/unload")
def unload():
global _model
with _lock:
_model = None
gc.collect()
torch.cuda.empty_cache()
return {"loaded": False}
@app.get("/health")
def health():
return {"loaded": _model is not None, "model_dir": MODEL_DIR}
@app.post("/tts")
def tts(req: TTSReq):
try:
audio, sr = _synth(req)
except Exception as e:
return JSONResponse(status_code=500, content={"detail": str(e)})
buf = io.BytesIO()
sf.write(buf, audio, sr, subtype="PCM_16", format="WAV")
return Response(content=buf.getvalue(), media_type="audio/wav",
headers={"X-Sample-Rate": str(sr)})
if __name__ == "__main__":
uvicorn.run(app, host="127.0.0.1", port=PORT, log_level="info")