#!/usr/bin/env python3 """ Pipeline 包装脚本 此脚本作为独立子进程运行，执行 TrainingPipeline 并将进度以 JSON 格式输出到 stdout。 主进程（AsyncTrainingManager）通过解析 stdout 来获取实时进度。 进度消息格式: ##PROGRESS##{"type": "progress", "stage": "...", ...}## Usage: python run_pipeline.py --config /path/to/config.json --task-id task-123 """ import argparse import json import sys import os import traceback from datetime import datetime from typing import Dict, Any # 确保可以导入项目模块（在导入其他模块之前） from pathlib import Path _SCRIPT_DIR = Path(__file__).parent.resolve() _API_SERVER_ROOT = _SCRIPT_DIR.parent.parent _PROJECT_ROOT = _API_SERVER_ROOT.parent sys.path.insert(0, str(_PROJECT_ROOT)) # 导入配置模块 from project_config import settings, PROJECT_ROOT, get_pythonpath # 进度消息前缀和后缀，用于主进程解析 PROGRESS_PREFIX = "##PROGRESS##" PROGRESS_SUFFIX = "##" def emit_progress(progress_info: Dict[str, Any]) -> None: """ 输出进度消息到 stdout Args: progress_info: 进度信息字典 """ # 确保有时间戳 if "timestamp" not in progress_info: progress_info["timestamp"] = datetime.utcnow().isoformat() json_str = json.dumps(progress_info, ensure_ascii=False) print(f"{PROGRESS_PREFIX}{json_str}{PROGRESS_SUFFIX}", flush=True) def emit_log(level: str, message: str, **extra) -> None: """ 输出日志消息 Args: level: 日志级别 (info, warning, error) message: 日志消息 **extra: 额外数据 """ emit_progress({ "type": "log", "level": level, "message": message, **extra }) def load_config(config_path: str) -> Dict[str, Any]: """ 加载配置文件 Args: config_path: 配置文件路径 Returns: 配置字典 """ with open(config_path, 'r', encoding='utf-8') as f: return json.load(f) def build_pipeline(config: Dict[str, Any]): """ 根据配置构建 TrainingPipeline Args: config: 配置字典，包含: - exp_name: 实验名称 - version: 模型版本 - stages: 要执行的阶段列表 - 各阶段的具体配置 Returns: TrainingPipeline 实例 """ from training_pipeline import ( TrainingPipeline, ModelVersion, # 配置类 AudioSliceConfig, ASRConfig, DenoiseConfig, FeatureExtractionConfig, SoVITSTrainConfig, GPTTrainConfig, InferenceConfig, # 阶段类 AudioSliceStage, ASRStage, DenoiseStage, TextFeatureStage, HuBERTFeatureStage, SemanticTokenStage, SoVITSTrainStage, GPTTrainStage, InferenceStage, ) pipeline = TrainingPipeline() exp_name = config["exp_name"] version_str = config.get("version", "v2") version = ModelVersion(version_str) if isinstance(version_str, str) else version_str # 通用配置参数 base_params = { "exp_name": exp_name, "exp_root": config.get("exp_root", "logs"), "gpu_numbers": config.get("gpu_numbers", "0"), "is_half": config.get("is_half", True), } # 阶段配置映射 stage_builders = { "audio_slice": lambda cfg: AudioSliceStage(AudioSliceConfig( **base_params, input_path=cfg.get("input_path", ""), threshold=cfg.get("threshold", -34), min_length=cfg.get("min_length", 4000), min_interval=cfg.get("min_interval", 300), hop_size=cfg.get("hop_size", 10), max_sil_kept=cfg.get("max_sil_kept", 500), max_amp=cfg.get("max_amp", 0.9), alpha=cfg.get("alpha", 0.25), n_parts=cfg.get("n_parts", 4), )), "asr": lambda cfg: ASRStage(ASRConfig( **base_params, model=cfg.get("model", "达摩 ASR (中文)"), model_size=cfg.get("model_size", "large"), language=cfg.get("language", "zh"), precision=cfg.get("precision", "float32"), )), "denoise": lambda cfg: DenoiseStage(DenoiseConfig( **base_params, input_dir=cfg.get("input_dir", ""), output_dir=cfg.get("output_dir", "output/denoise_opt"), )), "text_feature": lambda cfg: TextFeatureStage(FeatureExtractionConfig( **base_params, version=version, bert_pretrained_dir=cfg.get("bert_pretrained_dir", "GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large"), ssl_pretrained_dir=cfg.get("ssl_pretrained_dir", "GPT_SoVITS/pretrained_models/chinese-hubert-base"), pretrained_s2G=cfg.get("pretrained_s2G", "GPT_SoVITS/pretrained_models/gsv-v2final-pretrained/s2G2333k.pth"), )), "hubert_feature": lambda cfg: HuBERTFeatureStage(FeatureExtractionConfig( **base_params, version=version, bert_pretrained_dir=cfg.get("bert_pretrained_dir", "GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large"), ssl_pretrained_dir=cfg.get("ssl_pretrained_dir", "GPT_SoVITS/pretrained_models/chinese-hubert-base"), pretrained_s2G=cfg.get("pretrained_s2G", "GPT_SoVITS/pretrained_models/gsv-v2final-pretrained/s2G2333k.pth"), )), "semantic_token": lambda cfg: SemanticTokenStage(FeatureExtractionConfig( **base_params, version=version, bert_pretrained_dir=cfg.get("bert_pretrained_dir", "GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large"), ssl_pretrained_dir=cfg.get("ssl_pretrained_dir", "GPT_SoVITS/pretrained_models/chinese-hubert-base"), pretrained_s2G=cfg.get("pretrained_s2G", "GPT_SoVITS/pretrained_models/gsv-v2final-pretrained/s2G2333k.pth"), )), "sovits_train": lambda cfg: SoVITSTrainStage(SoVITSTrainConfig( **base_params, version=version, batch_size=cfg.get("batch_size", 4), total_epoch=cfg.get("total_epoch", 8), text_low_lr_rate=cfg.get("text_low_lr_rate", 0.4), save_every_epoch=cfg.get("save_every_epoch", 4), if_save_latest=cfg.get("if_save_latest", True), if_save_every_weights=cfg.get("if_save_every_weights", True), pretrained_s2G=cfg.get("pretrained_s2G", "GPT_SoVITS/pretrained_models/gsv-v2final-pretrained/s2G2333k.pth"), pretrained_s2D=cfg.get("pretrained_s2D", "GPT_SoVITS/pretrained_models/gsv-v2final-pretrained/s2D2333k.pth"), if_grad_ckpt=cfg.get("if_grad_ckpt", False), lora_rank=cfg.get("lora_rank", 32), )), "gpt_train": lambda cfg: GPTTrainStage(GPTTrainConfig( **base_params, version=version, batch_size=cfg.get("batch_size", 4), total_epoch=cfg.get("total_epoch", 15), save_every_epoch=cfg.get("save_every_epoch", 5), if_save_latest=cfg.get("if_save_latest", True), if_save_every_weights=cfg.get("if_save_every_weights", True), if_dpo=cfg.get("if_dpo", False), pretrained_s1=cfg.get("pretrained_s1", "GPT_SoVITS/pretrained_models/gsv-v2final-pretrained/s1bert25hz-5kh-longer-epoch=12-step=369668.ckpt"), )), "inference": lambda cfg: InferenceStage(InferenceConfig( **base_params, version=version, gpt_path=cfg.get("gpt_path", ""), sovits_path=cfg.get("sovits_path", ""), ref_text=cfg.get("ref_text", ""), ref_audio_path=cfg.get("ref_audio_path", ""), target_text=cfg.get("target_text", ""), text_split_method=cfg.get("text_split_method", "cut1"), )), } # 按顺序添加阶段 # stages 可以是: # 1. 字符串列表: ["audio_slice", "asr", ...] # 2. 字典列表: [{"type": "audio_slice", "threshold": -30}, ...] stages = config.get("stages", []) for stage_item in stages: # 判断是字符串还是字典 if isinstance(stage_item, str): stage_type = stage_item stage_config = config # 使用全局配置作为阶段配置 elif isinstance(stage_item, dict): stage_type = stage_item.get("type") # 合并全局配置和阶段特定配置 stage_config = {**config, **stage_item} else: emit_log("warning", f"无效的阶段配置类型: {type(stage_item)}") continue if stage_type in stage_builders: stage = stage_builders[stage_type](stage_config) pipeline.add_stage(stage) emit_log("info", f"已添加阶段: {stage.name}") else: emit_log("warning", f"未知阶段类型: {stage_type}") return pipeline def run_pipeline(config: Dict[str, Any], task_id: str) -> bool: """ 执行 Pipeline Args: config: 配置字典 task_id: 任务ID Returns: 是否成功完成 """ emit_progress({ "type": "progress", "status": "running", "message": "正在初始化训练流水线...", "task_id": task_id, "progress": 0.0, "overall_progress": 0.0, }) try: pipeline = build_pipeline(config) stages = pipeline.get_stages() if not stages: emit_progress({ "type": "progress", "status": "failed", "message": "没有配置任何训练阶段", "task_id": task_id, }) return False emit_log("info", f"训练流水线已初始化，共 {len(stages)} 个阶段") # 执行 Pipeline for progress in pipeline.run(): # 转换进度格式 emit_progress({ "type": "progress", "status": "running", "stage": progress.get("stage"), "stage_index": progress.get("stage_index"), "total_stages": progress.get("total_stages"), "progress": progress.get("progress", 0.0), "overall_progress": progress.get("overall_progress", 0.0), "message": progress.get("message"), "task_id": task_id, "data": progress.get("data", {}), }) # 检查是否失败 if progress.get("status") == "failed": emit_progress({ "type": "progress", "status": "failed", "stage": progress.get("stage"), "message": progress.get("message", "阶段执行失败"), "task_id": task_id, }) return False # 完成 emit_progress({ "type": "progress", "status": "completed", "message": "训练流水线执行完成", "task_id": task_id, "progress": 1.0, "overall_progress": 1.0, }) return True except Exception as e: error_msg = str(e) error_trace = traceback.format_exc() emit_progress({ "type": "progress", "status": "failed", "message": f"执行出错: {error_msg}", "error": error_msg, "traceback": error_trace, "task_id": task_id, }) return False def main(): """主函数""" parser = argparse.ArgumentParser(description="执行 GPT-SoVITS 训练流水线") parser.add_argument("--config", required=True, help="配置文件路径 (JSON)") parser.add_argument("--task-id", required=True, help="任务ID") args = parser.parse_args() emit_log("info", f"启动训练任务: {args.task_id}") emit_log("info", f"配置文件: {args.config}") try: config = load_config(args.config) except Exception as e: emit_progress({ "type": "progress", "status": "failed", "message": f"加载配置文件失败: {e}", "task_id": args.task_id, }) sys.exit(1) success = run_pipeline(config, args.task_id) sys.exit(0 if success else 1) if __name__ == "__main__": main()