inference/additional_entrypoints/generate_vstyle_query.remote.py

import os
import json
import re
import argparse
import numpy as np
import soundfile as sf
import librosa
import torch
from tqdm import tqdm
import random
from qwen_tts import Qwen3TTSModel

def remove_brackets_content_zh(text):
    # 使用正则表达式匹配【】及其内部的所有内容，并将其替换为空字符串
    # 问号 '?' 表示非贪婪匹配，确保每次只匹配一个完整的【xxxx】
    cleaned_text = re.sub(r'【.*?】', '', text)
    return cleaned_text

def remove_brackets_content_en(text):
    # 将【】及其内容替换为一个空格 ' '，以保证英文句子之间有正确的分隔
    cleaned_text = re.sub(r'【.*?】', ' ', text)
    
    # .strip() 会自动删除字符串最开头和最末尾的多余空格
    # re.sub(r'\s+', ' ', ...) 会把句子中间如果偶然出现的多个连续空格合并成一个
    cleaned_text = re.sub(r'\s+', ' ', cleaned_text).strip()
    
    return cleaned_text

def get_target_control(item, target_key):
    """
    获取目标的 control 字典。
    完美兼容“无后缀基础版”与“带数字后缀进化版”同时存在的情况。
    """
    max_idx = -1
    best_key = None
    
    # 1. 保底探测：先把 target_key 自身当做基线 (相当于 index = -1)
    if target_key in item:
        best_key = target_key
        
    # 2. 进化探测：去寻找 target_key_0, target_key_1... 找数字最大的覆盖基线
    pattern = re.compile(rf"{re.escape(target_key)}_(\d+)")
    for key in item.keys():
        match = pattern.fullmatch(key)
        if match:
            idx = int(match.group(1))
            if idx > max_idx:
                max_idx = idx
                best_key = key
                
    if best_key:
        return item.get(best_key), best_key
    return None, None

def read_jsonl(file_path):
    data = []
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            for line_number, line in enumerate(f, start=1):
                line = line.strip()
                if not line: continue
                try:
                    item = json.loads(line)
                    # Allow resume subsets to preserve the original full-jsonl line id.
                    item["line_idx"] = int(item.get("line_idx", line_number))
                    data.append(item)
                except json.JSONDecodeError as e:
                    print(f"[Warning] 第 {line_number} 行解析失败: {e}")
    except Exception as e:
        print(f"[Error] 读取文件异常: {e}")
    return data

def trim_silence(audio, top_db=45):
    if len(audio) == 0:
        return audio
    trimmed_audio, _ = librosa.effects.trim(audio, top_db=top_db)
    return trimmed_audio

def parse_requested_speakers(speakers_arg):
    return [x.strip() for x in speakers_arg.split(",") if x.strip()]

def get_nonempty_segment_indexes(control):
    return [
        idx for idx, part in enumerate(control["Control"])
        if part.get("sample_text", "").strip()
    ]

def write_item_metadata(item, control, out_sub_dir, line_idx):
    os.makedirs(out_sub_dir, exist_ok=True)
    control_json_path = os.path.join(out_sub_dir, "control.json")
    with open(control_json_path, "w", encoding="utf-8") as f:
        json.dump(control, f, ensure_ascii=False, indent=4)

    keys_to_keep = ["audio_content", "ability", "file_name", "instruct_id", "language"]
    instruct_data = {k: item[k] for k in keys_to_keep if k in item}
    instruct_id = item.get("instruct_id", line_idx)
    instruct_json_path = os.path.join(out_sub_dir, f"{instruct_id}_instruct.json")
    with open(instruct_json_path, "w", encoding="utf-8") as f:
        json.dump(instruct_data, f, ensure_ascii=False, indent=4)

def item_needs_generation(item, args):
    control = item["_parsed_control"]
    line_idx = item.get("line_idx")
    out_sub_dir = os.path.join(args.output_dir, str(line_idx))
    requested_speakers = parse_requested_speakers(args.speakers)

    voice_instruct_zh = control["Global"].get("instruct_zh", "")
    voice_instruct_en = control["Global"].get("instruct_en", "")
    if (
        not args.en_only
        and voice_instruct_zh
        and not os.path.exists(os.path.join(out_sub_dir, f"{line_idx}_vd_zh.wav"))
    ):
        return True
    if not args.zh_only and voice_instruct_en and not os.path.exists(os.path.join(out_sub_dir, f"{line_idx}_vd_en.wav")):
        return True

    # Without an explicit speaker subset, the supported speaker list is only
    # known after model load, so stay conservative.
    if not requested_speakers:
        return True

    segment_indexes = get_nonempty_segment_indexes(control)
    for speaker in requested_speakers:
        if not args.en_only:
            final_zh_path = os.path.join(out_sub_dir, f"{line_idx}_cv_{speaker}_zh.wav")
            if not os.path.exists(final_zh_path):
                return True
            for seg_idx in segment_indexes:
                seg_path = os.path.join(out_sub_dir, f"{line_idx}_cv_{speaker}_zh_{seg_idx}.wav")
                if not os.path.exists(seg_path):
                    return True

        if not args.zh_only:
            final_en_path = os.path.join(out_sub_dir, f"{line_idx}_cv_{speaker}_en.wav")
            if not os.path.exists(final_en_path):
                return True
            for seg_idx in segment_indexes:
                seg_path = os.path.join(out_sub_dir, f"{line_idx}_cv_{speaker}_en_{seg_idx}.wav")
                if not os.path.exists(seg_path):
                    return True

    write_item_metadata(item, control, out_sub_dir, line_idx)
    return False

