inference_mp.py

"""
YingMusicSinger 批量推理脚本
支持多卡多进程、进度条显示
输入支持 JSONL 文件 或 LyricEditBench 数据集

用法:
    # JSONL 输入，4卡
    python batch_infer.py \
        --input_type jsonl \
        --input_path /path/to/input.jsonl \
        --output_dir /path/to/output \
        --ckpt_path /path/to/ckpts \
        --num_gpus 4

    # LyricEditBench 输入
    python batch_infer.py \
        --input_type lyric_edit_bench \
        --output_dir /path/to/output \
        --ckpt_path /path/to/ckpts \
        --num_gpus 4
"""

import argparse
import json
import os
import sys
import traceback
from pathlib import Path

import torch
import torch.multiprocessing as mp
import torchaudio
from datasets import Audio, Dataset
from huggingface_hub import hf_hub_download
from tqdm import tqdm


def load_jsonl(path: str) -> list[dict]:
    items = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                items.append(json.loads(line))
    return items


def build_dataset_from_local(gtsinger_root: str):
    """
    Build LyricEditBench dataset using your local GTSinger directory.

    Args:
        gtsinger_root: Root directory of your local GTSinger dataset.
    """
    # Download the inherited metadata from HuggingFace
    json_path = hf_hub_download(
        repo_id="ASLP-lab/LyricEditBench",
        filename="GTSinger_Inherited.json",
        repo_type="dataset",
    )

    with open(json_path, "r") as f:
        data = json.load(f)

    gtsinger_root = str(Path(gtsinger_root).resolve())

    # Prepend local root to relative paths
    for item in data:
        item["melody_ref_path"] = os.path.join(gtsinger_root, item["melody_ref_path"])
        item["timbre_ref_path"] = os.path.join(gtsinger_root, item["timbre_ref_path"])
        # Set audio fields to the resolved file paths
        item["melody_ref_audio"] = item["melody_ref_path"]
        item["timbre_ref_audio"] = item["timbre_ref_path"]

    # Build HuggingFace Dataset with Audio features
    ds = Dataset.from_list(data)
    ds = ds.cast_column("melody_ref_audio", Audio())
    ds = ds.cast_column("timbre_ref_audio", Audio())

    return ds


def load_subset(data: list, subset_id: str) -> list:
    """Filter dataset by a subset ID list."""
    subset_path = hf_hub_download(
        repo_id="ASLP-lab/LyricEditBench",
        filename=f"id_lists/{subset_id}.txt",
        repo_type="dataset",
    )

    with open(subset_path, "r") as f:
        id_set = set(line.strip() for line in f if line.strip())

    return [item for item in data if item["id"] in id_set]


def load_lyric_edit_bench(input_type) -> list[dict]:
    # If you have GTsinger downloaded, use this:

    ds_full = build_dataset_from_local(
        "/user-fs/chenzihao/zhengjunjie/datas/Music/openvocaldata/GTSinger"
    )

    # else, you kan use this:
    # from datasets import load_dataset

    # ds_full = load_dataset("ASLP-lab/LyricEditBench", split="test")

    # ds_full loaded

    subset_1k = load_subset(ds_full, "1K")
    print(f"Loaded {len(subset_1k)} items")

    items = []
    for row in subset_1k:
        if input_type == "lyric_edit_bench_melody_control":
            items.append(
                {
                    "id": row.get("id", ""),
                    "melody_ref_path": row.get("melody_ref_path", ""),
                    "gen_text": row.get("gen_text", ""),
                    "timbre_ref_path": row.get("timbre_ref_path", ""),
                    "timbre_ref_text": row.get("timbre_ref_text", ""),
                }
            )
        elif input_type == "lyric_edit_bench_sing_edit":
            items.append(
                {
                    "id": row.get("id", ""),
                    "melody_ref_path": row.get("melody_ref_path", ""),
                    "gen_text": row.get("gen_text", ""),
                    "timbre_ref_path": row.get("melody_ref_path", ""),
                    "timbre_ref_text": row.get("melody_ref_text", ""),
                }
            )
        else:
            assert 0
    return items


