Merge pull request #6 from MinusZoneAI/dev

Dev
2024-07-05 00:07:44 +08:00
parent 17c3256867 29dde16432
commit 015267c8c0
15 changed files with 3957 additions and 55 deletions
@@ -50,8 +50,7 @@ NODE_CLASS_MAPPINGS["MZ_KohyaSSInitWorkspace"] = MZ_KohyaSSInitWorkspace
 NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSInitWorkspace"] = f"{AUTHOR_NAME} - KohyaSSInitWorkspace"


-class MZ_ImagesCopyWorkspace:
-
+class MZ_KohyaSSDatasetConfig:
    @classmethod
    def INPUT_TYPES(s):
        return {
@@ -59,21 +58,26 @@ class MZ_ImagesCopyWorkspace:
                "workspace_config": ("MZ_TT_SS_WorkspaceConfig",),
                "images": ("IMAGE",),
                "enable_bucket": (["enable", "disable"], {"default": "enable"}),
-                "resolution": ("INT", {"default": 512}),
+                "resolution": ("INT", {"default": 1024}),
                "num_repeats": ("INT", {"default": 1}),
                "batch_size": ("INT", {"default": 1}),
-                # "class_name": ("STRING", {"default": "girl", "dynamicPrompts": True}),
                "force_clear": (["enable", "disable"], {"default": "disable"}),
                "force_clear_only_images": (["enable", "disable"], {"default": "disable"}),
                "same_caption_generate": (["enable", "disable"], {"default": "disable"}),
                "same_caption": ("STRING", {"default": "", "dynamicPrompts": True, "multiline": True}),
            },
+            "optional": {
+                "conditioning_images": ("IMAGE",),
+            }
        }

    RETURN_TYPES = (f"STRING",)
    RETURN_NAMES = ("workspace_images_dir",)

    # OUTPUT_NODE = True
+    MZ_DESC = """
+如果训练类型是controlnet,必须传入预处理后的图片(conditioning_images)
+"""

    FUNCTION = "start"
    CATEGORY = CATEGORY_NAME
@@ -83,9 +87,13 @@ class MZ_ImagesCopyWorkspace:
        return mz_train_tools_core.MZ_ImageSelecter_call(kwargs)


-NODE_CLASS_MAPPINGS["MZ_ImagesCopyWorkspace"] = MZ_ImagesCopyWorkspace
+NODE_CLASS_MAPPINGS["MZ_ImagesCopyWorkspace"] = MZ_KohyaSSDatasetConfig
 NODE_DISPLAY_NAME_MAPPINGS["MZ_ImagesCopyWorkspace"] = f"{AUTHOR_NAME} - ImagesCopyWorkspace"

