#!/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()