diff --git a/__init__.py b/__init__.py index 3340171..628484b 100644 --- a/__init__.py +++ b/__init__.py @@ -24,7 +24,7 @@ class MZ_KohyaSSInitWorkspace: return { "required": { "workspace_name": ("STRING", {"default": ""}), - "branch": ("STRING", {"default": "main"}), + "branch": ("STRING", {"default": "71e2c91"}), "seed": ("INT", {"default": 0}), }, } @@ -213,7 +213,7 @@ class MZ_KohyaSSTrain: # 使用walk查询所有的workspace中的所有lora模型,lora存放在每个workspace的output目录下 workspaces_loras = [] - for root, dirs, files in os.walk(workspaces_dir): + for root, dirs, files in os.walk(workspaces_dir): if root.endswith("output"): for file in files: if file.endswith(".safetensors"): @@ -240,14 +240,16 @@ class MZ_KohyaSSTrain: return { "required": { "train_config": ("MZ_TT_SS_TrainConfig",), + "base_lora": (loras, {"default": "latest"}), + "sample_generate": (["enable", "disable"], {"default": "enable"}), + "sample_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}), }, "optional": { - "use_lora": (loras, {"default": "latest"}), }, } - RETURN_TYPES = (f"STRING",) - RETURN_NAMES = ("train_output_dir",) + RETURN_TYPES = () + RETURN_NAMES = () OUTPUT_NODE = True FUNCTION = "start" CATEGORY = CATEGORY_NAME diff --git a/hook_kohya_ss_run.py b/hook_kohya_ss_run.py index bb32f70..c8f3339 100644 --- a/hook_kohya_ss_run.py +++ b/hook_kohya_ss_run.py @@ -1,6 +1,7 @@ import os +import random -import torch +import torch import logging import sys import json @@ -33,46 +34,129 @@ import safetensors.torch import hook_kohya_ss_utils +other_config = {} +original_save_model = None + + +train_config_json = "{}" + +sample_images_pipe_class = None + + def sample_images(self, *args, **kwargs): # accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet + accelerator = args[0] + cmd_args = args[1] + epoch = args[2] global_step = args[3] - last_noise_pred = hook_kohya_ss_utils.running_info.get( - "last_noise_pred", None) - noise_pred_latent_path = None - if last_noise_pred is not None: - noise_pred_latent_path = os.path.join( - tempfile.gettempdir(), f"sample_images_{global_step % 5}.latent") + device = args[4] + vae = args[5] + tokenizer = args[6] + text_encoder = args[7] + unet = args[8] - output = {} - output["latent_tensor"] = last_noise_pred - output["latent_format_version_0"] = torch.tensor([]) - safetensors.torch.save_file(output, noise_pred_latent_path) + # print(f"sample_images: args = {args}") + # print(f"sample_images: kwargs = {kwargs}") + + # last_noise_pred = hook_kohya_ss_utils.running_info.get( + # "last_noise_pred", None) + + # noise_pred_latent_path = None + # if last_noise_pred is not None: + # noise_pred_latent_path = os.path.join( + # tempfile.gettempdir(), f"sample_images_{global_step % 5}.latent") + + # output = {} + # output["latent_tensor"] = last_noise_pred + # output["latent_format_version_0"] = torch.tensor([]) + # safetensors.torch.save_file(output, noise_pred_latent_path) + + if epoch is not None and cmd_args.save_every_n_epochs is not None and epoch % cmd_args.save_every_n_epochs == 0: + + prompt_dict_list = other_config.get("prompt_dict_list", []) + if len(prompt_dict_list) == 0: + seed = other_config.get("seed", 0) + + prompt_dict_list.append({ + "controlnet_image": None, + "prompt": other_config.get("sample_prompt", ""), + "seed": seed, + "negative_prompt": "", + "enum": 0, + "sample_sampler": "euler_a", + "sample_steps": 20, + "scale": 5.0, + }) + else: + for i, prompt_dict in enumerate(prompt_dict_list): + if prompt_dict.get("controlnet_image", None) is None: + prompt_dict["controlnet_image"] = None + if prompt_dict.get("seed", None) is None: + prompt_dict["seed"] = 0 + if prompt_dict.get("negative_prompt", None) is None: + prompt_dict["negative_prompt"] = "" + if prompt_dict.get("enum", None) is None: + prompt_dict["enum"] = i + hook_kohya_ss_utils.generate_image( + cmd_args=cmd_args, + accelerator=accelerator, + epoch=epoch, + text_encoder=text_encoder, + tokenizer=tokenizer, + unet=unet, + vae=vae, + pipe_class=sample_images_pipe_class, + prompt_dict_list=prompt_dict_list, + ) LOG({ "type": "sample_images", "global_step": global_step, - "latent": noise_pred_latent_path, + # "latent": noise_pred_latent_path, }) def run_lora_sd1_5(): hook_kohya_ss_utils.hook_kohya_ss() + + # 覆盖sample_images生成函数,包括进度条和生成图片功能 import train_network train_network.NetworkTrainer.sample_images = sample_images + + # 配置对应的pipeline + import library.train_util + global sample_images_pipe_class + sample_images_pipe_class = library.train_util.StableDiffusionLongPromptWeightingPipeline + trainer = train_network.NetworkTrainer() train_config = json.loads(train_config_json) train_args = config2args(train_network.setup_parser(), train_config) + LOG({ + "type": "start_train", + }) trainer.train(train_args) def run_lora_sdxl(): hook_kohya_ss_utils.hook_kohya_ss() + + # 覆盖sample_images生成函数,包括进度条和生成图片功能 import sdxl_train_network sdxl_train_network.SdxlNetworkTrainer.sample_images = sample_images + + # 配置对应的pipeline + import library.sdxl_train_util + global sample_images_pipe_class + sample_images_pipe_class = library.sdxl_train_util.SdxlStableDiffusionLongPromptWeightingPipeline + trainer = sdxl_train_network.SdxlNetworkTrainer() train_config = json.loads(train_config_json) train_args = config2args(sdxl_train_network.setup_parser(), train_config) + + LOG({ + "type": "start_train", + }) trainer.train(train_args) @@ -99,8 +183,12 @@ if __name__ == "__main__": parser.add_argument("--train_config_json", type=str, default="") parser.add_argument("--train_func", type=str, default="") parser.add_argument("--master_port", type=int, default=0) + parser.add_argument("--other_config_json", type=str, default="{}") args = parser.parse_args() + other_config_json = args.other_config_json + other_config = json.loads(other_config_json) + master_port = args.master_port print(f"master_port = {master_port}") diff --git a/hook_kohya_ss_utils.py b/hook_kohya_ss_utils.py index 8f37d45..e1f6031 100644 --- a/hook_kohya_ss_utils.py +++ b/hook_kohya_ss_utils.py @@ -348,3 +348,72 @@ def _load_target_model(args: argparse.Namespace, weight_dtype, device="cpu", une print("additional VAE loaded") return text_encoder, vae, unet, load_stable_diffusion_format + + +def generate_image(pipe_class, cmd_args, accelerator, vae, tokenizer, text_encoder, unet, epoch, prompt_dict_list): + import library.train_util + + # for multi gpu distributed inference. this is a singleton, so it's safe to use it here + distributed_state = library.train_util.PartialState() + org_vae_device = vae.device # CPUにいるはず + vae.to(distributed_state.device) + unet = accelerator.unwrap_model(unet) + + if isinstance(text_encoder, (list, tuple)): + text_encoder = [accelerator.unwrap_model(te) for te in text_encoder] + else: + text_encoder = accelerator.unwrap_model(text_encoder) + + default_scheduler = library.train_util.get_my_scheduler( + sample_sampler="k_euler", + v_parameterization=cmd_args.v_parameterization, + ) + + pipeline = pipe_class( + text_encoder=text_encoder, + vae=vae, + unet=unet, + tokenizer=tokenizer, + scheduler=default_scheduler, + safety_checker=None, + feature_extractor=None, + requires_safety_checker=False, + clip_skip=cmd_args.clip_skip, + ) + pipeline.to(distributed_state.device) + + workspaces_dir = os.path.dirname(cmd_args.dataset_config) + sample_images_path = os.path.join( + workspaces_dir, "sample_images") + + os.makedirs(sample_images_path, exist_ok=True) + + lora_output_name = cmd_args.output_name + + # 画像生成 + save_dir = sample_images_path + prompt_replacement = None + steps = 0 + controlnet = None + + # save random state to restore later + rng_state = torch.get_rng_state() + cuda_rng_state = None + try: + cuda_rng_state = torch.cuda.get_rng_state() if torch.cuda.is_available() else None + except Exception: + pass + + with torch.no_grad(): + for prompt_dict in prompt_dict_list: + library.train_util.sample_image_inference( + accelerator, cmd_args, pipeline, save_dir, prompt_dict, epoch, steps, prompt_replacement, controlnet=controlnet + ) + + del pipeline + library.train_util.clean_memory_on_device(accelerator.device) + + torch.set_rng_state(rng_state) + if cuda_rng_state is not None: + torch.cuda.set_rng_state(cuda_rng_state) + vae.to(org_vae_device) diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index 553dd9d..775b74e 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -260,17 +260,50 @@ def generate_toml_config(output_path, enable_bucket=True, resolution=512, batch_ toml.dump(config, f) -def latent2image(previews, latent_file): - samples = nodes.LoadLatent().load(latent_file)[0] - with torch.no_grad(): +# def latent2image(previews, latent_file): +# samples = nodes.LoadLatent().load(latent_file)[0] +# with torch.no_grad(): - samples = samples["samples"].to("cuda") - preview_bytes = previews.decode_latent_to_preview_image( - "", samples) - return preview_bytes[1] +# samples = samples["samples"].to("cuda") +# preview_bytes = previews.decode_latent_to_preview_image( +# "", samples) +# return preview_bytes[1] + +from PIL import Image -def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func): +def get_sample_images(train_config): + output_name = train_config.get("output_name") + sample_images_dir = os.path.join( + os.path.dirname(train_config.get("dataset_config")), "sample_images" + ) + pil_images = [] + pre_render_texts_x = [] + if os.path.exists(sample_images_dir): + image_files = os.listdir(sample_images_dir) + image_files = list( + filter(lambda x: x.endswith(".png"), image_files)) + # 筛选 output_name 前缀 + image_files = list( + filter(lambda x: x.startswith(output_name), image_files)) + + image_files = sorted(image_files, key=lambda x: x) + + for image_file in image_files: + pil_image = Image.open(os.path.join(sample_images_dir, image_file)) + pil_images.append([pil_image]) + pre_render_texts_x.append(image_file) + result = Utils.xy_image( + pre_render_images=pil_images, + pre_render_texts_x=pre_render_texts_x, + pre_render_texts_y=[""], + ) + return result + + +def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, other_config={}): + + other_config_str = json.dumps(other_config) exec_pyfile = os.path.join(os.path.dirname( __file__), "hook_kohya_ss_run.py",) @@ -303,12 +336,12 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func): resp = log if resp.get("type") == "sample_images": global_step = resp.get("global_step") - latent_path = resp.get("latent") - preview_bytes = None - if latent_path is not None: - preview_bytes = latent2image(previewer, latent_path) + xy_img = get_sample_images(train_config) + + max_side = max(xy_img.width, xy_img.height) + pb.update( - int(global_step), int(max_train_steps), preview_bytes) + int(global_step), int(max_train_steps), ("JPEG", xy_img, max_side)) else: print(f"LOG: {log}") except Exception as e: @@ -319,7 +352,7 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func): 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)], + "--train_config_json", train_config_str, "--train_func", trainer_func, "--master_port", str(port), "--other_config_json", other_config_str], check=True, ) stop_server() @@ -364,10 +397,10 @@ def MZ_KohyaSSTrain_call(args={}): sys.path.append(kohya_ss_tool_dir) check_install() - use_lora = args.get("use_lora", "empty") - if use_lora == "empty": + base_lora = args.get("base_lora", "empty") + if base_lora == "empty": pass - elif use_lora == "latest": + elif base_lora == "latest": workspace_lora_dir = os.path.join(workspace_dir, "output") if os.path.exists(workspace_lora_dir): workspace_lora_files = os.listdir(workspace_lora_dir) @@ -379,16 +412,16 @@ def MZ_KohyaSSTrain_call(args={}): 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( + base_lora = os.path.join( workspace_lora_dir, workspace_lora_files[0]) else: - use_lora = "empty" + base_lora = "empty" else: pass train_config = config.get("train_config") - if use_lora != "empty" and os.path.exists(use_lora): - train_config["network_weights"] = use_lora + if base_lora != "empty" and os.path.exists(base_lora): + train_config["network_weights"] = base_lora train_config["dim_from_weights"] = True if "network_dim" in train_config: @@ -406,13 +439,21 @@ def MZ_KohyaSSTrain_call(args={}): # train_args = config2args(train_network.setup_parser(), train_config) # trainer.train(train_args) # return + sample_generate = args.get("sample_generate", "enable") + sample_prompt = args.get("sample_prompt", "") + if sample_generate == "enable": + other_config = { + "sample_prompt": sample_prompt, + } + else: + other_config = {} if train_type == "lora_sd1_5": run_hook_kohya_ss_run_file( - kohya_ss_tool_dir, train_config, "run_lora_sd1_5") + kohya_ss_tool_dir, train_config, "run_lora_sd1_5", other_config) elif train_type == "lora_sdxl": run_hook_kohya_ss_run_file( - kohya_ss_tool_dir, train_config, "run_lora_sdxl") + kohya_ss_tool_dir, train_config, "run_lora_sdxl", other_config) else: raise Exception( f"暂时不支持的训练类型: {train_type}") diff --git a/mz_train_tools_utils.py b/mz_train_tools_utils.py index cba290b..b3f1098 100644 --- a/mz_train_tools_utils.py +++ b/mz_train_tools_utils.py @@ -713,9 +713,10 @@ class Utils: pil_img_info = ("JPEG", pil_img, 512) if pil_img is None: pil_img_info = None - if type(pil_img) == Tuple: + if type(pil_img) == Tuple or type(pil_img) == list or type(pil_img) == tuple: pil_img_info = pil_img try: + print("pil_img_info:", type(pil_img), pil_img_info) self.pbar.update_absolute( step, total_steps, pil_img_info) except Exception as e: @@ -770,9 +771,9 @@ class Utils: "===========================httpd.serve_forever() end=======================================") threading.Thread(target=serve_forever).start() - def check_server_stop(): + def check_server_stop(): while is_running["value"]: - print("is_running : ", is_running) + # print("is_running : ", is_running) time.sleep(1) httpd.shutdown() httpd.server_close() @@ -781,42 +782,132 @@ class Utils: return stop_server, port + def xy_image(pre_render_images, pre_render_texts_x, pre_render_texts_y): + # 去掉pre_render_texts_x中所有相同字符串前缀 -MODEL_ZOO = [ - { - "model": "ip-adapter-plus_sd15", - "find_path": [ - "ipadapter", - ], - "file_path": "ipadapter/ip-adapter-plus_sd15.safetensors", - "url": "https://www.modelscope.cn/api/v1/models/dannysun85/IP-Adapter/repo?Revision=master&FilePath=models%2Fip-adapter-plus_sd15.safetensors", - "SHA256": "a1c250be40455cc61a43da1201ec3f1edaea71214865fb47f57927e06cbe4996", - }, - { - "model": "ip-adapter-plus_sdxl_vit-h", - "find_path": [ - "ipadapter", - ], - "file_path": "ipadapter/ip-adapter-plus_sdxl_vit-h.safetensors", - "url": "https://www.modelscope.cn/api/v1/models/dannysun85/IP-Adapter/repo?Revision=master&FilePath=sdxl_models%2Fip-adapter-plus_sdxl_vit-h.safetensors", - "SHA256": "a1c250be40455cc61a43da1201ec3f1edaea71214865fb47f57927e06cbe4996", - }, - { - "model": "CLIP-ViT-H-14-laion2B-s32B-b79K", - "find_path": [ - "clip_vision", - ], - "file_path": "clip_vision/CLIP-ViT-H-14-laion2B-s32B-b79K.safetensors", - "url": "https://www.modelscope.cn/api/v1/models/dannysun85/IP-Adapter/repo?Revision=master&FilePath=models%2Fimage_encoder%2Fmodel.safetensors", - "SHA256": "6ca9667da1ca9e0b0f75e46bb030f7e011f44f86cbfb8d5a36590fcd7507b030", - }, - { - "model": "iclight_sd15_fc", - "find_path": [ - "unet", - ], - "file_path": "unet/iclight_sd15_fc.safetensors", - "url": "https://www.modelscope.cn/api/v1/models/AI-ModelScope/ic-light/repo?Revision=master&FilePath=iclight_sd15_fc.safetensors", - "SHA256": "a033fbaaa2f3f7859fa6a4477ee63ebbf9c116bf3569d5811856d2807f3468cd", - }, -] + def get_common_prefix(pre_render_texts_x): + common_prefix = "" + for i in range(len(pre_render_texts_x[0])): + c = pre_render_texts_x[0][i] + for j in range(len(pre_render_texts_x)): + if pre_render_texts_x[j][i] != c: + return common_prefix + else: + common_prefix += c + + return common_prefix + + common_prefix = get_common_prefix(pre_render_texts_x) + + if common_prefix != "": + pre_render_texts_x = [x.replace(common_prefix, "") for x in pre_render_texts_x] + + x_enable_num = len(pre_render_images) + y_enable_num = len(pre_render_images[0]) + max_width = 0 + max_height = 0 + for x in range(0, len(pre_render_images)): + for y in range(0, len(pre_render_images[x])): + if pre_render_images[x][y].width < 512: + org_width = pre_render_images[x][y].width + org_height = pre_render_images[x][y].height + pre_render_images[x][y] = pre_render_images[x][y].resize( + (512, int(org_height * 512 / org_width))) + + if pre_render_images[x][y].width > max_width: + max_width = pre_render_images[x][y].width + if pre_render_images[x][y].height > max_height: + max_height = pre_render_images[x][y].height + + image_xy_canvas = Image.new( + "RGB", (max_width * x_enable_num, max_height * y_enable_num)) + for i in range(x_enable_num): + for j in range(y_enable_num): + image_xy_canvas.paste( + pre_render_images[i][j], (max_width * i, max_height * j)) + + # draw axis + from PIL import ImageDraw, ImageFont + import PIL + # 获取PIL版本号 + pil_version = PIL.__version__ + if pil_version >= "10.0.0": + def textsize(self, text, font): + left, top, right, bottom = self.textbbox((0, 0), text, font) + return right - left, bottom - top + ImageDraw.ImageDraw.textsize = textsize + + full_padding = 260 + + full_canvas_width = max_width * x_enable_num + full_padding * 2 + full_canvas_height = max_height * y_enable_num + full_padding * 2 + full_canvas = Image.new( + "RGB", (full_canvas_width, full_canvas_height), "white") + + full_canvas.paste(image_xy_canvas, (full_padding, full_padding)) + + draw = ImageDraw.Draw(full_canvas) + + # top + draw.line((full_padding, full_padding, full_canvas_width - full_padding, + full_padding), fill="black", width=4) + + # left + draw.line((full_padding, full_padding, full_padding, + full_canvas_height - full_padding), fill="black", width=4) + + # bottom + draw.line((full_padding, full_canvas_height - full_padding, full_canvas_width - full_padding, + full_canvas_height - full_padding), fill="black", width=4) + + # right + draw.line((full_canvas_width - full_padding, full_padding, full_canvas_width - full_padding, + full_canvas_height - full_padding), fill="black", width=4) + + font_fullpath = Utils.download_model( + { + "url": "https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=font%2FAlibabaPuHuiTi-2-75-SemiBold.ttf", + "output": "font/AlibabaPuHuiTi-2-75-SemiBold.ttf", + } + ) + if os.path.exists(font_fullpath): + font = ImageFont.truetype(font_fullpath, size=32,) + else: + font = ImageFont.load_default() + for i in range(x_enable_num): + textwidth, textheight = draw.textsize( + pre_render_texts_x[i], font) + offset_x = max_width * i + full_padding + \ + ((max_width - textwidth) // 2) + offset_y = full_padding - textheight - 24 + + draw.text((offset_x, offset_y), + pre_render_texts_x[i], font=font, fill="black") + + for j in range(y_enable_num): + textwidth, textheight = draw.textsize( + pre_render_texts_y[j], font) + + label_text = pre_render_texts_y[j] + + def is_exceed_width(t): + textwidth, _ = draw.textsize( + t, font) + return textwidth > max_width - 24 + + if textwidth > full_padding - 24: + # 超过宽度就换行 + label_text = "" + for c in pre_render_texts_y[j]: + if is_exceed_width(label_text + c): + break + label_text += c + + offset_x = full_padding - textwidth - 24 + offset_y = max_height * j + full_padding + \ + ((max_height - textheight) // 2) + + draw.text((offset_x, offset_y), + label_text, font=font, fill="black") + + return full_canvas