diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index 814d057..0612e56 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -2,7 +2,6 @@ import argparse import hashlib -import io import os import shutil import sys @@ -38,9 +37,7 @@ def MZ_KohyaSSCloneRepo_call(args={}): try: if not os.path.exists(kohya_ss_lora_dir) or not os.path.exists(os.path.join(kohya_ss_lora_dir, ".git")): subprocess.run( - ["git", "clone", "--depth", "1", git_url, kohya_ss_lora_dir], check=True, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT,) + ["git", "clone", "--depth", "1", git_url, kohya_ss_lora_dir], check=True) # 切换远程分支 git remote set-branches origin 'main' branch = args.get("branch", "main") @@ -409,7 +406,8 @@ 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,32 +438,37 @@ 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, + ) + stop_server() + is_running = False + except subprocess.CalledProcessError as e: + stop_server() + is_running = False + stdout_str = "" + if e.stdout is not None: + stdout_str = e.stdout.decode("utf-8") + + stderr_str = "" + if e.stderr is not None: + stderr_str = e.stderr.decode("utf-8") - - 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") - # 清除制表符 - stdout_str = stdout_str.replace("\t", " ") - if stderr is not None: - stderr_str = stderr.decode("utf-8") - retcode = process.poll() - if retcode: - raise Exception(f"""请查看以上具体信息, -执行失败: -========================stdout======================== + raise Exception(f"""训练失败!!!具体报错信息请查看控制台... +=======================stdout======================= {stdout_str} -========================stderr======================== +=======================stderr======================= {stderr_str} -====================================================== -""") +==================================================== + """) + except Exception as e: + stop_server() + is_running = False + raise Exception(f"训练失败!!! 具体报错信息请查看控制台...") + def MZ_KohyaSSTrain_call(args={}): args = args.copy()