diff --git a/__init__.py b/__init__.py index 9c0c8c8..3340171 100644 --- a/__init__.py +++ b/__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",) diff --git a/hook_kohya_ss_utils.py b/hook_kohya_ss_utils.py index 33b8050..ca54f5b 100644 --- a/hook_kohya_ss_utils.py +++ b/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を読み込む diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index 3ce2286..77ffe9c 100644 --- a/mz_train_tools_core.py +++ b/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}")