'update'

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