# ================= 主程序 =================
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_jsonl", required=True, type=str)
    parser.add_argument("--output_dir", required=True, type=str)
    parser.add_argument("--custom_voice_path", required=True, type=str)
    parser.add_argument("--voice_design_path", required=True, type=str)
    parser.add_argument("--num_gpus", default=4, type=int, help="总分块数")
    parser.add_argument("--gpu_id", required=True, type=int, help="当前处理的块编号 (0 到 num_gpus-1)")
    parser.add_argument("--control_key", type=str, default=None, help="强制指定读取的 control 字段，如 final_generated_control")
    parser.add_argument("--speakers", type=str, default="", help="Comma-separated CustomVoice speaker subset. Default keeps all speakers.")
    parser.add_argument("--zh_only", action="store_true", help="Only generate zh VoiceDesign/CustomVoice prompt wavs.")
    parser.add_argument("--en_only", action="store_true", help="Only generate en VoiceDesign/CustomVoice prompt wavs.")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)
    
    # 1. 读取数据
    all_data = read_jsonl(args.input_jsonl)
    if not all_data:
        print("无有效数据，退出。")
        return

    # 2. 提前过滤出需要处理的有效数据，确保负载均衡
    valid_data = []
    for item in all_data:
        control, used_key = get_target_control(item, args.control_key)
        
        # 校验 control 是否存在以及格式是否完整
        if not control or "Global" not in control or "Control" not in control:
            continue
            
        # 校验是否包含有效文本
        texts = [c.get("sample_text", "") for c in control["Control"]]
        full_text = "".join(texts)
        if not full_text.strip():
            continue
            
        # 把提取好的 control 存回 item 中，后续直接用，避免二次解析
        item["_parsed_control"] = control
        valid_data.append(item)

    if not valid_data:
        print("未找到包含有效 control_key 的数据，退出。")
        return

    # 3. 对有效数据进行均衡分块
    chunk_size = (len(valid_data) + args.num_gpus - 1) // args.num_gpus
    chunks = [valid_data[i:i + chunk_size] for i in range(0, len(valid_data), chunk_size)]
    
    if args.gpu_id >= len(chunks):
        print(f"[Worker {args.gpu_id}] 没有分配到数据块，任务结束。")
        return
        
    my_chunk = chunks[args.gpu_id]
    my_chunk = [item for item in my_chunk if item_needs_generation(item, args)]
    if not my_chunk:
        print(f"[Worker {args.gpu_id}] 没有待生成音频，跳过模型加载。")
        return
    
    device = "cuda:0" 
    print(f"[Worker {args.gpu_id}] 启动，共需处理 {len(my_chunk)} 条有效数据 (总有效数据: {len(valid_data)})。")

    # 4. 加载模型
    print(f"[Worker {args.gpu_id}] 正在加载 VoiceDesign 模型...")
    vd_model = Qwen3TTSModel.from_pretrained(
        args.voice_design_path, device_map=device, dtype=torch.bfloat16, attn_implementation="flash_attention_2"
    )

    print(f"[Worker {args.gpu_id}] 正在加载 CustomVoice 模型...")
    cv_model = Qwen3TTSModel.from_pretrained(
        args.custom_voice_path, device_map=device, dtype=torch.bfloat16, attn_implementation="flash_attention_2"
    )
    
    supported_speakers = cv_model.get_supported_speakers()
    if args.speakers.strip():
        requested = [x.strip() for x in args.speakers.split(",") if x.strip()]
        supported_speakers = [x for x in requested if x in supported_speakers]
        if not supported_speakers:
            raise ValueError(f"No requested speakers are supported: {requested}")
    print(f"[Worker {args.gpu_id}] 支持的说话人: {supported_speakers}")
    random.seed(args.gpu_id)

    # 5. 遍历处理分配给当前 Worker 的数据
    for item in tqdm(my_chunk, desc=f"Worker {args.gpu_id} Progress"):

        current_seed = random.randint(0, 200)
        random.seed(current_seed)         # Python 原生随机库
        np.random.seed(current_seed)      # Numpy 随机库
        torch.manual_seed(current_seed)   # PyTorch CPU
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(current_seed)

        line_idx = item.get("line_idx")
        control = item["_parsed_control"]
            
        out_sub_dir = os.path.join(args.output_dir, str(line_idx))
        os.makedirs(out_sub_dir, exist_ok=True)

        voice_instruct_zh = control["Global"].get("instruct_zh", "")
        voice_instruct_en = control["Global"].get("instruct_en", "")

        expressive_instructs_zh = [c.get("instruct_zh", "") for c in control["Control"]]
        expressive_instructs_en = [c.get("instruct_en", "") for c in control["Control"]]
        texts = [c.get("sample_text", "") for c in control["Control"]]
        full_text = "".join(texts)

        # --- VoiceDesign 生成 ---
        try:
            vd_zh_path = os.path.join(out_sub_dir, f"{line_idx}_vd_zh.wav")
            if not args.en_only and voice_instruct_zh and not os.path.exists(vd_zh_path):
                wavs, sr = vd_model.generate_voice_design(text=full_text, language="Auto", instruct=remove_brackets_content_zh(voice_instruct_zh))
                sf.write(vd_zh_path, wavs[0], sr)
                
            if not args.zh_only:
                vd_en_path = os.path.join(out_sub_dir, f"{line_idx}_vd_en.wav")
                if voice_instruct_en and not os.path.exists(vd_en_path):
                    wavs, sr = vd_model.generate_voice_design(text=full_text, language="Auto", instruct=remove_brackets_content_en(voice_instruct_en))
                    sf.write(vd_en_path, wavs[0], sr)
        except Exception as e:
            print(f"[Worker {args.gpu_id}] 行号 {line_idx} VoiceDesign 生成失败: {e}")

        # --- CustomVoice 生成 ---
        needs_trimming = len(texts) >= 2
        for speaker in supported_speakers:
            try:
                # 中文部分
                cv_zh_segments = []
                final_sr_zh = 24000
                final_zh_path = os.path.join(out_sub_dir, f"{line_idx}_cv_{speaker}_zh.wav")
                
                if not args.en_only:
                    for seg_idx, (text_seg, inst_seg) in enumerate(zip(texts, expressive_instructs_zh)):
                        if not text_seg.strip(): continue
                        seg_wav_path = os.path.join(out_sub_dir, f"{line_idx}_cv_{speaker}_zh_{seg_idx}.wav")
                        
                        # 判断切片是否存在
                        if os.path.exists(seg_wav_path):
                            audio_data, sr = sf.read(seg_wav_path)
                        else:
                            wavs, sr = cv_model.generate_custom_voice(text=text_seg, language="Auto", speaker=speaker, instruct=remove_brackets_content_zh(inst_seg))
                            audio_data = trim_silence(wavs[0], top_db=45) if needs_trimming else wavs[0]
                            sf.write(seg_wav_path, audio_data, sr)
                            
                        cv_zh_segments.append(audio_data)
                        final_sr_zh = sr
                        
                    # 合并音频（仅当合并文件不存在时写入）
                    if cv_zh_segments and not os.path.exists(final_zh_path):
                        final_zh_audio = np.concatenate(cv_zh_segments)
                        sf.write(final_zh_path, final_zh_audio, final_sr_zh)

                if not args.zh_only:
                    # 英文部分
                    cv_en_segments = []
                    final_sr_en = 24000
                    final_en_path = os.path.join(out_sub_dir, f"{line_idx}_cv_{speaker}_en.wav")
                    
                    for seg_idx, (text_seg, inst_seg) in enumerate(zip(texts, expressive_instructs_en)):
                        if not text_seg.strip(): continue
                        seg_wav_path = os.path.join(out_sub_dir, f"{line_idx}_cv_{speaker}_en_{seg_idx}.wav")
                        
                        # 判断切片是否存在
                        if os.path.exists(seg_wav_path):
                            audio_data, sr = sf.read(seg_wav_path)
                        else:
                            wavs, sr = cv_model.generate_custom_voice(text=text_seg, language="Auto", speaker=speaker, instruct=remove_brackets_content_en(inst_seg))
                            audio_data = trim_silence(wavs[0], top_db=45) if needs_trimming else wavs[0]
                            sf.write(seg_wav_path, audio_data, sr)

                        cv_en_segments.append(audio_data)
                        final_sr_en = sr
                        
                    # 合并音频（仅当合并文件不存在时写入）
                    if cv_en_segments and not os.path.exists(final_en_path):
                        final_en_audio = np.concatenate(cv_en_segments)
                        sf.write(final_en_path, final_en_audio, final_sr_en)

            except Exception as e:
                print(f"[Worker {args.gpu_id}] 行号 {line_idx} Speaker {speaker} CustomVoice 生成失败: {e}")
        
        write_item_metadata(item, control, out_sub_dir, line_idx)

    print(f"[Worker {args.gpu_id}] 任务完成！")

if __name__ == "__main__":
    main()