+# 别名
+NODE_CLASS_MAPPINGS["MZ_KohyaSSDatasetConfig"] = MZ_KohyaSSDatasetConfig
+NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSDatasetConfig"] = f"{AUTHOR_NAME} - KohyaSSDatasetConfig"
+

 class MZ_KohyaSSUseConfig:
    train_config_template_dir = os.path.join(
@@ -94,7 +102,8 @@ class MZ_KohyaSSUseConfig:

    @classmethod
    def INPUT_TYPES(s):
-        train_config_templates = os.listdir(s.train_config_template_dir)
+        train_config_templates = Utils.listdir(s.train_config_template_dir)
+
        # 去掉json后缀
        train_config_templates = [os.path.splitext(x)[0]
                                  for x in train_config_templates]
@@ -217,7 +226,7 @@ class AlwaysEqualProxy(str):
        return False


-class MZ_KohyaSSTrain:
+class MZ_KohyaSSLoraTrain:

    @classmethod
    def INPUT_TYPES(s):
@@ -232,6 +241,9 @@ class MZ_KohyaSSTrain:
        # 使用walk查询所有的workspace中的所有lora模型,lora存放在每个workspace的output目录下
        workspaces_loras = []
        for root, dirs, files in os.walk(workspaces_dir):
+
+            # 排除隐藏文件夹
+            dirs[:] = [d for d in dirs if not d.startswith(".")]
            if root.endswith("output"):
                for file in files:
                    if file.endswith(".safetensors"):
@@ -278,8 +290,62 @@ class MZ_KohyaSSTrain:
        return mz_train_tools_core.MZ_KohyaSSTrain_call(kwargs)


-NODE_CLASS_MAPPINGS["MZ_KohyaSSTrain"] = MZ_KohyaSSTrain
-NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain"
+NODE_CLASS_MAPPINGS["MZ_KohyaSSTrain"] = MZ_KohyaSSLoraTrain
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_KohyaSSTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(old version)"
+
+NODE_CLASS_MAPPINGS["MZ_KohyaSSLoraTrain"] = MZ_KohyaSSLoraTrain
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_KohyaSSLoraTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(lora)"
+
+
+class MZ_KohyaSSControlnetTrain:
+
+    @classmethod
+    def INPUT_TYPES(s):
+        models = [
+            "latest",
+            "empty",
+        ]
+
+        comfyui_full_m_path = []
+        comfyui_basemodels = folder_paths.get_filename_list("controlnet")
+        for b_model in comfyui_basemodels:
+            m_path = folder_paths.get_full_path("controlnet", b_model)
+            comfyui_full_m_path.append(m_path)
+
+        # 按创建时间排序
+        comfyui_full_m_path = sorted(
+            comfyui_full_m_path, key=lambda x: os.path.getctime(x), reverse=True)
+
+        models = models + comfyui_full_m_path
+
+        return {
+            "required": {
+                "train_config": ("MZ_TT_SS_TrainConfig",),
+                "base_controlnet": (models, {"default": "latest"}),
+                "sample_generate": (["enable", "disable"], {"default": "enable"}),
+                "sample_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
+            },
+            "optional": {
+                "has_no_effect": (AlwaysEqualProxy("*"),),
+            },
+        }
+
+    RETURN_TYPES = ()
+    RETURN_NAMES = ()
+    OUTPUT_NODE = True
+    FUNCTION = "start"
+    CATEGORY = CATEGORY_NAME
+
+    def start(self, **kwargs):
+        importlib.reload(mz_train_tools_core)
+        return mz_train_tools_core.MZ_KohyaSSTrain_call(kwargs)
+
+
+NODE_CLASS_MAPPINGS["MZ_KohyaSSControlnetTrain"] = MZ_KohyaSSControlnetTrain
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_KohyaSSControlnetTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(controlnet)"


 class MZ_LoadImagesFromDirectoryPath:
@@ -304,9 +370,10 @@ class MZ_LoadImagesFromDirectoryPath:
        image_dir = kwargs["directory"]
        if not os.path.exists(image_dir):
            return (images,)
-        images = os.listdir(image_dir)
-        images = [x for x in images if x.endswith(
-            ".png") or x.endswith(".jpg")]
+        images = Utils.listdir(image_dir)
+
+        images = [x for x in images if x.lower().endswith(
+            ".png") or x.lower().endswith(".jpg")]
        images = [os.path.join(image_dir, x) for x in images]

        pil_images = []
@@ -317,9 +384,345 @@ class MZ_LoadImagesFromDirectoryPath:
        for pil_image in pil_images:
            tensor_images.append(Utils.pil2tensor(pil_image))

-        return (tensor_images,)
+        return (Utils.list_tensor2tensor(tensor_images),)


 NODE_CLASS_MAPPINGS["MZ_LoadImagesFromDirectoryPath"] = MZ_LoadImagesFromDirectoryPath
 NODE_DISPLAY_NAME_MAPPINGS[
    "MZ_LoadImagesFromDirectoryPath"] = f"{AUTHOR_NAME} - LoadImagesFromDirectoryPath"
+
+
+from . import mz_train_tools_core_HYDiT
+
+
+class MZ_HYDiTInitWorkspace:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "train_name": ("STRING", {"default": ""}),
+                "branch": ("STRING", {"default": "5657364143e44ac90f72aeb47b81bd505a95665d"}),
+                "source": ([
+                    "github",
+                    "githubfast",
+                    "521github",
+                    "kkgithub",
+                ], {"default": "github"}),
+                "seed": ("INT", {"default": 0}),
+            },
+        }
+
+    RETURN_TYPES = ("MZ_TT_HYDiT_WorkspaceConfig",)
+    RETURN_NAMES = ("workspace_config",)
+
+    FUNCTION = "start"
+
+    CATEGORY = CATEGORY_NAME
+
+    def start(self, **kwargs):
+        importlib.reload(mz_train_tools_core_HYDiT)
+        return mz_train_tools_core_HYDiT.MZ_HYDiTInitWorkspace_call(kwargs.copy())
+
+
+NODE_CLASS_MAPPINGS["MZ_HYDiTInitWorkspace"] = MZ_HYDiTInitWorkspace
+NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTInitWorkspace"] = f"{AUTHOR_NAME} - HYDiTInitWorkspace"
+
+
+class MZ_HYDiTDatasetConfig:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "workspace_config": ("MZ_TT_HYDiT_WorkspaceConfig",),
+                "resolution": ("INT", {"default": 1024}),
+                "force_clear": (["enable", "disable"], {"default": "disable"}),
+                "force_clear_only_images": (["enable", "disable"], {"default": "disable"}),
+                "same_caption_generate": (["enable", "disable"], {"default": "disable"}),
+                "same_caption": ("STRING", {"default": "", "dynamicPrompts": True, "multiline": True}),
+            },
+            "optional": {
+                "images": ("IMAGE",),
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("workspace_images_dir",)
+
+    FUNCTION = "start"
+
+    CATEGORY = CATEGORY_NAME
+
+    def start(self, **kwargs):
+        importlib.reload(mz_train_tools_core_HYDiT)
+        return mz_train_tools_core_HYDiT.MZ_HYDiTDatasetConfig_call(kwargs.copy())
+
+
+NODE_CLASS_MAPPINGS["MZ_HYDiTDatasetConfig"] = MZ_HYDiTDatasetConfig
+NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTDatasetConfig"] = f"{AUTHOR_NAME} - HYDiTDatasetConfig"
+
+
+# class MZ_HYDiTUseConfig:
+#     @classmethod
+#     def INPUT_TYPES(s):
+
+#         return {
+#             "required": {
+#                 "workspace_config": ("MZ_TT_HYDiT_WorkspaceConfig",),
+#                 "workspace_images_dir": ("STRING", {"forceInput": True}),
+#             },
+#             "optional": {
+#             }
+#         }
+
+#     RETURN_TYPES = ("MZ_TT_HYDiT_TrainConfig",)
+#     RETURN_NAMES = ("train_config",)
+
+#     FUNCTION = "start"
+
+#     CATEGORY = CATEGORY_NAME
+
+#     def start(self, **kwargs):
+#         importlib.reload(mz_train_tools_core_HYDiT)
+#         # kwargs["train_config_template_dir"] = self.train_config_template_dir
+#         return mz_train_tools_core_HYDiT.MZ_HYDiTUseConfig_call(kwargs)
+
+
+# NODE_CLASS_MAPPINGS["MZ_HYDiTUseConfig"] = MZ_HYDiTUseConfig
+# NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTUseConfig"] = f"{AUTHOR_NAME} - HYDiTUseConfig"
+
+
+class MZ_HYDiTAdvConfig:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "lr": ("STRING", {"default": "1e-5"}),
+                "rope_real": (["enable", "disable"], {"default": "enable"}),
+                # ['Wqkv', 'q_proj', 'kv_proj', 'out_proj']
+                "target_modules_Wqkv": (["enable", "disable"], {"default": "enable"}),
+                "target_modules_q_proj": (["enable", "disable"], {"default": "enable"}),
+                "target_modules_kv_proj": (["enable", "disable"], {"default": "enable"}),
+                "target_modules_out_proj": (["enable", "disable"], {"default": "enable"}),
+                "warmup_min_lr": ("STRING", {"default": "1e-6"}),
+                # parser.add_argument("--warmup-num-steps", type=float, default=0)
+                # parser.add_argument("--weight-decay", type=float, default=0, help="weight-decay in optimizer")
+                "weight_decay": ("FLOAT", {"default": 0}),
+                "warmup_num_steps": ("FLOAT", {"default": 0}),
+                # parser.add_argument("--uncond-p", type=float, default=0.2,
+                #                     help="The probability of dropping training text used for CLIP feature extraction")
+                # parser.add_argument("--uncond-p-t5", type=float, default=0.2,
+                #                     help="The probability of dropping training text used for mT5 feature extraction")
+                "uncond_p": ("FLOAT", {"default": 0.2}),
+                "uncond_p_t5": ("FLOAT", {"default": 0.2}),
+
+                # parser.add_argument("--use-flash-attn", action="store_true", help="During training, "
+                #                                                                 "flash attention is used to accelerate training.")
+                # parser.add_argument("--no-flash-attn", dest="use_flash_attn",
+                #                     action="store_false", help="During training, flash attention is not used to accelerate training.")
+                # parser.add_argument("--use-zero-stage", type=int, default=1, help="Use AngelPTM zero stage. Support 2 and 3")
+                # parser.add_argument("--grad-accu-steps", type=int, default=1, help="Gradient accumulation steps.")
+                "use_flash_attn": (["enable", "disable"], {"default": "disable"}),
+                "use_zero_stage": ("INT", {"default": 2}),
+                "grad_accu_steps": ("INT", {"default": 1}),
+                #  parser.add_argument("--extra-fp16", action="store_true", help="Use extra fp16 for vae and text_encoder.")
+                "extra_fp16": (["enable", "disable"], {"default": "enable"}),
+                # parser.add_argument("--qk-norm", action="store_true", help="Query Key normalization. See http://arxiv.org/abs/2302.05442 for details.")
+                "qk_norm": (["enable", "disable"], {"default": "enable"}),
+                # parser.add_argument("--norm", type=str, choices=["rms", "laryer"], default="layer", help="Normalization layer type")
+                "norm": (["rms", "layer"], {"default": "layer"}),
+            }
+        }
+
+    RETURN_TYPES = ("MZ_TT_HYDiT_AdvConfig",)
+    RETURN_NAMES = ("advanced_config",)
+
+    FUNCTION = "start"
+
+    CATEGORY = CATEGORY_NAME
+
+    def start(self, **kwargs):
+        return (kwargs.copy(),)
+
+
+NODE_CLASS_MAPPINGS["MZ_HYDiTAdvConfig"] = MZ_HYDiTAdvConfig
+NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTAdvConfig"] = f"{AUTHOR_NAME} - HYDiTAdvConfig"
+
+
+class MZ_HYDiTTrain:
+    @classmethod
+    def INPUT_TYPES(s):
+
+        hunyuan_models_path = os.path.join(
+            Utils.get_comfyui_models_path(), "hunyuan")
+        os.makedirs(hunyuan_models_path, exist_ok=True)
+
+        models = Utils.get_models_by_folder(hunyuan_models_path)
+
+        folders = Utils.get_folders_by_folder(hunyuan_models_path)
+
+        vae_models = Utils.get_models_by_folder(
+            os.path.join(Utils.get_comfyui_models_path(), "vae"))
+        unet_models = Utils.get_models_by_folder(
+            os.path.join(Utils.get_comfyui_models_path(), "unet"))
+
+        workspaces_root = os.path.join(
+            folder_paths.output_directory, "mz_train_workspaces")
+
+        loras = mz_train_tools_core_HYDiT.search_loras([
+            workspaces_root,
+            os.path.join(Utils.get_comfyui_models_path(), "loras"),
+        ])
+
+        return {
+            "required": {
+                "workspace_config": ("MZ_TT_HYDiT_WorkspaceConfig",),
+                "unet_path": (["auto"] + models + unet_models, {"default": "auto"}),
+                "ema_to_module": (["enable", "disable"], {"default": "enable"}),
+                "vae_ema_path": (["auto"] + folders + vae_models, {"default": "auto"}),
+                "text_encoder_path": (["auto"] + folders, {"default": "auto"}),
+                "tokenizer_path": (["auto"] + folders, {"default": "auto"}),
+                "t5_encoder_path": (["none", "auto"] + folders, {"default": "none"}),
+                "resolution": ("INT", {"default": 1024, "step": 16}),
+                "batch_size": ("INT", {"default": 1}),
+                "epochs": ("INT", {"default": 50}),
+                "ckpt_every": ("INT", {"default": 500}),
+                "rank": ("INT", {"default": 8}),
+                "base_lora": (["latest", "empty"] + loras, {"default": "latest"}),
+                "sample_generate": (["enable", "disable"], {"default": "enable"}),
+                "sample_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
+            },
+            "optional": {
+                "advanced_config": ("MZ_TT_HYDiT_AdvConfig",),
+                "workspace_images_dir": ("STRING", {"forceInput": True}),
+                "has_no_effect": (AlwaysEqualProxy("*"),),
+            }
+        }
+
+    RETURN_TYPES = ()
+    RETURN_NAMES = ()
+
+    FUNCTION = "start"
+
+    MZ_DESC = """
+base_size The base resolution (n, n) from which to create multiple resolutions | Recommended values: 256/512/1024
+"""
+
+    OUTPUT_NODE = True
+
+    CATEGORY = CATEGORY_NAME
+
+    def start(self, **kwargs):
+        importlib.reload(mz_train_tools_core_HYDiT)
+        return mz_train_tools_core_HYDiT.MZ_HYDiTTrain_call(kwargs.copy())
+
+
+NODE_CLASS_MAPPINGS["MZ_HYDiTTrain"] = MZ_HYDiTTrain
+NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTTrain"] = f"{AUTHOR_NAME} - HYDiTTrain"
+
+
+class MZ_HYDiTSimpleT2I:
+    @classmethod
+    def INPUT_TYPES(s):
+        hunyuan_models_path = os.path.join(
+            Utils.get_comfyui_models_path(), "hunyuan")
+        os.makedirs(hunyuan_models_path, exist_ok=True)
+
+        models = Utils.get_models_by_folder(hunyuan_models_path)
+        folders = Utils.get_folders_by_folder(hunyuan_models_path)
+
+        vae_models = Utils.get_models_by_folder(
+            os.path.join(Utils.get_comfyui_models_path(), "vae"))
+        unet_models = Utils.get_models_by_folder(
+            os.path.join(Utils.get_comfyui_models_path(), "unet"))
+
+        comfyui_full_loras = mz_train_tools_core_HYDiT.search_loras([
+            os.path.join(Utils.get_comfyui_models_path(), "loras"),
+        ])
+
+        return {
+            "required": {
+                "branch": ("STRING", {"default": "5657364143e44ac90f72aeb47b81bd505a95665d"}),
+                "source": ([
+                    "github",
+                    "githubfast",
+                    "521github",
+                    "kkgithub",
+                ], {"default": "github"}),
+                "unet_path": (["auto"] + models + unet_models, {"default": "auto"}),
+                "vae_ema_path": (["auto"] + folders + vae_models, {"default": "auto"}),
+                "text_encoder_path": (["auto"] + folders, {"default": "auto"}),
+                "tokenizer_path": (["auto"] + folders, {"default": "auto"}),
+                "t5_encoder_path": (["none", "auto"] + folders, {"default": "auto"}),
+                "lora_path": (["none"] + comfyui_full_loras, {"default": "none"}),
+                "seed": ("INT", {"default": 0}),
+                "steps": ("INT", {"default": 20}),
+                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
+                "scheduler": ([
+                    "ddpm", "ddim", "dpmms", "uni_pc", "dpmpp_2m_karras"
+                ], {"default": "ddpm"}),
+                "prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
+                "negative_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
+                "width": ("INT", {"default": 512, "max": 8192, "step": 16}),
+                "height": ("INT", {"default": 512, "max": 8192, "step": 16}),
+                "keep_device": (["enable", "disable"], {"default": "enable"}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+
+    RETURN_NAMES = ("images",)
+
+    FUNCTION = "start"
+
+    CATEGORY = CATEGORY_NAME
+
+    def start(self, **kwargs):
+        importlib.reload(mz_train_tools_core_HYDiT)
+        return mz_train_tools_core_HYDiT.MZ_HYDiTSimpleT2I_call(kwargs)
+
+
+NODE_CLASS_MAPPINGS["MZ_HYDiTSimpleT2I"] = MZ_HYDiTSimpleT2I
+NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTSimpleT2I"] = f"{AUTHOR_NAME} - HYDiTSimpleT2I"
+
+
+class MZ_TrainToolsDebug:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "object": (AlwaysEqualProxy("*"),),
+                "indent": ("INT", {"default": 2}),
+                "depth": ("INT", {"default": 5}),
+                "width": ("INT", {"default": 80}),
+                "compact": (["enable", "disable"], {"default": "enable"}),
+                "sort_keys": (["enable", "disable"], {"default": "enable"}),
+                "underscore_numbers": (["enable", "disable"], {"default": "enable"}),
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("debug",)
+
+    FUNCTION = "start"
+
+    CATEGORY = CATEGORY_NAME
+
+    def start(self, **kwargs):
+
+        from pprint import pprint, pp
+        object = kwargs["object"]
+        indent = kwargs["indent"]
+        depth = kwargs["depth"]
+        width = kwargs["width"]
+        compact = kwargs["compact"] == "enable"
+        sort_keys = kwargs["sort_keys"] == "enable"
+        underscore_numbers = kwargs["underscore_numbers"] == "enable"
+
+        debug = pp(object, stream=None, indent=indent, depth=depth, width=width,
+                   compact=compact, sort_dicts=sort_keys, underscore_numbers=underscore_numbers)
+
+        return (debug,)
+
+
+NODE_CLASS_MAPPINGS["MZ_TrainToolsDebug"] = MZ_TrainToolsDebug
+NODE_DISPLAY_NAME_MAPPINGS["MZ_TrainToolsDebug"] = f"{AUTHOR_NAME} - TrainToolsDebug"
@@ -0,0 +1,28 @@
+{
+  "metadata": {
+    "train_type": "controlnet_sd1_5"
+  },
+  "train_config": {
+    "pretrained_model_name_or_path": "",
+    "max_train_steps": "4500",
+    "xformers": true,
+    "sdpa": false,
+    "fp8_base": false,
+    "mixed_precision": "fp16",
+    "cache_latents": true,
+    "cache_latents_to_disk": true,
+    "learning_rate": "1e-5",
+    "lr_scheduler": "cosine_with_restarts",
+    "optimizer_type": "AdamW",
+    "save_every_n_epochs": "20",
+    "shuffle_caption": true,
+    "lr_warmup_steps": "0",
+    "save_precision": "fp16",
+    "lr_scheduler_num_cycles": "1",
+    "persistent_data_loader_workers": true,
+    "noise_offset": "0.1",
+    "output_dir": "",
+    "output_name": "",
+    "lowram": false
+  }
+}
@@ -0,0 +1,155 @@
+import argparse
+
+from index_kits.dataset.make_dataset_core import startup, make_multireso
+from index_kits.common import show_index_info
+from index_kits import __version__
+
+
+def common_args(parser):
+    parser.add_argument('-t', '--target', type=str,
+                        required=True, help='Save path')
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="""
+    IndexKits is a tool to build and manage index files for large-scale datasets. 
+    It supports both base index and multi-resolution index.
+    
+    Introduction
+    ------------
+    This command line tool provides the following functionalities:
+    1. Show index v2 information
+    2. Build base index v2
+    3. Build multi-resolution index v2
+    
+    Examples
+    --------
+    1. Show index v2 information
+        index_kits show /path/to/index.json
+    
+    2. Build base index v2
+        Default usage:
+        index_kits base -c /path/to/config.yaml -t /path/to/index.json
+        
+        Use multiple processes:
+        index_kits base -c /path/to/config.yaml -t /path/to/index.json -w 40
+    
+    3. Build multi-resolution index v2
+    
+        Build with a configuration file:
+        index_kits multireso -c /path/to/config.yaml -t /path/to/index_mb_gt512.json
+        
+        Build by specifying arguments without a configuration file:
+        index_kits multireso --src /path/to/index.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json 
+        
+        Build by specifying target-ratios:
+        index_kits multireso --src /path/to/index.json --base-size 512 --target-ratios 1:1 4:3 3:4 16:9 9:16 --min-size 512 -t /path/to/index_mb_gt512.json
+        
+        Build with multiple source index files.
+        index_kits multireso --src /path/to/index1.json /path/to/index2.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json
+    """, formatter_class=argparse.RawTextHelpFormatter)
+    sub_parsers = parser.add_subparsers(dest='task', required=True)
+
+    # Show index message
+    show_parser = sub_parsers.add_parser('show', description="""
+    Show base/multireso index v2 information.
+    
+    Example
+    -------
+    index_kits show /path/to/index.json
+    """, formatter_class=argparse.RawTextHelpFormatter)
+    show_parser.add_argument(
+        'src', type=str, help='Path to a base/multireso index file.')
+    show_parser.add_argument(
+        '--arrow-files', action='store_true', help='Show arrow files only.')
+    show_parser.add_argument('--depth', type=int, default=1,
+                             help='Arrow file depth. Default is 1, the level of last folder in the arrow file path. '
+                                  'Set it to 0 to show the full path including `xxx/last_folder/*.arrow`.')
+
+    # Single resolution bucket
+    base_parser = sub_parsers.add_parser('base', description="""
+    Build base index v2.
+    
+    Example
+    -------
+    index_kits base -c /path/to/config.yaml -t /path/to/index.json
+    """, formatter_class=argparse.RawTextHelpFormatter)
+    base_parser.add_argument('-c', '--config', type=str,
+                             required=True, help='Configuration file path')
+    common_args(base_parser)
+    base_parser.add_argument('-w', '--world-size', type=int, default=1)
+    base_parser.add_argument('--work-dir', type=str,
+                             default='.', help='Work directory')
+    base_parser.add_argument('--use-cache', action='store_true', help='Use cache to avoid reprocessing. '
+                                                                      'Perform merge pkl results directly.')
+
+    # Multi-resolution bucket
+    mo_parser = sub_parsers.add_parser('multireso', description="""
+    Build multi-resolution index v2
+    
+    Example
+    -------
+    Build with a configuration file:
+    index_kits multireso -c /path/to/config.yaml -t /path/to/index_mb_gt512.json
+    
+    Build by specifying arguments without a configuration file:
+    index_kits multireso --src /path/to/index.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json
+    
+    Build by specifying target-ratios:
+    index_kits multireso --src /path/to/index.json --base-size 512 --target-ratios 1:1 4:3 3:4 16:9 9:16 --min-size 512 -t /path/to/index_mb_gt512.json
+    
+    Build with multiple source index files.
+    index_kits multireso --src /path/to/index1.json /path/to/index2.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json
+    """, formatter_class=argparse.RawTextHelpFormatter)
+    mo_parser.add_argument('-c', '--config', type=str, default=None,
+                           help='Configuration file path in a yaml format. Either --config or --src must be provided.')
+    mo_parser.add_argument('-s', '--src', type=str, nargs='+', default=None,
+                           help='Source index files. Either --config or --src must be provided.')
+    common_args(mo_parser)
+    mo_parser.add_argument('--base-size', type=int, default=None,
+                           help="Base size. Typically set as 256/512/1024 according to image size you train model.")
+    mo_parser.add_argument('--reso-step', type=int, default=None,
+                           help="Resolution step. Either reso_step or target_ratios must be provided.")
+    mo_parser.add_argument('--target-ratios', type=str, nargs='+', default=None,
+                           help="Target ratios. Either reso_step or target_ratios must be provided.")
+    mo_parser.add_argument('--md5-file', type=str, default=None,
+                           help='You can provide an md5 to height and width file to accelerate the process. '
+                                'It is a pickle file that contains a dict, which maps md5 to (height, width) tuple.')
+    mo_parser.add_argument('--align', type=int, default=16,
+                           help="Used when --target-ratios is provided. Align size of source image height and width.")
+    mo_parser.add_argument('--min-size', type=int, default=0,
+                           help="Minimum size. Images smaller than this size will be ignored.")
+
+    # Common
+    parser.add_argument('-v', '--version', action='version',
+                        version=f'%(prog)s {__version__}')
+
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    args = get_args()
+    if args.task == 'show':
+        show_index_info(args.src,
+                        args.arrow_files,
+                        args.depth,
+                        )
+    elif args.task == 'base':
+        startup(args.config,
+                args.target,
+                args.world_size,
+                args.work_dir,
+                use_cache=args.use_cache,
+                )
+    elif args.task == 'multireso':
+        make_multireso(args.target,
+                       args.config,
+                       args.src,
+                       args.base_size,
+                       args.reso_step,
+                       args.target_ratios,
+                       args.align,
+                       args.min_size,
+                       args.md5_file,
+                       )
@@ -0,0 +1,595 @@
+import gc
+import json
+import os
+import random
+import sys
+import time
+from functools import partial
+from glob import glob
+from pathlib import Path
+import numpy as np
+import safetensors.torch
+
+import deepspeed
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+from torch.utils.data import DataLoader
+from torch.distributed.optim import ZeroRedundancyOptimizer
+from torchvision.transforms import functional as TF
+from diffusers.models import AutoencoderKL
+from transformers import BertModel, BertTokenizer, logging as tf_logging
+
+from hydit.config import get_args
+from hydit.lr_scheduler import WarmupLR
+from hydit.data_loader.arrow_load_stream import TextImageArrowStream
+from hydit.diffusion import create_diffusion
+from hydit.ds_config import deepspeed_config_from_args
+from hydit.modules.ema import EMA
+from hydit.modules.fp16_layers import Float16Module
+from hydit.modules.models import HUNYUAN_DIT_MODELS
+from hydit.modules.posemb_layers import init_image_posemb
+from hydit.utils.tools import create_logger, set_seeds, create_exp_folder, get_trainable_params
+from IndexKits.index_kits import ResolutionGroup
+from IndexKits.index_kits.sampler import DistributedSamplerWithStartIndex, BlockDistributedSampler
+from peft import LoraConfig, get_peft_model
+
+
+def deepspeed_initialize(args, logger, model, opt, deepspeed_config):
+    logger.info(f"Initialize deepspeed...")
+    logger.info(f"    Using deepspeed optimizer")
+
+    def get_learning_rate_scheduler(warmup_min_lr, lr, warmup_num_steps, opt):
+        return WarmupLR(opt, warmup_min_lr, lr, warmup_num_steps)
+
+    logger.info(
+        f"    Building scheduler with warmup_min_lr={args.warmup_min_lr}, warmup_num_steps={args.warmup_num_steps}")
+    logger.info(
+        f"    deepspeed_config={deepspeed_config}")
+
+    model, opt, _, scheduler = deepspeed.initialize(
+        model=model,
+        model_parameters=get_trainable_params(model),
+        config_params=deepspeed_config,
+        args=args,
+        lr_scheduler=partial(get_learning_rate_scheduler, args.warmup_min_lr,
+                             args.lr, args.warmup_num_steps) if args.warmup_num_steps > 0 else None,
+    )
+    return model, opt, scheduler
+
+
+def save_checkpoint(args, rank, logger, model, ema, epoch, train_steps, checkpoint_dir):
+    def save_lora_weight(checkpoint_dir, client_state, tag=f"{train_steps:07d}.pt"):
+        cur_ckpt_save_dir = f"{checkpoint_dir}/{tag}"
+        if rank == 0:
+            if args.use_fp16:
+                model.module.module.save_pretrained(cur_ckpt_save_dir)
+            else:
+                model.module.save_pretrained(cur_ckpt_save_dir)
+
+    checkpoint_path = "[Not rank 0. Disabled output.]"
+
+    client_state = {
+        "steps": train_steps,
+        "epoch": epoch,
+        "args": args
+    }
+    if ema is not None:
+        client_state['ema'] = ema.state_dict()
+
+    dst_paths = []
+    if train_steps % args.ckpt_every == 0:
+        checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt"
+        try:
+            if args.training_parts == "lora":
+                save_lora_weight(checkpoint_dir, client_state,
+                                 tag=f"{train_steps:07d}.pt")
+            else:
+                model.save_checkpoint(
+                    checkpoint_dir, client_state=client_state, tag=f"{train_steps:07d}.pt")
+            dst_paths.append(checkpoint_path)
+            logger.info(f"Saved checkpoint to {checkpoint_path}")
+        except:
+            logger.error(f"Saved failed to {checkpoint_path}")
+
+    if train_steps % args.ckpt_latest_every == 0 or train_steps == args.max_training_steps:
+        save_name = "latest.pt"
+        checkpoint_path = f"{checkpoint_dir}/{save_name}"
+        try:
+            if args.training_parts == "lora":
+                save_lora_weight(checkpoint_dir, client_state,
+                                 tag=f"{save_name}")
+            else:
+                model.save_checkpoint(
+                    checkpoint_dir, client_state=client_state, tag=f"{save_name}")
+            dst_paths.append(checkpoint_path)
+            logger.info(f"Saved checkpoint to {checkpoint_path}")
+        except:
+            logger.error(f"Saved failed to {checkpoint_path}")
+
+    dist.barrier()
+    if rank == 0 and len(dst_paths) > 0:
+        # Delete optimizer states to avoid occupying too much disk space.
+        for dst_path in dst_paths:
+            for opt_state_path in glob(f"{dst_path}/zero_dp_rank_*_tp_rank_00_pp_rank_00_optim_states.pt"):
+                os.remove(opt_state_path)
+
+    return checkpoint_path
+
+
+@torch.no_grad()
+def prepare_model_inputs(args, batch, device, vae, text_encoder, text_encoder_t5, freqs_cis_img):
+    try:
+        from .hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar
+    except:
+        from hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar
+
+    image, text_embedding, text_embedding_mask, text_embedding_t5, text_embedding_mask_t5, kwargs = batch
+
+    # clip & mT5 text embedding
+    text_embedding = text_embedding.to(device)
+    text_embedding_mask = text_embedding_mask.to(device)
+    encoder_hidden_states = text_encoder(
+        text_embedding.to(device),
+        attention_mask=text_embedding_mask.to(device),
+    )[0]
+    text_embedding_t5 = text_embedding_t5.to(device).squeeze(1)
+    text_embedding_mask_t5 = text_embedding_mask_t5.to(device).squeeze(1)
+    with torch.no_grad():
+        output_t5 = text_encoder_t5(
+            input_ids=text_embedding_t5,
+            attention_mask=text_embedding_mask_t5 if T5_ENCODER['attention_mask'] else None,
+            output_hidden_states=True
+        )
+        encoder_hidden_states_t5 = output_t5['hidden_states'][T5_ENCODER['layer_index']].detach(
+        )
+
+    # additional condition
+    image_meta_size = kwargs['image_meta_size'].to(device)
+    style = kwargs['style'].to(device)
+
+    if args.extra_fp16:
+        image = image.half()
+        image_meta_size = image_meta_size.half() if image_meta_size is not None else None
+
+    # Map input images to latent space + normalize latents:
+    image = image.to(device)
+    vae_scaling_factor = vae.config.scaling_factor
+    latents = vae.encode(image).latent_dist.sample().mul_(vae_scaling_factor)
+
+    # positional embedding
+    _, _, height, width = image.shape
+    reso = f"{height}x{width}"
+    cos_cis_img, sin_cis_img = freqs_cis_img[reso]
+
+    # Model conditions
+    model_kwargs = dict(
+        encoder_hidden_states=encoder_hidden_states,
+        text_embedding_mask=text_embedding_mask,
+        encoder_hidden_states_t5=encoder_hidden_states_t5,
+        text_embedding_mask_t5=text_embedding_mask_t5,
+        image_meta_size=image_meta_size,
+        style=style,
+        cos_cis_img=cos_cis_img,
+        sin_cis_img=sin_cis_img,
+    )
+
+    return latents, model_kwargs
+
+
+def Core(args):
+    if args.training_parts == "lora":
+        args.use_ema = False
+
+    assert torch.cuda.is_available(), "Training currently requires at least one GPU."
+
+    dist.init_process_group("nccl")
+    world_size = dist.get_world_size()
+    batch_size = args.batch_size
+    grad_accu_steps = args.grad_accu_steps
+    global_batch_size = world_size * batch_size * grad_accu_steps
+
+    rank = dist.get_rank()
+    device = rank % torch.cuda.device_count()
+    seed = args.global_seed * world_size + rank
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.cuda.set_device(device)
+    print(f"Starting rank={rank}, seed={seed}, world_size={world_size}.")
+    deepspeed_config = deepspeed_config_from_args(args, global_batch_size)
+
+    # Setup an experiment folder
+    experiment_dir, checkpoint_dir, logger = create_exp_folder(args, rank)
+
+    # Log all the arguments
+    logger.info(sys.argv)
+    logger.info(str(args))
+    # Save to a json file
+    args_dict = vars(args)
+    args_dict['world_size'] = world_size
+    with open(f"{experiment_dir}/args.json", 'w') as f:
+        json.dump(args_dict, f, indent=4)
+
+    # Disable the message "Some weights of the model checkpoint at ... were not used when initializing BertModel."
+    # If needed, just comment the following line.
+    tf_logging.set_verbosity_error()
+
+    # ===========================================================================
+    # Building HYDIT
+    # ===========================================================================
+
+    logger.info("Building HYDIT Model.")
+
+    # ---------------------------------------------------------------------------
+    #   Training sample base size, such as 256/512/1024. Notice that this size is
+    #   just a base size, not necessary the actual size of training samples. Actual
+    #   size of the training samples are correlated with `resolutions` when enabling
+    #   multi-resolution training.
+    # ---------------------------------------------------------------------------
+    image_size = args.image_size
+    if len(image_size) == 1:
+        image_size = [image_size[0], image_size[0]]
+    if len(image_size) != 2:
+        raise ValueError(f"Invalid image size: {args.image_size}")
+    assert image_size[0] % 8 == 0 and image_size[1] % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder). " \
+                                                              f"got {image_size}"
+    latent_size = [image_size[0] // 8, image_size[1] // 8]
+
+    # initialize model by deepspeed
+    assert args.deepspeed, f"Must enable deepspeed in this script: train_deepspeed.py"
+    with deepspeed.zero.Init(data_parallel_group=torch.distributed.group.WORLD,
+                             remote_device=None if args.remote_device == 'none' else args.remote_device,
+                             config_dict_or_path=deepspeed_config,
+                             mpu=None,
+                             enabled=args.zero_stage == 3):
+        model = HUNYUAN_DIT_MODELS[args.model](args,
+                                               input_size=latent_size,
+                                               log_fn=logger.info,
+                                               )
+    # Multi-resolution / Single-resolution training.
+    if args.multireso:
+        resolutions = ResolutionGroup(image_size[0],
+                                      align=16,
+                                      step=args.reso_step,
+                                      target_ratios=args.target_ratios).data
+    else:
+        resolutions = ResolutionGroup(image_size[0],
+                                      align=16,
+                                      target_ratios=['1:1']).data
+
+    freqs_cis_img = init_image_posemb(args.rope_img,
+                                      resolutions=resolutions,
+                                      patch_size=model.patch_size,
+                                      hidden_size=model.hidden_size,
+                                      num_heads=model.num_heads,
+                                      log_fn=logger.info,
+                                      rope_real=args.rope_real,
+                                      )
+
+    # Create EMA model and convert to fp16 if needed.
+    ema = None
+    if args.use_ema:
+        ema = EMA(args, model, device, logger)
+
+    # Setup FP16 main model:
+    if args.use_fp16:
+        model = Float16Module(model, args)
+    logger.info(
+        f"    Using main model with data type {'fp16' if args.use_fp16 else 'fp32'}")
+
+    diffusion = create_diffusion(
+        noise_schedule=args.noise_schedule,
+        predict_type=args.predict_type,
+        learn_sigma=args.learn_sigma,
+        mse_loss_weight_type=args.mse_loss_weight_type,
+        beta_start=args.beta_start,
+        beta_end=args.beta_end,
+        noise_offset=args.noise_offset,
+    )
+
+    try:
+        from .hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar, CustomizeEmbeds
+    except:
+        from hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar, CustomizeEmbeds
+
+    # Setup VAE
+    logger.info(f"    Loading vae from {VAE_EMA_PATH}")
+    vae = AutoencoderKL.from_pretrained(VAE_EMA_PATH)
+    # Setup BERT text encoder
+    logger.info(f"    Loading Bert text encoder from {TEXT_ENCODER}")
+    text_encoder = BertModel.from_pretrained(
+        TEXT_ENCODER, False, revision=None)
+    # Setup BERT tokenizer:
+    logger.info(f"    Loading Bert tokenizer from {TOKENIZER}")
+    tokenizer = BertTokenizer.from_pretrained(TOKENIZER)
+    # Setup T5 text encoder
+    from hydit.modules.text_encoder import MT5Embedder
+    mt5_path = T5_ENCODER['MT5']
+    if mt5_path is None:
+        embedder_t5 = CustomizeEmbeds()
+    else:
+        embedder_t5 = MT5Embedder(
+            mt5_path, torch_dtype=T5_ENCODER['torch_dtype'], max_length=args.text_len_t5)
+    tokenizer_t5 = embedder_t5.tokenizer
+    text_encoder_t5 = embedder_t5.model
+
+    if args.extra_fp16:
+        logger.info(f"    Using fp16 for extra modules: vae, text_encoder")
+        vae = vae.half().to(device)
+        text_encoder = text_encoder.half().to(device)
+        text_encoder_t5 = text_encoder_t5.half().to(device)
+    else:
+        vae = vae.to(device)
+        text_encoder = text_encoder.to(device)
+        text_encoder_t5 = text_encoder_t5.to(device)
+
+    logger.info(
+        f"    Optimizer parameters: lr={args.lr}, weight_decay={args.weight_decay}")
+    logger.info("    Using deepspeed optimizer")
+    opt = None
+
+    # ===========================================================================
+    # Building Dataset
+    # ===========================================================================
+
+    logger.info(f"Building Streaming Dataset.")
+    logger.info(f"    Loading index file {args.index_file} (v2)")
+
+    dataset = TextImageArrowStream(args=args,
+                                   resolution=image_size[0],
+                                   random_flip=args.random_flip,
+                                   log_fn=logger.info,
+                                   index_file=args.index_file,
+                                   multireso=args.multireso,
+                                   batch_size=batch_size,
+                                   world_size=world_size,
+                                   random_shrink_size_cond=args.random_shrink_size_cond,
+                                   merge_src_cond=args.merge_src_cond,
+                                   uncond_p=args.uncond_p,
+                                   text_ctx_len=args.text_len,
+                                   tokenizer=tokenizer,
+                                   uncond_p_t5=args.uncond_p_t5,
+                                   text_ctx_len_t5=args.text_len_t5,
+                                   tokenizer_t5=tokenizer_t5,
+                                   )
+    if args.multireso:
+        sampler = BlockDistributedSampler(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed,
+                                          shuffle=False, drop_last=True, batch_size=batch_size)
+    else:
+        sampler = DistributedSamplerWithStartIndex(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed,
+                                                   shuffle=False, drop_last=True)
+    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, sampler=sampler,
+                        num_workers=args.num_workers, pin_memory=True, drop_last=True)
+    logger.info(f"    Dataset contains {len(dataset):,} images.")
+    logger.info(f"    Index file: {args.index_file}.")
+    if args.multireso:
+        logger.info(f'    Using MultiResolutionBucketIndexV2 with step {dataset.index_manager.step} '
+                    f'and base size {dataset.index_manager.base_size}')
+        logger.info(f'\n  {dataset.index_manager.resolutions}')
+
+    # ===========================================================================
+    # Loading parameter
+    # ===========================================================================
+
+    logger.info(f"Loading parameter")
+    start_epoch = 0
+    start_epoch_step = 0
+    train_steps = 0
+    # Resume checkpoint if needed
+    # if args.resume is not None or len(args.resume) > 0:
+    if True:
+        model, ema, start_epoch, start_epoch_step, train_steps = model_resume(
+            args, model, ema, logger)
+ 
+    if args.training_parts == "lora":
+        lora_ckpt = args.lora_ckpt
+        if lora_ckpt is not None:
+            lastest_checkpoint = lora_ckpt
+            from peft.peft_model import PeftModel
+            print(f"Loading lora model from {lastest_checkpoint}")
+            if args.use_fp16:
+                model.module = PeftModel.from_pretrained(
+                    model.module, lastest_checkpoint, is_trainable=True)
+            else:
+                model = PeftModel.from_pretrained(
+                    model, lastest_checkpoint, is_trainable=True)
+        else:
+            loraconfig = LoraConfig(
+                r=args.rank,
+                lora_alpha=args.rank,
+                target_modules=args.target_modules
+            )
+
+            if args.use_fp16:
+                model.module = get_peft_model(model.module, loraconfig)
+            else:
+                model = get_peft_model(model, loraconfig)
+
+    logger.info(f"    Training parts: {args.training_parts}")
+
+    model, opt, scheduler = deepspeed_initialize(
+        args, logger, model, opt, deepspeed_config)
+
+    # ===========================================================================
+    # Training
+    # ===========================================================================
+
+    # print model structure
+    # for name, param in model.named_parameters():
+    #     print(name, param.size())
+    # raise Exception("stop")
+
+    model.train()
+    if args.use_ema:
+        ema.eval()
+
+    print(f"    Worker {rank} ready.")
+    dist.barrier()
+
+    iters_per_epoch = len(loader)
+    logger.info(
+        " ****************************** Running training ******************************")
+    logger.info(f"      Number GPUs:               {world_size}")
+    logger.info(f"      Number training samples:   {len(dataset):,}")
+    logger.info(
+        f"      Number parameters:         {sum(p.numel() for p in model.parameters()):,}")
+    logger.info(
+        f"      Number trainable params:   {sum(p.numel() for p in get_trainable_params(model)):,}")
+    logger.info(
+        "    ------------------------------------------------------------------------------")
+    logger.info(f"      Iters per epoch:           {iters_per_epoch:,}")
+    logger.info(f"      Batch size per device:     {batch_size}")
+    logger.info(
+        f"      Batch size all device:     {batch_size * world_size * grad_accu_steps:,} (world_size * batch_size * grad_accu_steps)")
+    logger.info(f"      Gradient Accu steps:       {args.grad_accu_steps}")
+    logger.info(
+        f"      Total optimization steps:  {args.epochs * iters_per_epoch // grad_accu_steps:,}")
+
+    logger.info(
+        f"      Training epochs:           {start_epoch}/{args.epochs}")
+    logger.info(
+        f"      Training epoch steps:      {start_epoch_step:,}/{iters_per_epoch:,}")
+    logger.info(
+        f"      Training total steps:      {train_steps:,}/{min(args.max_training_steps, args.epochs * iters_per_epoch):,}")
+    logger.info(
+        "    ------------------------------------------------------------------------------")
+    logger.info(f"      Noise schedule:            {args.noise_schedule}")
+    logger.info(
+        f"      Beta limits:               ({args.beta_start}, {args.beta_end})")
+    logger.info(f"      Learn sigma:               {args.learn_sigma}")
+    logger.info(f"      Prediction type:           {args.predict_type}")
+    logger.info(f"      Noise offset:              {args.noise_offset}")
+
+    logger.info(
+        "    ------------------------------------------------------------------------------")
+    logger.info(
+        f"      Using EMA model:           {args.use_ema} ({args.ema_dtype})")
+    if args.use_ema:
+        logger.info(
+            f"      Using EMA decay:           {ema.max_value if args.use_ema else None}")
+        logger.info(
+            f"      Using EMA warmup power:    {ema.power if args.use_ema else None}")
+    logger.info(f"      Using main model fp16:     {args.use_fp16}")
+    logger.info(f"      Using extra modules fp16:  {args.extra_fp16}")
+    logger.info(
+        "    ------------------------------------------------------------------------------")
+    logger.info(f"      Experiment directory:      {experiment_dir}")
+    logger.info(
+        "    *******************************************************************************")
+
+    if args.gc_interval > 0:
+        gc.disable()
+        gc.collect()
+
+    # Variables for monitoring/logging purposes:
+    log_steps = 0
+    running_loss = 0
+    start_time = time.time()
+
+    if args.async_ema:
+        ema_stream = torch.cuda.Stream()
+
+    easy_sample_images(args, vae, text_encoder, tokenizer, model, embedder_t5,
+                       target_height=768, target_width=1280, train_steps=0)
+    pbar = PBar(args.epochs * len(loader))
+
+    # Training loop
+    for epoch in range(start_epoch, args.epochs):
+        logger.info(f"    Start random shuffle with seed={seed}")
+        # Makesure all processors use the same seed to shuffle dataset.
+        dataset.shuffle(seed=args.global_seed + epoch, fast=True)
+        logger.info(f"    End of random shuffle")
+
+        # Move sampler to start_index
+        if not args.multireso:
+            start_index = start_epoch_step * world_size * batch_size
+            if start_index != sampler.start_index:
+                sampler.start_index = start_index
+                # Reset start_epoch_step to zero, to ensure next epoch will start from the beginning.
+                start_epoch_step = 0
+                logger.info(f"      Iters left this epoch: {len(loader):,}")
+
+        logger.info(f"    Beginning epoch {epoch}...")
+        step = 0
+        for batch in loader:
+            step += 1
+
+            latents, model_kwargs = prepare_model_inputs(
+                args, batch, device, vae, text_encoder, text_encoder_t5, freqs_cis_img)
+
+            # training model by deepspeed while use fp16
+            if args.use_fp16:
+                if args.use_ema and args.async_ema:
+                    with torch.cuda.stream(ema_stream):
+                        ema.update(model.module.module, step=step)
+                    torch.cuda.current_stream().wait_stream(ema_stream)
+
+            loss_dict = diffusion.training_losses(
+                model=model, x_start=latents, model_kwargs=model_kwargs)
+            loss = loss_dict["loss"].mean()
+            model.backward(loss)
+            last_batch_iteration = (
+                train_steps + 1) // (global_batch_size // (batch_size * world_size))
+            model.step(
+                lr_kwargs={'last_batch_iteration': last_batch_iteration})
+
+            if args.use_ema and not args.async_ema or (args.async_ema and step == len(loader) - 1):
+                if args.use_fp16:
+                    ema.update(model.module.module, step=step)
+                else:
+                    ema.update(model.module, step=step)
+
+            # ===========================================================================
+            # Log loss values:
+            # ===========================================================================
+            running_loss += loss.item()
+            log_steps += 1
+            train_steps += 1
+            if train_steps % args.log_every == 0:
+                # Measure training speed:
+                torch.cuda.synchronize()
+                end_time = time.time()
+                steps_per_sec = log_steps / (end_time - start_time)
+                # Reduce loss history over all processes:
+                avg_loss = torch.tensor(
+                    running_loss / log_steps, device=device)
+                dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM)
+                avg_loss = avg_loss.item() / world_size
+                # get lr from deepspeed fused optimizer
+                logger.info(f"(step={train_steps:07d}) " +
+                            (f"(update_step={train_steps // args.grad_accu_steps:07d}) " if args.grad_accu_steps > 1 else "") +
+                            f"Train Loss: {avg_loss:.4f}, "
+                            f"Lr: {opt.param_groups[0]['lr']:.6g}, "
+                            f"Steps/Sec: {steps_per_sec:.2f}, "
+                            f"Samples/Sec: {int(steps_per_sec * batch_size * world_size):d}")
+                # Reset monitoring variables:
+                running_loss = 0
+                log_steps = 0
+                start_time = time.time()
+
+            # collect gc:
+            if args.gc_interval > 0 and (step % args.gc_interval == 0):
+                gc.collect()
+
+            pbar.step(
+                f"Epoch {epoch}, step {step}, loss {loss.item():.4f}", args.epochs * len(loader), train_steps)
+
+            if (train_steps % args.ckpt_every == 0 or train_steps % args.ckpt_latest_every == 0  # or train_steps == args.max_training_steps
+                ) and train_steps > 0:
+                easy_sample_images(args, vae, text_encoder, tokenizer, model, embedder_t5,
+                                   target_height=768, target_width=1280, train_steps=train_steps)
+                save_checkpoint(args, rank, logger, model, ema,
+                                epoch, train_steps, checkpoint_dir)
+
+            if train_steps >= args.max_training_steps:
+                logger.info(f"Breaking step loop at {train_steps}.")
+                break
+
+        if train_steps >= args.max_training_steps:
+            logger.info(f"Breaking epoch loop at {epoch}.")
+            break
+
+    dist.destroy_process_group()
@@ -0,0 +1,175 @@
+
+import argparse
+import json
+import os
+import sys
+from types import SimpleNamespace
+
+import torch
+
+
+class SimpleNamespaceCNWarrper(SimpleNamespace):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.__dict__.update(kwargs)
+        self.__iter__ = lambda: iter(kwargs.keys())
+    # is not iterable
+
+    def __iter__(self):
+        return iter(self.__dict__.keys())
+    # object has no attribute 'num_attention_heads'
+
+    def __getattr__(self, name):
+        return self.__dict__.get(name, None)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        conflict_handler='resolve',
+    )
+    parser.add_argument("--sys_path", type=str, default="")
+    parser.add_argument("--train_config_file", type=str, default="")
+    parser.add_argument("--mz_master_port", type=int, default=0)
+    args = parser.parse_args()
+
+    master_port = args.mz_master_port
+
+    print(f"master_port = {master_port}")
+
+    try:
+        from . import hook_HYDiT_utils
+    except Exception as e:
+        import hook_HYDiT_utils
+
+    hook_HYDiT_utils.set_master_port(master_port)
+
+    sys_path = args.sys_path
+    if sys_path != "":
+        sys.path.append(sys_path)
+
+    print("HYDi run hook")
+
+    try:
+        from . import hook_HYDiT_main_train_deepspeed
+    except Exception as e:
+        import hook_HYDiT_main_train_deepspeed
+
+    import hydit.config
+
+    def _handle_conflict_error(self, *args, **kwargs):
+        pass
+
+    def parse_args(self, args=None, namespace=None):
+        args, argv = self.parse_known_args(args, namespace)
+        return args
+
+    train_config_file = args.train_config_file
+
+    if train_config_file == "":
+        raise ValueError("train_config_file is empty")
+
+    train_config = {}
+    with open(train_config_file, "r") as f:
+        train_config = json.load(f)
+
+    argparse.ArgumentParser._handle_conflict_error = _handle_conflict_error
+    argparse.ArgumentParser._handle_conflict_resolve = _handle_conflict_error
+    argparse.ArgumentParser.parse_args = parse_args
+    margs = hydit.config.get_args()
+    margs.model = train_config.get("model", "DiT-g/2")
+
+    margs.task_flag = train_config.get("task_flag")
+
+    margs.resume_split = train_config.get("resume_split", True)
+
+    margs.ema_to_module = train_config.get("ema_to_module", True)
+
+    margs.deepspeed = False
+
+    margs.predict_type = train_config.get("predict_type", "v_prediction")
+
+    margs.training_parts = train_config.get("training_parts", "lora")
+
+    margs.batch_size = train_config.get("batch_size", 1)
+
+    margs.grad_accu_steps = train_config.get("grad_accu_steps", 1)
+
+    margs.global_seed = train_config.get("global_seed", 0)
+
+    margs.use_flash_attn = train_config.get("use_flash_attn", False)
+
+    margs.use_fp16 = train_config.get("use_fp16", True)
+
+    margs.qk_norm = train_config.get("qk_norm", True)
+
+    margs.ema_dtype = train_config.get("ema_dtype", "fp32")
+
+    margs.async_ema = False
+
+    margs.ckpt_latest_every = 0x7fffffff
+
+    margs.multireso = train_config.get("multireso", True)
+
+    margs.epochs = train_config.get("epochs", 50)
+
+    margs.target_ratios = train_config.get(
+        "target_ratios", ['1:1', '3:4', '4:3', '16:9', '9:16'])
+    
+    margs.rope_img = train_config.get("rope_img", "base1024")
+
+    margs.image_size = train_config.get("image_size", 1024)
+
+    margs.rope_real = train_config.get("rope_real", True)
+
+    margs.index_file = train_config.get("index_file", None)
+
+    margs.lr = train_config.get("lr", 1e-5)
+
+    margs.rank = train_config.get("rank", 8)
+
+    margs.noise_offset = train_config.get("noise_offset", 0.0)
+
+    margs.log_every = train_config.get("log_every", 99999999999999)
+
+    margs.use_zero_stage = train_config.get("use_zero_stage", 2)
+
+    margs.global_batch_size = train_config.get("global_batch_size", 1)
+
+    margs.deepspeed = True
+
+    margs.results_dir = train_config.get("results_dir")
+
+    margs.mse_loss_weight_type = train_config.get(
+        "mse_loss_weight_type", "constant")
+
+    for k, v in train_config.items():
+        if hasattr(margs, k):
+            setattr(margs, k, v)
+
+    hook_HYDiT_utils.set_unet_path(
+        train_config.get("unet_path"))
+
+    hook_HYDiT_utils.set_vae_ema_path(
+        train_config.get("vae_ema_path"))
+
+    hook_HYDiT_utils.set_text_encoder_path(
+        train_config.get("text_encoder_path"))
+
+    hook_HYDiT_utils.set_tokenizer_path(
+        train_config.get("tokenizer_path"))
+
+    hook_HYDiT_utils.set_t5_encoder_path(
+        train_config.get("t5_encoder_path"))
+
+    hook_HYDiT_utils.set_train_config(train_config)
+
+    try:
+        # deepspeed/runtime/engine
+        import deepspeed.runtime.engine
+        deepspeed.runtime.engine.DeepSpeedEngine._do_sanity_check = lambda x: None
+    except Exception as e:
+        pass
+
+    if type(margs.image_size) == int:
+        margs.image_size = [margs.image_size, margs.image_size]
+    hook_HYDiT_main_train_deepspeed.Core(margs)
@@ -0,0 +1,379 @@
+
+
+import json
+import os
+import time
+import torch
+
+
+UNET_PATH = "ckpts/t2i/model/pytorch_model_ema.pt"
+
+
+def set_unet_path(path):
+    global UNET_PATH
+    UNET_PATH = path
+
+
+VAE_EMA_PATH = "ckpts/t2i/sdxl-vae-fp16-fix"
+
+
+def set_vae_ema_path(path):
+    global VAE_EMA_PATH
+    VAE_EMA_PATH = path
+
+
+TOKENIZER = "ckpts/t2i/tokenizer"
+
+
+def set_tokenizer_path(path):
+    global TOKENIZER
+    TOKENIZER = path
+
+
+TEXT_ENCODER = 'ckpts/t2i/clip_text_encoder'
+
+
+def set_text_encoder_path(path):
+    global TEXT_ENCODER
+    TEXT_ENCODER = path
+
+
+T5_ENCODER = {
+    'MT5': None,
+    'attention_mask': True,
+    'layer_index': -1,
+    'attention_pool': True,
+    'torch_dtype': torch.float16,
+    'learnable_replace': True
+}
+
+
+def set_t5_encoder_path(path):
+    global T5_ENCODER
+    T5_ENCODER['MT5'] = path
+
+
+global TRAIN_CONFIG
+
+
+def set_train_config(train_config):
+    global TRAIN_CONFIG
+    TRAIN_CONFIG = train_config
+
+
+def easy_sample_images(
+        args,
+        vae=None,
+        text_encoder=None,
+        tokenizer=None,
+        model=None,
+        embedder_t5=None,
+        target_height=768,
+        target_width=1280,
+        prompt="A photo of a girl with a hat on a sunny day",
+        negative_prompt="",
+        batch_size=1,
+        guidance_scale=2.0,
+        infer_steps=20,
+        sampler='dpmpp_2m_karras',
+        train_steps=0,
+        seed=0,
+):
+    from hydit.diffusion.pipeline import StableDiffusionPipeline
+    from diffusers import schedulers
+    from hydit.constants import SAMPLER_FACTORY
+    from hydit.modules.posemb_layers import get_fill_resize_and_crop, get_2d_rotary_pos_embed
+    from hydit.modules.models import HUNYUAN_DIT_CONFIG
+
+    import traceback
+    with torch.cuda.amp.autocast():
+
+        workspace_dir = TRAIN_CONFIG.get("workspace_dir")
+        sample_config_file = TRAIN_CONFIG.get("sample_config_file", None)
+        if sample_config_file is None:
+            print("sample_config_file is not set.")
+            return
+        try:
+            sample_config = json.load(open(sample_config_file, "r"))
+        except Exception as e:
+            print(f"Failed to load sample_config_file: {sample_config_file}")
+            return
+        sample_images_dir = os.path.join(workspace_dir, "sample_images")
+        os.makedirs(sample_images_dir, exist_ok=True)
+
+        sampler_factory = SAMPLER_FACTORY.copy()
+
+        sampler_factory["uni_pc"] = {
+            'scheduler': 'UniPCMultistepScheduler',
+            'name': 'UniPCMultistepScheduler',
+            'kwargs': {
+                'beta_schedule': 'scaled_linear',
+                'beta_start': 0.00085,
+                'beta_end': 0.03,
+                'prediction_type': 'v_prediction',
+                'trained_betas': None,
+                'solver_order': 2,
+            }
+        }
+        sampler_factory["dpmpp_2m_karras"] = {
+            'scheduler': 'DPMSolverMultistepScheduler',
+            'name': 'DPMSolverMultistepScheduler',
+            'kwargs': {
+                'beta_schedule': 'scaled_linear',
+                'beta_start': 0.00085,
+                'beta_end': 0.03,
+                'prediction_type': 'v_prediction',
+                'trained_betas': None,
+                'solver_order': 2,
+                'algorithm_type': 'dpmsolver++',
+                "use_karras_sigmas": True,
+            }
+        }
+
+        # Load sampler from factory
+        kwargs = sampler_factory[sampler]['kwargs']
+        scheduler = sampler_factory[sampler]['scheduler']
+
+        # Build scheduler according to the sampler.
+        scheduler_class = getattr(schedulers, scheduler)
+        scheduler = scheduler_class(**kwargs)
+
+        # Set timesteps for inference steps.
+        scheduler.set_timesteps(infer_steps, "cuda")
+
+        def calc_rope(height, width):
+            model_config = HUNYUAN_DIT_CONFIG["DiT-g/2"]
+            patch_size = model_config['patch_size']
+            head_size = model_config['hidden_size'] // model_config['num_heads']
+            th = height // 8 // patch_size
+            tw = width // 8 // patch_size
+            base_size = 512 // 8 // patch_size
+            start, stop = get_fill_resize_and_crop((th, tw), base_size)
+            sub_args = [start, stop, (th, tw)]
+            rope = get_2d_rotary_pos_embed(head_size, *sub_args)
+            return rope
+
+        pipeline = StableDiffusionPipeline(vae=vae,
+                                           text_encoder=text_encoder,
+                                           tokenizer=tokenizer,
+                                           unet=model.module,
+                                           scheduler=scheduler,
+                                           feature_extractor=None,
+                                           safety_checker=None,
+                                           requires_safety_checker=False,
+                                           embedder_t5=embedder_t5,
+                                           )
+        pipeline = pipeline.to("cuda")
+        # attr _execution_device is not defined
+
+        style = torch.as_tensor([0, 0] * batch_size, device="cuda")
+
+        src_size_cond = (target_width, target_height)
+        size_cond = list(src_size_cond) + [target_width, target_height, 0, 0]
+        image_meta_size = torch.as_tensor(
+            [size_cond] * 2 * batch_size, device="cuda",)
+
+        if type(sample_config) != list:
+            sample_config = [sample_config]
+
+        for i, sample in enumerate(sample_config):
+            prompt = sample.get("prompt", "")
+            negative_prompt = sample.get("negative_prompt", "")
+            guidance_scale = sample.get("cfg", guidance_scale)
+            infer_steps = sample.get("steps", infer_steps)
+            width = sample.get("width", target_width)
+            height = sample.get("height", target_height)
+
+            freqs_cis_img = calc_rope(height, width)
+
+            try:
+                generator = torch.Generator(device="cuda")
+                generator.manual_seed(seed)
+                samples = pipeline(
+                    height=height,
+                    width=width,
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    num_images_per_prompt=batch_size,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=infer_steps,
+                    style=style,
+                    return_dict=False,
+                    use_fp16=True,
+                    learn_sigma=args.learn_sigma,
+                    freqs_cis_img=freqs_cis_img,
+                    image_meta_size=image_meta_size,
+                    generator=generator,
+                )[0]
+
+                pass
+            except Exception as e:
+                print(f"Failed to sample images: {e} ")
+                # 打印堆栈信息
+                traceback.print_exc()
+                print(f"Failed to sample pipeline: {pipeline} ")
+
+            # print("samples:",type(samples),)
+            # input("Press Enter to continue...")
+            # print("samples:",samples,)
+
+            if type(samples) == list:
+                pil_image = samples[0]
+            else:
+                pil_image = samples
+
+            sample_filename = f"{args.task_flag}_train_steps_{train_steps:07d}.png"
+            sample_filename_path = os.path.join(
+                sample_images_dir, sample_filename)
+            pil_image.save(sample_filename_path)
+
+    return None
+
+
+def model_resume(args, model, ema, logger):
+    """
+    Load pretrained weights.
+    """
+    start_epoch = 0
+    start_epoch_step = 0
+    train_steps = 0
+    resume_path = UNET_PATH
+
+    logger.info(f"Resume from checkpoint {resume_path}")
+
+    if args.resume_split:
+        # Resume main model
+
+        resume_ckpt_module = torch.load(
+            resume_path, map_location=lambda storage, loc: storage)
+        model.load_state_dict(resume_ckpt_module, strict=False)
+
+        # Resume ema model
+        if args.use_ema:
+            if args.module_to_ema:
+                if "resume_ckpt_module" in locals():
+                    logger.info(f"    Resume ema model from main states.")
+                    ema.load_state_dict(resume_ckpt_module, strict=args.strict)
+                else:
+                    logger.info(f"    Resume ema model from module states.")
+                    resume_ckpt_module = torch.load(
+                        resume_path, map_location=lambda storage, loc: storage)
+                    ema.load_state_dict(resume_ckpt_module, strict=args.strict)
+            else:
+                if "resume_ckpt_ema" in locals():
+                    logger.info(f"    Resume ema model from EMA states.")
+                    ema.load_state_dict(resume_ckpt_ema, strict=args.strict)
+                else:
+                    logger.info(f"    Resume ema model from EMA states.")
+                    resume_ckpt_ema = torch.load(resume_path,
+                                                 map_location=lambda storage, loc: storage)
+                    ema.load_state_dict(resume_ckpt_ema, strict=args.strict)
+    else:
+        raise ValueError(
+            "    “If `resume` is True, then either `resume_split` must be true.”")
+
+    return model, ema, start_epoch, start_epoch_step, train_steps
+
+
+import tqdm
+
+import requests
+
+
+def set_master_port(port):
+    global master_port
+    master_port = port
+
+
+master_port = 0
+
+
+def LOG(log):
+    if master_port == 0:
+        raise Exception("master_port is 0")
+    # 发送http
+    try:
+        resp = requests.request("post", f"http://127.0.0.1:{master_port}/log", data=json.dumps(log), headers={
+                                "Content-Type": "application/json"})
+    except Exception as e:
+        return
+    if resp.status_code != 200:
+        raise Exception(f"LOG failed: {resp.text}")
+
+
+# with tqdm(total=total_steps, initial=train_steps) as pbar:
+# pbar.update(1)
+# pbar.set_description(
+#     f"Epoch {epoch}, step {step}, loss {loss.item():.4f}, mean_loss {mean_loss / step:.4f}")
+class PBar:
+    def __init__(self, total):
+        self.pbar = tqdm.tqdm(total=total)
+
+    def step(self, desc, total_steps, train_steps):
+        self.pbar.update(1)
+        self.pbar.set_description(desc)
+
+        LOG({
+            "type": "sample_images",
+            "global_step": train_steps,
+            "total_steps": total_steps,
+            # "latent": noise_pred_latent_path,
+        })
+
+
+from torch import nn
+
+
+class CustomizeEmbedsModel(nn.Module):
+    dtype = torch.float16
+    # x = torch.zeros(1, 1, 256, 2048)
+    x = None
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    def to(self, *args, **kwargs):
+        self.dtype = torch.float16
+        return self
+
+    def forward(self, *args, **kwargs):
+        input_ids = kwargs.get("input_ids", None)
+        if self.x is None:
+            if input_ids is None:
+                batch_size = 1
+            else:
+                batch_size = input_ids.shape[0]
+            self.x = torch.zeros(1, batch_size, 256, 2048, dtype=self.dtype)
+
+        if kwargs.get("output_hidden_states", False):
+            return {
+                "hidden_states": self.x.to("cuda"),
+                "input_ids": torch.zeros(1, 1),
+            }
+        return self.x
+
+
+class CustomizeTokenizer(dict):
+
+    added_tokens_encoder = []
+    input_ids = torch.zeros(1, 256)
+    attention_mask = torch.zeros(1, 256)
+
+    def __init__(self, *args, **kwargs):
+        self['added_tokens_encoder'] = self.added_tokens_encoder
+        self['input_ids'] = self.input_ids
+        self['attention_mask'] = self.attention_mask
+
+    def tokenize(self, text):
+        return text
+
+    def __call__(self, *args, **kwargs):
+        return self
+
+
+class CustomizeEmbeds():
+    def __init__(self):
+        super().__init__()
+        self.tokenizer = CustomizeTokenizer()
+        self.model = CustomizeEmbedsModel().to("cuda")
+        self.max_length = 256
@@ -48,6 +48,10 @@ train_config_json = "{}"
 sample_images_pipe_class = None


