blind/audio_1/scripts/02_gen_tts_cosyvoice2.py

#!/usr/bin/env python3
"""
步骤 2：TTS v2 — CosyVoice2-0.5B 批量生成
输出目录：audio_1/tts/v2_cosyvoice2/


使用方法：
  conda activate cosyvoice
  python audio_1/scripts/02_gen_tts_cosyvoice2.py
  python audio_1/scripts/02_gen_tts_cosyvoice2.py --dry-run
"""
import os, sys, glob, json, time, tempfile, subprocess, argparse
import numpy as np

PROJECT_DIR    = "/home/xsl/blind"
VOICE_DIR      = f"{PROJECT_DIR}/voice"
OUT_DIR        = f"{PROJECT_DIR}/audio_1/tts/v2_cosyvoice2"
REF_TEXTS      = f"{PROJECT_DIR}/audio_1/ref_texts.json"
MODEL_DIR      = "/home/xsl/tools/CosyVoice/pretrained_models/CosyVoice2-0.5B"
COSYVOICE_DIR  = "/home/xsl/tools/CosyVoice"

os.makedirs(OUT_DIR, exist_ok=True)

sys.path.insert(0, COSYVOICE_DIR)
sys.path.insert(0, f"{PROJECT_DIR}/audio")
from batch_tts_voxcpm import ROLES, LINES

with open(REF_TEXTS, encoding="utf-8") as f:
    REF_TEXT_MAP = json.load(f)

# ──────────────────────────────────────────────────────────────
def find_ref_audio(role_key):
    role_dir = os.path.join(VOICE_DIR, ROLES[role_key]["dir"])
    for ext in ("*.wav", "*.mp3", "*.flac"):
        files = sorted(glob.glob(os.path.join(role_dir, ext)))
        if files:
            return files[0]
    return None

def load_audio_16k(path):
    """加载音频，重采样到 16kHz，返回 numpy float32 array"""
    import librosa
    audio, _ = librosa.load(path, sr=16000, mono=True)
    return audio

def wav_to_mp3(src, dst):
    subprocess.run(["ffmpeg", "-y", "-i", src, "-codec:a", "libmp3lame",
                    "-b:a", "128k", "-ar", "24000", "-ac", "1", dst],
                   check=True, capture_output=True)

def save_wav(audio_tensor, sr, path):
    """保存 torch tensor 或 numpy array 到 WAV"""
    import soundfile as sf
    if hasattr(audio_tensor, 'numpy'):
        audio = audio_tensor.squeeze().numpy()
    else:
        audio = np.squeeze(audio_tensor)
    sf.write(path, audio, sr)

# ──────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--role", help="只合成指定角色")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--resume", action="store_true", default=True)
    args = parser.parse_args()

    lines = LINES
    if args.role:
        lines = [(f,r,t) for f,r,t in lines if r == args.role]

    print(f"\n{'='*60}")
    print(f" TTS v2 — CosyVoice2-0.5B  ({len(lines)} 条)")
    print(f" 输出: {OUT_DIR}")
    print(f"{'='*60}\n")

    # 预检参考音频
    ref_cache = {}
    for role_key in ROLES:
        ref_path = find_ref_audio(role_key)
        if ref_path:
            ref_cache[role_key] = {
                "path": ref_path,
                "audio16k": None,  # 懒加载
                "text": REF_TEXT_MAP.get(role_key, ""),
            }

    if args.dry_run:
        for f,r,t in lines:
            out = os.path.join(OUT_DIR, f"{f}_wav_v1.mp3")
            done = "✓" if os.path.exists(out) else "·"
            ref_ok = "✓" if r in ref_cache else "✗"
            print(f"  [{done}] ref={ref_ok} {f}  [{r}]  {t[:40]}")
        return

    # 加载模型
    print("加载 CosyVoice2 模型...", end=" ", flush=True)
    t0 = time.time()
    from cosyvoice.cli.cosyvoice import CosyVoice2
    cosyvoice = CosyVoice2(MODEL_DIR)
    print(f"就绪 ({time.time()-t0:.1f}s)  SR={cosyvoice.sample_rate}")

    tmp = tempfile.mkdtemp(prefix="cosyvoice2_")
    ok = fail = skip = 0

    for i, (fname, role, text) in enumerate(lines, 1):
        out_mp3 = os.path.join(OUT_DIR, f"{fname}_wav_v1.mp3")
        if args.resume and os.path.exists(out_mp3):
            print(f"  [{i:3d}/{len(lines)}] SKIP {fname}")
            skip += 1; continue

        if role not in ref_cache:
            print(f"  [{i:3d}/{len(lines)}] MISS {fname} (无参考音频)")
            fail += 1; continue

        speed = ROLES[role]["speed"]
        ref = ref_cache[role]
        print(f"  [{i:3d}/{len(lines)}] {fname}  [{role}]  {text[:30]}... ", end="", flush=True)
        t0 = time.time()

        try:
            # 懒加载参考音频
            if ref["audio16k"] is None:
                ref["audio16k"] = load_audio_16k(ref["path"])

            # 合成：zero-shot 语音克隆
            import torch
            prompt_speech = torch.from_numpy(ref["audio16k"]).unsqueeze(0)
            audio_chunks = []
            for result in cosyvoice.inference_zero_shot(
                tts_text=text,
                prompt_text=ref["text"],
                prompt_speech_16k=prompt_speech,
                speed=speed,
                stream=False,
            ):
                audio_chunks.append(result['tts_speech'].squeeze().numpy())

            # 合并并保存
            full_audio = np.concatenate(audio_chunks) if audio_chunks else np.array([])
            out_wav = os.path.join(tmp, f"{fname}.wav")
            import soundfile as sf
            sf.write(out_wav, full_audio, cosyvoice.sample_rate)
            wav_to_mp3(out_wav, out_mp3)
            print(f"ok ({time.time()-t0:.1f}s  {os.path.getsize(out_mp3)//1024}KB)")
            ok += 1
        except Exception as e:
            print(f"FAIL: {e}"); fail += 1

    print(f"\n完成: {ok} ok, {skip} skip, {fail} fail → {OUT_DIR}")

if __name__ == "__main__":
    main()