'update'

__init__.py

@@ -88,7 +88,6 @@ class MZ_KohyaSSUseConfig:
         # 去掉json后缀
         train_config_templates = [os.path.splitext(x)[0]
                                   for x in train_config_templates]
-
         return {
             "required": {
                 "workspace_config": ("MZ_TT_SS_WorkspaceConfig",),
@@ -101,7 +100,7 @@ class MZ_KohyaSSUseConfig:
                 "learning_rate": ("STRING", {"default": "1e-5"}),
             },
             "optional": {
-                "advanced_config": ("MZ_TT_SS_AdvConfig",),
+                "save_advanced_config": ("MZ_TT_SS_AdvConfig",),
             }
         }
 
@@ -204,10 +203,47 @@ class MZ_KohyaSSTrain:
 
     @classmethod
     def INPUT_TYPES(s):
+        loras = [
+            "latest",
+            "empty",
+        ]
+
+        workspaces_dir = os.path.join(
+            folder_paths.output_directory, "mz_train_workspaces")
+
+        # 使用walk查询所有的workspace中的所有lora模型,lora存放在每个workspace的output目录下
+        workspaces_loras = []
+        for root, dirs, files in os.walk(workspaces_dir): 
+            if root.endswith("output"):
+                for file in files:
+                    if file.endswith(".safetensors"):
+                        workspaces_loras.append(
+                            os.path.join(root, file)
+                        )
+
+        # 按创建时间排序
+        workspaces_loras = sorted(
+            workspaces_loras, key=lambda x: os.path.getctime(x), reverse=True)
+
+        comfyui_full_loras = []
+        comfyui_loras = folder_paths.get_filename_list("loras")
+        for lora in comfyui_loras:
+            lora_path = folder_paths.get_full_path("loras", lora)
+            comfyui_full_loras.append(lora_path)
+
+        # 按创建时间排序
+        comfyui_full_loras = sorted(
+            comfyui_full_loras, key=lambda x: os.path.getctime(x), reverse=True)
+
+        loras = loras + workspaces_loras + comfyui_full_loras
+
         return {
             "required": {
                 "train_config": ("MZ_TT_SS_TrainConfig",),
             },
+            "optional": {
+                "use_lora": (loras, {"default": "latest"}),
+            },
         }
 
     RETURN_TYPES = (f"STRING",)

hook_kohya_ss_utils.py

@@ -116,6 +116,7 @@ huggingface_hub.file_download._hf_hub_download_to_cache_dir = _hf_hub_download_t
 import diffusers.loaders.single_file
 original_snapshot_download = diffusers.loaders.single_file.snapshot_download
 
+
 def _snapshot_download(repo_id, *args, **kwargs):
     print(f"_snapshot_download: {repo_id}")
     if repo_id == "stabilityai/stable-diffusion-3-medium-diffusers":
@@ -124,6 +125,9 @@ def _snapshot_download(repo_id, *args, **kwargs):
     if repo_id == "runwayml/stable-diffusion-v1-5":
         return os.path.join(
             os.path.dirname(__file__), "configs", "models_config", "stable-diffusion-v1-5",)
+    if repo_id == "stabilityai/stable-diffusion-xl-base-1.0":
+        return os.path.join(
+            os.path.dirname(__file__), "configs", "models_config", "stable-diffusion-xl-base-1.0",)
     # return original_snapshot_download(repo_id, *args, **kwargs)
     raise NotImplementedError("_snapshot_download is not supported")
 
@@ -333,7 +337,6 @@ def _load_target_model(args: argparse.Namespace, weight_dtype, device="cpu", une
         )
         original_unet.load_state_dict(unet.state_dict())
         unet = original_unet
-        unet = unet.train()
         print("U-Net converted to original U-Net")
 
     # VAEを読み込む

mz_train_tools_core.py

@@ -164,7 +164,7 @@ def MZ_KohyaSSUseConfig_call(args={}):
         config["train_config"]["learning_rate"] = str(
             args.get("learning_rate"))
 
-        advanced_config = args.get("advanced_config", {}).copy()
+        advanced_config = args.get("save_advanced_config", {}).copy()
 
         for k in advanced_config:
             if type(advanced_config[k]) == str and advanced_config[k] == "":
@@ -327,7 +327,6 @@ def MZ_KohyaSSTrain_call(args={}):
 
     # pb.update(1, 1, img)
     # raise Exception(f"MZ_KohyaSSTrain_call: {args}")
-
     train_config = args.get("train_config", {})
     workspace_config = train_config.get("workspace_config", {})
     workspace_name = workspace_config.get("workspace_name", None)
@@ -353,6 +352,35 @@ 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":
+        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])
+    else:
+        pass
+
+    train_config = config.get("train_config")
+    if use_lora != "empty" and os.path.exists(use_lora):
+        train_config["network_weights"] = use_lora
+        train_config["dim_from_weights"] = True
+
+        if "network_dim" in train_config:
+            del train_config["network_dim"]
+        if "network_alpha" in train_config:
+            del train_config["network_alpha"]
+        if "network_dropout" in train_config:
+            del train_config["network_dropout"]
+
     train_type = config.get("metadata").get("train_type")
     # 在ComfyUI中无法运行, 梯度有问题, 原因不清楚
     # import train_network
@@ -363,11 +391,11 @@ def MZ_KohyaSSTrain_call(args={}):
     # return
 
     if train_type == "lora_sd1_5":
-        run_hook_kohya_ss_run_file(kohya_ss_tool_dir, config.get(
-            "train_config"), "run_lora_sd1_5")
+        run_hook_kohya_ss_run_file(
+            kohya_ss_tool_dir, train_config, "run_lora_sd1_5")
     elif train_type == "lora_sdxl":
-        run_hook_kohya_ss_run_file(kohya_ss_tool_dir, config.get(
-            "train_config"), "run_lora_sdxl")
+        run_hook_kohya_ss_run_file(
+            kohya_ss_tool_dir, train_config, "run_lora_sdxl")
     else:
         raise Exception(
             f"暂时不支持的训练类型: {train_type}")