def worker(
    rank: int,
    world_size: int,
    items: list[dict],
    output_dir: str,
    ckpt_path: str,
    args: argparse.Namespace,
):
    """每个 GPU 上运行的 worker 进程"""
    device = f"cuda:{rank}"
    torch.cuda.set_device(rank)

    # ---- 加载模型 ----
    from src.YingMusicSinger.infer.YingMusicSinger import YingMusicSinger

    model = YingMusicSinger.from_pretrained(ckpt_path)
    model.to(device)
    model.eval()

    # ---- 分片: 每个 worker 处理自己那份 ----
    shard = items[rank::world_size]

    # ---- 只在 rank 0 显示进度条 ----
    pbar = tqdm(
        shard,
        desc=f"[GPU {rank}]",
        position=rank,
        leave=True,
        disable=(rank != 0 and not args.show_all_progress),
    )

    success, fail = 0, 0
    for item in pbar:
        item_id = item.get("id", f"unknown_{success + fail}")
        out_path = os.path.join(output_dir, f"{item_id}.wav")

        # 跳过已存在的文件
        if os.path.exists(out_path) and not args.overwrite:
            success += 1
            pbar.set_postfix(ok=success, err=fail)
            continue

        try:
            with torch.no_grad():
                audio, sr = model(
                    ref_audio_path=item["timbre_ref_path"],
                    melody_audio_path=item["melody_ref_path"],
                    ref_text=item.get("timbre_ref_text", ""),
                    target_text=item.get("gen_text", ""),
                    lrc_align_mode=args.lrc_align_mode,
                    sil_len_to_end=args.sil_len_to_end,
                    t_shift=args.t_shift,
                    nfe_step=args.nfe_step,
                    cfg_strength=args.cfg_strength,
                    seed=args.seed
                    if args.seed != -1
                    else torch.randint(0, 2**32, (1,)).item(),
                )

            torchaudio.save(out_path, audio, sample_rate=sr)
            success += 1

        except Exception as e:
            fail += 1
            print(f"\n[GPU {rank}] ERROR on {item_id}: {e}", file=sys.stderr)
            if args.verbose:
                traceback.print_exc()

        pbar.set_postfix(ok=success, err=fail)

    pbar.close()
    print(f"[GPU {rank}] Done. success={success}, fail={fail}")


def main():
    parser = argparse.ArgumentParser(description="YingMusicSinger 批量推理")

    # ---- 输入 ----
    parser.add_argument(
        "--input_type",
        type=str,
        required=True,
        choices=[
            "jsonl",
            "lyric_edit_bench_melody_control",
            "lyric_edit_bench_sing_edit",
        ],
        help="输入类型: jsonl / lyric_edit_bench_melody_control 或 lyric_edit_bench_sing_edit",
    )
    parser.add_argument(
        "--input_path",
        type=str,
        default=None,
        help="JSONL 文件路径 (input_type=jsonl 时必填)",
    )

    # ---- 输出 ----
    parser.add_argument(
        "--output_dir",
        type=str,
        required=True,
        help="输出目录",
    )

    # ---- 模型 ----
    parser.add_argument(
        "--ckpt_path",
        type=str,
        required=False,
        help="模型 checkpoint 路径 (save_pretrained 保存的目录)",
        default=None,
    )

    # ---- 推理参数 ----
    parser.add_argument(
        "--num_gpus", type=int, default=None, help="使用 GPU 数量，默认全部"
    )
    parser.add_argument(
        "--lrc_align_mode",
        type=str,
        default="sentence_level",
        choices=["sentence_level"],
    )
    parser.add_argument("--sil_len_to_end", type=float, default=0.5)
    parser.add_argument("--t_shift", type=float, default=0.5)
    parser.add_argument("--nfe_step", type=int, default=32)
    parser.add_argument("--cfg_strength", type=float, default=3.0)
    parser.add_argument("--seed", type=int, default=-1)

    # ---- 其它 ----
    parser.add_argument("--overwrite", action="store_true", help="覆盖已有输出文件")
    parser.add_argument(
        "--show_all_progress", action="store_true", help="所有 GPU 都显示进度条"
    )
    parser.add_argument("--verbose", action="store_true", help="打印详细错误信息")

    args = parser.parse_args()

    # ---- 校验 ----
    if args.input_type == "jsonl":
        assert args.input_path is not None, "--input_path 是 jsonl 模式下必填的"
        assert os.path.isfile(args.input_path), f"文件不存在: {args.input_path}"

    # ---- 加载数据 ----
    print("加载数据...")
    if args.input_type == "jsonl":
        items = load_jsonl(args.input_path)
    else:
        items = load_lyric_edit_bench(args.input_type)
    print(f"共 {len(items)} 条数据")

    # ---- 确定 GPU 数量 ----
    available_gpus = torch.cuda.device_count()
    num_gpus = args.num_gpus or available_gpus
    num_gpus = min(num_gpus, available_gpus, len(items))
    assert num_gpus > 0, "没有可用的 GPU"
    print(f"使用 {num_gpus} 张 GPU")

    # ---- 创建输出目录 ----
    os.makedirs(args.output_dir, exist_ok=True)

    # ---- 启动多进程 ----
    if num_gpus == 1:
        # 单卡直接跑，不需要 spawn
        worker(0, 1, items, args.output_dir, args.ckpt_path, args)
    else:
        mp.set_start_method("spawn", force=True)
        processes = []
        for rank in range(num_gpus):
            p = mp.Process(
                target=worker,
                args=(rank, num_gpus, items, args.output_dir, args.ckpt_path, args),
            )
            p.start()
            processes.append(p)

        for p in processes:
            p.join()

    print(f"\n推理完成! 输出目录: {args.output_dir}")


if __name__ == "__main__":
    main()