'update'

hook_kohya_ss_run.py

@@ -1,8 +1,6 @@
 import os
 
-import torch
-os.environ["HTTP_PROXY"] = "http://127.0.0.1:0"
-os.environ["HTTPS_PROXY"] = "http://127.0.0.1:0"
+import torch 
 import logging
 import sys
 import json
@@ -50,11 +48,11 @@ def sample_images(self, *args, **kwargs):
         output["latent_format_version_0"] = torch.tensor([])
         safetensors.torch.save_file(output, noise_pred_latent_path)
 
-    LOG(json.dumps({
+    LOG({
         "type": "sample_images",
         "global_step": global_step,
         "latent": noise_pred_latent_path,
-    }))
+    })
 
 
 def run_lora_sd1_5():
@@ -84,23 +82,15 @@ func_map = {
 }
 
 
-import socket
-sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-udp_port = 0
-
-
-def is_connected():
-    try:
-        sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
-    except socket.error:
-        return False
-    return True
+import requests
 
 
 def LOG(log):
-    sock.sendto(log.encode(), ('127.0.0.1', udp_port))
-    if not is_connected():
-        raise Exception("Connection closed")
+    # 发送http
+    resp = requests.request("post", f"http://127.0.0.1:{master_port}/log", data=json.dumps(log), headers={
+                            "Content-Type": "application/json"})
+    if resp.status_code != 200:
+        raise Exception(f"LOG failed: {resp.text}")
 
 
 if __name__ == "__main__":
@@ -108,12 +98,12 @@ if __name__ == "__main__":
     parser.add_argument("--sys_path", type=str, default="")
     parser.add_argument("--train_config_json", type=str, default="")
     parser.add_argument("--train_func", type=str, default="")
-    parser.add_argument("--udp_port", type=int, default=0)
+    parser.add_argument("--master_port", type=int, default=0)
     args = parser.parse_args()
 
-    udp_port = args.udp_port
+    master_port = args.master_port
 
-    print(f"udp_port = {udp_port}")
+    print(f"master_port = {master_port}")
 
     sys_path = args.sys_path
     if sys_path != "":

hook_kohya_ss_utils.py

@@ -69,10 +69,13 @@ def request_wrapper(*args, **kwargs):
 from requests import api
 from requests import Session
 last_request = api.request
+original_session_request = Session.request
 api.request = request_wrapper
 
 
 def Session_request_wrapper(cls, method, url, **kwargs):
+    if url.startswith("http://127.0.0.1"):
+        return original_session_request(cls, method, url, **kwargs)
     # print(f"Session_request_wrapper requesting {url}")
     # print(f"Session_request_wrapper requesting kwargs: {kwargs}")
     if url in source_replacement_table:

mz_train_tools_core.py

@@ -157,6 +157,8 @@ def MZ_KohyaSSUseConfig_call(args={}):
 
         config["train_config"]["max_train_epochs"] = str(
             args.get("max_train_epochs"))
+        if config["train_config"]["max_train_epochs"] == "0":
+            config["train_config"]["max_train_epochs"] = False
 
         config["train_config"]["save_every_n_epochs"] = str(
             args.get("save_every_n_epochs"))
@@ -286,9 +288,19 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func):
     previewer = pb.get_previewer()
     import traceback
 
+    import comfy.model_management
+
+    stop_server = None
+
     def log_callback(log):
         try:
-            resp = json.loads(log)
+            comfy.model_management.throw_exception_if_processing_interrupted()
+        except Exception as e:
+            stop_server()
+            return is_running
+
+        try:
+            resp = log
             if resp.get("type") == "sample_images":
                 global_step = resp.get("global_step")
                 latent_path = resp.get("latent")
@@ -303,19 +315,19 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func):
             print(f"LOG: {log} e: {e} ")
             print(f"stack: {traceback.format_exc()}")
         return is_running
-    server, port = Utils.UDP_Server(log_callback)
+    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, "--udp_port", str(port)],
+                "--train_config_json", train_config_str, "--train_func", trainer_func, "--master_port", str(port)],
             check=True,
         )
+        stop_server()
         is_running = False
-        server.close()
     except Exception as e:
+        stop_server()
         is_running = False
-        server.close()
-        raise Exception(f"训练失败!!! 具体报错信息请查看控制台 {e}")
+        raise Exception(f"训练失败!!! 具体报错信息请查看控制台...")
 
 
 def MZ_KohyaSSTrain_call(args={}):
@@ -357,15 +369,20 @@ def MZ_KohyaSSTrain_call(args={}):
         pass
     elif use_lora == "latest":
         workspace_lora_dir = os.path.join(workspace_dir, "output")