+def utils_sample_images(*args, **kwargs):
+    return sample_images(None, *args, **kwargs)
+
+
 def sample_images(self, *args, **kwargs):
    #  accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet
    accelerator = args[0]
@@ -63,27 +67,15 @@ def sample_images(self, *args, **kwargs):
    # print(f"sample_images: args = {args}")
    # print(f"sample_images: kwargs = {kwargs}")

-    # last_noise_pred = hook_kohya_ss_utils.running_info.get(
-    #     "last_noise_pred", None)
-
-    # noise_pred_latent_path = None
-    # if last_noise_pred is not None:
-    #     noise_pred_latent_path = os.path.join(
-    #         tempfile.gettempdir(), f"sample_images_{global_step % 5}.latent")
-
-    #     output = {}
-    #     output["latent_tensor"] = last_noise_pred
-    #     output["latent_format_version_0"] = torch.tensor([])
-    #     safetensors.torch.save_file(output, noise_pred_latent_path)
+    controlnet = kwargs.get("controlnet", None)

    if epoch is not None and cmd_args.save_every_n_epochs is not None and epoch % cmd_args.save_every_n_epochs == 0:

        prompt_dict_list = other_config.get("prompt_dict_list", [])
        if len(prompt_dict_list) == 0:
            seed = other_config.get("seed", 0)
