diff --git a/hook_kohya_ss_run.py b/hook_kohya_ss_run.py index 988d861..bb32f70 100644 --- a/hook_kohya_ss_run.py +++ b/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 != "": diff --git a/hook_kohya_ss_utils.py b/hook_kohya_ss_utils.py index ca54f5b..8f37d45 100644 --- a/hook_kohya_ss_utils.py +++ b/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: diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index 77ffe9c..553dd9d 100644 --- a/mz_train_tools_core.py +++ b/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 diff --git a/mz_train_tools_utils.py b/mz_train_tools_utils.py index aed4ead..cba290b 100644 --- a/mz_train_tools_utils.py +++ b/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 = [