'update'

__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

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}")

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)

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}")

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