-
-            prompt_dict_list.append({
-                "controlnet_image": None,
+            prompt_dict = {
+                "controlnet_image": other_config.get("controlnet_image", None),
                "prompt": other_config.get("sample_prompt", ""),
                "seed": seed,
                "negative_prompt": "",
@@ -91,7 +83,9 @@ def sample_images(self, *args, **kwargs):
                "sample_sampler": "euler_a",
                "sample_steps": 20,
                "scale": 5.0,
-            })
+            }
+
+            prompt_dict_list.append(prompt_dict)
        else:
            for i, prompt_dict in enumerate(prompt_dict_list):
                if prompt_dict.get("controlnet_image", None) is None:
@@ -103,6 +97,7 @@ def sample_images(self, *args, **kwargs):
                if prompt_dict.get("enum", None) is None:
                    prompt_dict["enum"] = i
        hook_kohya_ss_utils.generate_image(
+            pipe_class=sample_images_pipe_class,
            cmd_args=cmd_args,
            accelerator=accelerator,
            epoch=epoch,
@@ -110,8 +105,8 @@ def sample_images(self, *args, **kwargs):
            tokenizer=tokenizer,
            unet=unet,
            vae=vae,
-            pipe_class=sample_images_pipe_class,
            prompt_dict_list=prompt_dict_list,
+            controlnet=controlnet,
        )

    LOG({
@@ -166,9 +161,55 @@ def run_lora_sdxl():
    trainer.train(train_args)


+from types import SimpleNamespace
+
+
+class SimpleNamespaceCNWarrper(SimpleNamespace):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.__dict__.update(kwargs)  # or self.__dict__ = kwargs
+        self.__dict__["mid_block_type"] = "UNetMidBlock2DCrossAttn"
+        self.__dict__["_diffusers_version"] = "0.6.0"
+        self.__iter__ = lambda: iter(kwargs.keys())
+    # is not iterable
+
+    def __iter__(self):
+        return iter(self.__dict__.keys())
+    # object has no attribute 'num_attention_heads'
+
+    def __getattr__(self, name):
+        return self.__dict__.get(name, None)
+
+
+def run_controlnet_sd1_5():
+    import types
+    types.SimpleNamespace = SimpleNamespaceCNWarrper
+    hook_kohya_ss_utils.hook_kohya_ss()
+    # 覆盖sample_images生成函数,包括进度条和生成图片功能
+
+    import train_controlnet
+
+    # 配置对应的pipeline
+    import library.train_util
+    library.train_util.sample_images = utils_sample_images
+
+    global sample_images_pipe_class
+    sample_images_pipe_class = library.train_util.StableDiffusionLongPromptWeightingPipeline
+
+    train_config = json.loads(train_config_json)
+    train_args = config2args(train_controlnet.setup_parser(), train_config)
+
+    LOG({
+        "type": "start_train",
+    })
+
+    train_controlnet.train(train_args)
+
+
 func_map = {
    "run_lora_sd1_5": run_lora_sd1_5,
    "run_lora_sdxl": run_lora_sdxl,
+    "run_controlnet_sd1_5": run_controlnet_sd1_5,
 }


@@ -93,6 +93,12 @@ def MZ_ImageSelecter_call(args={}):
    images = args.get("images")
    pil_images = Utils.tensors2pil_list(images)

+    conditioning_images = args.get("conditioning_images", None)
+    conditioning_pil_images = None
+    if conditioning_images is not None:
+        conditioning_pil_images = Utils.tensors2pil_list(
+            conditioning_images)
+
    resolution = args.get("resolution", 512)

    workspace_config = args.get("workspace_config", {})
@@ -106,12 +112,16 @@ def MZ_ImageSelecter_call(args={}):
    train_images_dir = os.path.join(workspace_dir, "train_images")
    os.makedirs(train_images_dir, exist_ok=True)

+    conditioning_images_dir = os.path.join(
+        workspace_dir, "conditioning_images")
+
    force_clear = args.get("force_clear") == "enable"
    force_clear_only_images = args.get("force_clear_only_images") == "enable"
    if force_clear:
        if force_clear_only_images:
-            for file in os.listdir(train_images_dir):
-                if file.endswith(".png") or file.endswith(".jpg"):
+            images_files = Utils.listdir(train_images_dir)
+            for file in images_files:
+                if file.lower().endswith(".png") or file.lower().endswith(".jpg"):
                    os.remove(os.path.join(train_images_dir, file))
        else:
            shutil.rmtree(train_images_dir)
@@ -120,10 +130,16 @@ def MZ_ImageSelecter_call(args={}):
    saved_images_path = []
    for i, pil_image in enumerate(pil_images):
        pil_image = Utils.resize_max(pil_image, resolution, resolution)
+        width, height = pil_image.size
        filename = hashlib.md5(pil_image.tobytes()).hexdigest() + ".png"
        pil_image.save(os.path.join(train_images_dir, filename))
        saved_images_path.append(filename)

+        if conditioning_pil_images is not None:
+            os.makedirs(conditioning_images_dir, exist_ok=True)
+            conditioning_pil_images[i].resize((width, height)).save(
+                os.path.join(conditioning_images_dir, filename))
+
    same_caption_generate = args.get("same_caption_generate") == "enable"
    if same_caption_generate:
        same_caption = args.get("same_caption").strip()
@@ -136,12 +152,17 @@ def MZ_ImageSelecter_call(args={}):
                    f.write(same_caption)

    dataset_config_path = os.path.join(workspace_dir, "dataset.toml")
+
+    if conditioning_images is None:
+        conditioning_images_dir = None
+
    generate_toml_config(
        dataset_config_path,
        enable_bucket=args.get("enable_bucket") == "enable",
        resolution=args.get("resolution"),
        batch_size=args.get("batch_size"),
        image_dir=train_images_dir,
+        conditioning_data_dir=conditioning_images_dir,
        caption_extension=".caption",
        num_repeats=args.get("num_repeats"),
    )
@@ -290,7 +311,7 @@ def check_install():
 import logging


-def generate_toml_config(output_path, enable_bucket=True, resolution=512, batch_size=1, image_dir=None, caption_extension=".caption", num_repeats=10, ):
+def generate_toml_config(output_path, enable_bucket=True, resolution=512, batch_size=1, image_dir=None, conditioning_data_dir=None, caption_extension=".caption", num_repeats=10, ):
    check_install()
    import toml
    config = {
@@ -304,6 +325,7 @@ def generate_toml_config(output_path, enable_bucket=True, resolution=512, batch_
                'subsets': [
                    {
                        'image_dir': image_dir,
+                        'conditioning_data_dir': conditioning_data_dir,
                        'caption_extension': caption_extension,
                        'num_repeats': num_repeats,
                    },
@@ -336,9 +358,9 @@ def get_sample_images(train_config):
    pil_images = []
    pre_render_texts_x = []
    if os.path.exists(sample_images_dir):
-        image_files = os.listdir(sample_images_dir)
+        image_files = Utils.listdir(sample_images_dir)
        image_files = list(
-            filter(lambda x: x.endswith(".png"), image_files))
+            filter(lambda x: x.lower().endswith(".png"), image_files))
        # 筛选 output_name 前缀
        image_files = list(
            filter(lambda x: x.startswith(output_name), image_files))
@@ -374,11 +396,8 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot
    if trainer_func.find("sdxl") != -1:
        taesd_type = "sdxl"

-
-    
- 
    pb = Utils.progress_bar(train_config.get("max_train_steps"), taesd_type)
- 
+
    import traceback

    import comfy.model_management
@@ -401,7 +420,6 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot
                max_side = max(xy_img.width, xy_img.height)
                # print(f"global_step: {global_step}, max_train_steps: {max_train_steps}")

-
                total_steps = resp.get("total_steps")
                pb.update(
                    int(global_step), int(total_steps), ("JPEG", xy_img, max_side))
@@ -466,7 +484,7 @@ def MZ_KohyaSSTrain_call(args={}):
    elif base_lora == "latest":
        workspace_lora_dir = os.path.join(workspace_dir, "output")
        if os.path.exists(workspace_lora_dir):
-            workspace_lora_files = os.listdir(workspace_lora_dir)
+            workspace_lora_files = Utils.listdir(workspace_lora_dir)
            workspace_lora_files = list(
                filter(lambda x: x.endswith(".safetensors"), workspace_lora_files))
            workspace_lora_files = list(
@@ -494,14 +512,38 @@ def MZ_KohyaSSTrain_call(args={}):
        if "network_dropout" in train_config:
            del train_config["network_dropout"]

+    
+
+
+    base_controlnet = args.get("base_controlnet", "empty")
+    if base_controlnet == "empty":
+        pass
+    elif base_controlnet == "latest":
+        workspace_controlnet_dir = os.path.join(workspace_dir, "output")
+        if os.path.exists(workspace_controlnet_dir):
+            workspace_controlnet_files = Utils.listdir(workspace_controlnet_dir)
+            workspace_controlnet_files = list(
+                filter(lambda x: x.endswith(".safetensors"), workspace_controlnet_files))
+            workspace_controlnet_files = list(
+                map(lambda x: os.path.join(workspace_controlnet_dir, x), workspace_controlnet_files))
+            # 排序
+            workspace_controlnet_files = sorted(
+                workspace_controlnet_files, key=lambda x: os.path.getctime(x), reverse=True)
+            if len(workspace_controlnet_files) > 0:
+                base_controlnet = os.path.join(
+                    workspace_controlnet_dir, workspace_controlnet_files[0])
+        else:
+            base_controlnet = "empty"
+    else:
+        pass
+
+    if base_controlnet != "empty" and os.path.exists(base_controlnet):
+        train_config["controlnet_model_name_or_path"] = base_controlnet 
+
+
    train_type = config.get("metadata").get("train_type")
-    # 在ComfyUI中无法运行, 梯度有问题, 原因不清楚
-    # import train_network
-    # train_network.logger = logging.getLogger()
-    # trainer = train_network.NetworkTrainer()
-    # train_args = config2args(train_network.setup_parser(), train_config)
-    # trainer.train(train_args)
-    # return
+
+
    sample_generate = args.get("sample_generate", "enable")
    sample_prompt = args.get("sample_prompt", "")
    if sample_generate == "enable":
@@ -517,6 +559,18 @@ def MZ_KohyaSSTrain_call(args={}):
    elif train_type == "lora_sdxl":
        run_hook_kohya_ss_run_file(
            kohya_ss_tool_dir, train_config, "run_lora_sdxl", other_config)
+    elif train_type == "controlnet_sd1_5":
+
+        conditioning_images_dir = os.path.join(
+            workspace_dir, "conditioning_images")
+        conditioning_images_onec = ""
+        if os.path.exists(conditioning_images_dir):
+            conditioning_images_onec = Utils.listdir(conditioning_images_dir)[0]
+            other_config["controlnet_image"] = os.path.join(
+                conditioning_images_dir, conditioning_images_onec)
+        
+        run_hook_kohya_ss_run_file(
+            kohya_ss_tool_dir, train_config, "run_controlnet_sd1_5", other_config)
    else:
        raise Exception(
            f"暂时不支持的训练类型: {train_type}")
@@ -67,10 +67,10 @@ class Utils:

    def listdir_png(path):
        try:
-            files = os.listdir(path)
+            files = Utils.listdir(path)
            new_files = []
            for file in files:
-                if file.endswith(".png"):
+                if file.lower().endswith(".png"):
                    new_files.append(file)
            files = new_files
            files.sort(key=lambda x: int(os.path.basename(x).split(".")[0]))
@@ -78,6 +78,15 @@ class Utils:
        except Exception as e:
            return []

+    def listdir(path):
+        try:
+            files = os.listdir(path)
+            # 排除.开头的文件
+            files = [file for file in files if not file.startswith(".")]
+            return files
+        except Exception as e:
+            return []
+
    def tensor2pil(image):
        return Image.fromarray(np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))

@@ -676,6 +685,8 @@ class Utils:

            if os.path.exists(find_fullpath):
                for root, dirs, files in os.walk(find_fullpath):
+                    # 排除隐藏文件夹
+                    dirs[:] = [d for d in dirs if not d.startswith('.')]
                    for file in files:
                        if target_sha256 == Utils.file_sha256(os.path.join(root, file)):
                            return os.path.join(root, file)
@@ -710,12 +721,12 @@ class Utils:
                return previewer

            def update(self, step, total_steps, pil_img=None):
-                pil_img_info = ("JPEG", pil_img, 512)
-                if pil_img is None:
-                    pil_img_info = None
-                if type(pil_img) == Tuple or type(pil_img) == list or type(pil_img) == tuple:
-                    pil_img_info = pil_img
                try:
+                    pil_img_info = ("JPEG", pil_img, 512)
+                    if pil_img is None:
+                        pil_img_info = None
+                    if type(pil_img) == Tuple or type(pil_img) == list or type(pil_img) == tuple:
+                        pil_img_info = pil_img
                    # print("pil_img_info:", type(pil_img), pil_img_info)
                    # print("step:", step, "total_steps:", total_steps)
                    self.pbar.update_absolute(
@@ -787,6 +798,8 @@ class Utils:
        # 去掉pre_render_texts_x中所有相同字符串前缀

        def get_common_prefix(pre_render_texts_x):
+            if len(pre_render_texts_x) == 0:
+                return ""
            common_prefix = ""
            for i in range(len(pre_render_texts_x[0])):
                c = pre_render_texts_x[0][i]
@@ -913,3 +926,23 @@ class Utils:
                      label_text, font=font, fill="black")

        return full_canvas
+
+    def get_models_by_folder(dir_path):
+        models = []
+        for root, dirs, files in os.walk(dir_path):
+            # 排除隐藏文件夹    
+            dirs[:] = [d for d in dirs if not d.startswith('.')]
+            for file in files:
+                if file.endswith(".pth") or file.endswith(".pt") or file.endswith(".pkl") or file.endswith(".onnx") or file.endswith(".safetensors"):
+                    models.append(os.path.join(root, file))
+        return models
+
+    def get_folders_by_folder(dir_path):
+        folders = []
+        for root, dirs, files in os.walk(dir_path):
+            
+            # 排除隐藏文件夹    
+            dirs[:] = [d for d in dirs if not d.startswith('.')]
+            for dir in dirs:
+                folders.append(os.path.join(root, dir))
+        return folders
@@ -0,0 +1,210 @@
+import json
+import os
+import shutil
+from PIL import Image
+import matplotlib.pyplot as plt
+
+from warnings import filterwarnings
+
+import numpy as np
+import torch
+import pytorch_lightning as pl
+import torch.nn as nn
+from PIL import Image
+
+
+from transformers import CLIPProcessor, CLIPModel
+import torch.nn.functional as F
+
+
+class MLP(pl.LightningModule):
+    def __init__(self, input_size, xcol='emb', ycol='avg_rating'):
+        super().__init__()
+        self.input_size = input_size
+        self.xcol = xcol
+        self.ycol = ycol
+        self.layers = nn.Sequential(
+            nn.Linear(self.input_size, 1024),
+            # nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(1024, 128),
+            # nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(128, 64),
+            # nn.ReLU(),
+            nn.Dropout(0.1),
+
+            nn.Linear(64, 16),
+            # nn.ReLU(),
+
+            nn.Linear(16, 1)
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+    def training_step(self, batch, batch_idx):
+        x = batch[self.xcol]
+        y = batch[self.ycol].reshape(-1, 1)
+        x_hat = self.layers(x)
+        loss = F.mse_loss(x_hat, y)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x = batch[self.xcol]
+        y = batch[self.ycol].reshape(-1, 1)
+        x_hat = self.layers(x)
+        loss = F.mse_loss(x_hat, y)
+        return loss
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+        return optimizer
+
+
+def normalized(a, axis=-1, order=2):
+    import numpy as np
+
+    l2 = np.atleast_1d(np.linalg.norm(a, order, axis))
+    l2[l2 == 0] = 1
+    return a / np.expand_dims(l2, axis)
+
+
+class Aestheic:
+    def __init__(self, clip_path, aes_model):
+
+        model = MLP(768)  # CLIP embedding dim is 768 for CLIP ViT L 14
+
+        # load the model you trained previously or the model available in this repo
+        s = torch.load(aes_model)
+
+        model.load_state_dict(s)
+
+        model.to("cuda")
+        model.eval()
+
+        self.model = model
+
+        self.device = "cuda"
+
+        model2 = CLIPModel.from_pretrained(
+            clip_path, local_files_only=True).to("cuda")
+        preprocess = CLIPProcessor.from_pretrained(
+            clip_path, local_files_only=True)
+
+        self.model2 = model2
+        self.preprocess = preprocess
+
+    def prediction(self, img_paths):
+        if not isinstance(img_paths, list):
+            img_paths = [img_paths]
+
+        pil_images = []
+        for img_path in img_paths:
+            pil_image = Image.open(img_path)
+            pil_images.append(pil_image)
+
+        inputs = self.preprocess(
+            images=pil_images, return_tensors="pt",
+        ).to("cuda")
+        with torch.no_grad():
+            image_features = self.model2.get_image_features(**inputs)
+            print(image_features.shape)
+
+        im_emb_arr = normalized(image_features.cpu().detach().numpy())
+
+        prediction = self.model(torch.from_numpy(
+            im_emb_arr).to(self.device).type(torch.cuda.FloatTensor))
+
+        prediction_score = prediction.cpu().detach().numpy()
+        prediction_score = prediction_score.flatten()
+
+        images_score = {
+            img_path: float(score) for img_path, score in zip(img_paths, prediction_score)
+        }
+
+        # 按value排序
+        images_score = dict(
+            sorted(images_score.items(), key=lambda x: x[1], reverse=True))
+
+        return images_score
+
+
+import argparse
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--input_dir", type=str, default="")
+    parser.add_argument("-o", "--output_dir", type=str, default="")
+    parser.add_argument("-c", "--clip_path", type=str,
+                        default="D:\\数据集\\clip-vit-large-patch14")
+    parser.add_argument("-m", "--aes_model", type=str,
+                        default="D:\\数据集\\sac+logos+ava1-l14-linearMSE.pth")
+    parser.add_argument("-s", "--score", type=float, default=5.9)
+
+    parser.add_argument("-b", "--batch_size", type=int, default=1)
+
+    args = parser.parse_args()
+    limitscore = args.score
+
+    aes = Aestheic(args.clip_path, args.aes_model)
+
+    input_dir = args.input_dir
+    output_dir = args.output_dir
+    batch_size = args.batch_size
+    total = 0
+    for root, dirs, files in os.walk(input_dir):
+        # 排除隐藏文件夹
+        dirs[:] = [d for d in dirs if not d.startswith('.')]
+        for file in files:
+            if not file.endswith(".png") and not file.endswith(".jpg"):
+                continue
+            total += 1
+
+    import tqdm
+    print(f"total: {total}")
+
+    images = []
+
+    with tqdm.tqdm(total=total) as pbar:
+        for root, dirs, files in os.walk(input_dir):
+            # 排除隐藏文件夹    
+            dirs[:] = [d for d in dirs if not d.startswith('.')]
+            for file in files:
+                if not file.endswith(".png") and not file.endswith(".jpg"):
+                    continue
+                file_path = os.path.join(root, file)
+                pbar.update(1)
+
+                pbar.set_description(
+                    f"Processing {os.path.basename(file_path)}")
+
+                if len(images) < batch_size:
+                    images.append(file_path)
+                else:
+                    image_scroe_map = aes.prediction(images)
+                    images = []
+                    # print(scroe)
+                    # print(json.dumps(scroe, indent=4, ensure_ascii=False))
+                    # exit(0)
+
+                    for img_path, score in image_scroe_map.items():
+                        print(f"{img_path}: {score}")
+
+                        if score > limitscore:
+                            if not os.path.exists(output_dir):
+                                os.makedirs(output_dir)
+                            shutil.copy(img_path, output_dir)
+
+    if len(images) > 0:
+        aes.prediction(images)
+        images = []
+
+        for img_path, score in image_scroe_map.items():
+            print(f"{img_path}: {score}")
+
+            if score > limitscore:
+                if not os.path.exists(output_dir):
+                    os.makedirs(output_dir)
+                shutil.copy(img_path, output_dir)
+
+                
@@ -0,0 +1,260 @@
+import json
+import os
+from collections import OrderedDict
+from os import PathLike
+import shutil
+from typing import Final
+
+import torch
+import torch.nn as nn
+from transformers import (
+    SiglipImageProcessor,
+    SiglipVisionConfig,
+    SiglipVisionModel,
+    logging,
+)
+from transformers.image_processing_utils import BatchFeature
+from transformers.modeling_outputs import ImageClassifierOutputWithNoAttention
+
+logging.set_verbosity_error()
+
+URL: Final[str] = (
+    "https://github.com/discus0434/aesthetic-predictor-v2-5/raw/main/models/aesthetic_predictor_v2_5.pth"
+)
+
+
+class AestheticPredictorV2_5Head(nn.Module):
+    def __init__(self, config: SiglipVisionConfig) -> None:
+        super().__init__()
+        self.scoring_head = nn.Sequential(
+            nn.Linear(config.hidden_size, 1024),
+            nn.Dropout(0.5),
+            nn.Linear(1024, 128),
+            nn.Dropout(0.5),
+            nn.Linear(128, 64),
+            nn.Dropout(0.5),
+            nn.Linear(64, 16),
+            nn.Dropout(0.2),
+            nn.Linear(16, 1),
+        )
+
+    def forward(self, image_embeds: torch.Tensor) -> torch.Tensor:
+        return self.scoring_head(image_embeds)
+
+
+class AestheticPredictorV2_5Model(SiglipVisionModel):
+    PATCH_SIZE = 14
+
+    def __init__(self, config: SiglipVisionConfig, *args, **kwargs) -> None:
+        super().__init__(config, *args, **kwargs)
+        self.layers = AestheticPredictorV2_5Head(config)
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor | None = None,
+        labels: torch.Tensor | None = None,
+        return_dict: bool | None = None,
+    ) -> tuple | ImageClassifierOutputWithNoAttention:
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        outputs = super().forward(
+            pixel_values=pixel_values,
+            return_dict=return_dict,
+        )
+        image_embeds = outputs.pooler_output
+        image_embeds_norm = image_embeds / \
+            image_embeds.norm(dim=-1, keepdim=True)
+        prediction = self.layers(image_embeds_norm)
+
+        loss = None
+        if labels is not None:
+            loss_fct = nn.MSELoss()
+            loss = loss_fct()
+
+        if not return_dict:
+            return (loss, prediction, image_embeds)
+
+        return ImageClassifierOutputWithNoAttention(
+            loss=loss,
+            logits=prediction,
+            hidden_states=image_embeds,
+        )
+
+
+class AestheticPredictorV2_5Processor(SiglipImageProcessor):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+    def __call__(self, *args, **kwargs) -> BatchFeature:
+        return super().__call__(*args, **kwargs)
+
+    @classmethod
+    def from_pretrained(
+        self,
+        pretrained_model_name_or_path: str
+        | PathLike = "google/siglip-so400m-patch14-384",
+        *args,
+        **kwargs,
+    ) -> "AestheticPredictorV2_5Processor":
+        return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+
+def convert_v2_5_from_siglip(
+    predictor_name_or_path: str | PathLike | None = None,
+    encoder_model_name: str = "google/siglip-so400m-patch14-384",
+    *args,
+    **kwargs,
+) -> tuple[AestheticPredictorV2_5Model, AestheticPredictorV2_5Processor]:
+    model = AestheticPredictorV2_5Model.from_pretrained(
+        encoder_model_name, *args, **kwargs
+    )
+
+    processor = AestheticPredictorV2_5Processor.from_pretrained(
+        encoder_model_name, *args, **kwargs
+    )
+
+    if predictor_name_or_path is None or not os.path.exists(predictor_name_or_path):
+        state_dict = torch.hub.load_state_dict_from_url(
+            URL, map_location="cpu")
+    else:
+        state_dict = torch.load(predictor_name_or_path, map_location="cpu")
+
+    assert isinstance(state_dict, OrderedDict)
+
+    model.layers.load_state_dict(state_dict)
+    model.eval()
+
+    return model, processor
+
+
+from PIL import Image
+
+
+class AestheticPredictor:
+    def __init__(self, clip_path, aesthetic_predictor_path):
+        # load model and preprocessor
+        self.model, self.preprocessor = convert_v2_5_from_siglip(
+            low_cpu_mem_usage=True,
+            trust_remote_code=True,
+            encoder_model_name=clip_path,
+            predictor_name_or_path=aesthetic_predictor_path,
+        )
+        if torch.cuda.is_available():
+            self.model = self.model.to(torch.bfloat16).cuda()
+
+    def inference(self, img_paths) -> float:
+        if not isinstance(img_paths, list):
+            img_paths = [img_paths]
+
+        pil_images = []
+        for img_path in img_paths:
+            pil_image = Image.open(img_path)
+            pil_images.append(pil_image)
+
+        # preprocess image
+        pixel_values = self.preprocessor(
+            images=pil_images, return_tensors="pt"
+        ).pixel_values
+
+        if torch.cuda.is_available():
+            pixel_values = pixel_values.to(torch.bfloat16).cuda()
+
+        # predict aesthetic score
+        with torch.inference_mode():
+            score = self.model(
+                pixel_values).logits.squeeze().float().cpu().numpy()
+        prediction_score = score.flatten()
+
+        images_score = {
+            img_path: float(score) for img_path, score in zip(img_paths, prediction_score)
+        }
+
+        # 按value排序
+        images_score = dict(
+            sorted(images_score.items(), key=lambda x: x[1], reverse=True))
+
+        return images_score
+
+
+
+import argparse
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--input_dir", type=str, default="")
+    parser.add_argument("-o", "--output_dir", type=str, default="")
+    parser.add_argument("-c", "--clip_path", type=str,
+                        default="D:\\数据集\\siglip-so400m-patch14-384")
+    parser.add_argument("-m", "--aes_model", type=str,
+                        default="D:\\数据集\\aesthetic_predictor_v2_5.pth")
+    parser.add_argument("-s", "--score", type=float, default=5.9)
+
+    parser.add_argument("-b", "--batch_size", type=int, default=1)
+
+    args = parser.parse_args()
+    limitscore = args.score
+
+    aes = AestheticPredictor(args.clip_path, args.aes_model)
+
+    input_dir = args.input_dir
+    output_dir = args.output_dir
+    batch_size = args.batch_size
+    total = 0
+    for root, dirs, files in os.walk(input_dir):
+        
+        # 排除隐藏文件夹    
+        dirs[:] = [d for d in dirs if not d.startswith('.')]
+        for file in files:
+            if not file.endswith(".png") and not file.endswith(".jpg"):
+                continue
+            total += 1
+
+    import tqdm
+    print(f"total: {total}")
+
+    images = []
+
+    with tqdm.tqdm(total=total) as pbar:
+        for root, dirs, files in os.walk(input_dir):
+            # 排除隐藏文件夹    
+            dirs[:] = [d for d in dirs if not d.startswith('.')]
+            for file in files:
+                if not file.endswith(".png") and not file.endswith(".jpg"):
+                    continue
+                file_path = os.path.join(root, file)
+                pbar.update(1)
+
+                pbar.set_description(
+                    f"Processing {os.path.basename(file_path)}")
+
+                if os.path.exists(os.path.join(output_dir, os.path.basename(file_path))):
+                    continue
+
+                if len(images) < batch_size:
+                    images.append(file_path)
+                else:
+                    image_scroe_map = aes.inference(images)
+                    images = []
+                    for img_path, score in image_scroe_map.items():
+                        print(f"{img_path}: {score}")
+
+                        if score > limitscore:
+                            if not os.path.exists(output_dir):
+                                os.makedirs(output_dir)
+                            shutil.copy(img_path, output_dir)
+
+    if len(images) > 0:
+        image_scroe_map = aes.inference(images)
+        images = []
+
+        for img_path, score in image_scroe_map.items():
+            print(f"{img_path}: {score}")
+
+            if score > limitscore:
+                if not os.path.exists(output_dir):
+                    os.makedirs(output_dir)
+                shutil.copy(img_path, output_dir)
+
+                
@@ -0,0 +1,200 @@
+import argparse
+import hashlib
+import io
+import os
+from PIL import Image
+import numpy as np
+import onnxruntime as ort
+import cv2
+
+
+class Utils:
+    def Md5(bytes):
+        return hashlib.md5(bytes).hexdigest()
+
+    def resize_by(image, percent):
+        # 判断类型是否为PIL
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+
+        width, height = image.size
+        new_width = int(width * percent)
+        new_height = int(height * percent)
+        return image.resize((new_width, new_height), Image.LANCZOS)
+
+    def resize_max(im, dst_w, dst_h):
+        src_w, src_h = im.size
+
+        if src_h > src_w:
+            newWidth = dst_w
+            newHeight = dst_w * src_h // src_w
+        else:
+            newWidth = dst_h * src_w // src_h
+            newHeight = dst_h
+
+        newHeight = newHeight // 8 * 8
+        newWidth = newWidth // 8 * 8
+
+        return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS)
+
+    def resize_min(im, dst_w, dst_h):
+        src_w, src_h = im.size
+
+        if src_h < src_w:
+            newWidth = dst_w
+            newHeight = dst_w * src_h // src_w
+        else:
+            newWidth = dst_h * src_w // src_h
+            newHeight = dst_h
+
+        newHeight = newHeight // 8 * 8
+        newWidth = newWidth // 8 * 8
+
+        return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS)
+
+    # 根据人脸进行剪裁
+    def resize_crop(im, dst_w, dst_h, focus_point):
+
+        src_w, src_h = im.size
+
+        if src_h > src_w:
+            newWidth = dst_w
+            newHeight = dst_w * src_h // src_w
+        else:
+            newWidth = dst_h * src_w // src_h
+            newHeight = dst_h
+
+        newHeight = int(newHeight // 8 * 8)
+        newWidth = int(newWidth // 8 * 8)
+
+        im = im.resize((newWidth, newHeight), Image.Resampling.LANCZOS)
+
+        left = (newWidth - dst_w) // 2
+        top = (newHeight - dst_h) // 2
+
+        right = left + dst_w
+        bottom = top + dst_h
+
+        focus_point_x, focus_point_y = focus_point
+        offset_w = dst_w // 8
+        offset_h = dst_h // 8
+        if focus_point_x < left + dst_w // 8:
+            left = 0
+            right = dst_w
+        elif focus_point_x > right - dst_w // 8:
+            right = newWidth
+            left = newWidth - dst_w
+        if focus_point_y < top + dst_h // 8:
+            top = 0
+            bottom = dst_h
+        elif focus_point_y > bottom - dst_h // 8:
+            bottom = newHeight
+            top = newHeight - dst_h
+
+        return im.crop((left, top, right, bottom))
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--input_dir", type=str, default="")
+    parser.add_argument("-o", "--output_dir", type=str, default="")
+    parser.add_argument("-m", "--face_model_path", type=str, default="")
+    parser.add_argument("-r", "--resize", type=int, default=1024)
+    parser.add_argument("-s", "--score", type=float, default=0.8)
+    parser.add_argument("-f", "--format", type=str, default="jpg")
+
+    args = parser.parse_args()
+    limitscore = args.score
+    sess = ort.InferenceSession(args.face_model_path)
+    sess.set_providers(['CUDAExecutionProvider'])
+    input_dir = args.input_dir
+    output_dir = args.output_dir
+    resize = args.resize
+    img_format = args.format
+    total = 0
+    for root, dirs, files in os.walk(input_dir):
+        # 排除隐藏文件夹
+        dirs[:] = [d for d in dirs if not d.startswith('.')]
+        for file in files:
+            if not file.endswith(".png") and not file.endswith(".jpg"):
+                continue
+            total += 1
+
+    import tqdm
+    print(f"total: {total}")
+
+    os.makedirs(output_dir, exist_ok=True)
+    with tqdm.tqdm(total=total) as pbar:
+
+        # 排除隐藏文件夹
+        dirs[:] = [d for d in dirs if not d.startswith('.')]
+        for root, dirs, files in os.walk(input_dir):
+            for file in files:
+                if not file.endswith(".png") and not file.endswith(".jpg"):
+                    continue
+                file_path = os.path.join(root, file)
+                pbar.update(1)
+
+                with open(file_path, "rb") as f:
+                    md5_str = Utils.Md5(f.read())
+                if os.path.exists(os.path.join(output_dir, md5_str + ".png")):
+                    continue
+
+                image = Image.open(file_path)
+                image = Utils.resize_min(image, resize, resize)
+
+                cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+
+                print(f"start check {file_path}")
+                scores, bboxes, keypoints, aligned_imgs, landmarks, affine_matrices = sess.run(
+                    None, {"input": cv2_img})
+                if len(bboxes) != 1:
+                    continue
+
+                bboxes = bboxes[0]
+                scores = scores[0]
+                if limitscore > scores:
+                    continue
+
+                affine_matrices = affine_matrices[0]
+                angle = np.degrees(np.arctan2(
+                    affine_matrices[1][0], affine_matrices[0][0]))
+
+                real_angle = 0
+                if angle < -45 and angle > -135:
+                    # 向左旋转90度
+                    real_angle = 90
+                elif angle > 45 and angle < 135:
+                    # 向右旋转90度
+                    real_angle = -90
+                elif angle < -135 or angle > 135:
+                    # 向左旋转180度
+                    real_angle = 180
+
+                if real_angle != 0:
+                    continue
+
+                # print(f"{file_path} bboxes: \n{bboxes} {scores} {real_angle}")
+
+                center_point = (bboxes[0] + bboxes[2]) / \
+                    2, (bboxes[1] + bboxes[3]) / 2
+                image = Utils.resize_crop(
+                    image, resize * 3 / 4, resize, center_point)
+                cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+                scores, bboxes, keypoints, aligned_imgs, landmarks, affine_matrices = sess.run(
+                    None, {"input": cv2_img})
+                if len(bboxes) != 1:
+                    continue
+                # cv2.imshow("image", cv2.cvtColor(
+                #     np.array(image), cv2.COLOR_RGB2BGR))
+                # cv2.waitKey(0)
+                if img_format == "png":
+                    output_path = os.path.join(output_dir, md5_str + ".png")
+                    image.save(output_path)
+                else:
+                    output_path = os.path.join(output_dir, md5_str + ".jpg")
+                    image.save(output_path, "jpeg", quality=95)
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,245 @@
+import argparse
+import hashlib
+import io
+import os
+from PIL import Image
+import numpy as np
+import onnxruntime as ort
+import cv2
+
+import piexif
+
+
+class Utils:
+    def Md5(bytes):
+        return hashlib.md5(bytes).hexdigest()
+
+    def resize_by(image, percent):
+        # 判断类型是否为PIL
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+
+        width, height = image.size
+        new_width = int(width * percent)
+        new_height = int(height * percent)
+        return image.resize((new_width, new_height), Image.LANCZOS)
+
+    def resize_max(im, dst_w, dst_h):
+        src_w, src_h = im.size
+
+        if src_h > src_w:
+            newWidth = dst_w
+            newHeight = dst_w * src_h // src_w
+        else:
+            newWidth = dst_h * src_w // src_h
+            newHeight = dst_h
+
+        newHeight = newHeight // 8 * 8
+        newWidth = newWidth // 8 * 8
+
+        return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS)
+
+    def resize_min(im, dst_w, dst_h):
+        src_w, src_h = im.size
+
+        if src_h < src_w:
+            newWidth = dst_w
+            newHeight = dst_w * src_h // src_w
+        else:
+            newWidth = dst_h * src_w // src_h
+            newHeight = dst_h
+
+        newHeight = newHeight // 8 * 8
+        newWidth = newWidth // 8 * 8
+
+        return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS)
+
+    # 根据人脸进行剪裁
+    def resize_crop(im, dst_w, dst_h, focus_point):
+
+        src_w, src_h = im.size
+
+        if src_h > src_w:
+            newWidth = dst_w
+            newHeight = dst_w * src_h // src_w
+        else:
+            newWidth = dst_h * src_w // src_h
+            newHeight = dst_h
+
+        newHeight = int(newHeight // 8 * 8)
+        newWidth = int(newWidth // 8 * 8)
+
+        im = im.resize((newWidth, newHeight), Image.Resampling.LANCZOS)
+
+        left = (newWidth - dst_w) // 2
+        top = (newHeight - dst_h) // 2
+
+        right = left + dst_w
+        bottom = top + dst_h
+
+        focus_point_x, focus_point_y = focus_point
+        offset_w = dst_w // 8
+        offset_h = dst_h // 8
+        if focus_point_x < left + dst_w // 8:
+            left = 0
+            right = dst_w
+        elif focus_point_x > right - dst_w // 8:
+            right = newWidth
+            left = newWidth - dst_w
+        if focus_point_y < top + dst_h // 8:
+            top = 0
+            bottom = dst_h
+        elif focus_point_y > bottom - dst_h // 8:
+            bottom = newHeight
+            top = newHeight - dst_h
+
+        return im.crop((left, top, right, bottom))
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--input_dir", type=str, default="")
+    parser.add_argument("-o", "--output_dir", type=str, default="")
+    parser.add_argument("-m", "--face_model_path", type=str, default="")
+    parser.add_argument("-r", "--resize", type=int, default=1024)
+    parser.add_argument("-p", "--face_proportion", type=float, default=0.5) 
+    parser.add_argument("-f", "--format", type=str, default="jpg")
+
+    args = parser.parse_args()
+    sess = ort.InferenceSession(args.face_model_path)
+    sess.set_providers(['CUDAExecutionProvider'])
+    input_dir = args.input_dir
+    output_dir = args.output_dir
+    resize = args.resize
+    img_format = args.format
+    total = 0
+    for root, dirs, files in os.walk(input_dir):
+        # 排除隐藏文件夹
+        dirs[:] = [d for d in dirs if not d.startswith('.')]
+        for file in files:
+            if not file.endswith(".png") and not file.endswith(".jpg"):
+                continue
+            total += 1
+
+    import tqdm
+    print(f"total: {total}")
+
+    os.makedirs(output_dir, exist_ok=True)
+    with tqdm.tqdm(total=total) as pbar:
+
+        # 排除隐藏文件夹
+        dirs[:] = [d for d in dirs if not d.startswith('.')]
+        for root, dirs, files in os.walk(input_dir):
+            for file in files:
+                if not file.endswith(".png") and not file.endswith(".jpg"):
+                    continue
+                file_path = os.path.join(root, file)
+                pbar.update(1)
+
+                with open(file_path, "rb") as f:
+                    md5_str = Utils.Md5(f.read())
+                if os.path.exists(os.path.join(output_dir, md5_str + ".png")):
+                    continue
+
+                image = Image.open(file_path)
+                max_image = image.copy()
+                image = Utils.resize_min(image, 1920, 1920)
+                zoom = image.size[0] / max_image.size[0]
+
+                cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+
+                print(f"start check {file_path}")
+                scores, bboxes, keypoints, aligned_imgs, landmarks, affine_matrices = sess.run(
+                    None, {"input": cv2_img})
+                if len(bboxes) != 1:
+                    continue
+
+                bboxes = bboxes[0]
+
+                affine_matrices = affine_matrices[0]
+                angle = np.degrees(np.arctan2(
+                    affine_matrices[1][0], affine_matrices[0][0]))
+
+                real_angle = 0
+                if angle < -45 and angle > -135:
+                    # 向左旋转90度
+                    real_angle = 90
+                elif angle > 45 and angle < 135:
+                    # 向右旋转90度
+                    real_angle = -90
+                elif angle < -135 or angle > 135:
+                    # 向左旋转180度
+                    real_angle = 180
+
+                if real_angle != 0:
+                    continue
+
+                # print(f"{file_path} bboxes: \n{bboxes} {scores} {real_angle}")
+
+                bboxes = bboxes / zoom
+
+                image = max_image
+
+                # 预期人脸高度占图片高度的一定比例
+                pr = args.face_proportion
+                face_width = bboxes[2] - bboxes[0]
+                pre_image_width = face_width / pr
+
+                face_center_x = (bboxes[2] + bboxes[0]) / 2
+                crop_x1 = face_center_x - pre_image_width / 2
+                crop_x2 = face_center_x + pre_image_width / 2
+
+                # 判断是否超出图片边界
+                if crop_x1 < 0 or crop_x2 > image.width:
+                    print(
+                        f"crop_x1: {crop_x1}, crop_x2: {crop_x2}, image.width: {image.width}")
+                    continue
+
+                face_center_y = (bboxes[3] + bboxes[1]) / 2
+                # 总高度和宽度的比例是4:3
+                pre_image_height = pre_image_width * 4 / 3
+
+                # 中点位于高度的0.382比例
+                crop_y1 = face_center_y - pre_image_height * 0.382
+                crop_y2 = crop_y1 + pre_image_height
+
+                # 判断是否超出图片边界
+                if crop_y1 < 0 or crop_y2 > image.height:
+                    print(
+                        f"crop_y1: {crop_y1}, crop_y2: {crop_y2}, image.height: {image.height}")
+                    continue
+
+                # 判断是否比resize小
+                if pre_image_width < resize or pre_image_height < resize:
+                    print(
+                        f"pre_image_width: {pre_image_width}, pre_image_height: {pre_image_height}")
+                    continue
+
+                # 进行裁剪
+                image = image.crop((crop_x1, crop_y1, crop_x2, crop_y2))
+
+                # 进行缩放并居中剪裁
+                image = image.resize((
+                    resize * 3 // 4, resize
+                ), Image.LANCZOS)
+
+                # cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+                # cv2.imshow("image", cv2_img)
+                # cv2.waitKey(0)
+
+                if img_format == "png":
+                    output_path = os.path.join(output_dir, md5_str + ".png")
+                    image.save(output_path)
+                else:
+                    output_path = os.path.join(output_dir, md5_str + ".jpg")
+                    image.save(output_path, "jpeg", quality=95)
+                    # 插入IPTC
+                    # piexif.insert(output_path, piexif.ImageIFD, {
+                    #     "CameraOwnerName": "CameraOwnerName",
+                    # })
+
+                print(f"save to {output_path}")
+
+
+if __name__ == "__main__":
+    main()