diff --git a/__init__.py b/__init__.py index c47b053..502d71d 100644 --- a/__init__.py +++ b/__init__.py @@ -50,8 +50,7 @@ NODE_CLASS_MAPPINGS["MZ_KohyaSSInitWorkspace"] = MZ_KohyaSSInitWorkspace NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSInitWorkspace"] = f"{AUTHOR_NAME} - KohyaSSInitWorkspace" -class MZ_ImagesCopyWorkspace: - +class MZ_KohyaSSDatasetConfig: @classmethod def INPUT_TYPES(s): return { @@ -59,21 +58,26 @@ class MZ_ImagesCopyWorkspace: "workspace_config": ("MZ_TT_SS_WorkspaceConfig",), "images": ("IMAGE",), "enable_bucket": (["enable", "disable"], {"default": "enable"}), - "resolution": ("INT", {"default": 512}), + "resolution": ("INT", {"default": 1024}), "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}), }, + "optional": { + "conditioning_images": ("IMAGE",), + } } RETURN_TYPES = (f"STRING",) RETURN_NAMES = ("workspace_images_dir",) # OUTPUT_NODE = True + MZ_DESC = """ +如果训练类型是controlnet,必须传入预处理后的图片(conditioning_images) +""" FUNCTION = "start" CATEGORY = CATEGORY_NAME @@ -83,9 +87,13 @@ class MZ_ImagesCopyWorkspace: return mz_train_tools_core.MZ_ImageSelecter_call(kwargs) -NODE_CLASS_MAPPINGS["MZ_ImagesCopyWorkspace"] = MZ_ImagesCopyWorkspace +NODE_CLASS_MAPPINGS["MZ_ImagesCopyWorkspace"] = MZ_KohyaSSDatasetConfig NODE_DISPLAY_NAME_MAPPINGS["MZ_ImagesCopyWorkspace"] = f"{AUTHOR_NAME} - ImagesCopyWorkspace" +# 别名 +NODE_CLASS_MAPPINGS["MZ_KohyaSSDatasetConfig"] = MZ_KohyaSSDatasetConfig +NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSDatasetConfig"] = f"{AUTHOR_NAME} - KohyaSSDatasetConfig" + class MZ_KohyaSSUseConfig: train_config_template_dir = os.path.join( @@ -94,7 +102,8 @@ class MZ_KohyaSSUseConfig: @classmethod def INPUT_TYPES(s): - train_config_templates = os.listdir(s.train_config_template_dir) + train_config_templates = Utils.listdir(s.train_config_template_dir) + # 去掉json后缀 train_config_templates = [os.path.splitext(x)[0] for x in train_config_templates] @@ -217,7 +226,7 @@ class AlwaysEqualProxy(str): return False -class MZ_KohyaSSTrain: +class MZ_KohyaSSLoraTrain: @classmethod def INPUT_TYPES(s): @@ -232,6 +241,9 @@ class MZ_KohyaSSTrain: # 使用walk查询所有的workspace中的所有lora模型,lora存放在每个workspace的output目录下 workspaces_loras = [] for root, dirs, files in os.walk(workspaces_dir): + + # 排除隐藏文件夹 + dirs[:] = [d for d in dirs if not d.startswith(".")] if root.endswith("output"): for file in files: if file.endswith(".safetensors"): @@ -278,8 +290,62 @@ 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: @@ -304,9 +370,10 @@ class MZ_LoadImagesFromDirectoryPath: image_dir = kwargs["directory"] if not os.path.exists(image_dir): return (images,) - images = os.listdir(image_dir) - images = [x for x in images if x.endswith( - ".png") or x.endswith(".jpg")] + images = Utils.listdir(image_dir) + + images = [x for x in images if x.lower().endswith( + ".png") or x.lower().endswith(".jpg")] images = [os.path.join(image_dir, x) for x in images] pil_images = [] @@ -317,9 +384,345 @@ class MZ_LoadImagesFromDirectoryPath: for pil_image in pil_images: tensor_images.append(Utils.pil2tensor(pil_image)) - return (tensor_images,) + return (Utils.list_tensor2tensor(tensor_images),) NODE_CLASS_MAPPINGS["MZ_LoadImagesFromDirectoryPath"] = MZ_LoadImagesFromDirectoryPath NODE_DISPLAY_NAME_MAPPINGS[ "MZ_LoadImagesFromDirectoryPath"] = f"{AUTHOR_NAME} - LoadImagesFromDirectoryPath" + + +from . import mz_train_tools_core_HYDiT + + +class MZ_HYDiTInitWorkspace: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "train_name": ("STRING", {"default": ""}), + "branch": ("STRING", {"default": "5657364143e44ac90f72aeb47b81bd505a95665d"}), + "source": ([ + "github", + "githubfast", + "521github", + "kkgithub", + ], {"default": "github"}), + "seed": ("INT", {"default": 0}), + }, + } + + RETURN_TYPES = ("MZ_TT_HYDiT_WorkspaceConfig",) + RETURN_NAMES = ("workspace_config",) + + FUNCTION = "start" + + CATEGORY = CATEGORY_NAME + + def start(self, **kwargs): + importlib.reload(mz_train_tools_core_HYDiT) + return mz_train_tools_core_HYDiT.MZ_HYDiTInitWorkspace_call(kwargs.copy()) + + +NODE_CLASS_MAPPINGS["MZ_HYDiTInitWorkspace"] = MZ_HYDiTInitWorkspace +NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTInitWorkspace"] = f"{AUTHOR_NAME} - HYDiTInitWorkspace" + + +class MZ_HYDiTDatasetConfig: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "workspace_config": ("MZ_TT_HYDiT_WorkspaceConfig",), + "resolution": ("INT", {"default": 1024}), + "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}), + }, + "optional": { + "images": ("IMAGE",), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("workspace_images_dir",) + + FUNCTION = "start" + + CATEGORY = CATEGORY_NAME + + def start(self, **kwargs): + importlib.reload(mz_train_tools_core_HYDiT) + return mz_train_tools_core_HYDiT.MZ_HYDiTDatasetConfig_call(kwargs.copy()) + + +NODE_CLASS_MAPPINGS["MZ_HYDiTDatasetConfig"] = MZ_HYDiTDatasetConfig +NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTDatasetConfig"] = f"{AUTHOR_NAME} - HYDiTDatasetConfig" + + +# class MZ_HYDiTUseConfig: +# @classmethod +# def INPUT_TYPES(s): + +# return { +# "required": { +# "workspace_config": ("MZ_TT_HYDiT_WorkspaceConfig",), +# "workspace_images_dir": ("STRING", {"forceInput": True}), +# }, +# "optional": { +# } +# } + +# RETURN_TYPES = ("MZ_TT_HYDiT_TrainConfig",) +# RETURN_NAMES = ("train_config",) + +# FUNCTION = "start" + +# CATEGORY = CATEGORY_NAME + +# def start(self, **kwargs): +# importlib.reload(mz_train_tools_core_HYDiT) +# # kwargs["train_config_template_dir"] = self.train_config_template_dir +# return mz_train_tools_core_HYDiT.MZ_HYDiTUseConfig_call(kwargs) + + +# NODE_CLASS_MAPPINGS["MZ_HYDiTUseConfig"] = MZ_HYDiTUseConfig +# NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTUseConfig"] = f"{AUTHOR_NAME} - HYDiTUseConfig" + + +class MZ_HYDiTAdvConfig: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "lr": ("STRING", {"default": "1e-5"}), + "rope_real": (["enable", "disable"], {"default": "enable"}), + # ['Wqkv', 'q_proj', 'kv_proj', 'out_proj'] + "target_modules_Wqkv": (["enable", "disable"], {"default": "enable"}), + "target_modules_q_proj": (["enable", "disable"], {"default": "enable"}), + "target_modules_kv_proj": (["enable", "disable"], {"default": "enable"}), + "target_modules_out_proj": (["enable", "disable"], {"default": "enable"}), + "warmup_min_lr": ("STRING", {"default": "1e-6"}), + # parser.add_argument("--warmup-num-steps", type=float, default=0) + # parser.add_argument("--weight-decay", type=float, default=0, help="weight-decay in optimizer") + "weight_decay": ("FLOAT", {"default": 0}), + "warmup_num_steps": ("FLOAT", {"default": 0}), + # parser.add_argument("--uncond-p", type=float, default=0.2, + # help="The probability of dropping training text used for CLIP feature extraction") + # parser.add_argument("--uncond-p-t5", type=float, default=0.2, + # help="The probability of dropping training text used for mT5 feature extraction") + "uncond_p": ("FLOAT", {"default": 0.2}), + "uncond_p_t5": ("FLOAT", {"default": 0.2}), + + # parser.add_argument("--use-flash-attn", action="store_true", help="During training, " + # "flash attention is used to accelerate training.") + # parser.add_argument("--no-flash-attn", dest="use_flash_attn", + # action="store_false", help="During training, flash attention is not used to accelerate training.") + # parser.add_argument("--use-zero-stage", type=int, default=1, help="Use AngelPTM zero stage. Support 2 and 3") + # parser.add_argument("--grad-accu-steps", type=int, default=1, help="Gradient accumulation steps.") + "use_flash_attn": (["enable", "disable"], {"default": "disable"}), + "use_zero_stage": ("INT", {"default": 2}), + "grad_accu_steps": ("INT", {"default": 1}), + # parser.add_argument("--extra-fp16", action="store_true", help="Use extra fp16 for vae and text_encoder.") + "extra_fp16": (["enable", "disable"], {"default": "enable"}), + # parser.add_argument("--qk-norm", action="store_true", help="Query Key normalization. See http://arxiv.org/abs/2302.05442 for details.") + "qk_norm": (["enable", "disable"], {"default": "enable"}), + # parser.add_argument("--norm", type=str, choices=["rms", "laryer"], default="layer", help="Normalization layer type") + "norm": (["rms", "layer"], {"default": "layer"}), + } + } + + RETURN_TYPES = ("MZ_TT_HYDiT_AdvConfig",) + RETURN_NAMES = ("advanced_config",) + + FUNCTION = "start" + + CATEGORY = CATEGORY_NAME + + def start(self, **kwargs): + return (kwargs.copy(),) + + +NODE_CLASS_MAPPINGS["MZ_HYDiTAdvConfig"] = MZ_HYDiTAdvConfig +NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTAdvConfig"] = f"{AUTHOR_NAME} - HYDiTAdvConfig" + + +class MZ_HYDiTTrain: + @classmethod + def INPUT_TYPES(s): + + hunyuan_models_path = os.path.join( + Utils.get_comfyui_models_path(), "hunyuan") + os.makedirs(hunyuan_models_path, exist_ok=True) + + models = Utils.get_models_by_folder(hunyuan_models_path) + + folders = Utils.get_folders_by_folder(hunyuan_models_path) + + vae_models = Utils.get_models_by_folder( + os.path.join(Utils.get_comfyui_models_path(), "vae")) + unet_models = Utils.get_models_by_folder( + os.path.join(Utils.get_comfyui_models_path(), "unet")) + + workspaces_root = os.path.join( + folder_paths.output_directory, "mz_train_workspaces") + + loras = mz_train_tools_core_HYDiT.search_loras([ + workspaces_root, + os.path.join(Utils.get_comfyui_models_path(), "loras"), + ]) + + return { + "required": { + "workspace_config": ("MZ_TT_HYDiT_WorkspaceConfig",), + "unet_path": (["auto"] + models + unet_models, {"default": "auto"}), + "ema_to_module": (["enable", "disable"], {"default": "enable"}), + "vae_ema_path": (["auto"] + folders + vae_models, {"default": "auto"}), + "text_encoder_path": (["auto"] + folders, {"default": "auto"}), + "tokenizer_path": (["auto"] + folders, {"default": "auto"}), + "t5_encoder_path": (["none", "auto"] + folders, {"default": "none"}), + "resolution": ("INT", {"default": 1024, "step": 16}), + "batch_size": ("INT", {"default": 1}), + "epochs": ("INT", {"default": 50}), + "ckpt_every": ("INT", {"default": 500}), + "rank": ("INT", {"default": 8}), + "base_lora": (["latest", "empty"] + loras, {"default": "latest"}), + "sample_generate": (["enable", "disable"], {"default": "enable"}), + "sample_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}), + }, + "optional": { + "advanced_config": ("MZ_TT_HYDiT_AdvConfig",), + "workspace_images_dir": ("STRING", {"forceInput": True}), + "has_no_effect": (AlwaysEqualProxy("*"),), + } + } + + RETURN_TYPES = () + RETURN_NAMES = () + + FUNCTION = "start" + + MZ_DESC = """ +base_size The base resolution (n, n) from which to create multiple resolutions | Recommended values: 256/512/1024 +""" + + OUTPUT_NODE = True + + CATEGORY = CATEGORY_NAME + + def start(self, **kwargs): + importlib.reload(mz_train_tools_core_HYDiT) + return mz_train_tools_core_HYDiT.MZ_HYDiTTrain_call(kwargs.copy()) + + +NODE_CLASS_MAPPINGS["MZ_HYDiTTrain"] = MZ_HYDiTTrain +NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTTrain"] = f"{AUTHOR_NAME} - HYDiTTrain" + + +class MZ_HYDiTSimpleT2I: + @classmethod + def INPUT_TYPES(s): + hunyuan_models_path = os.path.join( + Utils.get_comfyui_models_path(), "hunyuan") + os.makedirs(hunyuan_models_path, exist_ok=True) + + models = Utils.get_models_by_folder(hunyuan_models_path) + folders = Utils.get_folders_by_folder(hunyuan_models_path) + + vae_models = Utils.get_models_by_folder( + os.path.join(Utils.get_comfyui_models_path(), "vae")) + unet_models = Utils.get_models_by_folder( + os.path.join(Utils.get_comfyui_models_path(), "unet")) + + comfyui_full_loras = mz_train_tools_core_HYDiT.search_loras([ + os.path.join(Utils.get_comfyui_models_path(), "loras"), + ]) + + return { + "required": { + "branch": ("STRING", {"default": "5657364143e44ac90f72aeb47b81bd505a95665d"}), + "source": ([ + "github", + "githubfast", + "521github", + "kkgithub", + ], {"default": "github"}), + "unet_path": (["auto"] + models + unet_models, {"default": "auto"}), + "vae_ema_path": (["auto"] + folders + vae_models, {"default": "auto"}), + "text_encoder_path": (["auto"] + folders, {"default": "auto"}), + "tokenizer_path": (["auto"] + folders, {"default": "auto"}), + "t5_encoder_path": (["none", "auto"] + folders, {"default": "auto"}), + "lora_path": (["none"] + comfyui_full_loras, {"default": "none"}), + "seed": ("INT", {"default": 0}), + "steps": ("INT", {"default": 20}), + "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}), + "scheduler": ([ + "ddpm", "ddim", "dpmms", "uni_pc", "dpmpp_2m_karras" + ], {"default": "ddpm"}), + "prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}), + "negative_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}), + "width": ("INT", {"default": 512, "max": 8192, "step": 16}), + "height": ("INT", {"default": 512, "max": 8192, "step": 16}), + "keep_device": (["enable", "disable"], {"default": "enable"}), + } + } + + RETURN_TYPES = ("IMAGE",) + + RETURN_NAMES = ("images",) + + FUNCTION = "start" + + CATEGORY = CATEGORY_NAME + + def start(self, **kwargs): + importlib.reload(mz_train_tools_core_HYDiT) + return mz_train_tools_core_HYDiT.MZ_HYDiTSimpleT2I_call(kwargs) + + +NODE_CLASS_MAPPINGS["MZ_HYDiTSimpleT2I"] = MZ_HYDiTSimpleT2I +NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTSimpleT2I"] = f"{AUTHOR_NAME} - HYDiTSimpleT2I" + + +class MZ_TrainToolsDebug: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "object": (AlwaysEqualProxy("*"),), + "indent": ("INT", {"default": 2}), + "depth": ("INT", {"default": 5}), + "width": ("INT", {"default": 80}), + "compact": (["enable", "disable"], {"default": "enable"}), + "sort_keys": (["enable", "disable"], {"default": "enable"}), + "underscore_numbers": (["enable", "disable"], {"default": "enable"}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("debug",) + + FUNCTION = "start" + + CATEGORY = CATEGORY_NAME + + def start(self, **kwargs): + + from pprint import pprint, pp + object = kwargs["object"] + indent = kwargs["indent"] + depth = kwargs["depth"] + width = kwargs["width"] + compact = kwargs["compact"] == "enable" + sort_keys = kwargs["sort_keys"] == "enable" + underscore_numbers = kwargs["underscore_numbers"] == "enable" + + debug = pp(object, stream=None, indent=indent, depth=depth, width=width, + compact=compact, sort_dicts=sort_keys, underscore_numbers=underscore_numbers) + + return (debug,) + + +NODE_CLASS_MAPPINGS["MZ_TrainToolsDebug"] = MZ_TrainToolsDebug +NODE_DISPLAY_NAME_MAPPINGS["MZ_TrainToolsDebug"] = f"{AUTHOR_NAME} - TrainToolsDebug" 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/examples/captioner.png b/examples/captioner.png new file mode 100644 index 0000000..53d5c3b Binary files /dev/null and b/examples/captioner.png differ diff --git a/hook_HYDiT_idk_run.py b/hook_HYDiT_idk_run.py new file mode 100644 index 0000000..50b8f27 --- /dev/null +++ b/hook_HYDiT_idk_run.py @@ -0,0 +1,155 @@ +import argparse + +from index_kits.dataset.make_dataset_core import startup, make_multireso +from index_kits.common import show_index_info +from index_kits import __version__ + + +def common_args(parser): + parser.add_argument('-t', '--target', type=str, + required=True, help='Save path') + + +def get_args(): + parser = argparse.ArgumentParser(description=""" + IndexKits is a tool to build and manage index files for large-scale datasets. + It supports both base index and multi-resolution index. + + Introduction + ------------ + This command line tool provides the following functionalities: + 1. Show index v2 information + 2. Build base index v2 + 3. Build multi-resolution index v2 + + Examples + -------- + 1. Show index v2 information + index_kits show /path/to/index.json + + 2. Build base index v2 + Default usage: + index_kits base -c /path/to/config.yaml -t /path/to/index.json + + Use multiple processes: + index_kits base -c /path/to/config.yaml -t /path/to/index.json -w 40 + + 3. Build multi-resolution index v2 + + Build with a configuration file: + index_kits multireso -c /path/to/config.yaml -t /path/to/index_mb_gt512.json + + Build by specifying arguments without a configuration file: + index_kits multireso --src /path/to/index.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json + + Build by specifying target-ratios: + index_kits multireso --src /path/to/index.json --base-size 512 --target-ratios 1:1 4:3 3:4 16:9 9:16 --min-size 512 -t /path/to/index_mb_gt512.json + + Build with multiple source index files. + index_kits multireso --src /path/to/index1.json /path/to/index2.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json + """, formatter_class=argparse.RawTextHelpFormatter) + sub_parsers = parser.add_subparsers(dest='task', required=True) + + # Show index message + show_parser = sub_parsers.add_parser('show', description=""" + Show base/multireso index v2 information. + + Example + ------- + index_kits show /path/to/index.json + """, formatter_class=argparse.RawTextHelpFormatter) + show_parser.add_argument( + 'src', type=str, help='Path to a base/multireso index file.') + show_parser.add_argument( + '--arrow-files', action='store_true', help='Show arrow files only.') + show_parser.add_argument('--depth', type=int, default=1, + help='Arrow file depth. Default is 1, the level of last folder in the arrow file path. ' + 'Set it to 0 to show the full path including `xxx/last_folder/*.arrow`.') + + # Single resolution bucket + base_parser = sub_parsers.add_parser('base', description=""" + Build base index v2. + + Example + ------- + index_kits base -c /path/to/config.yaml -t /path/to/index.json + """, formatter_class=argparse.RawTextHelpFormatter) + base_parser.add_argument('-c', '--config', type=str, + required=True, help='Configuration file path') + common_args(base_parser) + base_parser.add_argument('-w', '--world-size', type=int, default=1) + base_parser.add_argument('--work-dir', type=str, + default='.', help='Work directory') + base_parser.add_argument('--use-cache', action='store_true', help='Use cache to avoid reprocessing. ' + 'Perform merge pkl results directly.') + + # Multi-resolution bucket + mo_parser = sub_parsers.add_parser('multireso', description=""" + Build multi-resolution index v2 + + Example + ------- + Build with a configuration file: + index_kits multireso -c /path/to/config.yaml -t /path/to/index_mb_gt512.json + + Build by specifying arguments without a configuration file: + index_kits multireso --src /path/to/index.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json + + Build by specifying target-ratios: + index_kits multireso --src /path/to/index.json --base-size 512 --target-ratios 1:1 4:3 3:4 16:9 9:16 --min-size 512 -t /path/to/index_mb_gt512.json + + Build with multiple source index files. + index_kits multireso --src /path/to/index1.json /path/to/index2.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json + """, formatter_class=argparse.RawTextHelpFormatter) + mo_parser.add_argument('-c', '--config', type=str, default=None, + help='Configuration file path in a yaml format. Either --config or --src must be provided.') + mo_parser.add_argument('-s', '--src', type=str, nargs='+', default=None, + help='Source index files. Either --config or --src must be provided.') + common_args(mo_parser) + mo_parser.add_argument('--base-size', type=int, default=None, + help="Base size. Typically set as 256/512/1024 according to image size you train model.") + mo_parser.add_argument('--reso-step', type=int, default=None, + help="Resolution step. Either reso_step or target_ratios must be provided.") + mo_parser.add_argument('--target-ratios', type=str, nargs='+', default=None, + help="Target ratios. Either reso_step or target_ratios must be provided.") + mo_parser.add_argument('--md5-file', type=str, default=None, + help='You can provide an md5 to height and width file to accelerate the process. ' + 'It is a pickle file that contains a dict, which maps md5 to (height, width) tuple.') + mo_parser.add_argument('--align', type=int, default=16, + help="Used when --target-ratios is provided. Align size of source image height and width.") + mo_parser.add_argument('--min-size', type=int, default=0, + help="Minimum size. Images smaller than this size will be ignored.") + + # Common + parser.add_argument('-v', '--version', action='version', + version=f'%(prog)s {__version__}') + + args = parser.parse_args() + return args + + +if __name__ == '__main__': + args = get_args() + if args.task == 'show': + show_index_info(args.src, + args.arrow_files, + args.depth, + ) + elif args.task == 'base': + startup(args.config, + args.target, + args.world_size, + args.work_dir, + use_cache=args.use_cache, + ) + elif args.task == 'multireso': + make_multireso(args.target, + args.config, + args.src, + args.base_size, + args.reso_step, + args.target_ratios, + args.align, + args.min_size, + args.md5_file, + ) diff --git a/hook_HYDiT_main_train_deepspeed.py b/hook_HYDiT_main_train_deepspeed.py new file mode 100644 index 0000000..c25fa9d --- /dev/null +++ b/hook_HYDiT_main_train_deepspeed.py @@ -0,0 +1,595 @@ +import gc +import json +import os +import random +import sys +import time +from functools import partial +from glob import glob +from pathlib import Path +import numpy as np +import safetensors.torch + +import deepspeed +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.distributed as dist +from torch.utils.data import DataLoader +from torch.distributed.optim import ZeroRedundancyOptimizer +from torchvision.transforms import functional as TF +from diffusers.models import AutoencoderKL +from transformers import BertModel, BertTokenizer, logging as tf_logging + +from hydit.config import get_args +from hydit.lr_scheduler import WarmupLR +from hydit.data_loader.arrow_load_stream import TextImageArrowStream +from hydit.diffusion import create_diffusion +from hydit.ds_config import deepspeed_config_from_args +from hydit.modules.ema import EMA +from hydit.modules.fp16_layers import Float16Module +from hydit.modules.models import HUNYUAN_DIT_MODELS +from hydit.modules.posemb_layers import init_image_posemb +from hydit.utils.tools import create_logger, set_seeds, create_exp_folder, get_trainable_params +from IndexKits.index_kits import ResolutionGroup +from IndexKits.index_kits.sampler import DistributedSamplerWithStartIndex, BlockDistributedSampler +from peft import LoraConfig, get_peft_model + + +def deepspeed_initialize(args, logger, model, opt, deepspeed_config): + logger.info(f"Initialize deepspeed...") + logger.info(f" Using deepspeed optimizer") + + def get_learning_rate_scheduler(warmup_min_lr, lr, warmup_num_steps, opt): + return WarmupLR(opt, warmup_min_lr, lr, warmup_num_steps) + + logger.info( + f" Building scheduler with warmup_min_lr={args.warmup_min_lr}, warmup_num_steps={args.warmup_num_steps}") + logger.info( + f" deepspeed_config={deepspeed_config}") + + model, opt, _, scheduler = deepspeed.initialize( + model=model, + model_parameters=get_trainable_params(model), + config_params=deepspeed_config, + args=args, + lr_scheduler=partial(get_learning_rate_scheduler, args.warmup_min_lr, + args.lr, args.warmup_num_steps) if args.warmup_num_steps > 0 else None, + ) + return model, opt, scheduler + + +def save_checkpoint(args, rank, logger, model, ema, epoch, train_steps, checkpoint_dir): + def save_lora_weight(checkpoint_dir, client_state, tag=f"{train_steps:07d}.pt"): + cur_ckpt_save_dir = f"{checkpoint_dir}/{tag}" + if rank == 0: + if args.use_fp16: + model.module.module.save_pretrained(cur_ckpt_save_dir) + else: + model.module.save_pretrained(cur_ckpt_save_dir) + + checkpoint_path = "[Not rank 0. Disabled output.]" + + client_state = { + "steps": train_steps, + "epoch": epoch, + "args": args + } + if ema is not None: + client_state['ema'] = ema.state_dict() + + dst_paths = [] + if train_steps % args.ckpt_every == 0: + checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt" + try: + if args.training_parts == "lora": + save_lora_weight(checkpoint_dir, client_state, + tag=f"{train_steps:07d}.pt") + else: + model.save_checkpoint( + checkpoint_dir, client_state=client_state, tag=f"{train_steps:07d}.pt") + dst_paths.append(checkpoint_path) + logger.info(f"Saved checkpoint to {checkpoint_path}") + except: + logger.error(f"Saved failed to {checkpoint_path}") + + if train_steps % args.ckpt_latest_every == 0 or train_steps == args.max_training_steps: + save_name = "latest.pt" + checkpoint_path = f"{checkpoint_dir}/{save_name}" + try: + if args.training_parts == "lora": + save_lora_weight(checkpoint_dir, client_state, + tag=f"{save_name}") + else: + model.save_checkpoint( + checkpoint_dir, client_state=client_state, tag=f"{save_name}") + dst_paths.append(checkpoint_path) + logger.info(f"Saved checkpoint to {checkpoint_path}") + except: + logger.error(f"Saved failed to {checkpoint_path}") + + dist.barrier() + if rank == 0 and len(dst_paths) > 0: + # Delete optimizer states to avoid occupying too much disk space. + for dst_path in dst_paths: + for opt_state_path in glob(f"{dst_path}/zero_dp_rank_*_tp_rank_00_pp_rank_00_optim_states.pt"): + os.remove(opt_state_path) + + return checkpoint_path + + +@torch.no_grad() +def prepare_model_inputs(args, batch, device, vae, text_encoder, text_encoder_t5, freqs_cis_img): + try: + from .hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar + except: + from hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar + + image, text_embedding, text_embedding_mask, text_embedding_t5, text_embedding_mask_t5, kwargs = batch + + # clip & mT5 text embedding + text_embedding = text_embedding.to(device) + text_embedding_mask = text_embedding_mask.to(device) + encoder_hidden_states = text_encoder( + text_embedding.to(device), + attention_mask=text_embedding_mask.to(device), + )[0] + text_embedding_t5 = text_embedding_t5.to(device).squeeze(1) + text_embedding_mask_t5 = text_embedding_mask_t5.to(device).squeeze(1) + with torch.no_grad(): + output_t5 = text_encoder_t5( + input_ids=text_embedding_t5, + attention_mask=text_embedding_mask_t5 if T5_ENCODER['attention_mask'] else None, + output_hidden_states=True + ) + encoder_hidden_states_t5 = output_t5['hidden_states'][T5_ENCODER['layer_index']].detach( + ) + + # additional condition + image_meta_size = kwargs['image_meta_size'].to(device) + style = kwargs['style'].to(device) + + if args.extra_fp16: + image = image.half() + image_meta_size = image_meta_size.half() if image_meta_size is not None else None + + # Map input images to latent space + normalize latents: + image = image.to(device) + vae_scaling_factor = vae.config.scaling_factor + latents = vae.encode(image).latent_dist.sample().mul_(vae_scaling_factor) + + # positional embedding + _, _, height, width = image.shape + reso = f"{height}x{width}" + cos_cis_img, sin_cis_img = freqs_cis_img[reso] + + # Model conditions + model_kwargs = dict( + encoder_hidden_states=encoder_hidden_states, + text_embedding_mask=text_embedding_mask, + encoder_hidden_states_t5=encoder_hidden_states_t5, + text_embedding_mask_t5=text_embedding_mask_t5, + image_meta_size=image_meta_size, + style=style, + cos_cis_img=cos_cis_img, + sin_cis_img=sin_cis_img, + ) + + return latents, model_kwargs + + +def Core(args): + if args.training_parts == "lora": + args.use_ema = False + + assert torch.cuda.is_available(), "Training currently requires at least one GPU." + + dist.init_process_group("nccl") + world_size = dist.get_world_size() + batch_size = args.batch_size + grad_accu_steps = args.grad_accu_steps + global_batch_size = world_size * batch_size * grad_accu_steps + + rank = dist.get_rank() + device = rank % torch.cuda.device_count() + seed = args.global_seed * world_size + rank + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.cuda.set_device(device) + print(f"Starting rank={rank}, seed={seed}, world_size={world_size}.") + deepspeed_config = deepspeed_config_from_args(args, global_batch_size) + + # Setup an experiment folder + experiment_dir, checkpoint_dir, logger = create_exp_folder(args, rank) + + # Log all the arguments + logger.info(sys.argv) + logger.info(str(args)) + # Save to a json file + args_dict = vars(args) + args_dict['world_size'] = world_size + with open(f"{experiment_dir}/args.json", 'w') as f: + json.dump(args_dict, f, indent=4) + + # Disable the message "Some weights of the model checkpoint at ... were not used when initializing BertModel." + # If needed, just comment the following line. + tf_logging.set_verbosity_error() + + # =========================================================================== + # Building HYDIT + # =========================================================================== + + logger.info("Building HYDIT Model.") + + # --------------------------------------------------------------------------- + # Training sample base size, such as 256/512/1024. Notice that this size is + # just a base size, not necessary the actual size of training samples. Actual + # size of the training samples are correlated with `resolutions` when enabling + # multi-resolution training. + # --------------------------------------------------------------------------- + image_size = args.image_size + if len(image_size) == 1: + image_size = [image_size[0], image_size[0]] + if len(image_size) != 2: + raise ValueError(f"Invalid image size: {args.image_size}") + assert image_size[0] % 8 == 0 and image_size[1] % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder). " \ + f"got {image_size}" + latent_size = [image_size[0] // 8, image_size[1] // 8] + + # initialize model by deepspeed + assert args.deepspeed, f"Must enable deepspeed in this script: train_deepspeed.py" + with deepspeed.zero.Init(data_parallel_group=torch.distributed.group.WORLD, + remote_device=None if args.remote_device == 'none' else args.remote_device, + config_dict_or_path=deepspeed_config, + mpu=None, + enabled=args.zero_stage == 3): + model = HUNYUAN_DIT_MODELS[args.model](args, + input_size=latent_size, + log_fn=logger.info, + ) + # Multi-resolution / Single-resolution training. + if args.multireso: + resolutions = ResolutionGroup(image_size[0], + align=16, + step=args.reso_step, + target_ratios=args.target_ratios).data + else: + resolutions = ResolutionGroup(image_size[0], + align=16, + target_ratios=['1:1']).data + + freqs_cis_img = init_image_posemb(args.rope_img, + resolutions=resolutions, + patch_size=model.patch_size, + hidden_size=model.hidden_size, + num_heads=model.num_heads, + log_fn=logger.info, + rope_real=args.rope_real, + ) + + # Create EMA model and convert to fp16 if needed. + ema = None + if args.use_ema: + ema = EMA(args, model, device, logger) + + # Setup FP16 main model: + if args.use_fp16: + model = Float16Module(model, args) + logger.info( + f" Using main model with data type {'fp16' if args.use_fp16 else 'fp32'}") + + diffusion = create_diffusion( + noise_schedule=args.noise_schedule, + predict_type=args.predict_type, + learn_sigma=args.learn_sigma, + mse_loss_weight_type=args.mse_loss_weight_type, + beta_start=args.beta_start, + beta_end=args.beta_end, + noise_offset=args.noise_offset, + ) + + try: + from .hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar, CustomizeEmbeds + except: + from hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar, CustomizeEmbeds + + # Setup VAE + logger.info(f" Loading vae from {VAE_EMA_PATH}") + vae = AutoencoderKL.from_pretrained(VAE_EMA_PATH) + # Setup BERT text encoder + logger.info(f" Loading Bert text encoder from {TEXT_ENCODER}") + text_encoder = BertModel.from_pretrained( + TEXT_ENCODER, False, revision=None) + # Setup BERT tokenizer: + logger.info(f" Loading Bert tokenizer from {TOKENIZER}") + tokenizer = BertTokenizer.from_pretrained(TOKENIZER) + # Setup T5 text encoder + from hydit.modules.text_encoder import MT5Embedder + mt5_path = T5_ENCODER['MT5'] + if mt5_path is None: + embedder_t5 = CustomizeEmbeds() + else: + embedder_t5 = MT5Embedder( + mt5_path, torch_dtype=T5_ENCODER['torch_dtype'], max_length=args.text_len_t5) + tokenizer_t5 = embedder_t5.tokenizer + text_encoder_t5 = embedder_t5.model + + if args.extra_fp16: + logger.info(f" Using fp16 for extra modules: vae, text_encoder") + vae = vae.half().to(device) + text_encoder = text_encoder.half().to(device) + text_encoder_t5 = text_encoder_t5.half().to(device) + else: + vae = vae.to(device) + text_encoder = text_encoder.to(device) + text_encoder_t5 = text_encoder_t5.to(device) + + logger.info( + f" Optimizer parameters: lr={args.lr}, weight_decay={args.weight_decay}") + logger.info(" Using deepspeed optimizer") + opt = None + + # =========================================================================== + # Building Dataset + # =========================================================================== + + logger.info(f"Building Streaming Dataset.") + logger.info(f" Loading index file {args.index_file} (v2)") + + dataset = TextImageArrowStream(args=args, + resolution=image_size[0], + random_flip=args.random_flip, + log_fn=logger.info, + index_file=args.index_file, + multireso=args.multireso, + batch_size=batch_size, + world_size=world_size, + random_shrink_size_cond=args.random_shrink_size_cond, + merge_src_cond=args.merge_src_cond, + uncond_p=args.uncond_p, + text_ctx_len=args.text_len, + tokenizer=tokenizer, + uncond_p_t5=args.uncond_p_t5, + text_ctx_len_t5=args.text_len_t5, + tokenizer_t5=tokenizer_t5, + ) + if args.multireso: + sampler = BlockDistributedSampler(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed, + shuffle=False, drop_last=True, batch_size=batch_size) + else: + sampler = DistributedSamplerWithStartIndex(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed, + shuffle=False, drop_last=True) + loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, sampler=sampler, + num_workers=args.num_workers, pin_memory=True, drop_last=True) + logger.info(f" Dataset contains {len(dataset):,} images.") + logger.info(f" Index file: {args.index_file}.") + if args.multireso: + logger.info(f' Using MultiResolutionBucketIndexV2 with step {dataset.index_manager.step} ' + f'and base size {dataset.index_manager.base_size}') + logger.info(f'\n {dataset.index_manager.resolutions}') + + # =========================================================================== + # Loading parameter + # =========================================================================== + + logger.info(f"Loading parameter") + start_epoch = 0 + start_epoch_step = 0 + train_steps = 0 + # Resume checkpoint if needed + # if args.resume is not None or len(args.resume) > 0: + if True: + model, ema, start_epoch, start_epoch_step, train_steps = model_resume( + args, model, ema, logger) + + if args.training_parts == "lora": + lora_ckpt = args.lora_ckpt + if lora_ckpt is not None: + lastest_checkpoint = lora_ckpt + from peft.peft_model import PeftModel + print(f"Loading lora model from {lastest_checkpoint}") + if args.use_fp16: + model.module = PeftModel.from_pretrained( + model.module, lastest_checkpoint, is_trainable=True) + else: + model = PeftModel.from_pretrained( + model, lastest_checkpoint, is_trainable=True) + else: + loraconfig = LoraConfig( + r=args.rank, + lora_alpha=args.rank, + target_modules=args.target_modules + ) + + if args.use_fp16: + model.module = get_peft_model(model.module, loraconfig) + else: + model = get_peft_model(model, loraconfig) + + logger.info(f" Training parts: {args.training_parts}") + + model, opt, scheduler = deepspeed_initialize( + args, logger, model, opt, deepspeed_config) + + # =========================================================================== + # Training + # =========================================================================== + + # print model structure + # for name, param in model.named_parameters(): + # print(name, param.size()) + # raise Exception("stop") + + model.train() + if args.use_ema: + ema.eval() + + print(f" Worker {rank} ready.") + dist.barrier() + + iters_per_epoch = len(loader) + logger.info( + " ****************************** Running training ******************************") + logger.info(f" Number GPUs: {world_size}") + logger.info(f" Number training samples: {len(dataset):,}") + logger.info( + f" Number parameters: {sum(p.numel() for p in model.parameters()):,}") + logger.info( + f" Number trainable params: {sum(p.numel() for p in get_trainable_params(model)):,}") + logger.info( + " ------------------------------------------------------------------------------") + logger.info(f" Iters per epoch: {iters_per_epoch:,}") + logger.info(f" Batch size per device: {batch_size}") + logger.info( + f" Batch size all device: {batch_size * world_size * grad_accu_steps:,} (world_size * batch_size * grad_accu_steps)") + logger.info(f" Gradient Accu steps: {args.grad_accu_steps}") + logger.info( + f" Total optimization steps: {args.epochs * iters_per_epoch // grad_accu_steps:,}") + + logger.info( + f" Training epochs: {start_epoch}/{args.epochs}") + logger.info( + f" Training epoch steps: {start_epoch_step:,}/{iters_per_epoch:,}") + logger.info( + f" Training total steps: {train_steps:,}/{min(args.max_training_steps, args.epochs * iters_per_epoch):,}") + logger.info( + " ------------------------------------------------------------------------------") + logger.info(f" Noise schedule: {args.noise_schedule}") + logger.info( + f" Beta limits: ({args.beta_start}, {args.beta_end})") + logger.info(f" Learn sigma: {args.learn_sigma}") + logger.info(f" Prediction type: {args.predict_type}") + logger.info(f" Noise offset: {args.noise_offset}") + + logger.info( + " ------------------------------------------------------------------------------") + logger.info( + f" Using EMA model: {args.use_ema} ({args.ema_dtype})") + if args.use_ema: + logger.info( + f" Using EMA decay: {ema.max_value if args.use_ema else None}") + logger.info( + f" Using EMA warmup power: {ema.power if args.use_ema else None}") + logger.info(f" Using main model fp16: {args.use_fp16}") + logger.info(f" Using extra modules fp16: {args.extra_fp16}") + logger.info( + " ------------------------------------------------------------------------------") + logger.info(f" Experiment directory: {experiment_dir}") + logger.info( + " *******************************************************************************") + + if args.gc_interval > 0: + gc.disable() + gc.collect() + + # Variables for monitoring/logging purposes: + log_steps = 0 + running_loss = 0 + start_time = time.time() + + if args.async_ema: + ema_stream = torch.cuda.Stream() + + easy_sample_images(args, vae, text_encoder, tokenizer, model, embedder_t5, + target_height=768, target_width=1280, train_steps=0) + pbar = PBar(args.epochs * len(loader)) + + # Training loop + for epoch in range(start_epoch, args.epochs): + logger.info(f" Start random shuffle with seed={seed}") + # Makesure all processors use the same seed to shuffle dataset. + dataset.shuffle(seed=args.global_seed + epoch, fast=True) + logger.info(f" End of random shuffle") + + # Move sampler to start_index + if not args.multireso: + start_index = start_epoch_step * world_size * batch_size + if start_index != sampler.start_index: + sampler.start_index = start_index + # Reset start_epoch_step to zero, to ensure next epoch will start from the beginning. + start_epoch_step = 0 + logger.info(f" Iters left this epoch: {len(loader):,}") + + logger.info(f" Beginning epoch {epoch}...") + step = 0 + for batch in loader: + step += 1 + + latents, model_kwargs = prepare_model_inputs( + args, batch, device, vae, text_encoder, text_encoder_t5, freqs_cis_img) + + # training model by deepspeed while use fp16 + if args.use_fp16: + if args.use_ema and args.async_ema: + with torch.cuda.stream(ema_stream): + ema.update(model.module.module, step=step) + torch.cuda.current_stream().wait_stream(ema_stream) + + loss_dict = diffusion.training_losses( + model=model, x_start=latents, model_kwargs=model_kwargs) + loss = loss_dict["loss"].mean() + model.backward(loss) + last_batch_iteration = ( + train_steps + 1) // (global_batch_size // (batch_size * world_size)) + model.step( + lr_kwargs={'last_batch_iteration': last_batch_iteration}) + + if args.use_ema and not args.async_ema or (args.async_ema and step == len(loader) - 1): + if args.use_fp16: + ema.update(model.module.module, step=step) + else: + ema.update(model.module, step=step) + + # =========================================================================== + # Log loss values: + # =========================================================================== + running_loss += loss.item() + log_steps += 1 + train_steps += 1 + if train_steps % args.log_every == 0: + # Measure training speed: + torch.cuda.synchronize() + end_time = time.time() + steps_per_sec = log_steps / (end_time - start_time) + # Reduce loss history over all processes: + avg_loss = torch.tensor( + running_loss / log_steps, device=device) + dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM) + avg_loss = avg_loss.item() / world_size + # get lr from deepspeed fused optimizer + logger.info(f"(step={train_steps:07d}) " + + (f"(update_step={train_steps // args.grad_accu_steps:07d}) " if args.grad_accu_steps > 1 else "") + + f"Train Loss: {avg_loss:.4f}, " + f"Lr: {opt.param_groups[0]['lr']:.6g}, " + f"Steps/Sec: {steps_per_sec:.2f}, " + f"Samples/Sec: {int(steps_per_sec * batch_size * world_size):d}") + # Reset monitoring variables: + running_loss = 0 + log_steps = 0 + start_time = time.time() + + # collect gc: + if args.gc_interval > 0 and (step % args.gc_interval == 0): + gc.collect() + + pbar.step( + f"Epoch {epoch}, step {step}, loss {loss.item():.4f}", args.epochs * len(loader), train_steps) + + if (train_steps % args.ckpt_every == 0 or train_steps % args.ckpt_latest_every == 0 # or train_steps == args.max_training_steps + ) and train_steps > 0: + easy_sample_images(args, vae, text_encoder, tokenizer, model, embedder_t5, + target_height=768, target_width=1280, train_steps=train_steps) + save_checkpoint(args, rank, logger, model, ema, + epoch, train_steps, checkpoint_dir) + + if train_steps >= args.max_training_steps: + logger.info(f"Breaking step loop at {train_steps}.") + break + + if train_steps >= args.max_training_steps: + logger.info(f"Breaking epoch loop at {epoch}.") + break + + dist.destroy_process_group() diff --git a/hook_HYDiT_run.py b/hook_HYDiT_run.py new file mode 100644 index 0000000..1fe4222 --- /dev/null +++ b/hook_HYDiT_run.py @@ -0,0 +1,175 @@ + +import argparse +import json +import os +import sys +from types import SimpleNamespace + +import torch + + +class SimpleNamespaceCNWarrper(SimpleNamespace): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.__dict__.update(kwargs) + 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) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + conflict_handler='resolve', + ) + parser.add_argument("--sys_path", type=str, default="") + parser.add_argument("--train_config_file", type=str, default="") + parser.add_argument("--mz_master_port", type=int, default=0) + args = parser.parse_args() + + master_port = args.mz_master_port + + print(f"master_port = {master_port}") + + try: + from . import hook_HYDiT_utils + except Exception as e: + import hook_HYDiT_utils + + hook_HYDiT_utils.set_master_port(master_port) + + sys_path = args.sys_path + if sys_path != "": + sys.path.append(sys_path) + + print("HYDi run hook") + + try: + from . import hook_HYDiT_main_train_deepspeed + except Exception as e: + import hook_HYDiT_main_train_deepspeed + + import hydit.config + + def _handle_conflict_error(self, *args, **kwargs): + pass + + def parse_args(self, args=None, namespace=None): + args, argv = self.parse_known_args(args, namespace) + return args + + train_config_file = args.train_config_file + + if train_config_file == "": + raise ValueError("train_config_file is empty") + + train_config = {} + with open(train_config_file, "r") as f: + train_config = json.load(f) + + argparse.ArgumentParser._handle_conflict_error = _handle_conflict_error + argparse.ArgumentParser._handle_conflict_resolve = _handle_conflict_error + argparse.ArgumentParser.parse_args = parse_args + margs = hydit.config.get_args() + margs.model = train_config.get("model", "DiT-g/2") + + margs.task_flag = train_config.get("task_flag") + + margs.resume_split = train_config.get("resume_split", True) + + margs.ema_to_module = train_config.get("ema_to_module", True) + + margs.deepspeed = False + + margs.predict_type = train_config.get("predict_type", "v_prediction") + + margs.training_parts = train_config.get("training_parts", "lora") + + margs.batch_size = train_config.get("batch_size", 1) + + margs.grad_accu_steps = train_config.get("grad_accu_steps", 1) + + margs.global_seed = train_config.get("global_seed", 0) + + margs.use_flash_attn = train_config.get("use_flash_attn", False) + + margs.use_fp16 = train_config.get("use_fp16", True) + + margs.qk_norm = train_config.get("qk_norm", True) + + margs.ema_dtype = train_config.get("ema_dtype", "fp32") + + margs.async_ema = False + + margs.ckpt_latest_every = 0x7fffffff + + margs.multireso = train_config.get("multireso", True) + + margs.epochs = train_config.get("epochs", 50) + + margs.target_ratios = train_config.get( + "target_ratios", ['1:1', '3:4', '4:3', '16:9', '9:16']) + + margs.rope_img = train_config.get("rope_img", "base1024") + + margs.image_size = train_config.get("image_size", 1024) + + margs.rope_real = train_config.get("rope_real", True) + + margs.index_file = train_config.get("index_file", None) + + margs.lr = train_config.get("lr", 1e-5) + + margs.rank = train_config.get("rank", 8) + + margs.noise_offset = train_config.get("noise_offset", 0.0) + + margs.log_every = train_config.get("log_every", 99999999999999) + + margs.use_zero_stage = train_config.get("use_zero_stage", 2) + + margs.global_batch_size = train_config.get("global_batch_size", 1) + + margs.deepspeed = True + + margs.results_dir = train_config.get("results_dir") + + margs.mse_loss_weight_type = train_config.get( + "mse_loss_weight_type", "constant") + + for k, v in train_config.items(): + if hasattr(margs, k): + setattr(margs, k, v) + + hook_HYDiT_utils.set_unet_path( + train_config.get("unet_path")) + + hook_HYDiT_utils.set_vae_ema_path( + train_config.get("vae_ema_path")) + + hook_HYDiT_utils.set_text_encoder_path( + train_config.get("text_encoder_path")) + + hook_HYDiT_utils.set_tokenizer_path( + train_config.get("tokenizer_path")) + + hook_HYDiT_utils.set_t5_encoder_path( + train_config.get("t5_encoder_path")) + + hook_HYDiT_utils.set_train_config(train_config) + + try: + # deepspeed/runtime/engine + import deepspeed.runtime.engine + deepspeed.runtime.engine.DeepSpeedEngine._do_sanity_check = lambda x: None + except Exception as e: + pass + + if type(margs.image_size) == int: + margs.image_size = [margs.image_size, margs.image_size] + hook_HYDiT_main_train_deepspeed.Core(margs) diff --git a/hook_HYDiT_utils.py b/hook_HYDiT_utils.py new file mode 100644 index 0000000..d5ac04b --- /dev/null +++ b/hook_HYDiT_utils.py @@ -0,0 +1,379 @@ + + +import json +import os +import time +import torch + + +UNET_PATH = "ckpts/t2i/model/pytorch_model_ema.pt" + + +def set_unet_path(path): + global UNET_PATH + UNET_PATH = path + + +VAE_EMA_PATH = "ckpts/t2i/sdxl-vae-fp16-fix" + + +def set_vae_ema_path(path): + global VAE_EMA_PATH + VAE_EMA_PATH = path + + +TOKENIZER = "ckpts/t2i/tokenizer" + + +def set_tokenizer_path(path): + global TOKENIZER + TOKENIZER = path + + +TEXT_ENCODER = 'ckpts/t2i/clip_text_encoder' + + +def set_text_encoder_path(path): + global TEXT_ENCODER + TEXT_ENCODER = path + + +T5_ENCODER = { + 'MT5': None, + 'attention_mask': True, + 'layer_index': -1, + 'attention_pool': True, + 'torch_dtype': torch.float16, + 'learnable_replace': True +} + + +def set_t5_encoder_path(path): + global T5_ENCODER + T5_ENCODER['MT5'] = path + + +global TRAIN_CONFIG + + +def set_train_config(train_config): + global TRAIN_CONFIG + TRAIN_CONFIG = train_config + + +def easy_sample_images( + args, + vae=None, + text_encoder=None, + tokenizer=None, + model=None, + embedder_t5=None, + target_height=768, + target_width=1280, + prompt="A photo of a girl with a hat on a sunny day", + negative_prompt="", + batch_size=1, + guidance_scale=2.0, + infer_steps=20, + sampler='dpmpp_2m_karras', + train_steps=0, + seed=0, +): + from hydit.diffusion.pipeline import StableDiffusionPipeline + from diffusers import schedulers + from hydit.constants import SAMPLER_FACTORY + from hydit.modules.posemb_layers import get_fill_resize_and_crop, get_2d_rotary_pos_embed + from hydit.modules.models import HUNYUAN_DIT_CONFIG + + import traceback + with torch.cuda.amp.autocast(): + + workspace_dir = TRAIN_CONFIG.get("workspace_dir") + sample_config_file = TRAIN_CONFIG.get("sample_config_file", None) + if sample_config_file is None: + print("sample_config_file is not set.") + return + try: + sample_config = json.load(open(sample_config_file, "r")) + except Exception as e: + print(f"Failed to load sample_config_file: {sample_config_file}") + return + sample_images_dir = os.path.join(workspace_dir, "sample_images") + os.makedirs(sample_images_dir, exist_ok=True) + + sampler_factory = SAMPLER_FACTORY.copy() + + sampler_factory["uni_pc"] = { + 'scheduler': 'UniPCMultistepScheduler', + 'name': 'UniPCMultistepScheduler', + 'kwargs': { + 'beta_schedule': 'scaled_linear', + 'beta_start': 0.00085, + 'beta_end': 0.03, + 'prediction_type': 'v_prediction', + 'trained_betas': None, + 'solver_order': 2, + } + } + sampler_factory["dpmpp_2m_karras"] = { + 'scheduler': 'DPMSolverMultistepScheduler', + 'name': 'DPMSolverMultistepScheduler', + 'kwargs': { + 'beta_schedule': 'scaled_linear', + 'beta_start': 0.00085, + 'beta_end': 0.03, + 'prediction_type': 'v_prediction', + 'trained_betas': None, + 'solver_order': 2, + 'algorithm_type': 'dpmsolver++', + "use_karras_sigmas": True, + } + } + + # Load sampler from factory + kwargs = sampler_factory[sampler]['kwargs'] + scheduler = sampler_factory[sampler]['scheduler'] + + # Build scheduler according to the sampler. + scheduler_class = getattr(schedulers, scheduler) + scheduler = scheduler_class(**kwargs) + + # Set timesteps for inference steps. + scheduler.set_timesteps(infer_steps, "cuda") + + def calc_rope(height, width): + model_config = HUNYUAN_DIT_CONFIG["DiT-g/2"] + patch_size = model_config['patch_size'] + head_size = model_config['hidden_size'] // model_config['num_heads'] + th = height // 8 // patch_size + tw = width // 8 // patch_size + base_size = 512 // 8 // patch_size + start, stop = get_fill_resize_and_crop((th, tw), base_size) + sub_args = [start, stop, (th, tw)] + rope = get_2d_rotary_pos_embed(head_size, *sub_args) + return rope + + pipeline = StableDiffusionPipeline(vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + unet=model.module, + scheduler=scheduler, + feature_extractor=None, + safety_checker=None, + requires_safety_checker=False, + embedder_t5=embedder_t5, + ) + pipeline = pipeline.to("cuda") + # attr _execution_device is not defined + + style = torch.as_tensor([0, 0] * batch_size, device="cuda") + + src_size_cond = (target_width, target_height) + size_cond = list(src_size_cond) + [target_width, target_height, 0, 0] + image_meta_size = torch.as_tensor( + [size_cond] * 2 * batch_size, device="cuda",) + + if type(sample_config) != list: + sample_config = [sample_config] + + for i, sample in enumerate(sample_config): + prompt = sample.get("prompt", "") + negative_prompt = sample.get("negative_prompt", "") + guidance_scale = sample.get("cfg", guidance_scale) + infer_steps = sample.get("steps", infer_steps) + width = sample.get("width", target_width) + height = sample.get("height", target_height) + + freqs_cis_img = calc_rope(height, width) + + try: + generator = torch.Generator(device="cuda") + generator.manual_seed(seed) + samples = pipeline( + height=height, + width=width, + prompt=prompt, + negative_prompt=negative_prompt, + num_images_per_prompt=batch_size, + guidance_scale=guidance_scale, + num_inference_steps=infer_steps, + style=style, + return_dict=False, + use_fp16=True, + learn_sigma=args.learn_sigma, + freqs_cis_img=freqs_cis_img, + image_meta_size=image_meta_size, + generator=generator, + )[0] + + pass + except Exception as e: + print(f"Failed to sample images: {e} ") + # 打印堆栈信息 + traceback.print_exc() + print(f"Failed to sample pipeline: {pipeline} ") + + # print("samples:",type(samples),) + # input("Press Enter to continue...") + # print("samples:",samples,) + + if type(samples) == list: + pil_image = samples[0] + else: + pil_image = samples + + sample_filename = f"{args.task_flag}_train_steps_{train_steps:07d}.png" + sample_filename_path = os.path.join( + sample_images_dir, sample_filename) + pil_image.save(sample_filename_path) + + return None + + +def model_resume(args, model, ema, logger): + """ + Load pretrained weights. + """ + start_epoch = 0 + start_epoch_step = 0 + train_steps = 0 + resume_path = UNET_PATH + + logger.info(f"Resume from checkpoint {resume_path}") + + if args.resume_split: + # Resume main model + + resume_ckpt_module = torch.load( + resume_path, map_location=lambda storage, loc: storage) + model.load_state_dict(resume_ckpt_module, strict=False) + + # Resume ema model + if args.use_ema: + if args.module_to_ema: + if "resume_ckpt_module" in locals(): + logger.info(f" Resume ema model from main states.") + ema.load_state_dict(resume_ckpt_module, strict=args.strict) + else: + logger.info(f" Resume ema model from module states.") + resume_ckpt_module = torch.load( + resume_path, map_location=lambda storage, loc: storage) + ema.load_state_dict(resume_ckpt_module, strict=args.strict) + else: + if "resume_ckpt_ema" in locals(): + logger.info(f" Resume ema model from EMA states.") + ema.load_state_dict(resume_ckpt_ema, strict=args.strict) + else: + logger.info(f" Resume ema model from EMA states.") + resume_ckpt_ema = torch.load(resume_path, + map_location=lambda storage, loc: storage) + ema.load_state_dict(resume_ckpt_ema, strict=args.strict) + else: + raise ValueError( + " “If `resume` is True, then either `resume_split` must be true.”") + + return model, ema, start_epoch, start_epoch_step, train_steps + + +import tqdm + +import requests + + +def set_master_port(port): + global master_port + master_port = port + + +master_port = 0 + + +def LOG(log): + if master_port == 0: + raise Exception("master_port is 0") + # 发送http + try: + resp = requests.request("post", f"http://127.0.0.1:{master_port}/log", data=json.dumps(log), headers={ + "Content-Type": "application/json"}) + except Exception as e: + return + if resp.status_code != 200: + raise Exception(f"LOG failed: {resp.text}") + + +# with tqdm(total=total_steps, initial=train_steps) as pbar: +# pbar.update(1) +# pbar.set_description( +# f"Epoch {epoch}, step {step}, loss {loss.item():.4f}, mean_loss {mean_loss / step:.4f}") +class PBar: + def __init__(self, total): + self.pbar = tqdm.tqdm(total=total) + + def step(self, desc, total_steps, train_steps): + self.pbar.update(1) + self.pbar.set_description(desc) + + LOG({ + "type": "sample_images", + "global_step": train_steps, + "total_steps": total_steps, + # "latent": noise_pred_latent_path, + }) + + +from torch import nn + + +class CustomizeEmbedsModel(nn.Module): + dtype = torch.float16 + # x = torch.zeros(1, 1, 256, 2048) + x = None + + def __init__(self, *args, **kwargs): + super().__init__() + + def to(self, *args, **kwargs): + self.dtype = torch.float16 + return self + + def forward(self, *args, **kwargs): + input_ids = kwargs.get("input_ids", None) + if self.x is None: + if input_ids is None: + batch_size = 1 + else: + batch_size = input_ids.shape[0] + self.x = torch.zeros(1, batch_size, 256, 2048, dtype=self.dtype) + + if kwargs.get("output_hidden_states", False): + return { + "hidden_states": self.x.to("cuda"), + "input_ids": torch.zeros(1, 1), + } + return self.x + + +class CustomizeTokenizer(dict): + + added_tokens_encoder = [] + input_ids = torch.zeros(1, 256) + attention_mask = torch.zeros(1, 256) + + def __init__(self, *args, **kwargs): + self['added_tokens_encoder'] = self.added_tokens_encoder + self['input_ids'] = self.input_ids + self['attention_mask'] = self.attention_mask + + def tokenize(self, text): + return text + + def __call__(self, *args, **kwargs): + return self + + +class CustomizeEmbeds(): + def __init__(self): + super().__init__() + self.tokenizer = CustomizeTokenizer() + self.model = CustomizeEmbedsModel().to("cuda") + self.max_length = 256 diff --git a/hook_kohya_ss_run.py b/hook_kohya_ss_run.py index 85d995b..a61439b 100644 --- a/hook_kohya_ss_run.py +++ b/hook_kohya_ss_run.py @@ -48,6 +48,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] @@ -63,27 +67,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": "", @@ -91,7 +83,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: @@ -103,6 +97,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, @@ -110,8 +105,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({ @@ -166,9 +161,55 @@ 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.__dict__["_diffusers_version"] = "0.6.0" + 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 + library.train_util.sample_images = utils_sample_images + + 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..db1a3e0 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,12 +112,16 @@ 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: if force_clear_only_images: - for file in os.listdir(train_images_dir): - if file.endswith(".png") or file.endswith(".jpg"): + images_files = Utils.listdir(train_images_dir) + for file in images_files: + if file.lower().endswith(".png") or file.lower().endswith(".jpg"): os.remove(os.path.join(train_images_dir, file)) else: shutil.rmtree(train_images_dir) @@ -120,10 +130,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 +152,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 +311,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 +325,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, }, @@ -336,9 +358,9 @@ def get_sample_images(train_config): pil_images = [] pre_render_texts_x = [] if os.path.exists(sample_images_dir): - image_files = os.listdir(sample_images_dir) + image_files = Utils.listdir(sample_images_dir) image_files = list( - filter(lambda x: x.endswith(".png"), image_files)) + filter(lambda x: x.lower().endswith(".png"), image_files)) # 筛选 output_name 前缀 image_files = list( filter(lambda x: x.startswith(output_name), image_files)) @@ -374,11 +396,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 +420,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)) @@ -466,7 +484,7 @@ def MZ_KohyaSSTrain_call(args={}): 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) + workspace_lora_files = Utils.listdir(workspace_lora_dir) workspace_lora_files = list( filter(lambda x: x.endswith(".safetensors"), workspace_lora_files)) workspace_lora_files = list( @@ -494,14 +512,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 = Utils.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": @@ -517,6 +559,18 @@ 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": + + conditioning_images_dir = os.path.join( + workspace_dir, "conditioning_images") + conditioning_images_onec = "" + if os.path.exists(conditioning_images_dir): + conditioning_images_onec = Utils.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: raise Exception( f"暂时不支持的训练类型: {train_type}") diff --git a/mz_train_tools_core_HYDiT.py b/mz_train_tools_core_HYDiT.py new file mode 100644 index 0000000..90359a8 --- /dev/null +++ b/mz_train_tools_core_HYDiT.py @@ -0,0 +1,1124 @@ + + +import argparse +import csv +import hashlib +import os +import shutil +import sys +import json +import subprocess +import time + +import safetensors.torch +import torch + + +from .mz_train_tools_utils import Utils +import folder_paths +import nodes + + +git_accelerate_urls = { + "githubfast": "githubfast.com", + "521github": "521github.com", + "kkgithub": "kkgithub.com", +} + +# 初始化工具仓库和工作区 + + +def MZ_HYDiTInitClone_call(args={}): + mz_dir = Utils.get_minus_zone_models_path() + git_url = "https://github.com/Tencent/HunyuanDiT" + source = args.get("source", "github") + hunyuan_lora_dir = os.path.join(mz_dir, "train_tools", "HunyuanDiT") + if git_accelerate_urls.get(source, None) is not None: + git_url = f"https://{git_accelerate_urls[source]}/Tencent/HunyuanDiT" + try: + if not os.path.exists(hunyuan_lora_dir) or not os.path.exists(os.path.join(hunyuan_lora_dir, ".git")): + subprocess.run( + ["git", "clone", "--depth", "1", git_url, hunyuan_lora_dir], check=True) + + # 切换远程分支 git remote set-branches origin 'main' + branch = args.get("branch", "main") + + # 查看本地分支是否一致 + short_result = subprocess.run( + ["git", "rev-parse", "--abbrev-ref", "HEAD"], cwd=hunyuan_lora_dir, stdout=subprocess.PIPE, check=True) + result = subprocess.run( + ["git", "rev-parse", "HEAD"], cwd=hunyuan_lora_dir, stdout=subprocess.PIPE, check=True) + + short_current_branch = short_result.stdout.decode().strip() + long_current_branch = result.stdout.decode().strip() + print( + f"当前分支(current branch): {long_current_branch}({short_current_branch})") + print(f"目标分支(target branch): {branch}") + time.sleep(1) + if branch != long_current_branch and branch != short_current_branch: + subprocess.run( + ["git", "remote", "set-branches", "origin", branch], cwd=hunyuan_lora_dir, check=True) + subprocess.run( + ["git", "fetch", "--depth", "1", "origin", branch], cwd=hunyuan_lora_dir, check=True) + + # 恢复所有文件 + subprocess.run( + ["git", "checkout", "."], cwd=hunyuan_lora_dir, check=True) + + subprocess.run( + ["git", "checkout", branch], cwd=hunyuan_lora_dir, check=True) + + content = None + with open(os.path.join(hunyuan_lora_dir, "hydit/diffusion/pipeline.py"), "r", encoding="utf-8") as f: + pre_replace = "device = self._execution_device" + content = f.read() + content = content.replace( + pre_replace, "device = torch.device('cuda')") + with open(os.path.join(hunyuan_lora_dir, "hydit/diffusion/pipeline.py"), "w", encoding="utf-8") as f: + f.write(content) + + except Exception as e: + raise Exception(f"克隆kohya-ss/sd-scripts或者切换分支时出现异常,详细信息请查看控制台...") + + +def MZ_HYDiTInitWorkspace_call(args={}): + MZ_HYDiTInitClone_call(args) + + workspace_name = args.get("train_name", None) + workspace_name = workspace_name.strip() + + if workspace_name is None or workspace_name == "": + raise Exception("训练名称不能为空(train_name is required)") + + args["workspace_name"] = workspace_name + workspaces_dir = os.path.join( + folder_paths.output_directory, "mz_train_workspaces") + + os.makedirs(workspaces_dir, exist_ok=True) + + workspace_dir = os.path.join(workspaces_dir, workspace_name) + if not os.path.exists(workspace_dir): + os.makedirs(workspace_dir) + + args["workspace_dir"] = workspace_dir + + return ( + args, + ) + + +def MZ_HYDiTDatasetConfig_call(args={}): + workspace_config = args.get("workspace_config", {}) + workspace_name = workspace_config.get("workspace_name", None) + workspace_dir = os.path.join( + folder_paths.output_directory, "mz_train_workspaces", workspace_name) + train_images_dir = os.path.join(workspace_dir, "train_images") + os.makedirs(train_images_dir, exist_ok=True) + + images = args.get("images", None) + + if images is None: + print(f"训练数据位于:{train_images_dir},未检测到传入图片默认已有训练数据,将直接跳过数据准备步骤...") + same_caption_generate = args.get("same_caption_generate") == "enable" + if same_caption_generate: + saved_images_path = Utils.listdir(train_images_dir) + same_caption = args.get("same_caption").strip() + if same_caption != "": + # 循环已经保存的图片 + for i, filename in enumerate(saved_images_path): + if filename.lower().endswith(".png") or filename.lower().endswith(".jpg"): + base_filename = os.path.splitext(filename)[0] + caption_filename = base_filename + ".caption" + with open(os.path.join(train_images_dir, caption_filename), "w", encoding="utf-8") as f: + f.write(same_caption) + + + return ( + train_images_dir, + ) + + pil_images = Utils.tensors2pil_list(images) + + resolution = args.get("resolution", 1024) + + if workspace_name is None or workspace_name == "": + raise Exception("训练名称不能为空(workspace_name is required)") + + force_clear = args.get("force_clear") == "enable" + force_clear_only_images = args.get("force_clear_only_images") == "enable" + if force_clear: + if force_clear_only_images: + images_files = Utils.listdir(train_images_dir) + for file in images_files: + if file.lower().endswith(".png") or file.lower().endswith(".jpg"): + os.remove(os.path.join(train_images_dir, file)) + else: + shutil.rmtree(train_images_dir) + os.makedirs(train_images_dir, exist_ok=True) + + 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) + + same_caption_generate = args.get("same_caption_generate") == "enable" + if same_caption_generate: + same_caption = args.get("same_caption").strip() + if same_caption != "": + # 循环已经保存的图片 + for i, filename in enumerate(saved_images_path): + base_filename = os.path.splitext(filename)[0] + caption_filename = base_filename + ".caption" + with open(os.path.join(train_images_dir, caption_filename), "w", encoding="utf-8") as f: + f.write(same_caption) + + return ( + train_images_dir, + ) + + +HYDiT_MODEL = { + "HunyuanDiT/t2i/model/pytorch_model_ema.pt": { + "url": "https://www.modelscope.cn/models/wailovet/hy1.1/resolve/master/t2i/model/pytorch_model_ema.pt", + }, + "HunyuanDiT/t2i/model/pytorch_model_module.pt": { + "url": "https://www.modelscope.cn/models/wailovet/hy1.1/resolve/master/t2i/model/pytorch_model_module.pt", + }, + "HunyuanDiT/t2i/sdxl-vae-fp16-fix/diffusion_pytorch_model.safetensors": { + "url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fsdxl-vae-fp16-fix%2Fdiffusion_pytorch_model.safetensors", + }, + "HunyuanDiT/t2i/sdxl-vae-fp16-fix/config.json": { + "url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fsdxl-vae-fp16-fix%2Fconfig.json", + }, + "HunyuanDiT/t2i/clip_text_encoder/pytorch_model.bin": { + "url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fclip_text_encoder%2Fpytorch_model.bin", + }, + "HunyuanDiT/t2i/clip_text_encoder/config.json": { + "url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fclip_text_encoder%2Fconfig.json", + }, + "HunyuanDiT/t2i/tokenizer/special_tokens_map.json": { + "url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Ftokenizer%2Fspecial_tokens_map.json", + }, + "HunyuanDiT/t2i/tokenizer/tokenizer_config.json": { + "url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Ftokenizer%2Ftokenizer_config.json", + }, + "HunyuanDiT/t2i/tokenizer/vocab.txt": { + "url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Ftokenizer%2Fvocab.txt", + }, + "HunyuanDiT/t2i/tokenizer/vocab_org.txt": { + "url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Ftokenizer%2Fvocab_org.txt", + }, + "HunyuanDiT/t2i/mt5/config.json": { + "url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fmt5%2Fconfig.json", + }, + "HunyuanDiT/t2i/mt5/generation_config.json": { + "url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fmt5%2Fgeneration_config.json", + }, + "HunyuanDiT/t2i/mt5/special_tokens_map.json": { + "url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fmt5%2Fspecial_tokens_map.json", + }, + "HunyuanDiT/t2i/mt5/spiece.model": { + "url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fmt5%2Fspiece.model", + }, + "HunyuanDiT/t2i/mt5/tokenizer_config.json": { + "url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fmt5%2Ftokenizer_config.json", + }, + "HunyuanDiT/t2i/mt5/pytorch_model.bin": { + "url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fmt5%2Fpytorch_model.bin", + }, +} + + +def check_model_auto_download(args): + unet_path = args.get("unet_path", "auto") + hunyuan_base_path = os.path.join( + Utils.get_comfyui_models_path(), "hunyuan") + if unet_path == "auto": + ema_to_module = args.get("ema_to_module", "enable") == "enable" + if ema_to_module: + download_file = "HunyuanDiT/t2i/model/pytorch_model_ema.pt" + else: + download_file = "HunyuanDiT/t2i/model/pytorch_model_module.pt" + download_fullpath = os.path.join( + hunyuan_base_path, download_file) + if not os.path.exists(download_fullpath): + os.makedirs(os.path.dirname(download_fullpath), exist_ok=True) + + if os.path.exists(download_fullpath): + args["unet_path"] = download_fullpath + else: + success_path = Utils.download_file( + HYDiT_MODEL[download_file]["url"], + download_fullpath, + ) + if os.path.exists(success_path): + args["unet_path"] = download_fullpath + + vae_ema_path = args.get("vae_ema_path", "auto") + if vae_ema_path == "auto": + download_files = [] + prefix = "HunyuanDiT/t2i/sdxl-vae-fp16-fix/" + for key in HYDiT_MODEL: + if key.startswith(prefix): + download_files.append(key) + + for download_file in download_files: + download_fullpath = os.path.join( + hunyuan_base_path, download_file) + if not os.path.exists(download_fullpath): + os.makedirs(os.path.dirname(download_fullpath), exist_ok=True) + + if not os.path.exists(download_fullpath): + Utils.download_file( + HYDiT_MODEL[download_file]["url"], + download_fullpath, + ) + + args["vae_ema_path"] = os.path.join( + hunyuan_base_path, prefix) + + text_encoder_path = args.get("text_encoder_path", "auto") + if text_encoder_path == "auto": + download_files = [] + prefix = "HunyuanDiT/t2i/clip_text_encoder/" + for key in HYDiT_MODEL: + if key.startswith(prefix): + download_files.append(key) + + for download_file in download_files: + download_fullpath = os.path.join( + hunyuan_base_path, download_file) + if not os.path.exists(download_fullpath): + os.makedirs(os.path.dirname(download_fullpath), exist_ok=True) + + if not os.path.exists(download_fullpath): + Utils.download_file( + HYDiT_MODEL[download_file]["url"], + download_fullpath, + ) + + args["text_encoder_path"] = os.path.join( + hunyuan_base_path, prefix) + + tokenizer_path = args.get("tokenizer_path", "auto") + if tokenizer_path == "auto": + download_files = [] + prefix = "HunyuanDiT/t2i/tokenizer/" + for key in HYDiT_MODEL: + if key.startswith(prefix): + download_files.append(key) + + for download_file in download_files: + download_fullpath = os.path.join( + hunyuan_base_path, download_file) + if not os.path.exists(download_fullpath): + os.makedirs(os.path.dirname(download_fullpath), exist_ok=True) + + if not os.path.exists(download_fullpath): + Utils.download_file( + HYDiT_MODEL[download_file]["url"], + download_fullpath, + ) + + args["tokenizer_path"] = os.path.join( + hunyuan_base_path, prefix) + + t5_encoder_path = args.get("t5_encoder_path", "auto") + if t5_encoder_path == "auto": + download_files = [] + prefix = "HunyuanDiT/t2i/mt5/" + for key in HYDiT_MODEL: + if key.startswith(prefix): + download_files.append(key) + + for download_file in download_files: + download_fullpath = os.path.join( + hunyuan_base_path, download_file) + if not os.path.exists(download_fullpath): + os.makedirs(os.path.dirname(download_fullpath), exist_ok=True) + + if not os.path.exists(download_fullpath): + Utils.download_file( + HYDiT_MODEL[download_file]["url"], + download_fullpath, + ) + + args["t5_encoder_path"] = os.path.join( + hunyuan_base_path, prefix) + + return args + + +def check_required(): + packages = [ + "pandas", "pyarrow", "diffusers", "transformers", + "timm", "peft", "accelerate", "loguru", "einops", "sentencepiece", + "polygraphy", "protobuf" + ] + for package in packages: + try: + __import__(package) + except ImportError: + subprocess.run([sys.executable, "-m", "pip", + "install", package], check=True) + + try: + import deepspeed + except ImportError: + raise Exception("请手动安装合适的deepspeed版本") + + +def MZ_HYDiTTrain_call(args={}): + + check_required() + + advanced_config = args.get("advanced_config", {}) + fix_float_keys = ["lr", "warmup_min_lr", + "warmup_max_lr", "weight_decay", "warmup_num_steps",] + for key in fix_float_keys: + if key in advanced_config: + if isinstance(advanced_config[key], str): + advanced_config[key] = float(advanced_config[key]) + + args = check_model_auto_download(args) + # raise Exception(args) + resolution = args.get("resolution") + + workspace_config = args.get("workspace_config", {}) + workspace_name = workspace_config.get("workspace_name", None) + workspace_dir = workspace_config.get("workspace_dir", None) + + workspace_images_dir = args.get("workspace_images_dir", None) + if workspace_images_dir is None: + workspace_images_dir = os.path.join(workspace_dir, "train_images") + + full_filenames = Utils.listdir(workspace_images_dir) + + # 创建image_text.csv文件 + csv_filename = os.path.join(workspace_dir, "image_text.csv") + + if os.path.exists(csv_filename): + os.remove(csv_filename) + + max_width, max_height, min_width, min_height = 0, 0, 99999, 99999 + + from PIL import Image + + with open(csv_filename, "w", encoding="utf-8", newline="") as file: + writer = csv.writer(file) + + writer.writerow(["image_path", "text_zh"]) + for filename in full_filenames: + print(f"处理文件: {filename}") + if filename.lower().endswith(".png") or filename.lower().endswith(".jpg"): + + image_path = os.path.join(workspace_images_dir, filename) + pil_image = Image.open(image_path) + width, height = pil_image.size + if width > max_width: + max_width = width + if height > max_height: + max_height = height + if width < min_width: + min_width = width + if height < min_height: + min_height = height + + text_zh = "" + + # 去除后缀 + _caption_filename = filename.split(".")[0] + caption_filename = "" + if os.path.exists(os.path.join(workspace_images_dir, _caption_filename + ".txt")): + caption_filename = _caption_filename + ".txt" + elif os.path.exists(os.path.join(workspace_images_dir, _caption_filename + ".caption")): + caption_filename = _caption_filename + ".caption" + else: + caption_filename = "" + if caption_filename != "": + with open(os.path.join(workspace_images_dir, caption_filename), "r") as f: + caption = f.read() + text_zh = caption + + writer.writerow([image_path, text_zh]) + + if resolution > max_width and resolution > max_height: + print("resolution 大于所有图片的宽高,强制使用最大宽高作为resolution") + resolution = max(max_width, max_height) // 8 * 8 + + print(f"csv文件已生成: {csv_filename}") + + HYDiT_tool_dir = os.path.join( + Utils.get_minus_zone_models_path(), "train_tools", "HunyuanDiT") + + # python ./hydit/data_loader/csv2arrow.py ./dataset/porcelain/csvfile/image_text.csv ./dataset/porcelain/arrows + arrows_dir = os.path.join(workspace_dir, "arrows") + if os.path.exists(arrows_dir): + shutil.rmtree(arrows_dir) + + os.makedirs(arrows_dir, exist_ok=True) + csv2arrow_exec = os.path.join( + HYDiT_tool_dir, "hydit", "data_loader", "csv2arrow.py") + try: + subprocess.run([ + sys.executable, csv2arrow_exec, csv_filename, arrows_dir + ], check=True) + except Exception as e: + raise Exception(f"生成arrow文件时出现异常,详细信息请查看控制台...") + + arrows_files = Utils.listdir(arrows_dir) + + dataset_yaml_path = os.path.join(workspace_dir, "dataset.yaml") + + import yaml + with open(dataset_yaml_path, "w", encoding="utf-8") as f: + yaml.dump({ + "source": [os.path.join(arrows_dir, arrow_file) for arrow_file in arrows_files], + }, f) + + print(f"dataset.yaml文件已生成: {dataset_yaml_path}") + + idk_exec = os.path.join( + os.path.dirname(__file__), "hook_HYDiT_idk_run.py") + + try: + from index_kits import __version__ + except Exception as e: + subprocess.run([ + sys.executable, "-m", "pip", "install", "-e", os.path.join( + HYDiT_tool_dir, "IndexKits") + ], check=True) + # idk base -c dataset/yamls/porcelain.yaml -t dataset/porcelain/jsons/porcelain.json + + dataset_json = os.path.join(workspace_dir, "dataset.json") + try: + subprocess.run([ + sys.executable, idk_exec, "base", "-c", dataset_yaml_path, "-t", dataset_json, + ], check=True) + except Exception as e: + raise Exception(f"生成index文件时出现异常,详细信息请查看控制台...") + + print(f"index文件已生成: {dataset_json}") + + dataset_mt_yaml_path = os.path.join(workspace_dir, "dataset_mt.yaml") + + with open(dataset_mt_yaml_path, "w", encoding="utf-8") as f: + yaml.dump({ + "src": [dataset_json], + "base_size": resolution, + "target_ratios": ["1:1", "3:4", "4:3", "16:9", "9:16"], + }, f) + + print(f"dataset_mt.yaml文件已生成: {dataset_mt_yaml_path}") + + # idk multireso -c dataset/yamls/porcelain_mt.yaml -t dataset/porcelain/jsons/porcelain_mt.json + + dataset_mt_json = os.path.join(workspace_dir, "dataset_mt.json") + try: + subprocess.run([ + sys.executable, idk_exec, "multireso", "-c", dataset_mt_yaml_path, "-t", dataset_mt_json, + ], check=True) + except Exception as e: + raise Exception(f"生成多分辨率index文件时出现异常,详细信息请查看控制台...") + + print(f"多分辨率index文件已生成: {dataset_mt_json}") + + train_config = {} + + train_config_path = os.path.join(workspace_dir, "mz_train_config.json") + + if os.path.exists(train_config_path): + with open(train_config_path, "r", encoding="utf-8") as f: + train_config = json.load(f) + +# "unet": (["auto"] + models + unet_models, {"default": "auto"}), +# "vae": (["auto"] + folders + vae_models, {"default": "auto"}), +# "text_encoder": (["auto"] + folders, {"default": "auto"}), +# "tokenizer": (["auto"] + folders, {"default": "auto"}), +# "t5_encoder": (["auto"] + folders, {"default": "auto"}), + + epochs = args.get("epochs", 50) + output_dir = os.path.join(workspace_dir, "output") + + datetime = time.strftime("%Y%m%d_%H%M%S", time.localtime()) + output_name = f"{workspace_name}_{datetime}" + + lora_ckpt = args.get("base_lora") + if lora_ckpt == "empty": + lora_ckpt = None + elif lora_ckpt == "latest": + lora_ckpt = None + loras = search_loras([workspace_dir]) + if len(loras) > 0: + lora_ckpt = loras[0] + elif not os.path.exists(lora_ckpt): + raise Exception(f"未找到指定的lora文件: {lora_ckpt}") + + print(f"使用指定的lora文件: {lora_ckpt}") + train_config.update({ + "workspace_name": workspace_name, + "workspace_dir": workspace_dir, + "batch_size": args.get("batch_size", 1), + "ema_to_module": args.get("ema_to_module", "enable") == "enable", + "target_ratios": train_config.get("target_ratios", ["1:1", "3:4", "4:3", "16:9", "9:16"]), + "index_file": dataset_mt_json, + "epochs": epochs, + "ckpt_every": int(args.get("ckpt_every")), + "rope_img": f"base{resolution}", + "rope_real": args.get("rope_real", "enable") == "enable", + "image_size": resolution, + "unet_path": args.get("unet_path"), + "vae_ema_path": args.get("vae_ema_path"), + "text_encoder_path": args.get("text_encoder_path"), + "tokenizer_path": args.get("tokenizer_path"), + "t5_encoder_path": args.get("t5_encoder_path") if args.get("t5_encoder_path") != "none" else None, + "results_dir": output_dir, + "task_flag": output_name, + "lora_ckpt": lora_ckpt, + "rank": int(args.get("rank")), + }) + + if "target_modules" in train_config: + del train_config["target_modules"] + for key in advanced_config: + if key.startswith("target_modules_"): + if advanced_config[key] == "enable": + if "target_modules" not in train_config: + train_config["target_modules"] = [] + train_config["target_modules"].append( + key.replace("target_modules_", "")) + elif advanced_config[key] == "enable" or advanced_config[key] == "disable": + train_config[key] = advanced_config[key] == "enable" + else: + train_config[key] = advanced_config[key] + + sample_generate = args.get("sample_generate", "enable") + sample_prompt = args.get("sample_prompt", "") + sample_config_file = os.path.join(workspace_dir, "sample_config.json") + if sample_generate == "enable": + resolution = args.get("resolution") + sample_config = { + "prompt": sample_prompt, + "negative_prompt": "", + "cfg": 5.0, + "steps": 20, + "width": resolution, + "height": resolution, + } + with open(sample_config_file, "w", encoding="utf-8") as f: + json.dump(sample_config, f, indent=4, ensure_ascii=False) + + train_config["sample_config_file"] = sample_config_file + else: + train_config["sample_config_file"] = None + + with open(train_config_path, "w", encoding="utf-8") as f: + json.dump(train_config, f, indent=4, ensure_ascii=False) + + run_hook_HYDiT_pyexec(train_config_path, train_config, HYDiT_tool_dir) + return ( + args, + ) + + +def run_hook_HYDiT_pyexec(train_config_path, train_config, HYDiT_tool_dir): + pb = Utils.progress_bar(0, "sdxl") + import traceback + import comfy.model_management + stop_server = None + + is_running = True + + def log_callback(log): + try: + comfy.model_management.throw_exception_if_processing_interrupted() + except Exception as e: + stop_server() + return is_running + + try: + resp = log + if resp.get("type") == "sample_images": + total_steps = resp.get("total_steps") + global_step = resp.get("global_step") + xy_img = get_sample_images(train_config) + if xy_img is None: + pb.update( + int(global_step), int(total_steps), None) + return is_running + max_side = max(xy_img.width, xy_img.height) + + pb.update( + int(global_step), int(total_steps), ("JPEG", xy_img, max_side)) + else: + print(f"LOG: {log}") + except Exception as e: + print(f"LOG: {log} e: {e} ") + print(f"stack: {traceback.format_exc()}") + return is_running + stop_server, port = Utils.Simple_Server(log_callback) + + exec_pyfile = os.path.join(os.path.dirname( + __file__), "hook_HYDiT_run.py",) + + try: + subprocess.run( + [sys.executable, "-m", "torch.distributed.launch", "--use_env", exec_pyfile, "--sys_path", HYDiT_tool_dir, + "--mz_master_port", str(port), "--train_config_file", train_config_path], + check=True, + ) + stop_server() + is_running = False + except Exception as e: + stop_server() + is_running = False + raise Exception(f"训练失败!!! 具体报错信息请查看控制台...") + + +def get_sample_images(train_config): + if train_config.get("sample_config_file") is None: + return None + + from PIL import Image + output_name = train_config.get("task_flag") + + workspace_dir = train_config.get("workspace_dir") + + sample_images_dir = os.path.join(workspace_dir, "sample_images") + + pil_images = [] + pre_render_texts_x = [] + if os.path.exists(sample_images_dir): + image_files = Utils.listdir(sample_images_dir) + + image_files = list( + filter(lambda x: x.lower().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 + return None + + +def get_HunYuanDiT_model_from_path(model_path, lora_path, width, height): + + from hydit.modules.models import HUNYUAN_DIT_CONFIG + import importlib + import hydit.modules.models + importlib.reload(hydit.modules.models) + HunYuanDiT = hydit.modules.models.HunYuanDiT + + from types import SimpleNamespace + + # self.text_states_dim = args.text_states_dim + # self.text_states_dim_t5 = args.text_states_dim_t5 + # self.text_len = args.text_len + # self.text_len_t5 = args.text_len_t5 + # self.norm = args.norm + # use_flash_attn = args.infer_mode == 'fa' or args.use_flash_attn + + args = SimpleNamespace() + args.learn_sigma = True + args.text_states_dim = 1024 + args.text_states_dim_t5 = 2048 + args.text_len = 77 + args.text_len_t5 = 256 + args.norm = "layer" + args.infer_mode = "torch" + args.use_fp16 = True + args.qk_norm = True + args.use_flash_attn = False + + model_config = HUNYUAN_DIT_CONFIG["DiT-g/2"] + + latent_size = (width // 8, height // 8) + + model = HunYuanDiT(args, + input_size=latent_size, + **model_config, + ).half().to("cuda") + state_dict = torch.load( + model_path, map_location=lambda storage, loc: storage) + model.load_state_dict(state_dict, strict=False) + + if lora_path is not None: + + lora_path_files = Utils.listdir(lora_path) + safetensors_file = None + config_json_file = None + for file in lora_path_files: + if file.lower().endswith(".safetensors"): + safetensors_file = os.path.join(lora_path, file) + if file.lower().endswith(".json"): + config_json_file = os.path.join(lora_path, file) + + if safetensors_file is None or config_json_file is None: + raise Exception("未找到对应的lora文件") + + from peft import PeftConfig + + # 查询后缀 + from safetensors.torch import load_file + adapter_state_dict = load_file(safetensors_file) + + peft_config = PeftConfig() + + with open(config_json_file, "r", encoding="utf-8") as f: + loaded_attributes = json.load(f) + + for key, value in loaded_attributes.items(): + setattr(peft_config, key, value) + + model.load_adapter( + adapter_state_dict=adapter_state_dict, + peft_config=peft_config, + ) + model.merge_and_unload() + + return model + + +def clean_unet(): + unet_data = Utils.cache_get("HYDiT_UNET") + if unet_data is not None: + unet_data["model"].cpu() + Utils.cache_set("HYDiT_UNET", None) + torch.cuda.empty_cache() + + +def clean_vae(): + vae_data = Utils.cache_get("HYDiT_VAE_EMA") + if vae_data is not None: + vae_data["model"].cpu() + Utils.cache_set("HYDiT_VAE_EM", None) + torch.cuda.empty_cache() + + +def clean_clip_text_encoder(): + clip_text_encoder = Utils.cache_get("HYDiT_CLIP_TEXT_ENCODER") + if clip_text_encoder is not None: + clip_text_encoder["model"].cpu() + Utils.cache_set("HYDiT_CLIP_TEXT_ENCODER", None) + torch.cuda.empty_cache() + + +def clean_t5_encoder(): + t5_encoder = Utils.cache_get("HYDiT_T5_ENCODER") + if t5_encoder is not None: + t5_encoder["model"].cpu() + Utils.cache_set("HYDiT_T5_ENCODER", None) + torch.cuda.empty_cache() + + +from torch import nn + + +class CustomizeEmbedsModel(nn.Module): + dtype = torch.float16 + x = torch.zeros(1, 1, 256, 2048) + + def __init__(self, *args, **kwargs): + super().__init__() + + def to(self, *args, **kwargs): + self.dtype = torch.float16 + return self + + def forward(self, *args, **kwargs): + if kwargs.get("output_hidden_states", False): + return { + "hidden_states": self.x.to("cuda"), + "input_ids": torch.zeros(1, 1), + } + return self.x + + +class CustomizeTokenizer(dict): + + added_tokens_encoder = [] + input_ids = torch.zeros(1, 256) + attention_mask = torch.zeros(1, 256) + + def __init__(self): + self['added_tokens_encoder'] = self.added_tokens_encoder + self['input_ids'] = self.input_ids + self['attention_mask'] = self.attention_mask + + def tokenize(self, text): + return text + + def __call__(self, *args, **kwargs): + return self + + +class CustomizeEmbeds(): + def __init__(self): + super().__init__() + self.tokenizer = CustomizeTokenizer() + self.model = CustomizeEmbedsModel() + self.max_length = 256 + + +def MZ_HYDiTSimpleT2I_call(args={}): + check_required() + MZ_HYDiTInitClone_call(args) + args = check_model_auto_download(args) + + HYDiT_tool_dir = os.path.join( + Utils.get_minus_zone_models_path(), "train_tools", "HunyuanDiT") + if HYDiT_tool_dir not in sys.path: + print(f"add {HYDiT_tool_dir} to sys.path") + sys.path.append(HYDiT_tool_dir) + + from hydit.inference import SAMPLER_FACTORY, HUNYUAN_DIT_CONFIG, get_fill_resize_and_crop, get_2d_rotary_pos_embed + from hydit.modules.text_encoder import MT5Embedder + from hydit.utils.tools import set_seeds + from PIL import Image + from diffusers import schedulers, AutoencoderKL + from transformers import BertModel, BertTokenizer + + SAMPLER_FACTORY["uni_pc"] = { + 'scheduler': 'UniPCMultistepScheduler', + 'name': 'UniPCMultistepScheduler', + 'kwargs': { + 'beta_schedule': 'scaled_linear', + 'beta_start': 0.00085, + 'beta_end': 0.03, + 'prediction_type': 'v_prediction', + 'trained_betas': None, + 'solver_order': 2, + # 'algorithm_type': 'dpmsolver++', + } + } + SAMPLER_FACTORY["dpmpp_2m_karras"] = { + 'scheduler': 'DPMSolverMultistepScheduler', + 'name': 'DPMSolverMultistepScheduler', + 'kwargs': { + 'beta_schedule': 'scaled_linear', + 'beta_start': 0.00085, + 'beta_end': 0.03, + 'prediction_type': 'v_prediction', + 'trained_betas': None, + 'solver_order': 2, + 'algorithm_type': 'dpmsolver++', + "use_karras_sigmas": True, + } + } + + width = args.get("width") + height = args.get("height") + + unet_path = args.get("unet_path") + lora_path = args.get("lora_path") + unet_cache_key = f"HYDiT_UNET" + unet_data = Utils.cache_get(unet_cache_key) + if unet_data is None or unet_data.get("unet_path") != unet_path or unet_data.get("lora_path") != lora_path: + clean_unet() + _lora_path = lora_path if lora_path != "none" else None + unet = get_HunYuanDiT_model_from_path( + args.get("unet_path"), _lora_path, width, height) + Utils.cache_set(unet_cache_key, { + "model": unet, + "unet_path": unet_path, + "lora_path": lora_path, + }) + else: + unet = unet_data["model"] + + text_encoder_path = args.get("text_encoder_path") + clip_text_encoder_data = Utils.cache_get("HYDiT_CLIP_TEXT_ENCODER") + + if clip_text_encoder_data is None or clip_text_encoder_data.get("text_encoder_path") != text_encoder_path: + clean_clip_text_encoder() + clip_text_encoder = BertModel.from_pretrained( + text_encoder_path, False, revision=None, local_files_only=True).to("cuda") + Utils.cache_set("HYDiT_CLIP_TEXT_ENCODER", { + "model": clip_text_encoder, + "text_encoder_path": text_encoder_path, + }) + else: + clip_text_encoder = clip_text_encoder_data["model"] + + tokenizer_path = args.get("tokenizer_path") + text_tokenizer = BertTokenizer.from_pretrained( + tokenizer_path, revision=None, local_files_only=True) + + vae_ema_path = args.get("vae_ema_path") + vae_ema_data = Utils.cache_get("HYDiT_VAE_EMA") + if vae_ema_data is None or vae_ema_data.get("vae_ema_path") != vae_ema_path: + clean_vae() + if os.path.isdir(vae_ema_path): + vae = AutoencoderKL.from_pretrained( + vae_ema_path, local_files_only=True).to("cuda") + else: + vae = AutoencoderKL.from_single_file( + vae_ema_path, local_files_only=True).to("cuda") + + Utils.cache_set("HYDiT_VAE_EMA", { + "model": vae, + "vae_ema_path": vae_ema_path, + }) + else: + vae = vae_ema_data["model"] + + t5_encoder_path = args.get("t5_encoder_path") + if t5_encoder_path == "none": + t5_encoder = CustomizeEmbeds() + else: + t5_encoder_data = Utils.cache_get("HYDiT_T5_ENCODER") + if t5_encoder_data is None or t5_encoder_data.get("t5_encoder_path") != t5_encoder_path: + clean_t5_encoder() + t5_encoder = MT5Embedder( + t5_encoder_path, torch_dtype=torch.float16, max_length=256) + Utils.cache_set("HYDiT_T5_ENCODER", { + "model": t5_encoder, + "t5_encoder_path": t5_encoder_path, + }) + else: + t5_encoder = t5_encoder_data["model"] + + sampler = args.get("scheduler") + scheduler = SAMPLER_FACTORY[sampler]['scheduler'] + kwargs = SAMPLER_FACTORY[sampler]['kwargs'] + # Build scheduler according to the sampler. + scheduler_class = getattr(schedulers, scheduler) + scheduler = scheduler_class(**kwargs) + + scheduler.set_timesteps(args.get("steps"), "cuda") + + from hydit.diffusion.pipeline import StableDiffusionPipeline + pipeline = StableDiffusionPipeline(vae=vae, + text_encoder=clip_text_encoder, + tokenizer=text_tokenizer, + unet=unet, + scheduler=scheduler, + feature_extractor=None, + safety_checker=None, + requires_safety_checker=False, + progress_bar_config=None, + embedder_t5=t5_encoder, + infer_mode="torch", + ) + + pipeline = pipeline.to("cuda") + + seed = args.get("seed", 0) + target_height = int((height // 16) * 16) + target_width = int((width // 16) * 16) + prompt = args.get("prompt", "").strip() + negative_prompt = args.get("negative_prompt", "").strip() + + batch_size = 1 + style = torch.as_tensor([0, 0] * batch_size, device="cuda") + + src_size_cond = (target_width, target_height) + size_cond = list(src_size_cond) + [target_width, target_height, 0, 0] + image_meta_size = torch.as_tensor( + [size_cond] * 2 * batch_size, device="cuda",) + + def calc_rope(height, width): + model_config = HUNYUAN_DIT_CONFIG["DiT-g/2"] + patch_size = model_config['patch_size'] + head_size = model_config['hidden_size'] // model_config['num_heads'] + th = height // 8 // patch_size + tw = width // 8 // patch_size + base_size = 512 // 8 // patch_size + start, stop = get_fill_resize_and_crop((th, tw), base_size) + sub_args = [start, stop, (th, tw)] + rope = get_2d_rotary_pos_embed(head_size, *sub_args) + return rope + + freqs_cis_img = calc_rope(target_height, target_width) + guidance_scale = args.get("cfg", 7.0) + infer_steps = args.get("steps", 20) + + with torch.inference_mode(): + with torch.cuda.amp.autocast(): + + pbar = Utils.progress_bar(0, "sdxl") + preview = pbar.get_previewer() + + def callback(step: int, timestep: int, latents: torch.FloatTensor, pred_x0): + try: + pil_img = preview.decode_latent_to_preview_image( + None, + latents, + )[1] + pbar.update(step, infer_steps, pil_img) + except Exception as e: + print("decode_tensors error:", e) + generator = set_seeds(seed, device="cuda") + samples = pipeline( + height=height, + width=width, + prompt=prompt, + negative_prompt=negative_prompt, + num_images_per_prompt=batch_size, + guidance_scale=guidance_scale, + num_inference_steps=infer_steps, + style=style, + return_dict=False, + generator=generator, + use_fp16=True, + learn_sigma=True, + freqs_cis_img=freqs_cis_img, + image_meta_size=image_meta_size, + callback=callback, + callback_steps=2, + )[0] + + if type(samples) == list: + pil_image = samples[0] + else: + pil_image = samples + + tensor_image = Utils.pil2tensor(pil_image) + + keep_device = args.get("keep_device", "disable") == "enable" + if not keep_device: + del pipeline + del unet + del clip_text_encoder + del vae + del t5_encoder + clean_unet() + clean_clip_text_encoder() + clean_vae() + torch.cuda.empty_cache() + + return (Utils.list_tensor2tensor([tensor_image]),) + + +def search_loras(wdirs): + loras = [] + for wdir in wdirs: + for root, dirs, files in os.walk(wdir): + # 排除隐藏文件夹 + dirs[:] = [d for d in dirs if not d.startswith('.')] + # 存在adapter_config.json文件并且存在adapter_model.safetensors文件 + for file in files: + if file == "adapter_config.json": + lora_path = os.path.join(root, "adapter_model.safetensors") + if os.path.exists(lora_path): + loras.append(root) + + # 按时间倒序 + loras = sorted(loras, key=lambda x: os.path.getmtime(x), reverse=True) + return loras diff --git a/mz_train_tools_utils.py b/mz_train_tools_utils.py index 49ac5eb..04c099d 100644 --- a/mz_train_tools_utils.py +++ b/mz_train_tools_utils.py @@ -67,10 +67,10 @@ class Utils: def listdir_png(path): try: - files = os.listdir(path) + files = Utils.listdir(path) new_files = [] for file in files: - if file.endswith(".png"): + if file.lower().endswith(".png"): new_files.append(file) files = new_files files.sort(key=lambda x: int(os.path.basename(x).split(".")[0])) @@ -78,6 +78,15 @@ class Utils: except Exception as e: return [] + def listdir(path): + try: + files = os.listdir(path) + # 排除.开头的文件 + files = [file for file in files if not file.startswith(".")] + return files + except Exception as e: + return [] + def tensor2pil(image): return Image.fromarray(np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) @@ -676,6 +685,8 @@ class Utils: if os.path.exists(find_fullpath): for root, dirs, files in os.walk(find_fullpath): + # 排除隐藏文件夹 + dirs[:] = [d for d in dirs if not d.startswith('.')] for file in files: if target_sha256 == Utils.file_sha256(os.path.join(root, file)): return os.path.join(root, file) @@ -710,12 +721,12 @@ class Utils: return previewer def update(self, step, total_steps, pil_img=None): - pil_img_info = ("JPEG", pil_img, 512) - if pil_img is None: - pil_img_info = None - if type(pil_img) == Tuple or type(pil_img) == list or type(pil_img) == tuple: - pil_img_info = pil_img try: + pil_img_info = ("JPEG", pil_img, 512) + if pil_img is None: + pil_img_info = None + if type(pil_img) == Tuple or type(pil_img) == list or type(pil_img) == tuple: + pil_img_info = pil_img # print("pil_img_info:", type(pil_img), pil_img_info) # print("step:", step, "total_steps:", total_steps) self.pbar.update_absolute( @@ -787,6 +798,8 @@ class Utils: # 去掉pre_render_texts_x中所有相同字符串前缀 def get_common_prefix(pre_render_texts_x): + if len(pre_render_texts_x) == 0: + return "" common_prefix = "" for i in range(len(pre_render_texts_x[0])): c = pre_render_texts_x[0][i] @@ -913,3 +926,23 @@ class Utils: label_text, font=font, fill="black") return full_canvas + + def get_models_by_folder(dir_path): + models = [] + for root, dirs, files in os.walk(dir_path): + # 排除隐藏文件夹 + dirs[:] = [d for d in dirs if not d.startswith('.')] + for file in files: + if file.endswith(".pth") or file.endswith(".pt") or file.endswith(".pkl") or file.endswith(".onnx") or file.endswith(".safetensors"): + models.append(os.path.join(root, file)) + return models + + def get_folders_by_folder(dir_path): + folders = [] + for root, dirs, files in os.walk(dir_path): + + # 排除隐藏文件夹 + dirs[:] = [d for d in dirs if not d.startswith('.')] + for dir in dirs: + folders.append(os.path.join(root, dir)) + return folders diff --git a/tools/aesthetic.py b/tools/aesthetic.py new file mode 100644 index 0000000..ee3891d --- /dev/null +++ b/tools/aesthetic.py @@ -0,0 +1,210 @@ +import json +import os +import shutil +from PIL import Image +import matplotlib.pyplot as plt + +from warnings import filterwarnings + +import numpy as np +import torch +import pytorch_lightning as pl +import torch.nn as nn +from PIL import Image + + +from transformers import CLIPProcessor, CLIPModel +import torch.nn.functional as F + + +class MLP(pl.LightningModule): + def __init__(self, input_size, xcol='emb', ycol='avg_rating'): + super().__init__() + self.input_size = input_size + self.xcol = xcol + self.ycol = ycol + self.layers = nn.Sequential( + nn.Linear(self.input_size, 1024), + # nn.ReLU(), + nn.Dropout(0.2), + nn.Linear(1024, 128), + # nn.ReLU(), + nn.Dropout(0.2), + nn.Linear(128, 64), + # nn.ReLU(), + nn.Dropout(0.1), + + nn.Linear(64, 16), + # nn.ReLU(), + + nn.Linear(16, 1) + ) + + def forward(self, x): + return self.layers(x) + + def training_step(self, batch, batch_idx): + x = batch[self.xcol] + y = batch[self.ycol].reshape(-1, 1) + x_hat = self.layers(x) + loss = F.mse_loss(x_hat, y) + return loss + + def validation_step(self, batch, batch_idx): + x = batch[self.xcol] + y = batch[self.ycol].reshape(-1, 1) + x_hat = self.layers(x) + loss = F.mse_loss(x_hat, y) + return loss + + def configure_optimizers(self): + optimizer = torch.optim.Adam(self.parameters(), lr=1e-3) + return optimizer + + +def normalized(a, axis=-1, order=2): + import numpy as np + + l2 = np.atleast_1d(np.linalg.norm(a, order, axis)) + l2[l2 == 0] = 1 + return a / np.expand_dims(l2, axis) + + +class Aestheic: + def __init__(self, clip_path, aes_model): + + model = MLP(768) # CLIP embedding dim is 768 for CLIP ViT L 14 + + # load the model you trained previously or the model available in this repo + s = torch.load(aes_model) + + model.load_state_dict(s) + + model.to("cuda") + model.eval() + + self.model = model + + self.device = "cuda" + + model2 = CLIPModel.from_pretrained( + clip_path, local_files_only=True).to("cuda") + preprocess = CLIPProcessor.from_pretrained( + clip_path, local_files_only=True) + + self.model2 = model2 + self.preprocess = preprocess + + def prediction(self, img_paths): + if not isinstance(img_paths, list): + img_paths = [img_paths] + + pil_images = [] + for img_path in img_paths: + pil_image = Image.open(img_path) + pil_images.append(pil_image) + + inputs = self.preprocess( + images=pil_images, return_tensors="pt", + ).to("cuda") + with torch.no_grad(): + image_features = self.model2.get_image_features(**inputs) + print(image_features.shape) + + im_emb_arr = normalized(image_features.cpu().detach().numpy()) + + prediction = self.model(torch.from_numpy( + im_emb_arr).to(self.device).type(torch.cuda.FloatTensor)) + + prediction_score = prediction.cpu().detach().numpy() + prediction_score = prediction_score.flatten() + + images_score = { + img_path: float(score) for img_path, score in zip(img_paths, prediction_score) + } + + # 按value排序 + images_score = dict( + sorted(images_score.items(), key=lambda x: x[1], reverse=True)) + + return images_score + + +import argparse +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("-i", "--input_dir", type=str, default="") + parser.add_argument("-o", "--output_dir", type=str, default="") + parser.add_argument("-c", "--clip_path", type=str, + default="D:\\数据集\\clip-vit-large-patch14") + parser.add_argument("-m", "--aes_model", type=str, + default="D:\\数据集\\sac+logos+ava1-l14-linearMSE.pth") + parser.add_argument("-s", "--score", type=float, default=5.9) + + parser.add_argument("-b", "--batch_size", type=int, default=1) + + args = parser.parse_args() + limitscore = args.score + + aes = Aestheic(args.clip_path, args.aes_model) + + input_dir = args.input_dir + output_dir = args.output_dir + batch_size = args.batch_size + total = 0 + for root, dirs, files in os.walk(input_dir): + # 排除隐藏文件夹 + dirs[:] = [d for d in dirs if not d.startswith('.')] + for file in files: + if not file.endswith(".png") and not file.endswith(".jpg"): + continue + total += 1 + + import tqdm + print(f"total: {total}") + + images = [] + + with tqdm.tqdm(total=total) as pbar: + for root, dirs, files in os.walk(input_dir): + # 排除隐藏文件夹 + dirs[:] = [d for d in dirs if not d.startswith('.')] + for file in files: + if not file.endswith(".png") and not file.endswith(".jpg"): + continue + file_path = os.path.join(root, file) + pbar.update(1) + + pbar.set_description( + f"Processing {os.path.basename(file_path)}") + + if len(images) < batch_size: + images.append(file_path) + else: + image_scroe_map = aes.prediction(images) + images = [] + # print(scroe) + # print(json.dumps(scroe, indent=4, ensure_ascii=False)) + # exit(0) + + for img_path, score in image_scroe_map.items(): + print(f"{img_path}: {score}") + + if score > limitscore: + if not os.path.exists(output_dir): + os.makedirs(output_dir) + shutil.copy(img_path, output_dir) + + if len(images) > 0: + aes.prediction(images) + images = [] + + for img_path, score in image_scroe_map.items(): + print(f"{img_path}: {score}") + + if score > limitscore: + if not os.path.exists(output_dir): + os.makedirs(output_dir) + shutil.copy(img_path, output_dir) + + diff --git a/tools/aesthetic2_5.py b/tools/aesthetic2_5.py new file mode 100644 index 0000000..d51f377 --- /dev/null +++ b/tools/aesthetic2_5.py @@ -0,0 +1,260 @@ +import json +import os +from collections import OrderedDict +from os import PathLike +import shutil +from typing import Final + +import torch +import torch.nn as nn +from transformers import ( + SiglipImageProcessor, + SiglipVisionConfig, + SiglipVisionModel, + logging, +) +from transformers.image_processing_utils import BatchFeature +from transformers.modeling_outputs import ImageClassifierOutputWithNoAttention + +logging.set_verbosity_error() + +URL: Final[str] = ( + "https://github.com/discus0434/aesthetic-predictor-v2-5/raw/main/models/aesthetic_predictor_v2_5.pth" +) + + +class AestheticPredictorV2_5Head(nn.Module): + def __init__(self, config: SiglipVisionConfig) -> None: + super().__init__() + self.scoring_head = nn.Sequential( + nn.Linear(config.hidden_size, 1024), + nn.Dropout(0.5), + nn.Linear(1024, 128), + nn.Dropout(0.5), + nn.Linear(128, 64), + nn.Dropout(0.5), + nn.Linear(64, 16), + nn.Dropout(0.2), + nn.Linear(16, 1), + ) + + def forward(self, image_embeds: torch.Tensor) -> torch.Tensor: + return self.scoring_head(image_embeds) + + +class AestheticPredictorV2_5Model(SiglipVisionModel): + PATCH_SIZE = 14 + + def __init__(self, config: SiglipVisionConfig, *args, **kwargs) -> None: + super().__init__(config, *args, **kwargs) + self.layers = AestheticPredictorV2_5Head(config) + self.post_init() + + def forward( + self, + pixel_values: torch.FloatTensor | None = None, + labels: torch.Tensor | None = None, + return_dict: bool | None = None, + ) -> tuple | ImageClassifierOutputWithNoAttention: + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = super().forward( + pixel_values=pixel_values, + return_dict=return_dict, + ) + image_embeds = outputs.pooler_output + image_embeds_norm = image_embeds / \ + image_embeds.norm(dim=-1, keepdim=True) + prediction = self.layers(image_embeds_norm) + + loss = None + if labels is not None: + loss_fct = nn.MSELoss() + loss = loss_fct() + + if not return_dict: + return (loss, prediction, image_embeds) + + return ImageClassifierOutputWithNoAttention( + loss=loss, + logits=prediction, + hidden_states=image_embeds, + ) + + +class AestheticPredictorV2_5Processor(SiglipImageProcessor): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + + def __call__(self, *args, **kwargs) -> BatchFeature: + return super().__call__(*args, **kwargs) + + @classmethod + def from_pretrained( + self, + pretrained_model_name_or_path: str + | PathLike = "google/siglip-so400m-patch14-384", + *args, + **kwargs, + ) -> "AestheticPredictorV2_5Processor": + return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) + + +def convert_v2_5_from_siglip( + predictor_name_or_path: str | PathLike | None = None, + encoder_model_name: str = "google/siglip-so400m-patch14-384", + *args, + **kwargs, +) -> tuple[AestheticPredictorV2_5Model, AestheticPredictorV2_5Processor]: + model = AestheticPredictorV2_5Model.from_pretrained( + encoder_model_name, *args, **kwargs + ) + + processor = AestheticPredictorV2_5Processor.from_pretrained( + encoder_model_name, *args, **kwargs + ) + + if predictor_name_or_path is None or not os.path.exists(predictor_name_or_path): + state_dict = torch.hub.load_state_dict_from_url( + URL, map_location="cpu") + else: + state_dict = torch.load(predictor_name_or_path, map_location="cpu") + + assert isinstance(state_dict, OrderedDict) + + model.layers.load_state_dict(state_dict) + model.eval() + + return model, processor + + +from PIL import Image + + +class AestheticPredictor: + def __init__(self, clip_path, aesthetic_predictor_path): + # load model and preprocessor + self.model, self.preprocessor = convert_v2_5_from_siglip( + low_cpu_mem_usage=True, + trust_remote_code=True, + encoder_model_name=clip_path, + predictor_name_or_path=aesthetic_predictor_path, + ) + if torch.cuda.is_available(): + self.model = self.model.to(torch.bfloat16).cuda() + + def inference(self, img_paths) -> float: + if not isinstance(img_paths, list): + img_paths = [img_paths] + + pil_images = [] + for img_path in img_paths: + pil_image = Image.open(img_path) + pil_images.append(pil_image) + + # preprocess image + pixel_values = self.preprocessor( + images=pil_images, return_tensors="pt" + ).pixel_values + + if torch.cuda.is_available(): + pixel_values = pixel_values.to(torch.bfloat16).cuda() + + # predict aesthetic score + with torch.inference_mode(): + score = self.model( + pixel_values).logits.squeeze().float().cpu().numpy() + prediction_score = score.flatten() + + images_score = { + img_path: float(score) for img_path, score in zip(img_paths, prediction_score) + } + + # 按value排序 + images_score = dict( + sorted(images_score.items(), key=lambda x: x[1], reverse=True)) + + return images_score + + + +import argparse +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("-i", "--input_dir", type=str, default="") + parser.add_argument("-o", "--output_dir", type=str, default="") + parser.add_argument("-c", "--clip_path", type=str, + default="D:\\数据集\\siglip-so400m-patch14-384") + parser.add_argument("-m", "--aes_model", type=str, + default="D:\\数据集\\aesthetic_predictor_v2_5.pth") + parser.add_argument("-s", "--score", type=float, default=5.9) + + parser.add_argument("-b", "--batch_size", type=int, default=1) + + args = parser.parse_args() + limitscore = args.score + + aes = AestheticPredictor(args.clip_path, args.aes_model) + + input_dir = args.input_dir + output_dir = args.output_dir + batch_size = args.batch_size + total = 0 + for root, dirs, files in os.walk(input_dir): + + # 排除隐藏文件夹 + dirs[:] = [d for d in dirs if not d.startswith('.')] + for file in files: + if not file.endswith(".png") and not file.endswith(".jpg"): + continue + total += 1 + + import tqdm + print(f"total: {total}") + + images = [] + + with tqdm.tqdm(total=total) as pbar: + for root, dirs, files in os.walk(input_dir): + # 排除隐藏文件夹 + dirs[:] = [d for d in dirs if not d.startswith('.')] + for file in files: + if not file.endswith(".png") and not file.endswith(".jpg"): + continue + file_path = os.path.join(root, file) + pbar.update(1) + + pbar.set_description( + f"Processing {os.path.basename(file_path)}") + + if os.path.exists(os.path.join(output_dir, os.path.basename(file_path))): + continue + + if len(images) < batch_size: + images.append(file_path) + else: + image_scroe_map = aes.inference(images) + images = [] + for img_path, score in image_scroe_map.items(): + print(f"{img_path}: {score}") + + if score > limitscore: + if not os.path.exists(output_dir): + os.makedirs(output_dir) + shutil.copy(img_path, output_dir) + + if len(images) > 0: + image_scroe_map = aes.inference(images) + images = [] + + for img_path, score in image_scroe_map.items(): + print(f"{img_path}: {score}") + + if score > limitscore: + if not os.path.exists(output_dir): + os.makedirs(output_dir) + shutil.copy(img_path, output_dir) + + diff --git a/tools/face_resize_copy_dir.py b/tools/face_resize_copy_dir.py new file mode 100644 index 0000000..03eacf5 --- /dev/null +++ b/tools/face_resize_copy_dir.py @@ -0,0 +1,200 @@ +import argparse +import hashlib +import io +import os +from PIL import Image +import numpy as np +import onnxruntime as ort +import cv2 + + +class Utils: + def Md5(bytes): + return hashlib.md5(bytes).hexdigest() + + def resize_by(image, percent): + # 判断类型是否为PIL + if not isinstance(image, Image.Image): + image = Image.fromarray(image) + + width, height = image.size + new_width = int(width * percent) + new_height = int(height * percent) + return image.resize((new_width, new_height), Image.LANCZOS) + + def resize_max(im, dst_w, dst_h): + src_w, src_h = im.size + + if src_h > src_w: + newWidth = dst_w + newHeight = dst_w * src_h // src_w + else: + newWidth = dst_h * src_w // src_h + newHeight = dst_h + + newHeight = newHeight // 8 * 8 + newWidth = newWidth // 8 * 8 + + return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS) + + def resize_min(im, dst_w, dst_h): + src_w, src_h = im.size + + if src_h < src_w: + newWidth = dst_w + newHeight = dst_w * src_h // src_w + else: + newWidth = dst_h * src_w // src_h + newHeight = dst_h + + newHeight = newHeight // 8 * 8 + newWidth = newWidth // 8 * 8 + + return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS) + + # 根据人脸进行剪裁 + def resize_crop(im, dst_w, dst_h, focus_point): + + src_w, src_h = im.size + + if src_h > src_w: + newWidth = dst_w + newHeight = dst_w * src_h // src_w + else: + newWidth = dst_h * src_w // src_h + newHeight = dst_h + + newHeight = int(newHeight // 8 * 8) + newWidth = int(newWidth // 8 * 8) + + im = im.resize((newWidth, newHeight), Image.Resampling.LANCZOS) + + left = (newWidth - dst_w) // 2 + top = (newHeight - dst_h) // 2 + + right = left + dst_w + bottom = top + dst_h + + focus_point_x, focus_point_y = focus_point + offset_w = dst_w // 8 + offset_h = dst_h // 8 + if focus_point_x < left + dst_w // 8: + left = 0 + right = dst_w + elif focus_point_x > right - dst_w // 8: + right = newWidth + left = newWidth - dst_w + if focus_point_y < top + dst_h // 8: + top = 0 + bottom = dst_h + elif focus_point_y > bottom - dst_h // 8: + bottom = newHeight + top = newHeight - dst_h + + return im.crop((left, top, right, bottom)) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("-i", "--input_dir", type=str, default="") + parser.add_argument("-o", "--output_dir", type=str, default="") + parser.add_argument("-m", "--face_model_path", type=str, default="") + parser.add_argument("-r", "--resize", type=int, default=1024) + parser.add_argument("-s", "--score", type=float, default=0.8) + parser.add_argument("-f", "--format", type=str, default="jpg") + + args = parser.parse_args() + limitscore = args.score + sess = ort.InferenceSession(args.face_model_path) + sess.set_providers(['CUDAExecutionProvider']) + input_dir = args.input_dir + output_dir = args.output_dir + resize = args.resize + img_format = args.format + total = 0 + for root, dirs, files in os.walk(input_dir): + # 排除隐藏文件夹 + dirs[:] = [d for d in dirs if not d.startswith('.')] + for file in files: + if not file.endswith(".png") and not file.endswith(".jpg"): + continue + total += 1 + + import tqdm + print(f"total: {total}") + + os.makedirs(output_dir, exist_ok=True) + with tqdm.tqdm(total=total) as pbar: + + # 排除隐藏文件夹 + dirs[:] = [d for d in dirs if not d.startswith('.')] + for root, dirs, files in os.walk(input_dir): + for file in files: + if not file.endswith(".png") and not file.endswith(".jpg"): + continue + file_path = os.path.join(root, file) + pbar.update(1) + + with open(file_path, "rb") as f: + md5_str = Utils.Md5(f.read()) + if os.path.exists(os.path.join(output_dir, md5_str + ".png")): + continue + + image = Image.open(file_path) + image = Utils.resize_min(image, resize, resize) + + cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) + + print(f"start check {file_path}") + scores, bboxes, keypoints, aligned_imgs, landmarks, affine_matrices = sess.run( + None, {"input": cv2_img}) + if len(bboxes) != 1: + continue + + bboxes = bboxes[0] + scores = scores[0] + if limitscore > scores: + continue + + affine_matrices = affine_matrices[0] + angle = np.degrees(np.arctan2( + affine_matrices[1][0], affine_matrices[0][0])) + + real_angle = 0 + if angle < -45 and angle > -135: + # 向左旋转90度 + real_angle = 90 + elif angle > 45 and angle < 135: + # 向右旋转90度 + real_angle = -90 + elif angle < -135 or angle > 135: + # 向左旋转180度 + real_angle = 180 + + if real_angle != 0: + continue + + # print(f"{file_path} bboxes: \n{bboxes} {scores} {real_angle}") + + center_point = (bboxes[0] + bboxes[2]) / \ + 2, (bboxes[1] + bboxes[3]) / 2 + image = Utils.resize_crop( + image, resize * 3 / 4, resize, center_point) + cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) + scores, bboxes, keypoints, aligned_imgs, landmarks, affine_matrices = sess.run( + None, {"input": cv2_img}) + if len(bboxes) != 1: + continue + # cv2.imshow("image", cv2.cvtColor( + # np.array(image), cv2.COLOR_RGB2BGR)) + # cv2.waitKey(0) + if img_format == "png": + output_path = os.path.join(output_dir, md5_str + ".png") + image.save(output_path) + else: + output_path = os.path.join(output_dir, md5_str + ".jpg") + image.save(output_path, "jpeg", quality=95) + + +if __name__ == "__main__": + main() diff --git a/tools/face_resize_copy_dir2.py b/tools/face_resize_copy_dir2.py new file mode 100644 index 0000000..7a78823 --- /dev/null +++ b/tools/face_resize_copy_dir2.py @@ -0,0 +1,245 @@ +import argparse +import hashlib +import io +import os +from PIL import Image +import numpy as np +import onnxruntime as ort +import cv2 + +import piexif + + +class Utils: + def Md5(bytes): + return hashlib.md5(bytes).hexdigest() + + def resize_by(image, percent): + # 判断类型是否为PIL + if not isinstance(image, Image.Image): + image = Image.fromarray(image) + + width, height = image.size + new_width = int(width * percent) + new_height = int(height * percent) + return image.resize((new_width, new_height), Image.LANCZOS) + + def resize_max(im, dst_w, dst_h): + src_w, src_h = im.size + + if src_h > src_w: + newWidth = dst_w + newHeight = dst_w * src_h // src_w + else: + newWidth = dst_h * src_w // src_h + newHeight = dst_h + + newHeight = newHeight // 8 * 8 + newWidth = newWidth // 8 * 8 + + return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS) + + def resize_min(im, dst_w, dst_h): + src_w, src_h = im.size + + if src_h < src_w: + newWidth = dst_w + newHeight = dst_w * src_h // src_w + else: + newWidth = dst_h * src_w // src_h + newHeight = dst_h + + newHeight = newHeight // 8 * 8 + newWidth = newWidth // 8 * 8 + + return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS) + + # 根据人脸进行剪裁 + def resize_crop(im, dst_w, dst_h, focus_point): + + src_w, src_h = im.size + + if src_h > src_w: + newWidth = dst_w + newHeight = dst_w * src_h // src_w + else: + newWidth = dst_h * src_w // src_h + newHeight = dst_h + + newHeight = int(newHeight // 8 * 8) + newWidth = int(newWidth // 8 * 8) + + im = im.resize((newWidth, newHeight), Image.Resampling.LANCZOS) + + left = (newWidth - dst_w) // 2 + top = (newHeight - dst_h) // 2 + + right = left + dst_w + bottom = top + dst_h + + focus_point_x, focus_point_y = focus_point + offset_w = dst_w // 8 + offset_h = dst_h // 8 + if focus_point_x < left + dst_w // 8: + left = 0 + right = dst_w + elif focus_point_x > right - dst_w // 8: + right = newWidth + left = newWidth - dst_w + if focus_point_y < top + dst_h // 8: + top = 0 + bottom = dst_h + elif focus_point_y > bottom - dst_h // 8: + bottom = newHeight + top = newHeight - dst_h + + return im.crop((left, top, right, bottom)) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("-i", "--input_dir", type=str, default="") + parser.add_argument("-o", "--output_dir", type=str, default="") + parser.add_argument("-m", "--face_model_path", type=str, default="") + parser.add_argument("-r", "--resize", type=int, default=1024) + parser.add_argument("-p", "--face_proportion", type=float, default=0.5) + parser.add_argument("-f", "--format", type=str, default="jpg") + + args = parser.parse_args() + sess = ort.InferenceSession(args.face_model_path) + sess.set_providers(['CUDAExecutionProvider']) + input_dir = args.input_dir + output_dir = args.output_dir + resize = args.resize + img_format = args.format + total = 0 + for root, dirs, files in os.walk(input_dir): + # 排除隐藏文件夹 + dirs[:] = [d for d in dirs if not d.startswith('.')] + for file in files: + if not file.endswith(".png") and not file.endswith(".jpg"): + continue + total += 1 + + import tqdm + print(f"total: {total}") + + os.makedirs(output_dir, exist_ok=True) + with tqdm.tqdm(total=total) as pbar: + + # 排除隐藏文件夹 + dirs[:] = [d for d in dirs if not d.startswith('.')] + for root, dirs, files in os.walk(input_dir): + for file in files: + if not file.endswith(".png") and not file.endswith(".jpg"): + continue + file_path = os.path.join(root, file) + pbar.update(1) + + with open(file_path, "rb") as f: + md5_str = Utils.Md5(f.read()) + if os.path.exists(os.path.join(output_dir, md5_str + ".png")): + continue + + image = Image.open(file_path) + max_image = image.copy() + image = Utils.resize_min(image, 1920, 1920) + zoom = image.size[0] / max_image.size[0] + + cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) + + print(f"start check {file_path}") + scores, bboxes, keypoints, aligned_imgs, landmarks, affine_matrices = sess.run( + None, {"input": cv2_img}) + if len(bboxes) != 1: + continue + + bboxes = bboxes[0] + + affine_matrices = affine_matrices[0] + angle = np.degrees(np.arctan2( + affine_matrices[1][0], affine_matrices[0][0])) + + real_angle = 0 + if angle < -45 and angle > -135: + # 向左旋转90度 + real_angle = 90 + elif angle > 45 and angle < 135: + # 向右旋转90度 + real_angle = -90 + elif angle < -135 or angle > 135: + # 向左旋转180度 + real_angle = 180 + + if real_angle != 0: + continue + + # print(f"{file_path} bboxes: \n{bboxes} {scores} {real_angle}") + + bboxes = bboxes / zoom + + image = max_image + + # 预期人脸高度占图片高度的一定比例 + pr = args.face_proportion + face_width = bboxes[2] - bboxes[0] + pre_image_width = face_width / pr + + face_center_x = (bboxes[2] + bboxes[0]) / 2 + crop_x1 = face_center_x - pre_image_width / 2 + crop_x2 = face_center_x + pre_image_width / 2 + + # 判断是否超出图片边界 + if crop_x1 < 0 or crop_x2 > image.width: + print( + f"crop_x1: {crop_x1}, crop_x2: {crop_x2}, image.width: {image.width}") + continue + + face_center_y = (bboxes[3] + bboxes[1]) / 2 + # 总高度和宽度的比例是4:3 + pre_image_height = pre_image_width * 4 / 3 + + # 中点位于高度的0.382比例 + crop_y1 = face_center_y - pre_image_height * 0.382 + crop_y2 = crop_y1 + pre_image_height + + # 判断是否超出图片边界 + if crop_y1 < 0 or crop_y2 > image.height: + print( + f"crop_y1: {crop_y1}, crop_y2: {crop_y2}, image.height: {image.height}") + continue + + # 判断是否比resize小 + if pre_image_width < resize or pre_image_height < resize: + print( + f"pre_image_width: {pre_image_width}, pre_image_height: {pre_image_height}") + continue + + # 进行裁剪 + image = image.crop((crop_x1, crop_y1, crop_x2, crop_y2)) + + # 进行缩放并居中剪裁 + image = image.resize(( + resize * 3 // 4, resize + ), Image.LANCZOS) + + # cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) + # cv2.imshow("image", cv2_img) + # cv2.waitKey(0) + + if img_format == "png": + output_path = os.path.join(output_dir, md5_str + ".png") + image.save(output_path) + else: + output_path = os.path.join(output_dir, md5_str + ".jpg") + image.save(output_path, "jpeg", quality=95) + # 插入IPTC + # piexif.insert(output_path, piexif.ImageIFD, { + # "CameraOwnerName": "CameraOwnerName", + # }) + + print(f"save to {output_path}") + + +if __name__ == "__main__": + main()