diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index b2abceb..b460c03 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -2,6 +2,7 @@ import argparse import hashlib +import io import os import shutil import sys @@ -408,8 +409,7 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot pb = Utils.progress_bar(train_config.get("max_train_steps"), taesd_type) - import traceback - + import traceback import comfy.model_management stop_server = None @@ -440,21 +440,30 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot print(f"stack: {traceback.format_exc()}") return is_running stop_server, port = Utils.Simple_Server(log_callback) - try: - subprocess.run( - [sys.executable, exec_pyfile, "--sys_path", kohya_ss_tool_dir, - "--train_config_json", train_config_str, "--train_func", trainer_func, "--master_port", str(port), "--other_config_json", other_config_str], - check=True, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - ) - stop_server() - is_running = False - except Exception as e: - stop_server() - is_running = False - raise Exception(f"训练失败!!! 具体报错信息请查看控制台...") + + process = subprocess.Popen( + [sys.executable, exec_pyfile, "--sys_path", kohya_ss_tool_dir, + "--train_config_json", train_config_str, "--train_func", trainer_func, "--master_port", str(port), "--other_config_json", other_config_str], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + ) + stdout, stderr = process.communicate() + stdout_str = "" + stderr_str = "" + if stdout is not None: + stdout_str = stdout.decode("utf-8") + if stderr is not None: + stderr_str = stderr.decode("utf-8") + retcode = process.poll() + if retcode: + raise Exception(f"""请查看以上具体信息, +执行失败: +========================stdout======================== +{stdout_str} +========================stderr======================== +{stderr_str} +====================================================== +""") def MZ_KohyaSSTrain_call(args={}): args = args.copy()