blind/audio_2/scripts/01_gen_tts_styled.py

#!/usr/bin/env python3
"""
《林夏》TTS v2 — VoxCPM2 Style Control 批量生成
输出：audio_2/tts/v2_styled/

用法:
  VOXCPM_PORT=8002 bash /home/xsl/tts-server/start-voxcpm.sh
  VOXCPM_PORT=8002 python audio_2/scripts/01_gen_tts_styled.py
  python audio_2/scripts/01_gen_tts_styled.py --dry-run
  python audio_2/scripts/01_gen_tts_styled.py --role linxia --test
"""
import argparse
import base64
import glob
import os
import subprocess
import sys
import tempfile
import time

import requests

PROJECT_DIR = "/home/xsl/blind"
VOICE_DIR = f"{PROJECT_DIR}/voice"
OUT_DIR = f"{PROJECT_DIR}/audio_2/tts/v2_styled"
HOST = os.environ.get("VOXCPM_HOST", "127.0.0.1")
PORT = os.environ.get("VOXCPM_PORT", "8002")
BASE_URL = f"http://{HOST}:{PORT}"

sys.path.insert(0, f"{PROJECT_DIR}/audio")
sys.path.insert(0, f"{PROJECT_DIR}/audio_2")

from batch_tts_voxcpm import ROLES, LINES  # noqa: E402
from voice_style_prompts import get_style  # noqa: E402

os.makedirs(OUT_DIR, exist_ok=True)


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 convert_to_wav(src, dst):
    subprocess.run(
        ["ffmpeg", "-y", "-i", src, "-ar", "16000", "-ac", "1", "-acodec", "pcm_s16le", dst],
        check=True,
        capture_output=True,
    )


def wav_to_base64(path):
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode()


def apply_speed(src, speed, dst):
    if abs(speed - 1.0) < 0.01:
        import shutil

        shutil.copy2(src, dst)
        return
    subprocess.run(
        ["ffmpeg", "-y", "-i", src, "-filter:a", f"atempo={speed}", "-acodec", "pcm_s16le", dst],
        check=True,
        capture_output=True,
    )


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 check_server():
    try:
        r = requests.get(f"{BASE_URL}/health", timeout=5)
        if r.status_code != 200:
            return False
        data = r.json()
        return data.get("status") == "ok"
    except Exception:
        return False


def register_voice(wav_path):
    r = requests.post(
        f"{BASE_URL}/v1/voices",
        json={"wav_base64": wav_to_base64(wav_path)},
        timeout=60,
    )
    r.raise_for_status()
    return r.json()["voice_id"]


def synthesize_styled(voice_id, text, style, cfg=2.0, steps=10):
    r = requests.post(
        f"{BASE_URL}/v1/speech/styled",
        json={
            "text": text,
            "style": style,
            "voice_id": voice_id,
            "cfg_value": cfg,
            "inference_timesteps": steps,
        },
        timeout=180,
    )
    r.raise_for_status()
    return r.content


def main():
    parser = argparse.ArgumentParser(description="《林夏》Style Control 批量 TTS → audio_2")
    parser.add_argument("--role", help="只合成指定角色")
    parser.add_argument("--test", action="store_true", help="只合成第一条")
    parser.add_argument("--dry-run", action="store_true", help="只打印 style 分配")
    parser.add_argument("--force", action="store_true", help="覆盖已存在文件")
    parser.add_argument("--cfg", type=float, default=2.0)
    parser.add_argument("--steps", type=int, default=10)
    args = parser.parse_args()

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

    print(f"\n{'=' * 60}")
    print(" 《林夏》TTS v2 — Style Control")
    print(f" 服务: {BASE_URL}")
    print(f" 输出: {OUT_DIR}")
    print(f" 共 {len(lines)} 条")
    print(f"{'=' * 60}\n")

    if args.dry_run:
        for fname, role, text in lines:
            style = get_style(fname, role, text)
            out = os.path.join(OUT_DIR, f"{fname}_wav_v2.mp3")
            mark = "✓" if os.path.exists(out) else "·"
            print(f"  [{mark}] {fname}  [{role}]")
            print(f"       style: {style}")
            print(f"       text:  {text[:50]}{'...' if len(text) > 50 else ''}\n")
        return

    if not check_server():
        print(f"[错误] VoxCPM2 服务未就绪: {BASE_URL}/health")
        print("请先启动:")
        print(f"  VOXCPM_PORT={PORT} bash /home/xsl/tts-server/start-voxcpm.sh")
        sys.exit(1)

    voice_ids = {}
    tmp = tempfile.mkdtemp(prefix="voxcpm2_")
    for role_key in sorted(set(r for _, r, _ in lines)):
        ref = find_ref_audio(role_key)
        if not ref:
            print(f"  [跳过] {role_key}: 未找到参考音频")
            continue
        wav = ref if ref.endswith(".wav") else os.path.join(tmp, f"{role_key}.wav")
        if not ref.endswith(".wav"):
            convert_to_wav(ref, wav)
        print(f"  注册音色: {role_key} ... ", end="", flush=True)
        voice_ids[role_key] = register_voice(wav)
        print(f"ok ({voice_ids[role_key][:10]}...)")

    ok = skip = fail = 0
    manifest_path = os.path.join(OUT_DIR, "style_manifest.tsv")
    manifest_lines = ["filename\trole\tstyle\ttext\n"]

    for i, (fname, role, text) in enumerate(lines, 1):
        out_mp3 = os.path.join(OUT_DIR, f"{fname}_wav_v2.mp3")
        style = get_style(fname, role, text)
        manifest_lines.append(
            f"{fname}\t{role}\t{style.replace(chr(9), ' ')}\t{text.replace(chr(9), ' ')}\n"
        )

        if os.path.exists(out_mp3) and not args.force:
            print(f"  [{i:3d}/{len(lines)}] SKIP {fname}")
            skip += 1
            continue

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

        print(f"  [{i:3d}/{len(lines)}] {fname} ... ", end="", flush=True)
        t0 = time.time()
        try:
            wav_bytes = synthesize_styled(
                voice_ids[role], text, style, cfg=args.cfg, steps=args.steps
            )
            raw_wav = os.path.join(tmp, f"{fname}_raw.wav")
            spd_wav = os.path.join(tmp, f"{fname}_spd.wav")
            with open(raw_wav, "wb") as f:
                f.write(wav_bytes)
            apply_speed(raw_wav, ROLES[role]["speed"], spd_wav)
            wav_to_mp3(spd_wav, out_mp3)
            kb = os.path.getsize(out_mp3) // 1024
            print(f"ok ({time.time() - t0:.1f}s, {kb}KB)")
            ok += 1
        except Exception as e:
            print(f"FAIL: {e}")
            fail += 1

    with open(manifest_path, "w", encoding="utf-8") as f:
        f.writelines(manifest_lines)

    print(f"\n完成: {ok} ok, {skip} skip, {fail} fail")
    print(f"输出: {OUT_DIR}")
    print(f"清单: {manifest_path}")


if __name__ == "__main__":
    main()