67 lines
2.3 KiB
Python
67 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
|
"""CosyVoice2 性能对比：非流式 / 流式 / 流式+jit
|
|
|
|
用法: py -3 src/voice/cosyvoice-opt-test.py [--jit] [--stream]
|
|
"""
|
|
import os
|
|
import sys
|
|
import time
|
|
import subprocess
|
|
|
|
BASE = os.path.dirname(os.path.abspath(__file__))
|
|
MODEL_DIR = os.path.join(BASE, "models", "CosyVoice2-0.5B")
|
|
REF_WAV = os.path.join(BASE, "models", "melo-tts", "ref_female.wav")
|
|
REF_TEXT = "你好，我是小雅，今天想和你聊聊天。"
|
|
|
|
DO_JIT = "--jit" in sys.argv
|
|
DO_STREAM = "--stream" in sys.argv
|
|
|
|
sys.path.insert(0, os.path.join(BASE, "cosyvoice-src"))
|
|
sys.path.insert(0, os.path.join(BASE, "cosyvoice-src", "third_party", "Matcha-TTS"))
|
|
|
|
def gpu_mem():
|
|
try:
|
|
out = subprocess.run(
|
|
["nvidia-smi", "--query-gpu=memory.used,memory.total", "--format=csv,noheader,nounits"],
|
|
capture_output=True, text=True, timeout=5,
|
|
).stdout.strip()
|
|
return out
|
|
except Exception:
|
|
return "N/A"
|
|
|
|
def main():
|
|
tag = "stream" if DO_STREAM else "batch"
|
|
if DO_JIT:
|
|
tag += "+jit"
|
|
print(f"[opt] 配置: {tag} | GPU: {gpu_mem()} MiB", flush=True)
|
|
t0 = time.time()
|
|
|
|
from cosyvoice.cli.cosyvoice import AutoModel
|
|
cosyvoice = AutoModel(model_dir=MODEL_DIR, fp16=True, load_jit=DO_JIT)
|
|
print(f"[opt] 模型加载: {time.time()-t0:.1f}s | GPU: {gpu_mem()} MiB", flush=True)
|
|
|
|
text = "你好，我是小雅，很高兴认识你。今天我们聊聊白龙马的奇妙旅程吧。"
|
|
|
|
import torch
|
|
t1 = time.time()
|
|
first_chunk_t = None
|
|
chunks = []
|
|
gen = cosyvoice.inference_zero_shot(text, REF_TEXT, REF_WAV, stream=DO_STREAM)
|
|
for i, j in enumerate(gen):
|
|
if i == 0:
|
|
first_chunk_t = time.time() - t1
|
|
print(f"[opt] 首包: {first_chunk_t:.1f}s | GPU: {gpu_mem()} MiB", flush=True)
|
|
chunks.append(j['tts_speech'])
|
|
elapsed = time.time() - t1
|
|
wav = torch.cat(chunks, dim=1)
|
|
duration = wav.shape[1] / cosyvoice.sample_rate
|
|
print(f"[opt] 合成: {elapsed:.1f}s | 音频 {duration:.1f}s | RTF {elapsed/duration:.2f}"
|
|
f" | 首包 {first_chunk_t:.1f}s | GPU: {gpu_mem()} MiB", flush=True)
|
|
|
|
import soundfile as sf
|
|
sf.write(f"cosyvoice_{tag}.wav", wav.squeeze(0).cpu().numpy(), cosyvoice.sample_rate)
|
|
print(f"[opt] 已保存: cosyvoice_{tag}.wav", flush=True)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|