-        workspace_lora_files = os.listdir(workspace_lora_dir)
-        workspace_lora_files = list(
-            filter(lambda x: x.endswith(".safetensors"), workspace_lora_files))
-        # 排序
-        workspace_lora_files = sorted(
-            workspace_lora_files, key=lambda x: os.path.getctime(x), reverse=True)
-        if len(workspace_lora_files) > 0:
-            use_lora = os.path.join(
-                workspace_lora_dir, workspace_lora_files[0])
+        if os.path.exists(workspace_lora_dir):
+            workspace_lora_files = os.listdir(workspace_lora_dir)
+            workspace_lora_files = list(
+                filter(lambda x: x.endswith(".safetensors"), workspace_lora_files))
+            workspace_lora_files = list(
+                map(lambda x: os.path.join(workspace_lora_dir, x), workspace_lora_files))
+            # 排序
+            workspace_lora_files = sorted(
+                workspace_lora_files, key=lambda x: os.path.getctime(x), reverse=True)
+            if len(workspace_lora_files) > 0:
+                use_lora = os.path.join(
+                    workspace_lora_dir, workspace_lora_files[0])
+        else:
+            use_lora = "empty"
     else:
         pass
 

mz_train_tools_utils.py

@@ -1,10 +1,12 @@
 
+import json
 import os
 import shutil
 import subprocess
 import sys
 import threading
 import time
+import traceback
 from typing import Tuple
 import warnings
 import numpy as np
@@ -688,7 +690,7 @@ class Utils:
         type_str = str(type(model.model.model_config).__name__)
         return "SDXL" in type_str
 
-    def progress_bar(steps,taesd_type="sd1_5"):
+    def progress_bar(steps, taesd_type="sd1_5"):
         class pb:
             if taesd_type == "sd1_5":
                 taesd_decoder_name = "taesd_decoder"
@@ -696,6 +698,7 @@ class Utils:
                 taesd_decoder_name = "taesdxl_decoder"
             latent_rgb_factors = None
             latent_channels = 4
+
             def __init__(self, steps):
                 self.steps = steps
                 self.pbar = comfy.utils.ProgressBar(steps)
@@ -707,7 +710,6 @@ class Utils:
                 return previewer
 
             def update(self, step, total_steps, pil_img=None):
-                print(f"step: {step}/{total_steps}")
                 pil_img_info = ("JPEG", pil_img, 512)
                 if pil_img is None:
                     pil_img_info = None
@@ -721,28 +723,63 @@ class Utils:
 
         return pb(steps)
 
-    def UDP_Server(reader):
+    def get_free_port():
         import socket
-        import threading
-        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-        s.bind(("127.0.0.1", 0))
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.bind(('localhost', 0))
         port = s.getsockname()[1]
+        s.close()
+        return port
 
-        def recv():
-            while True:
-                data, addr = s.recvfrom(1024)
-                if s.fileno() == -1:
-                    break
-                try:
-                    data = str(data, encoding="utf-8")
-                except:
-                    pass
+    def Simple_Server(reader):
+        import threading
+
+        from http.server import HTTPServer, BaseHTTPRequestHandler
+
+        port = Utils.get_free_port()
+        is_running = {"value": True}
+
+        def stop_server():
+            is_running.update({"value": False})
+
+        httpd = None
+
+        class SimpleHTTPRequestHandler(BaseHTTPRequestHandler):
+            def do_POST(self):
+                self.send_response(200)
+                self.end_headers()
+                content_length = int(self.headers['Content-Length'])
+                post_data = self.rfile.read(content_length)
+                data = json.loads(post_data)
                 if reader(data) == False:
-                    break
+                    stop_server()
 
-        threading.Thread(target=recv).start()
+            def log_message(self, format, *args):
+                pass
 
-        return s, port
+        httpd = HTTPServer(('localhost', port), SimpleHTTPRequestHandler)
+
+        def serve_forever():
+            try:
+                print(
+                    "===========================httpd.serve_forever() start=======================================")
+                httpd.serve_forever()
+            except Exception as e:
+                pass
+            print(
+                "===========================httpd.serve_forever() end=======================================")
+        threading.Thread(target=serve_forever).start()
+
+        def check_server_stop(): 
+            while is_running["value"]:
+                print("is_running : ", is_running)
+                time.sleep(1)
+            httpd.shutdown()
+            httpd.server_close()
+
+        threading.Thread(target=check_server_stop).start()
+
+        return stop_server, port
 
 
 MODEL_ZOO = [