From ad50d8a1f1c032c57d0da5cc9642a40b17628dd5 Mon Sep 17 00:00:00 2001 From: wailovet Date: Fri, 16 Aug 2024 03:20:40 +0800 Subject: [PATCH] 'update' --- __init__.py | 649 +--------------- hook_HYDiT_idk_run.py | 155 ---- hook_HYDiT_main_train_deepspeed.py | 595 --------------- hook_HYDiT_run.py | 175 ----- hook_HYDiT_utils.py | 379 ---------- hook_kohya_ss_hunyuan_pipe.py | 326 -------- hook_kohya_ss_run.py | 183 +---- mz_train_tools_core.py | 489 +++--------- mz_train_tools_core_HYDiT.py | 1137 ---------------------------- mz_train_tools_utils.py | 110 ++- 10 files changed, 200 insertions(+), 3998 deletions(-) delete mode 100644 hook_HYDiT_idk_run.py delete mode 100644 hook_HYDiT_main_train_deepspeed.py delete mode 100644 hook_HYDiT_run.py delete mode 100644 hook_HYDiT_utils.py delete mode 100644 hook_kohya_ss_hunyuan_pipe.py delete mode 100644 mz_train_tools_core_HYDiT.py diff --git a/__init__.py b/__init__.py index f7f0d11..ef0ae08 100644 --- a/__init__.py +++ b/__init__.py @@ -67,6 +67,7 @@ class MZ_KohyaSSDatasetConfig: "same_caption": ("STRING", {"default": "", "dynamicPrompts": True, "multiline": True}), "image_format": (["png", "jpg", "webp"], {"default": "webp"}), "dataset_config_extension": ([".toml", ".json"], {"default": ".json"}), + "auto_reg": (["enable", "disable"], {"default": "enable"}), }, "optional": { "conditioning_images": ("IMAGE",), @@ -76,17 +77,12 @@ class MZ_KohyaSSDatasetConfig: RETURN_TYPES = (f"STRING",) RETURN_NAMES = ("workspace_images_dir",) - # OUTPUT_NODE = True - MZ_DESC = """ -如果训练类型是controlnet,必须传入预处理后的图片(conditioning_images) -""" - FUNCTION = "start" CATEGORY = CATEGORY_NAME + "/kohya_ss" def start(self, **kwargs): importlib.reload(mz_train_tools_core) - return mz_train_tools_core.MZ_ImageSelecter_call(kwargs) + return mz_train_tools_core.MZ_KohyaSSDatasetConfig_call(kwargs) NODE_CLASS_MAPPINGS["MZ_ImagesCopyWorkspace"] = MZ_KohyaSSDatasetConfig @@ -97,51 +93,6 @@ NODE_CLASS_MAPPINGS["MZ_KohyaSSDatasetConfig"] = MZ_KohyaSSDatasetConfig NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSDatasetConfig"] = f"{AUTHOR_NAME} - KohyaSSDatasetConfig" -class MZ_KohyaSSUseConfig_oldversion: - train_config_template_dir = os.path.join( - os.path.dirname(__file__), "configs", "kohya_ss_lora" - ) - - @classmethod - def INPUT_TYPES(s): - 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] - return { - "required": { - "workspace_config": ("MZ_TT_SS_WorkspaceConfig",), - "workspace_images_dir": ("STRING", {"forceInput": True}), - "train_config_template": (train_config_templates,), - "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), - "max_train_steps": ("INT", {"default": 0, "min": 0, "max": 0x7fffffff}), - "max_train_epochs": ("INT", {"default": 100, "min": 0, "max": 0x7fffffff}), - "save_every_n_epochs": ("INT", {"default": 10}), - "learning_rate": ("STRING", {"default": "1e-5"}), - }, - "optional": { - "save_advanced_config": ("MZ_TT_SS_AdvConfig",), - } - } - - RETURN_TYPES = (f"MZ_TT_SS_TrainConfig",) - RETURN_NAMES = ("train_config",) - - FUNCTION = "start" - CATEGORY = CATEGORY_NAME + "/kohya_ss" + "/v1" - - def start(self, **kwargs): - importlib.reload(mz_train_tools_core) - kwargs["train_config_template_dir"] = self.train_config_template_dir - return mz_train_tools_core.MZ_KohyaSSUseConfig_call(kwargs) - - -NODE_CLASS_MAPPINGS["MZ_KohyaSSUseConfig"] = MZ_KohyaSSUseConfig_oldversion -NODE_DISPLAY_NAME_MAPPINGS[ - "MZ_KohyaSSUseConfig"] = f"{AUTHOR_NAME} - KohyaSSUseConfig(old version)" - - class MZ_KohyaSSAdvConfig: @classmethod def INPUT_TYPES(s): @@ -232,75 +183,6 @@ class AlwaysEqualProxy(str): return False -class MZ_KohyaSSTrain_oldversion: - - @classmethod - def INPUT_TYPES(s): - loras = [ - "latest", - "empty", - ] - - workspaces_dir = os.path.join( - folder_paths.output_directory, "mz_train_workspaces") - - # 使用walk查询所有的workspace中的所有lora模型,lora存放在每个workspace的output目录下 - workspaces_loras = [] - for root, dirs, files in os.walk(workspaces_dir): - - # 排除隐藏文件夹 - dirs[:] = [d for d in dirs if not d.startswith(".")] - if root.endswith("output"): - for file in files: - if file.endswith(".safetensors"): - workspaces_loras.append( - os.path.join(root, file) - ) - - # 按创建时间排序 - workspaces_loras = sorted( - workspaces_loras, key=lambda x: os.path.getctime(x), reverse=True) - - comfyui_full_loras = [] - comfyui_loras = folder_paths.get_filename_list("loras") - for lora in comfyui_loras: - lora_path = folder_paths.get_full_path("loras", lora) - comfyui_full_loras.append(lora_path) - - # 按创建时间排序 - comfyui_full_loras = sorted( - comfyui_full_loras, key=lambda x: os.path.getctime(x), reverse=True) - - loras = loras + workspaces_loras + comfyui_full_loras - - return { - "required": { - "train_config": ("MZ_TT_SS_TrainConfig",), - "base_lora": (loras, {"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 + "/kohya_ss" + "/v1" - - def start(self, **kwargs): - importlib.reload(mz_train_tools_core) - return mz_train_tools_core.MZ_KohyaSSTrain_call(kwargs) - - -NODE_CLASS_MAPPINGS["MZ_KohyaSSTrain"] = MZ_KohyaSSTrain_oldversion -NODE_DISPLAY_NAME_MAPPINGS[ - "MZ_KohyaSSTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(old version)" - - class MZ_KohyaSSLoraTrain: train_config_template_dir = os.path.join( os.path.dirname(__file__), "configs", "kohya_ss_lora" @@ -367,7 +249,6 @@ class MZ_KohyaSSLoraTrain: return { "required": { "workspace_config": ("MZ_TT_SS_WorkspaceConfig",), - "train_config_template": (train_config_templates,), "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), "max_train_steps": ("INT", {"default": 0, "min": 0, "max": 0x7fffffff}), "max_train_epochs": ("INT", {"default": 100, "min": 0, "max": 0x7fffffff}), @@ -393,25 +274,6 @@ class MZ_KohyaSSLoraTrain: def start(self, **kwargs): importlib.reload(mz_train_tools_core) - # train_config = { - # "workspace_config": kwargs["workspace_config"], - # "train_config_template": kwargs["train_config_template"], - # "ckpt_name": kwargs["ckpt_name"], - # "max_train_steps": kwargs["max_train_steps"], - # "max_train_epochs": kwargs["max_train_epochs"], - # "save_every_n_epochs": kwargs["save_every_n_epochs"], - # "learning_rate": kwargs["learning_rate"], - # } - - # train_config["train_config_template_dir"] = self.train_config_template_dir - - # advanced_config = kwargs.get("advanced_config", None) - - # if advanced_config is not None: - # for k, v in advanced_config.items(): - # train_config[k] = v - - # kwargs["train_config"] = train_config return mz_train_tools_core.MZ_KohyaSSTrain_call(kwargs) @@ -420,212 +282,6 @@ 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 + "/kohya_ss" - - 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_KohyaSS_KohakuBlueleaf_HYHiDLoraTrain: - KohyaSSLoraTrain = MZ_KohyaSSLoraTrain() - - @classmethod - def INPUT_TYPES(s): - models, folders, vae_models, unet_models, loras = HYDiT_paths() - _kohya_input_types = s.KohyaSSLoraTrain.INPUT_TYPES() - - # "workspace_config": kwargs["workspace_config"], - # "train_config_template": kwargs["train_config_template"], - # "ckpt_name": kwargs["ckpt_name"], - # "max_train_steps": kwargs["max_train_steps"], - # "max_train_epochs": kwargs["max_train_epochs"], - # "save_every_n_epochs": kwargs["save_every_n_epochs"], - # "learning_rate": kwargs["learning_rate"], - - kohya_input_types = {"required": {}, "optional": {}} - kohya_input_types["required"]["unet_path"] = ( - ["auto"] + models + unet_models, {"default": "auto"}) - kohya_input_types["required"]["vae_ema_path"] = ( - ["auto"] + folders + vae_models, {"default": "auto"}) - kohya_input_types["required"]["text_encoder_path"] = ( - ["auto"] + folders, {"default": "auto"}) - kohya_input_types["required"]["tokenizer_path"] = ( - ["auto"] + folders, {"default": "auto"}) - kohya_input_types["required"]["t5_encoder_path"] = ( - ["none", "auto"] + folders, {"default": "none"}) - - for k, v in _kohya_input_types["required"].items(): - if k == "ckpt_name": - continue - if k == "ema_to_module": - continue - kohya_input_types["required"][k] = v - - for k, v in _kohya_input_types["optional"].items(): - kohya_input_types["optional"][k] = v - - return kohya_input_types - - RETURN_TYPES = () - RETURN_NAMES = () - OUTPUT_NODE = True - FUNCTION = "start" - CATEGORY = CATEGORY_NAME + "/kohya_ss_HYDiT_kohakublueleaf" - - def start(self, **kwargs): - kwargs["hunyuan_models_config"] = { - "unet_path": kwargs["unet_path"], - "vae_ema_path": kwargs["vae_ema_path"], - "text_encoder_path": kwargs["text_encoder_path"], - "tokenizer_path": kwargs["tokenizer_path"], - "t5_encoder_path": kwargs["t5_encoder_path"], - } - kwargs["ckpt_name"] = None - return self.KohyaSSLoraTrain.start(**kwargs) - - -NODE_CLASS_MAPPINGS["MZ_KohyaSS_KohakuBlueleaf_HYHiDLoraTrain"] = MZ_KohyaSS_KohakuBlueleaf_HYHiDLoraTrain -NODE_DISPLAY_NAME_MAPPINGS[ - "MZ_KohyaSS_KohakuBlueleaf_HYHiDLoraTrain"] = f"{AUTHOR_NAME} - KohyaSS_KohakuBlueleaf_HYHiDLoraTrain" - - -class MZ_KohyaSS_KohakuBlueleaf_HYHiDInitWorkspace: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "lora_name": ("STRING", {"default": ""}), - "branch": ("STRING", {"default": "0dc79edc01f2000de1dad5ad6d20d8b099bfafe2"}), - "source": ([ - "github", - "githubfast", - "521github", - "kkgithub", - ], {"default": "github"}), - "seed": ("INT", {"default": 0}), - }, - } - - RETURN_TYPES = ("MZ_TT_SS_WorkspaceConfig",) - RETURN_NAMES = ("workspace_config",) - - FUNCTION = "start" - CATEGORY = CATEGORY_NAME + "/kohya_ss_HYDiT_kohakublueleaf" - - def start(self, **kwargs): - importlib.reload(mz_train_tools_core) - kwargs["branch_repoid"] = "KohakuBlueleaf/sd-scripts" - kwargs["branch_local_name"] = "KohakuBlueleaf_kohya_ss_lora" - return mz_train_tools_core.MZ_KohyaSSInitWorkspace_call(kwargs) - - -NODE_CLASS_MAPPINGS["MZ_KohyaSS_KohakuBlueleaf_HYHiDInitWorkspace"] = MZ_KohyaSS_KohakuBlueleaf_HYHiDInitWorkspace -NODE_DISPLAY_NAME_MAPPINGS[ - "MZ_KohyaSS_KohakuBlueleaf_HYHiDInitWorkspace"] = f"{AUTHOR_NAME} - KohyaSS_KohakuBlueleaf_HYHiDInitWorkspace" - - -class MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I: - @classmethod - def INPUT_TYPES(s): - models, folders, vae_models, unet_models, _ = HYDiT_paths() - comfyui_full_loras = [] - comfyui_loras = folder_paths.get_filename_list("loras") - for lora in comfyui_loras: - lora_path = folder_paths.get_full_path("loras", lora) - comfyui_full_loras.append(lora_path) - return { - "required": { - "branch": ("STRING", {"default": "0dc79edc01f2000de1dad5ad6d20d8b099bfafe2"}), - "source": ([ - "github", - "githubfast", - "521github", - "kkgithub", - ], {"default": "github"}), - "version": (["1.1", "1.2"], {"default": "1.2"}), - "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": "none"}), - "lora_path": (["none"] + comfyui_full_loras, {"default": "none"}), - "seed": ("INT", {"default": 0}), - "steps": ("INT", {"default": 20}), - "cfg": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}), - "scheduler": ([ - "euler_ancestral", "dpmpp_2m_sde" - ], {"default": "dpmpp_2m_sde"}), - "prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}), - "negative_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}), - "width": ("INT", {"default": 1024, "max": 8192, "step": 16}), - "height": ("INT", {"default": 1024, "max": 8192, "step": 16}), - "keep_device": (["enable", "disable"], {"default": "enable"}), - }, - "optional": { - } - } - - RETURN_TYPES = ("IMAGE",) - RETURN_NAMES = ("images",) - - FUNCTION = "start" - CATEGORY = CATEGORY_NAME + "/kohya_ss_HYDiT_kohakublueleaf" - - def start(self, **kwargs): - importlib.reload(mz_train_tools_core) - kwargs["branch_repoid"] = "KohakuBlueleaf/sd-scripts" - kwargs["branch_local_name"] = "KohakuBlueleaf_kohya_ss_lora" - return mz_train_tools_core.MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I_call(kwargs) - - -NODE_CLASS_MAPPINGS["MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I"] = MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I -NODE_DISPLAY_NAME_MAPPINGS[ - "MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I"] = f"{AUTHOR_NAME} - KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I" - - class MZ_LoadImagesFromDirectoryPath: @classmethod def INPUT_TYPES(s): @@ -668,304 +324,3 @@ class MZ_LoadImagesFromDirectoryPath: 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 + "/HYDiT_native" - - 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 + "/HYDiT_native" - - 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_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 + "/HYDiT_native" - - def start(self, **kwargs): - return (kwargs.copy(),) - - -NODE_CLASS_MAPPINGS["MZ_HYDiTAdvConfig"] = MZ_HYDiTAdvConfig -NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTAdvConfig"] = f"{AUTHOR_NAME} - HYDiTAdvConfig" - - -def HYDiT_paths(): - 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 models, folders, vae_models, unet_models, loras - - -class MZ_HYDiTTrain: - @classmethod - def INPUT_TYPES(s): - - models, folders, vae_models, unet_models, loras = HYDiT_paths() - - 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 + "/HYDiT_native" - - 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 + "/HYDiT_native" - - 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"}), - "index": ("STRING", {"default": ""}), - }, - } - - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("debug",) - - FUNCTION = "start" - - CATEGORY = CATEGORY_NAME - - def start(self, **kwargs): - importlib.reload(mz_train_tools_core) - return mz_train_tools_core.MZ_TrainToolsDebug_call(kwargs.copy()) - - -NODE_CLASS_MAPPINGS["MZ_TrainToolsDebug"] = MZ_TrainToolsDebug -NODE_DISPLAY_NAME_MAPPINGS["MZ_TrainToolsDebug"] = f"{AUTHOR_NAME} - TrainToolsDebug" diff --git a/hook_HYDiT_idk_run.py b/hook_HYDiT_idk_run.py deleted file mode 100644 index 50b8f27..0000000 --- a/hook_HYDiT_idk_run.py +++ /dev/null @@ -1,155 +0,0 @@ -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 deleted file mode 100644 index c25fa9d..0000000 --- a/hook_HYDiT_main_train_deepspeed.py +++ /dev/null @@ -1,595 +0,0 @@ -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 deleted file mode 100644 index 1fe4222..0000000 --- a/hook_HYDiT_run.py +++ /dev/null @@ -1,175 +0,0 @@ - -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 deleted file mode 100644 index d5ac04b..0000000 --- a/hook_HYDiT_utils.py +++ /dev/null @@ -1,379 +0,0 @@ - - -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_hunyuan_pipe.py b/hook_kohya_ss_hunyuan_pipe.py deleted file mode 100644 index 8d5868e..0000000 --- a/hook_kohya_ss_hunyuan_pipe.py +++ /dev/null @@ -1,326 +0,0 @@ - -import inspect -import re -from typing import Callable, List, Optional, Union - -import numpy as np -import PIL.Image -import torch -from packaging import version -from tqdm import tqdm -from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer - -from diffusers import SchedulerMixin, StableDiffusionPipeline -from diffusers.models import AutoencoderKL, UNet2DConditionModel -from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput, StableDiffusionSafetyChecker -from diffusers.utils import logging -from PIL import Image -from library import sdxl_model_util, sdxl_train_util, train_util - -from library.sdxl_lpw_stable_diffusion import * - - -from library.hunyuan_models import * - - -def load_scheduler_sigmas(beta_start=0.00085, beta_end=0.018, num_train_timesteps=1000): - betas = torch.linspace(beta_start**0.5, beta_end**0.5, - num_train_timesteps, dtype=torch.float32) ** 2 - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - - sigmas = np.array(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5) - sigmas = np.concatenate([sigmas[::-1], [0.0]]).astype(np.float32) - sigmas = torch.from_numpy(sigmas) - return alphas_cumprod, sigmas - - -ATTN_MODE = "xformers" -CLIP_TOKENS = 75 * 2 + 2 -DEVICE = "cuda" -try: - from k_diffusion.external import DiscreteVDDPMDenoiser - from k_diffusion.sampling import sample_euler_ancestral, get_sigmas_exponential, sample_dpmpp_2m_sde -except ImportError: - import subprocess - import sys - - subprocess.check_call( - [sys.executable, "-m", "pip", "install", "k-diffusion"]) - - from k_diffusion.external import DiscreteVDDPMDenoiser - from k_diffusion.sampling import sample_euler_ancestral, get_sigmas_exponential, sample_dpmpp_2m_sde - - -from library.hunyuan_utils import get_cond, calc_rope - - -class HuanYuanDiffusionLongPromptWeightingPipeline: - - def __init__( - self, - vae: AutoencoderKL, - text_encoder, - tokenizer, - unet, - scheduler, - # clip_skip: int, - safety_checker, - feature_extractor, - requires_safety_checker=False, - clip_skip=0, - ): - # clip skip is ignored currently - # print("tokenizer: ", tokenizer) - # print("text_encoder: ", text_encoder) - self.unet = unet - self.scheduler = scheduler - self.safety_checker = safety_checker - self.feature_extractor = feature_extractor - self.requires_safety_checker = requires_safety_checker - self.vae = vae - self.vae_scale_factor = 2 ** ( - len(self.vae.config.block_out_channels) - 1) - self.progress_bar = lambda x: tqdm(x, leave=False) - - self.clip_skip = clip_skip - self.tokenizers = tokenizer - self.text_encoders = text_encoder - - def to(self, device=None, dtype=None): - if device is not None: - self.device = device - # self.vae.to(device=self.device) - if dtype is not None: - self.dtype = dtype - - @property - def _execution_device(self): - r""" - Returns the device on which the pipeline's models will be executed. After calling - `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module - hooks. - """ - if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"): - return self.device - for module in self.unet.modules(): - if ( - hasattr(module, "_hf_hook") - and hasattr(module._hf_hook, "execution_device") - and module._hf_hook.execution_device is not None - ): - return torch.device(module._hf_hook.execution_device) - return self.device - - def check_inputs(self, prompt, height, width, strength, callback_steps): - pass - - def get_timesteps(self, num_inference_steps, strength, device, is_text2img): - if is_text2img: - return self.scheduler.timesteps.to(device), num_inference_steps - else: - # get the original timestep using init_timestep - offset = self.scheduler.config.get("steps_offset", 0) - init_timestep = int(num_inference_steps * strength) + offset - init_timestep = min(init_timestep, num_inference_steps) - - t_start = max(num_inference_steps - init_timestep + offset, 0) - timesteps = self.scheduler.timesteps[t_start:].to(device) - return timesteps, num_inference_steps - t_start - - def run_safety_checker(self, image, device, dtype): - return image, None - - def decode_latents(self, latents): - with torch.no_grad(): - latents = 1 / sdxl_model_util.VAE_SCALE_FACTOR * latents - - image = self.vae.decode(latents.to(self.vae.dtype)).sample - image = (image / 2 + 0.5).clamp(0, 1) - # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 - image = image.cpu().permute(0, 2, 3, 1).float().numpy() - return image - - def prepare_extra_step_kwargs(self, generator, eta): - return {} - - @torch.no_grad() - def __call__( - self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, - image: Union[torch.FloatTensor, PIL.Image.Image] = None, - mask_image: Union[torch.FloatTensor, PIL.Image.Image] = None, - height: int = 512, - width: int = 512, - num_inference_steps: int = 50, - guidance_scale: float = 7.5, - strength: float = 0.8, - num_images_per_prompt: Optional[int] = 1, - eta: float = 0.0, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.FloatTensor] = None, - max_embeddings_multiples: Optional[int] = 3, - output_type: Optional[str] = "pil", - return_dict: bool = True, - controlnet=None, - controlnet_image=None, - callback: Optional[Callable[[ - int, int, torch.FloatTensor], None]] = None, - is_cancelled_callback: Optional[Callable[[], bool]] = None, - callback_steps: int = 1, - ): - - BETA_END = 0.018 - CFG_SCALE = guidance_scale - STEPS = num_inference_steps - patch_size = 2 - num_heads = 88 - - alphas, sigmas = load_scheduler_sigmas(beta_end=BETA_END) - denoiser = self.unet - clip_tokenizer = self.tokenizers[0] - clip_encoder = self.text_encoders[0] - - mt5_embedder = self.text_encoders[1] - - vae = self.vae - denoiser.eval() - denoiser.disable_fp32_silu() - denoiser.disable_fp32_layer_norm() - denoiser.set_attn_mode(ATTN_MODE) - vae.requires_grad_(False) - mt5_embedder.to(torch.float16) - - with torch.autocast("cuda"): - clip_h, clip_m, mt5_h, mt5_m = get_cond( - prompt, - mt5_embedder, - clip_tokenizer, - clip_encoder, - # Should be same as original implementation with max_length_clip=77 - # Support 75*n + 2 - max_length_clip=CLIP_TOKENS, - ) - neg_clip_h, neg_clip_m, neg_mt5_h, neg_mt5_m = get_cond( - negative_prompt, - mt5_embedder, - clip_tokenizer, - clip_encoder, - max_length_clip=CLIP_TOKENS, - ) - clip_h = torch.concat([clip_h, neg_clip_h], dim=0) - clip_m = torch.concat([clip_m, neg_clip_m], dim=0) - mt5_h = torch.concat([mt5_h, neg_mt5_h], dim=0) - mt5_m = torch.concat([mt5_m, neg_mt5_m], dim=0) - torch.cuda.empty_cache() - - style = torch.as_tensor([0] * 2, device=DEVICE) - # src hw, dst hw, 0, 0 - size_cond = [height, width, height, width, 0, 0] - image_meta_size = torch.as_tensor([size_cond] * 2, device=DEVICE) - - freqs_cis_img = calc_rope(height, width, patch_size, num_heads) - - denoiser_wrapper = DiscreteVDDPMDenoiser( - # A quick patch for learn_sigma - lambda *args, **kwargs: denoiser(* \ - args, **kwargs).chunk(2, dim=1)[0], - alphas, - False, - ).to(DEVICE) - - def cfg_denoise_func(x, sigma): - cond, uncond = denoiser_wrapper( - x.repeat(2, 1, 1, 1), - sigma.repeat(2), - encoder_hidden_states=clip_h, - text_embedding_mask=clip_m, - encoder_hidden_states_t5=mt5_h, - text_embedding_mask_t5=mt5_m, - image_meta_size=image_meta_size, - style=style, - cos_cis_img=freqs_cis_img[0], - sin_cis_img=freqs_cis_img[1], - ).chunk(2, dim=0) - return uncond + (cond - uncond) * CFG_SCALE - - sigmas = denoiser_wrapper.get_sigmas(STEPS).to(DEVICE) - sigmas = get_sigmas_exponential( - STEPS, denoiser_wrapper.sigma_min, denoiser_wrapper.sigma_max, DEVICE - ) - x1 = torch.randn(1, 4, height // 8, width // 8, - dtype=torch.float16, device=DEVICE) - - with torch.autocast("cuda"): - sample = sample_dpmpp_2m_sde( - cfg_denoise_func, - x1 * sigmas[0], - sigmas, - ) - torch.cuda.empty_cache() - latents = sample - return latents - - def text2img( - self, - prompt: Union[str, List[str]], - negative_prompt: Optional[Union[str, List[str]]] = None, - height: int = 512, - width: int = 512, - num_inference_steps: int = 50, - guidance_scale: float = 7.5, - num_images_per_prompt: Optional[int] = 1, - eta: float = 0.0, - generator: Optional[torch.Generator] = None, - latents: Optional[torch.FloatTensor] = None, - max_embeddings_multiples: Optional[int] = 3, - output_type: Optional[str] = "pil", - return_dict: bool = True, - callback: Optional[Callable[[ - int, int, torch.FloatTensor], None]] = None, - is_cancelled_callback: Optional[Callable[[], bool]] = None, - callback_steps: int = 1, - ): - - return self.__call__( - prompt=prompt, - negative_prompt=negative_prompt, - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - num_images_per_prompt=num_images_per_prompt, - eta=eta, - generator=generator, - latents=latents, - max_embeddings_multiples=max_embeddings_multiples, - output_type=output_type, - return_dict=return_dict, - callback=callback, - is_cancelled_callback=is_cancelled_callback, - callback_steps=callback_steps, - ) - - def latents_to_image(self, latents): - # 9. Post-processing - image = self.decode_latents(latents.to(self.vae.dtype)) - image = self.numpy_to_pil(image) - return image - - # copy from pil_utils.py - def numpy_to_pil(self, images: np.ndarray) -> Image.Image: - """ - Convert a numpy image or a batch of images to a PIL image. - """ - if images.ndim == 3: - images = images[None, ...] - images = (images * 255).round().astype("uint8") - if images.shape[-1] == 1: - # special case for grayscale (single channel) images - pil_images = [Image.fromarray( - image.squeeze(), mode="L") for image in images] - else: - pil_images = [Image.fromarray(image) for image in images] - - return pil_images diff --git a/hook_kohya_ss_run.py b/hook_kohya_ss_run.py index b6364e7..7fa68f1 100644 --- a/hook_kohya_ss_run.py +++ b/hook_kohya_ss_run.py @@ -27,14 +27,6 @@ def config2args(train_parser: argparse.ArgumentParser, config): return args -from PIL import Image - - -import numpy as np -import tempfile -import safetensors.torch - - import sys sys.path.append(os.path.dirname(__file__)) try: @@ -205,182 +197,9 @@ class SimpleNamespaceCNWarrper(SimpleNamespace): 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_args = config2args(train_controlnet.setup_parser(), train_config) - - LOG({ - "type": "start_train", - }) - - train_controlnet.train(train_args) - - -def run_lora_hunyuan1_2(): - hook_kohya_ss_utils.hook_kohya_ss() - - # 覆盖sample_images生成函数,包括进度条和生成图片功能 - import hunyuan_train_network - - # 还未实现 - hunyuan_train_network.HunYuanNetworkTrainer.sample_images = sample_images - - # def empty_sample_images(*args, **kwargs): - # pass - # sample_images_pipe_class = empty_sample_images - - import hook_kohya_ss_hunyuan_pipe - global sample_images_pipe_class - sample_images_pipe_class = hook_kohya_ss_hunyuan_pipe.HuanYuanDiffusionLongPromptWeightingPipeline - # 配置对应的pipeline - import library.hunyuan_utils - - print(json.dumps(other_config, indent=4)) - hunyuan_models_config = other_config.get("hunyuan_models_config", None) - - from transformers import ( - AutoTokenizer, - T5Tokenizer, - BertModel, - BertTokenizer, - ) - from diffusers import AutoencoderKL, LMSDiscreteScheduler - - def hunyuan_load_tokenizers(): - tokenizer = AutoTokenizer.from_pretrained( - hunyuan_models_config["tokenizer_path"], - local_files_only=True, - ) - tokenizer.eos_token_id = tokenizer.sep_token_id - t5_encoder_path = hunyuan_models_config.get("t5_encoder_path", None) - if t5_encoder_path == "none": - t5_encoder_path = None - tokenizer2 = None - if t5_encoder_path is not None: - tokenizer2 = T5Tokenizer.from_pretrained( - t5_encoder_path, - local_files_only=True, - ) - - return [tokenizer, tokenizer2] - - library.hunyuan_utils.load_tokenizers = hunyuan_load_tokenizers - - def hunyuan_load_model(model_path: str, dtype=torch.float16, device="cuda", use_extra_cond=False, dit_path=None): - - dit_path = hunyuan_models_config.get("unet_path", None) - - import library.hunyuan_models - # from hunyuan_models import MT5Embedder, HunYuanDiT, BertModel, DiT_g_2 - MT5Embedder = library.hunyuan_models.MT5Embedder - HunYuanDiT = library.hunyuan_models.HunYuanDiT - BertModel = library.hunyuan_models.BertModel - DiT_g_2 = library.hunyuan_models.DiT_g_2 - - denoiser, patch_size, head_dim = DiT_g_2( - input_size=(128, 128), use_extra_cond=use_extra_cond) - if dit_path is not None: - state_dict = torch.load(dit_path) - if 'state_dict' in state_dict: - state_dict = state_dict['state_dict'] - else: - state_dict = torch.load(os.path.join( - model_path, "denoiser/pytorch_model_module.pt")) - denoiser.load_state_dict(state_dict) - denoiser.to(device).to(dtype) - - clip_tokenizer = AutoTokenizer.from_pretrained( - hunyuan_models_config["tokenizer_path"], - local_files_only=True, - ) - clip_tokenizer.eos_token_id = 2 - clip_encoder = ( - BertModel.from_pretrained( - hunyuan_models_config["text_encoder_path"], - local_files_only=True, - ).to(device).to(dtype) - ) - - t5_encoder_path = hunyuan_models_config.get("t5_encoder_path", None) - if t5_encoder_path == "none": - t5_encoder_path = None - mt5_embedder = None - if t5_encoder_path is not None: - mt5_embedder = ( - MT5Embedder( - model_dir=hunyuan_models_config["t5_encoder_path"], - torch_dtype=dtype, - max_length=256) - .to(device) - .to(dtype) - ) - else: - - batch_size = train_args.train_batch_size - import library.config_util - user_config = library.config_util.load_user_config( - train_args.dataset_config) - datasets = user_config.get("datasets", []) - if len(datasets) > 0: - batch_size = datasets[0].get("batch_size", batch_size) - mt5_embedder = ( - hook_kohya_ss_utils.CustomizeMT5Embedder( - batch_size=batch_size, - ) - .to(device) - .to(dtype) - ) - - vae = ( - AutoencoderKL.from_pretrained( - hunyuan_models_config["vae_ema_path"], - local_files_only=True, - ) - .to(device) - .to(dtype) - ) - vae.requires_grad_(False) - return ( - denoiser, - patch_size, - head_dim, - clip_tokenizer, - clip_encoder, - mt5_embedder, - vae, - ) - - library.hunyuan_utils.load_model = hunyuan_load_model - - trainer = hunyuan_train_network.HunYuanNetworkTrainer() - train_args = config2args( - hunyuan_train_network.setup_parser(), train_config) - print(f"train_args = {train_args}") - - LOG({ - "type": "start_train", - }) - trainer.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, - "run_lora_hunyuan1_2": run_lora_hunyuan1_2, } @@ -450,4 +269,4 @@ if __name__ == "__main__": except Exception as e: print(f"Exception: {e}") if sys.platform == "win32": - input("Press Enter to continue...") \ No newline at end of file + input("Press Enter to continue...") diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index fe4e468..5f12574 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -16,65 +16,11 @@ import folder_paths import nodes -git_accelerate_urls = { - "githubfast": "githubfast.com", - "521github": "521github.com", - "kkgithub": "kkgithub.com", -} - - -def MZ_KohyaSSCloneRepo_call(args={}): - mz_dir = Utils.get_minus_zone_models_path() - - branch_repoid = args.get("branch_repoid", "kohya-ss/sd-scripts") - branch_local_name = args.get("branch_local_name", "kohya_ss_lora") - - git_url = f"https://github.com/{branch_repoid}" - source = args.get("source", "github") - kohya_ss_lora_dir = os.path.join(mz_dir, "train_tools", branch_local_name) - if git_accelerate_urls.get(source, None) is not None: - git_url = f"https://{git_accelerate_urls[source]}/{branch_repoid}" - try: - if not os.path.exists(kohya_ss_lora_dir) or not os.path.exists(os.path.join(kohya_ss_lora_dir, ".git")): - subprocess.run( - ["git", "clone", "--depth", "1", git_url, kohya_ss_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=kohya_ss_lora_dir, stdout=subprocess.PIPE, check=True) - result = subprocess.run( - ["git", "rev-parse", "HEAD"], cwd=kohya_ss_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}") - - if branch != long_current_branch and branch != short_current_branch: - subprocess.run( - ["git", "remote", "set-branches", "origin", branch], cwd=kohya_ss_lora_dir, check=True) - subprocess.run( - ["git", "fetch", "--depth", "1", "origin", branch], cwd=kohya_ss_lora_dir, check=True) - - # 恢复所有文件 - subprocess.run( - ["git", "checkout", "."], cwd=kohya_ss_lora_dir, check=True) - - subprocess.run( - ["git", "checkout", branch], cwd=kohya_ss_lora_dir, check=True) - - except Exception as e: - raise Exception(f"克隆kohya-ss/sd-scripts或者切换分支时出现异常,详细信息请查看控制台...") - # 初始化工具仓库和工作区 def MZ_KohyaSSInitWorkspace_call(args={}): - MZ_KohyaSSCloneRepo_call(args) + Utils.clone_repo(args) workspace_name = args.get("lora_name", None) workspace_name = workspace_name.strip() @@ -97,7 +43,7 @@ def MZ_KohyaSSInitWorkspace_call(args={}): ) -def MZ_ImageSelecter_call(args={}): +def MZ_KohyaSSDatasetConfig_call(args={}): images = args.get("images") pil_images = Utils.tensors2pil_list(images) @@ -175,7 +121,9 @@ def MZ_ImageSelecter_call(args={}): dataset_config_extension = args.get("dataset_config_extension") generate_dataset_config( - os.path.join(workspace_dir, "dataset" + dataset_config_extension), + workspace_dir=workspace_dir, + output_path=os.path.join( + workspace_dir, "dataset" + dataset_config_extension), enable_bucket=args.get("enable_bucket") == "enable", resolution=args.get("resolution"), batch_size=args.get("batch_size"), @@ -183,6 +131,7 @@ def MZ_ImageSelecter_call(args={}): conditioning_data_dir=conditioning_images_dir, caption_extension=caption_extension, num_repeats=args.get("num_repeats"), + is_reg=args.get("auto_reg") == "enable", ) return ( train_images_dir, @@ -343,7 +292,7 @@ def check_install(): import logging -def generate_dataset_config(output_path, enable_bucket=True, resolution=512, batch_size=1, image_dir=None, conditioning_data_dir=None, caption_extension=".caption", num_repeats=10, ): +def generate_dataset_config(workspace_dir, output_path, enable_bucket=True, resolution=512, batch_size=1, image_dir=None, conditioning_data_dir=None, caption_extension=".caption", num_repeats=10, is_reg=False): config = { 'general': { @@ -367,6 +316,13 @@ def generate_dataset_config(output_path, enable_bucket=True, resolution=512, bat if conditioning_data_dir is not None: config["datasets"][0]["subsets"][0]["conditioning_data_dir"] = conditioning_data_dir + if is_reg: + config["datasets"][0]["subsets"].append({ + "is_reg": True, + "image_dir": os.path.join(workspace_dir, "reg_images"), + "num_repeats": 1, + }) + if output_path.endswith(".toml"): check_install() import toml @@ -479,12 +435,8 @@ def run_hook_kohya_ss_run_file(workspace_dir, output_name, kohya_ss_tool_dir, tr from .mz_train_tools_utils import HSubprocess - screen_name = None - if use_screen: - screen_name = "mz_train_tools_core" - process_instance = HSubprocess( - cmd_list, screen_name=screen_name) + cmd_list) process_instance.wait() stop_server() @@ -495,6 +447,24 @@ def run_hook_kohya_ss_run_file(workspace_dir, output_name, kohya_ss_tool_dir, tr raise Exception(f"训练失败!!! 具体报错信息请查看控制台...") +def get_models_type(model_path): + + from comfy.model_detection import model_config_from_unet, unet_prefix_from_state_dict, detect_unet_config + from comfy.supported_models import SDXL, SD15 + import safetensors.torch + sd = safetensors.torch.load_file(model_path, device="cpu") + unet_prefix = unet_prefix_from_state_dict(sd) + + model_type = model_config_from_unet(sd, unet_prefix) + + if isinstance(model_type, SD15): + return "lora_sd1_5" + if isinstance(model_type, SDXL): + return "lora_sdxl" + + raise Exception(f"不支持的模型类型: {model_type}") + + def generate_kohya_ss_config(args): args = args.copy() workspace_config = args.get("workspace_config", {}).copy() @@ -512,7 +482,11 @@ def generate_kohya_ss_config(args): workspace_config_file = os.path.join(workspace_dir, "config.json") - train_config_template = args.get("train_config_template", None) + ckpt_name = args.get("ckpt_name", "") + ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + + train_config_template = get_models_type(ckpt_path) + train_config_template_dir = os.path.join( os.path.dirname(__file__), "configs", "kohya_ss_lora") train_config_template_file = os.path.join( @@ -525,7 +499,6 @@ def generate_kohya_ss_config(args): config["metadata"]["train_type"] = train_config_template ckpt_name = args.get("ckpt_name", "") if ckpt_name != "" and ckpt_name is not None: - ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) config["train_config"]["pretrained_model_name_or_path"] = ckpt_path # output_dir @@ -582,6 +555,41 @@ def generate_kohya_ss_config(args): return config +def generate_regularization_folder(ckpt_name, workspace_dir, caption_extension, + width=512, height=512): + train_images_dir = os.path.join(workspace_dir, "train_images") + saved_images_path = os.listdir(train_images_dir) + + reg_dir = os.path.join(workspace_dir, "reg_images") + + caption_files = [] + caption_value = [] + for i, filename in enumerate(saved_images_path): + if filename.lower().endswith(caption_extension): + caption_files.append(filename) + with open(os.path.join(train_images_dir, filename), "r", encoding="utf-8") as f: + caption_value.append(f.read()) + + if len(caption_files) == 0: + return + os.makedirs(reg_dir, exist_ok=True) + for filename, caption in zip(caption_files, caption_value): + if caption == "": + continue + + with open(os.path.join(reg_dir, filename), "w", encoding="utf-8") as f: + f.write(caption) + + image_filename = filename.replace(caption_extension, ".webp") + Utils.comfyui_fast_generate_picture( + ckpt_name=ckpt_name, + prompt=caption, + output=os.path.join(reg_dir, image_filename), + width=width, + height=height, + ) + + def MZ_KohyaSSTrain_call(args={}): args = args.copy() workspace_config = args.get("workspace_config").copy() @@ -678,8 +686,36 @@ def MZ_KohyaSSTrain_call(args={}): with open(workspace_config_file, "w", encoding="utf-8") as f: json.dump(config, f, indent=4, ensure_ascii=False) + dataset_config_file = config["train_config"]["dataset_config"] + with open(dataset_config_file, "r", encoding="utf-8") as f: + dataset_config = json.load(f) + + if dataset_config is None: + raise Exception("dataset_config is None") + + subsets_list = dataset_config.get("datasets")[0].get("subsets") + is_reg = True if len(subsets_list) == 2 and subsets_list[1].get( + "is_reg") else False + if is_reg: + caption_extension = subsets_list[0].get("caption_extension") + ckpt_name = args.get("ckpt_name") + resolution = dataset_config.get("datasets")[0].get("resolution") + if f"{resolution}".find("x") != -1: + resolutionsp = resolution.split("x") + width = int(resolutionsp[0]) + height = int(resolutionsp[1]) + elif f"{resolution}".find("*") != -1: + resolutionsp = resolution.split("*") + width = int(resolutionsp[0]) + height = int(resolutionsp[1]) + else: + width = height = int(resolution) + + generate_regularization_folder( + ckpt_name, workspace_dir, caption_extension, width=width, height=height) + # raise Exception( - # f"config: {json.dumps(config, indent=4, ensure_ascii=False)}") + # f"config: {json.dumps(dataset_config, indent=4, ensure_ascii=False)}") output_name = config["train_config"].get("output_name") @@ -704,21 +740,6 @@ def MZ_KohyaSSTrain_call(args={}): run_hook_kohya_ss_run_file( workspace_dir, output_name, kohya_ss_tool_dir, "run_controlnet_sd1_5") - elif train_type == "lora_hunyuan1_2" or train_type == "lora_hunyuan1_1": - hunyuan_models_config = args.get( - "hunyuan_models_config", {}) - from .mz_train_tools_core_HYDiT import check_model_auto_download - - hunyuan_models_config["version"] = config.get( - "metadata").get("version") - config["other_config"]["hunyuan_models_config"] = check_model_auto_download( - hunyuan_models_config) - - with open(workspace_config_file, "w", encoding="utf-8") as f: - json.dump(config, f, indent=4, ensure_ascii=False) - - run_hook_kohya_ss_run_file( - workspace_dir, output_name, kohya_ss_tool_dir, "run_lora_hunyuan1_2") else: raise Exception( f"暂时不支持的训练类型: {train_type}") @@ -726,309 +747,3 @@ def MZ_KohyaSSTrain_call(args={}): return ( "训练完成", ) - - -def MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I_call(args={}): - args = args.copy() - MZ_KohyaSSCloneRepo_call(args) - from .mz_train_tools_core_HYDiT import check_model_auto_download - args = check_model_auto_download(args) - import numpy as np - import torch - seed = args.get("seed", 0) - torch.manual_seed(seed) - from packaging import version - from transformers import AutoTokenizer, BertModel - from diffusers.models import AutoencoderKL - try: - from k_diffusion.external import DiscreteVDDPMDenoiser - from k_diffusion.sampling import sample_euler_ancestral, get_sigmas_exponential, sample_dpmpp_2m_sde - except ImportError: - subprocess.check_call( - [sys.executable, "-m", "pip", "install", "k-diffusion"]) - from k_diffusion.external import DiscreteVDDPMDenoiser - from k_diffusion.sampling import sample_euler_ancestral, get_sigmas_exponential, sample_dpmpp_2m_sde - - branch_local_name = args.get("branch_local_name") - kohya_ss_tool_dir = os.path.join( - Utils.get_minus_zone_models_path(), "train_tools", branch_local_name) - if kohya_ss_tool_dir not in sys.path: - sys.path.append(kohya_ss_tool_dir) - from library.hunyuan_models import DiT_g_2, MT5Embedder - from library.hunyuan_utils import get_cond, calc_rope - from networks.lora import create_network_from_weights - - def load_scheduler_sigmas(beta_start=0.00085, beta_end=0.018, num_train_timesteps=1000): - betas = torch.linspace(beta_start**0.5, beta_end ** - 0.5, num_train_timesteps, dtype=torch.float32) ** 2 - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - - sigmas = np.array(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5) - sigmas = np.concatenate([sigmas[::-1], [0.0]]).astype(np.float32) - sigmas = torch.from_numpy(sigmas) - return alphas_cumprod, sigmas - - version = args.get("version") - BETA_END = None - USE_EXTRA_COND = None - if version == "1.1": - BETA_END = 0.03 - USE_EXTRA_COND = True - else: - BETA_END = 0.018 - USE_EXTRA_COND = False - - ATTN_MODE = "xformers" - CLIP_TOKENS = 75 * 2 + 2 - dtype = DTYPE = torch.float16 - device = DEVICE = "cuda" - - image = None - - with torch.inference_mode(True), torch.no_grad(): - alphas, sigmas = load_scheduler_sigmas(beta_end=BETA_END) - - tokenizer_path = args.get("tokenizer_path") - clip_tokenizer = AutoTokenizer.from_pretrained( - tokenizer_path, local_files_only=True) - clip_tokenizer.eos_token_id = 2 - - text_encoder_path = args.get("text_encoder_path") - - clip_encoder = Utils.model_cache_get( - model_type="HYDiT_clip_encoder", model_path=text_encoder_path,) - if clip_encoder is None: - clip_encoder = ( - BertModel.from_pretrained( - text_encoder_path, local_files_only=True).to(device).to(dtype) - ) - Utils.model_cache_set( - model_type="HYDiT_clip_encoder", model_path=text_encoder_path, model=clip_encoder) - - t5_encoder_path = args.get("t5_encoder_path") - if t5_encoder_path != "none" and os.path.exists(t5_encoder_path): - mt5_embedder = Utils.model_cache_get( - model_type="HYDiT_mt5_embedder", model_path=t5_encoder_path,) - if mt5_embedder is None: - mt5_embedder = ( - MT5Embedder(t5_encoder_path, torch_dtype=dtype, - max_length=256,).to(device).to(dtype) - ) - Utils.model_cache_set( - model_type="HYDiT_mt5_embedder", model_path=t5_encoder_path, model=mt5_embedder) - else: - from .mz_train_tools_utils import CustomizeMT5Embedder - mt5_embedder = ( - CustomizeMT5Embedder( - batch_size=1, - ) - .to(device) - .to(dtype) - ) - - vae_ema_path = args.get("vae_ema_path") - vae = Utils.model_cache_get( - model_type="HYDiT_vae", model_path=vae_ema_path,) - if vae is None: - vae = ( - AutoencoderKL.from_pretrained( - vae_ema_path, local_files_only=True) - .to(device) - .to(dtype) - ) - Utils.model_cache_set( - model_type="HYDiT_vae", model_path=vae_ema_path, model=vae) - - unet_path = args.get("unet_path") - lora_path = args.get("lora_path") - - denoiser_args = Utils.model_cache_get( - model_type="HYDiT_unet_merge_lora", model_path=f"{unet_path}_{lora_path}") - - if denoiser_args is None: - denoiser, patch_size, head_dim = DiT_g_2( - input_size=(128, 128), use_extra_cond=USE_EXTRA_COND) - state_dict = torch.load(unet_path) - denoiser.load_state_dict(state_dict) - denoiser.to(device).to(dtype) - denoiser.eval() - denoiser.disable_fp32_silu() - denoiser.disable_fp32_layer_norm() - denoiser.set_attn_mode(ATTN_MODE) - - if lora_path is not None and lora_path != "none": - if not os.path.exists(lora_path): - raise Exception(f"lora_path: {lora_path} 不存在") - lora_net, state_dict = create_network_from_weights( - multiplier=1.0, - file=lora_path, - vae=vae, - text_encoder=[clip_encoder, mt5_embedder], - unet=denoiser, - ) - lora_net.apply_to( - text_encoder=[clip_encoder, mt5_embedder], - unet=denoiser, - ) - lora_net.load_state_dict(state_dict) - lora_net = lora_net.to(DEVICE, dtype=DTYPE) - - Utils.model_cache_set( - model_type="HYDiT_unet_merge_lora", model_path=f"{unet_path}_{lora_path}", model=(denoiser, patch_size, head_dim)) - else: - denoiser, patch_size, head_dim = denoiser_args - - vae.requires_grad_(False) - mt5_embedder.to(torch.float16) - prompt = args.get("prompt") - negative_prompt = args.get("negative_prompt") - with torch.autocast("cuda"): - clip_h, clip_m, mt5_h, mt5_m = get_cond( - prompt, - mt5_embedder, - clip_tokenizer, - clip_encoder, - # Should be same as original implementation with max_length_clip=77 - # Support 75*n + 2 - max_length_clip=CLIP_TOKENS, - ) - neg_clip_h, neg_clip_m, neg_mt5_h, neg_mt5_m = get_cond( - negative_prompt, - mt5_embedder, - clip_tokenizer, - clip_encoder, - max_length_clip=CLIP_TOKENS, - ) - clip_h = torch.concat([clip_h, neg_clip_h], dim=0) - clip_m = torch.concat([clip_m, neg_clip_m], dim=0) - mt5_h = torch.concat([mt5_h, neg_mt5_h], dim=0) - mt5_m = torch.concat([mt5_m, neg_mt5_m], dim=0) - torch.cuda.empty_cache() - - style = torch.as_tensor([0] * 2, device=DEVICE) - W = args.get("width") - H = args.get("height") - - size_cond = [H, W, H, W, 0, 0] - image_meta_size = torch.as_tensor([size_cond] * 2, device=DEVICE) - freqs_cis_img = calc_rope(H, W, patch_size, head_dim) - - denoiser_wrapper = DiscreteVDDPMDenoiser( - # A quick patch for learn_sigma - lambda *args, **kwargs: denoiser(* \ - args, **kwargs).chunk(2, dim=1)[0], - alphas, - False, - ).to(DEVICE) - - CFG_SCALE = cfg = args.get("cfg", 5.0) - STEPS = steps = args.get("steps", 25) - - def cfg_denoise_func(x, sigma): - cond, uncond = denoiser_wrapper( - x.repeat(2, 1, 1, 1), - sigma.repeat(2), - encoder_hidden_states=clip_h, - text_embedding_mask=clip_m, - encoder_hidden_states_t5=mt5_h, - text_embedding_mask_t5=mt5_m, - image_meta_size=image_meta_size, - style=style, - cos_cis_img=freqs_cis_img[0], - sin_cis_img=freqs_cis_img[1], - ).chunk(2, dim=0) - return uncond + (cond - uncond) * CFG_SCALE - - sigmas = denoiser_wrapper.get_sigmas(STEPS).to(DEVICE) - sigmas = get_sigmas_exponential( - STEPS, denoiser_wrapper.sigma_min, denoiser_wrapper.sigma_max, DEVICE - ) - x1 = torch.randn(1, 4, H // 8, W // 8, - dtype=torch.float16, device=DEVICE) - - pbar = Utils.progress_bar(STEPS, "sdxl") - preview = pbar.get_previewer() - - def generate_callback(args): - try: - i = args.get("i") - latents = args.get("denoised") - # 判断是否存在decode_latent_to_preview_image - if hasattr(preview, "decode_latent_to_preview_image"): - pil_img = preview.decode_latent_to_preview_image( - None, - latents, - )[1] - else: - pil_img = None - pbar.update(i, STEPS, pil_img) - except Exception as e: - print("generate_callback error:", e) - raise e - - with torch.autocast("cuda"): - scheduler = args.get("scheduler") - if scheduler == "euler_ancestral": - sample = sample_euler_ancestral( - cfg_denoise_func, - x1 * sigmas[0], - sigmas, - callback=generate_callback, - ) - else: - sample = sample_dpmpp_2m_sde( - cfg_denoise_func, - x1 * sigmas[0], - sigmas, - callback=generate_callback, - ) - torch.cuda.empty_cache() - with torch.no_grad(): - latent = sample / 0.13025 - image = vae.decode(latent).sample - image = (image / 2 + 0.5).clamp(0, 1) - image = image.permute(0, 2, 3, 1) - - keep_device = args.get("keep_device", "enable") - if keep_device == "disable": - Utils.model_cache_clean( - model_type="HYDiT_clip_encoder") - Utils.model_cache_clean( - model_type="HYDiT_mt5_embedder") - Utils.model_cache_clean( - model_type="HYDiT_vae") - Utils.model_cache_clean( - model_type="HYDiT_unet_merge_lora") - - return (image,) - - -def MZ_TrainToolsDebug_call(kwargs): - from pprint import pprint, pformat - 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" - - index = kwargs["index"] - # 用点号分割 - index = index.split(".") - for i in index: - if i == "": - continue - if isinstance(object, (list, tuple)): - object = object[int(i)] - elif isinstance(object, dict): - object = object.get(i) - elif hasattr(object, i): - object = getattr(object, i) - else: - object = object[i] - - debug = pformat(object, indent=indent, depth=depth, width=width, - compact=compact, sort_dicts=sort_keys, underscore_numbers=underscore_numbers) - return (debug,) diff --git a/mz_train_tools_core_HYDiT.py b/mz_train_tools_core_HYDiT.py deleted file mode 100644 index 4147e61..0000000 --- a/mz_train_tools_core_HYDiT.py +++ /dev/null @@ -1,1137 +0,0 @@ - - -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") or file.lower().endswith(".webp"): - 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") or file.lower().endswith(".webp"): - 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-v1.2/t2i/model/pytorch_model_module.pt": { - "url": "https://www.modelscope.cn/models/wailovet/hy1.2/resolve/master/t2i/model/pytorch_model_module.pt", - }, - "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": - version = args.get("version", None) - if version is None: - 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" - else: - if version == "1.2": - download_file = "HunyuanDiT-v1.2/t2i/model/pytorch_model_module.pt" - elif version == "1.1": - download_file = "HunyuanDiT/t2i/model/pytorch_model_module.pt" - else: - raise Exception(f"未知的版本号: {version}") - - 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") or file.lower().endswith(".webp"): - - 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 15aad24..53a758e 100644 --- a/mz_train_tools_utils.py +++ b/mz_train_tools_utils.py @@ -29,6 +29,90 @@ CACHE_POOL = {} class Utils: + def comfyui_fast_generate_picture(ckpt_name, prompt, output, width=512, height=512): + if os.path.exists(output): + return + import nodes + import node_helpers + model, clip, vae = nodes.CheckpointLoaderSimple().load_checkpoint( + ckpt_name=ckpt_name) + seed = 66 + steps = 20 + cfg = 7 + sampler_name = "euler_ancestral" + scheduler = "normal" + + # EmptyLatentImage + latent_image = nodes.EmptyLatentImage().generate(width, height)[0] + + # CLIPTextEncode + positive = nodes.CLIPTextEncode().encode(clip, prompt)[0] + negative = nodes.CLIPTextEncode().encode(clip, "")[0] + + # KSampler + latent = nodes.KSampler().sample( + model=model, + cfg=cfg, seed=seed, steps=steps, sampler_name=sampler_name, scheduler=scheduler, latent_image=latent_image, positive=positive, negative=negative)[0] + + # VAEDecode + image = nodes.VAEDecode().decode(vae, latent)[0] + + # SaveImage + Utils.tensor2pil(image).save(output) + + def clone_repo(args={}): + git_accelerate_urls = { + "githubfast": "githubfast.com", + "521github": "521github.com", + "kkgithub": "kkgithub.com", + } + mz_dir = Utils.get_minus_zone_models_path() + + branch_repoid = args.get("branch_repoid", "kohya-ss/sd-scripts") + branch_local_name = args.get("branch_local_name", "kohya_ss_lora") + + git_url = f"https://github.com/{branch_repoid}" + source = args.get("source", "github") + kohya_ss_lora_dir = os.path.join( + mz_dir, "train_tools", branch_local_name) + if git_accelerate_urls.get(source, None) is not None: + git_url = f"https://{git_accelerate_urls[source]}/{branch_repoid}" + try: + if not os.path.exists(kohya_ss_lora_dir) or not os.path.exists(os.path.join(kohya_ss_lora_dir, ".git")): + subprocess.run( + ["git", "clone", "--depth", "1", git_url, kohya_ss_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=kohya_ss_lora_dir, stdout=subprocess.PIPE, check=True) + result = subprocess.run( + ["git", "rev-parse", "HEAD"], cwd=kohya_ss_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}") + + if branch != long_current_branch and branch != short_current_branch: + subprocess.run( + ["git", "remote", "set-branches", "origin", branch], cwd=kohya_ss_lora_dir, check=True) + subprocess.run( + ["git", "fetch", "--depth", "1", "origin", branch], cwd=kohya_ss_lora_dir, check=True) + + # 恢复所有文件 + subprocess.run( + ["git", "checkout", "."], cwd=kohya_ss_lora_dir, check=True) + + subprocess.run( + ["git", "checkout", branch], cwd=kohya_ss_lora_dir, check=True) + + except Exception as e: + raise Exception(f"克隆kohya-ss/sd-scripts或者切换分支时出现异常,详细信息请查看控制台...") + def Md5(str): return hashlib.md5(str.encode('utf-8')).hexdigest() @@ -1143,33 +1227,20 @@ class HSubprocess: screen_name = None _mswindows = False - def __init__(self, args, screen_name=None): + def __init__(self, args): self._mswindows = (sys.platform == "win32") if self._mswindows: self.args = ["start", "/w"] self.args.extend(args) return - self.screen_name = screen_name - if screen_name is not None: - try: - subprocess.check_call(["screen", "-v"]) - except Exception as e: - raise Exception("Please install screen first.") - - screen_cmd = ["screen", "-R", screen_name, "-m", ] - screen_cmd.extend(args) - - self.args = screen_cmd - else: - self.args = args + self.args = args def stop(self): if self._mswindows: print('taskkill /F /FI "WINDOWTITLE eq hook_kohya_ss_run" /T') os.system( f'taskkill /F /FI "WINDOWTITLE eq hook_kohya_ss_run" /T') - if self.process_instance is not None: if self.screen_name is not None: @@ -1193,6 +1264,15 @@ class HSubprocess: self.process_instance_pid = None def wait(self): + if self._mswindows: + self.wait_win() + else: + self.wait_linux() + + def wait_linux(self): + subprocess.run(self.args, check=True) + + def wait_win(self): with subprocess.Popen( self.args, stdin=subprocess.PIPE,