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()