'update'

__init__.py

@@ -221,7 +221,7 @@ class AlwaysEqualProxy(str):
         return False
 
 
-class MZ_KohyaSSTrain:
+class MZ_KohyaSSLoraTrain:
 
     @classmethod
     def INPUT_TYPES(s):
@@ -282,8 +282,61 @@ class MZ_KohyaSSTrain:
         return mz_train_tools_core.MZ_KohyaSSTrain_call(kwargs)
 
 
-NODE_CLASS_MAPPINGS["MZ_KohyaSSTrain"] = MZ_KohyaSSTrain
-NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain"
+NODE_CLASS_MAPPINGS["MZ_KohyaSSTrain"] = MZ_KohyaSSLoraTrain
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_KohyaSSTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(old version)"
+
+NODE_CLASS_MAPPINGS["MZ_KohyaSSLoraTrain"] = MZ_KohyaSSLoraTrain
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_KohyaSSLoraTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(lora)"
+
+class MZ_KohyaSSControlnetTrain:
+
+    @classmethod
+    def INPUT_TYPES(s):
+        models = [
+            "latest",
+            "empty",
+        ]
+
+        comfyui_full_m_path = []
+        comfyui_basemodels = folder_paths.get_filename_list("controlnet")
+        for b_model in comfyui_basemodels:
+            m_path = folder_paths.get_full_path("controlnet", b_model)
+            comfyui_full_m_path.append(m_path)
+
+        # 按创建时间排序
+        comfyui_full_m_path = sorted(
+            comfyui_full_m_path, key=lambda x: os.path.getctime(x), reverse=True)
+
+        models = models + comfyui_full_m_path
+
+        return {
+            "required": {
+                "train_config": ("MZ_TT_SS_TrainConfig",),
+                "base_controlnet": (models, {"default": "latest"}),
+                "sample_generate": (["enable", "disable"], {"default": "enable"}),
+                "sample_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
+            },
+            "optional": {
+                "has_no_effect": (AlwaysEqualProxy("*"),),
+            },
+        }
+
+    RETURN_TYPES = ()
+    RETURN_NAMES = ()
+    OUTPUT_NODE = True
+    FUNCTION = "start"
+    CATEGORY = CATEGORY_NAME
+
+    def start(self, **kwargs):
+        importlib.reload(mz_train_tools_core)
+        return mz_train_tools_core.MZ_KohyaSSTrain_call(kwargs)
+
+
+NODE_CLASS_MAPPINGS["MZ_KohyaSSControlnetTrain"] = MZ_KohyaSSControlnetTrain
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_KohyaSSControlnetTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(controlnet)"
 
 
 class MZ_LoadImagesFromDirectoryPath:

hook_kohya_ss_run.py

@@ -43,6 +43,10 @@ train_config_json = "{}"
 sample_images_pipe_class = None
 
 
+def utils_sample_images(*args, **kwargs):
+    return sample_images(None, *args, **kwargs)
+
+
 def sample_images(self, *args, **kwargs):
     #  accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet
     accelerator = args[0]
@@ -58,27 +62,15 @@ def sample_images(self, *args, **kwargs):
     # 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)
+    controlnet = kwargs.get("controlnet", None)
 
     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_dict = {
+                "controlnet_image": other_config.get("controlnet_image", None),
                 "prompt": other_config.get("sample_prompt", ""),
                 "seed": seed,
                 "negative_prompt": "",
@@ -86,7 +78,9 @@ def sample_images(self, *args, **kwargs):
                 "sample_sampler": "euler_a",
                 "sample_steps": 20,
                 "scale": 5.0,
-            })
+            }
+
+            prompt_dict_list.append(prompt_dict)
         else:
             for i, prompt_dict in enumerate(prompt_dict_list):
                 if prompt_dict.get("controlnet_image", None) is None:
@@ -98,6 +92,7 @@ def sample_images(self, *args, **kwargs):
                 if prompt_dict.get("enum", None) is None:
                     prompt_dict["enum"] = i
         hook_kohya_ss_utils.generate_image(
+            pipe_class=sample_images_pipe_class,
             cmd_args=cmd_args,
             accelerator=accelerator,
             epoch=epoch,
@@ -105,8 +100,8 @@ def sample_images(self, *args, **kwargs):
             tokenizer=tokenizer,
             unet=unet,
             vae=vae,
-            pipe_class=sample_images_pipe_class,
             prompt_dict_list=prompt_dict_list,
+            controlnet=controlnet,
         )
 
     LOG({
@@ -169,6 +164,7 @@ class SimpleNamespaceCNWarrper(SimpleNamespace):
         super().__init__(*args, **kwargs)
         self.__dict__.update(kwargs)  # or self.__dict__ = kwargs
         self.__dict__["mid_block_type"] = "UNetMidBlock2DCrossAttn"
+        self.__dict__["_diffusers_version"] = "0.6.0"
         self.__iter__ = lambda: iter(kwargs.keys())
     # is not iterable
 
@@ -185,10 +181,13 @@ def run_controlnet_sd1_5():
     types.SimpleNamespace = SimpleNamespaceCNWarrper
     hook_kohya_ss_utils.hook_kohya_ss()
     # 覆盖sample_images生成函数,包括进度条和生成图片功能
+
     import train_controlnet
 
     # 配置对应的pipeline
     import library.train_util
+    library.train_util.sample_images = utils_sample_images
+
     global sample_images_pipe_class
     sample_images_pipe_class = library.train_util.StableDiffusionLongPromptWeightingPipeline
 

hook_kohya_ss_utils.py

@@ -216,6 +216,10 @@ def hook_kohya_ss():
     import library.utils
     import library.train_util
     import library.sdxl_train_util
+    import library.model_util
+
+    library.model_util.convert_controlnet_state_dict_to_diffusers = convert_controlnet_state_dict_to_diffusers
+
     library.utils.setup_logging = setup_logging
 
     library.train_util.load_tokenizer = load_tokenizers
@@ -350,7 +354,7 @@ def _load_target_model(args: argparse.Namespace, weight_dtype, device="cpu", une
     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):
+def generate_image(pipe_class, cmd_args, accelerator, vae, tokenizer, text_encoder, unet, epoch, prompt_dict_list, controlnet=None):
     import library.train_util
 
     # for multi gpu distributed inference. this is a singleton, so it's safe to use it here
@@ -369,6 +373,10 @@ def generate_image(pipe_class, cmd_args, accelerator, vae, tokenizer, text_encod
         v_parameterization=cmd_args.v_parameterization,
     )
 
+    # print("unet.config._diffusers_version:", unet.config._diffusers_version)
+    # print("unet.config:", unet.config)
+    # input("press enter to continue...")
+
     pipeline = pipe_class(
         text_encoder=text_encoder,
         vae=vae,
@@ -394,7 +402,6 @@ def generate_image(pipe_class, cmd_args, accelerator, vae, tokenizer, text_encod
     save_dir = sample_images_path
     prompt_replacement = None
     steps = 0
-    controlnet = None
 
     # save random state to restore later
     rng_state = torch.get_rng_state()
@@ -406,6 +413,9 @@ def generate_image(pipe_class, cmd_args, accelerator, vae, tokenizer, text_encod
 
     with torch.no_grad():
         for prompt_dict in prompt_dict_list:
+            print("prompt_dict:", json.dumps(
+                prompt_dict, indent=4, ensure_ascii=False))
+
             library.train_util.sample_image_inference(
                 accelerator, cmd_args, pipeline, save_dir, prompt_dict, epoch, steps, prompt_replacement, controlnet=controlnet
             )
@@ -417,3 +427,42 @@ def generate_image(pipe_class, cmd_args, accelerator, vae, tokenizer, text_encod
     if cuda_rng_state is not None:
         torch.cuda.set_rng_state(cuda_rng_state)
     vae.to(org_vae_device)
+
+
+def convert_controlnet_state_dict_to_diffusers(controlnet_state_dict):
+    import library.model_util as model_util
+    unet_conversion_map, unet_conversion_map_resnet, unet_conversion_map_layer = model_util.controlnet_conversion_map()
+
+    mapping = {k: k for k in controlnet_state_dict.keys()}
+
+    # print("controlnet_state_dict:", controlnet_state_dict.keys())
+    keys_first = list(controlnet_state_dict.keys())[0]
+
+    if keys_first.startswith("control_model."):
+        new_mapping = {}
+        for k, v in mapping.items():
+            new_mapping[k.replace("control_model.", "")] = v.replace(
+                "control_model.", "")
+        mapping = new_mapping
+
+    for sd_name, diffusers_name in unet_conversion_map:
+        mapping[sd_name] = diffusers_name
+    for k, v in mapping.items():
+        for sd_part, diffusers_part in unet_conversion_map_layer:
+            v = v.replace(sd_part, diffusers_part)
+        mapping[k] = v
+    for k, v in mapping.items():
+        if "resnets" in v:
+            for sd_part, diffusers_part in unet_conversion_map_resnet:
+                v = v.replace(sd_part, diffusers_part)
+            mapping[k] = v
+
+    if keys_first.startswith("control_model."):
+        new_state_dict = {
+            v: controlnet_state_dict[
+                f"control_model.{k}"] for k, v in mapping.items()}
+    else:
+        new_state_dict = {
+            v: controlnet_state_dict[k] for k, v in mapping.items()}
+
+    return new_state_dict

mz_train_tools_core.py

@@ -511,14 +511,38 @@ def MZ_KohyaSSTrain_call(args={}):
         if "network_dropout" in train_config:
             del train_config["network_dropout"]
 
+    
+
+
+    base_controlnet = args.get("base_controlnet", "empty")
+    if base_controlnet == "empty":
+        pass
+    elif base_controlnet == "latest":
+        workspace_controlnet_dir = os.path.join(workspace_dir, "output")
+        if os.path.exists(workspace_controlnet_dir):
+            workspace_controlnet_files = os.listdir(workspace_controlnet_dir)
+            workspace_controlnet_files = list(
+                filter(lambda x: x.endswith(".safetensors"), workspace_controlnet_files))
+            workspace_controlnet_files = list(
+                map(lambda x: os.path.join(workspace_controlnet_dir, x), workspace_controlnet_files))
+            # 排序
+            workspace_controlnet_files = sorted(
+                workspace_controlnet_files, key=lambda x: os.path.getctime(x), reverse=True)
+            if len(workspace_controlnet_files) > 0:
+                base_controlnet = os.path.join(
+                    workspace_controlnet_dir, workspace_controlnet_files[0])
+        else:
+            base_controlnet = "empty"
+    else:
+        pass
+
+    if base_controlnet != "empty" and os.path.exists(base_controlnet):
+        train_config["controlnet_model_name_or_path"] = base_controlnet 
+
+
     train_type = config.get("metadata").get("train_type")
-    # 在ComfyUI中无法运行, 梯度有问题, 原因不清楚
-    # import train_network
-    # train_network.logger = logging.getLogger()
-    # trainer = train_network.NetworkTrainer()
-    # 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":
@@ -535,6 +559,15 @@ def MZ_KohyaSSTrain_call(args={}):
         run_hook_kohya_ss_run_file(
             kohya_ss_tool_dir, train_config, "run_lora_sdxl", other_config)
     elif train_type == "controlnet_sd1_5":
+
+        conditioning_images_dir = os.path.join(
+            workspace_dir, "conditioning_images")
+        conditioning_images_onec = ""
+        if os.path.exists(conditioning_images_dir):
+            conditioning_images_onec = os.listdir(conditioning_images_dir)[0]
+            other_config["controlnet_image"] = os.path.join(
+                conditioning_images_dir, conditioning_images_onec)
+        
         run_hook_kohya_ss_run_file(
             kohya_ss_tool_dir, train_config, "run_controlnet_sd1_5", other_config)
     else: