diff --git a/__init__.py b/__init__.py index c47b053..471ff52 100644 --- a/__init__.py +++ b/__init__.py @@ -51,7 +51,6 @@ NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSInitWorkspace"] = f"{AUTHOR_NAME} - KohyaS class MZ_ImagesCopyWorkspace: - @classmethod def INPUT_TYPES(s): return { @@ -62,12 +61,17 @@ class MZ_ImagesCopyWorkspace: "resolution": ("INT", {"default": 512}), "num_repeats": ("INT", {"default": 1}), "batch_size": ("INT", {"default": 1}), - # "class_name": ("STRING", {"default": "girl", "dynamicPrompts": True}), "force_clear": (["enable", "disable"], {"default": "disable"}), "force_clear_only_images": (["enable", "disable"], {"default": "disable"}), "same_caption_generate": (["enable", "disable"], {"default": "disable"}), "same_caption": ("STRING", {"default": "", "dynamicPrompts": True, "multiline": True}), + "description": ("STRING", {"default": """ +如果训练类型是controlnet,必须传入预处理后的图片(conditioning_images) +""", "multiline": True}), }, + "optional": { + "conditioning_images": ("IMAGE",), + } } RETURN_TYPES = (f"STRING",) @@ -323,3 +327,8 @@ class MZ_LoadImagesFromDirectoryPath: NODE_CLASS_MAPPINGS["MZ_LoadImagesFromDirectoryPath"] = MZ_LoadImagesFromDirectoryPath NODE_DISPLAY_NAME_MAPPINGS[ "MZ_LoadImagesFromDirectoryPath"] = f"{AUTHOR_NAME} - LoadImagesFromDirectoryPath" + + +# 别名 +NODE_CLASS_MAPPINGS["MZ_DatasetConfig"] = MZ_ImagesCopyWorkspace +NODE_DISPLAY_NAME_MAPPINGS["MZ_DatasetConfig"] = f"{AUTHOR_NAME} - DatasetConfig" diff --git a/configs/kohya_ss_lora/controlnet_sd1_5.json b/configs/kohya_ss_lora/controlnet_sd1_5.json new file mode 100644 index 0000000..04ad320 --- /dev/null +++ b/configs/kohya_ss_lora/controlnet_sd1_5.json @@ -0,0 +1,28 @@ +{ + "metadata": { + "train_type": "controlnet_sd1_5" + }, + "train_config": { + "pretrained_model_name_or_path": "", + "max_train_steps": "4500", + "xformers": true, + "sdpa": false, + "fp8_base": false, + "mixed_precision": "fp16", + "cache_latents": true, + "cache_latents_to_disk": true, + "learning_rate": "1e-5", + "lr_scheduler": "cosine_with_restarts", + "optimizer_type": "AdamW", + "save_every_n_epochs": "20", + "shuffle_caption": true, + "lr_warmup_steps": "0", + "save_precision": "fp16", + "lr_scheduler_num_cycles": "1", + "persistent_data_loader_workers": true, + "noise_offset": "0.1", + "output_dir": "", + "output_name": "", + "lowram": false + } +} diff --git a/hook_kohya_ss_run.py b/hook_kohya_ss_run.py index d0488e6..0454550 100644 --- a/hook_kohya_ss_run.py +++ b/hook_kohya_ss_run.py @@ -161,9 +161,51 @@ def run_lora_sdxl(): trainer.train(train_args) +from types import SimpleNamespace + + +class SimpleNamespaceCNWarrper(SimpleNamespace): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.__dict__.update(kwargs) # or self.__dict__ = kwargs + self.__dict__["mid_block_type"] = "UNetMidBlock2DCrossAttn" + self.__iter__ = lambda: iter(kwargs.keys()) + # is not iterable + + def __iter__(self): + return iter(self.__dict__.keys()) + # object has no attribute 'num_attention_heads' + + def __getattr__(self, name): + return self.__dict__.get(name, None) + + +def run_controlnet_sd1_5(): + import types + types.SimpleNamespace = SimpleNamespaceCNWarrper + hook_kohya_ss_utils.hook_kohya_ss() + # 覆盖sample_images生成函数,包括进度条和生成图片功能 + import train_controlnet + + # 配置对应的pipeline + import library.train_util + global sample_images_pipe_class + sample_images_pipe_class = library.train_util.StableDiffusionLongPromptWeightingPipeline + + train_config = json.loads(train_config_json) + train_args = config2args(train_controlnet.setup_parser(), train_config) + + LOG({ + "type": "start_train", + }) + + train_controlnet.train(train_args) + + func_map = { "run_lora_sd1_5": run_lora_sd1_5, "run_lora_sdxl": run_lora_sdxl, + "run_controlnet_sd1_5": run_controlnet_sd1_5, } diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index 89bd55d..dd6d7c5 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -93,6 +93,12 @@ def MZ_ImageSelecter_call(args={}): images = args.get("images") pil_images = Utils.tensors2pil_list(images) + conditioning_images = args.get("conditioning_images", None) + conditioning_pil_images = None + if conditioning_images is not None: + conditioning_pil_images = Utils.tensors2pil_list( + conditioning_images) + resolution = args.get("resolution", 512) workspace_config = args.get("workspace_config", {}) @@ -106,6 +112,9 @@ def MZ_ImageSelecter_call(args={}): train_images_dir = os.path.join(workspace_dir, "train_images") os.makedirs(train_images_dir, exist_ok=True) + conditioning_images_dir = os.path.join( + workspace_dir, "conditioning_images") + force_clear = args.get("force_clear") == "enable" force_clear_only_images = args.get("force_clear_only_images") == "enable" if force_clear: @@ -120,10 +129,16 @@ def MZ_ImageSelecter_call(args={}): saved_images_path = [] for i, pil_image in enumerate(pil_images): pil_image = Utils.resize_max(pil_image, resolution, resolution) + width, height = pil_image.size filename = hashlib.md5(pil_image.tobytes()).hexdigest() + ".png" pil_image.save(os.path.join(train_images_dir, filename)) saved_images_path.append(filename) + if conditioning_pil_images is not None: + os.makedirs(conditioning_images_dir, exist_ok=True) + conditioning_pil_images[i].resize((width, height)).save( + os.path.join(conditioning_images_dir, filename)) + same_caption_generate = args.get("same_caption_generate") == "enable" if same_caption_generate: same_caption = args.get("same_caption").strip() @@ -136,12 +151,17 @@ def MZ_ImageSelecter_call(args={}): f.write(same_caption) dataset_config_path = os.path.join(workspace_dir, "dataset.toml") + + if conditioning_images is None: + conditioning_images_dir = None + generate_toml_config( dataset_config_path, enable_bucket=args.get("enable_bucket") == "enable", resolution=args.get("resolution"), batch_size=args.get("batch_size"), image_dir=train_images_dir, + conditioning_data_dir=conditioning_images_dir, caption_extension=".caption", num_repeats=args.get("num_repeats"), ) @@ -290,7 +310,7 @@ def check_install(): import logging -def generate_toml_config(output_path, enable_bucket=True, resolution=512, batch_size=1, image_dir=None, caption_extension=".caption", num_repeats=10, ): +def generate_toml_config(output_path, enable_bucket=True, resolution=512, batch_size=1, image_dir=None, conditioning_data_dir=None, caption_extension=".caption", num_repeats=10, ): check_install() import toml config = { @@ -304,6 +324,7 @@ def generate_toml_config(output_path, enable_bucket=True, resolution=512, batch_ 'subsets': [ { 'image_dir': image_dir, + 'conditioning_data_dir': conditioning_data_dir, 'caption_extension': caption_extension, 'num_repeats': num_repeats, }, @@ -374,11 +395,8 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot if trainer_func.find("sdxl") != -1: taesd_type = "sdxl" - - - pb = Utils.progress_bar(train_config.get("max_train_steps"), taesd_type) - + import traceback import comfy.model_management @@ -401,7 +419,6 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot max_side = max(xy_img.width, xy_img.height) # print(f"global_step: {global_step}, max_train_steps: {max_train_steps}") - total_steps = resp.get("total_steps") pb.update( int(global_step), int(total_steps), ("JPEG", xy_img, max_side)) @@ -517,6 +534,9 @@ def MZ_KohyaSSTrain_call(args={}): elif train_type == "lora_sdxl": run_hook_kohya_ss_run_file( kohya_ss_tool_dir, train_config, "run_lora_sdxl", other_config) + elif train_type == "controlnet_sd1_5": + run_hook_kohya_ss_run_file( + kohya_ss_tool_dir, train_config, "run_controlnet_sd1_5", other_config) else: raise Exception( f"暂时不支持的训练类型: {train_type}")