'update'

__init__.py

@@ -416,7 +416,7 @@ class MZ_HYDiTInitWorkspace:
 
     def start(self, **kwargs):
         importlib.reload(mz_train_tools_core_HYDiT)
-        return mz_train_tools_core_HYDiT.MZ_HYDiTInitWorkspace_call(kwargs)
+        return mz_train_tools_core_HYDiT.MZ_HYDiTInitWorkspace_call(kwargs.copy())
 
 
 NODE_CLASS_MAPPINGS["MZ_HYDiTInitWorkspace"] = MZ_HYDiTInitWorkspace
@@ -449,7 +449,7 @@ class MZ_HYDiTDatasetConfig:
 
     def start(self, **kwargs):
         importlib.reload(mz_train_tools_core_HYDiT)
-        return mz_train_tools_core_HYDiT.MZ_HYDiTDatasetConfig_call(kwargs)
+        return mz_train_tools_core_HYDiT.MZ_HYDiTDatasetConfig_call(kwargs.copy())
 
 
 NODE_CLASS_MAPPINGS["MZ_HYDiTDatasetConfig"] = MZ_HYDiTDatasetConfig
@@ -491,14 +491,40 @@ class MZ_HYDiTAdvConfig:
     def INPUT_TYPES(s):
         return {
             "required": {
-                "learn_sigma": (["enable", "disable"], {"default": "enable"}),
-                "text_states_dim": ("INT", {"default": 1024}),
-                "text_states_dim_t5": ("INT", {"default": 2048}),
-                "text_len": ("INT", {"default": 77}),
-                "text_len_t5": ("INT", {"default": 256}),
-                "norm": (["layer", "instance"], {"default": "layer"}),
-                "infer_mode": (["torch", "deepspeed"], {"default": "torch"}),
-                "use_fp16": (["enable", "disable"], {"default": "enable"}),
+                "lr": ("FLOAT", {"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": ("FLOAT", {"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"}),
             }
         }
 
@@ -510,7 +536,7 @@ class MZ_HYDiTAdvConfig:
     CATEGORY = CATEGORY_NAME
 
     def start(self, **kwargs):
-        return (kwargs,)
+        return (kwargs.copy(),)
 
 
 NODE_CLASS_MAPPINGS["MZ_HYDiTAdvConfig"] = MZ_HYDiTAdvConfig
@@ -532,24 +558,35 @@ class MZ_HYDiTTrain:
             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": "auto"}),
+                "t5_encoder_path": (["none", "auto"] + folders, {"default": "none"}),
                 "resolution": ("INT", {"default": 1024}),
                 "batch_size": ("INT", {"default": 1}),
                 "epochs": ("INT", {"default": 50}),
-                "ckpt_every": ("INT", {"default": 10}),
-                "ema_to_module": (["enable", "disable"], {"default": "enable"}),
-                "base_lora": (["latest", "empty"], {"default": "latest"}),
+                "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("*"),),
             }
@@ -570,7 +607,7 @@ base_size The base resolution (n, n) from which to create multiple resolutions |
 
     def start(self, **kwargs):
         importlib.reload(mz_train_tools_core_HYDiT)
-        return mz_train_tools_core_HYDiT.MZ_HYDiTTrain_call(kwargs)
+        return mz_train_tools_core_HYDiT.MZ_HYDiTTrain_call(kwargs.copy())
 
 
 NODE_CLASS_MAPPINGS["MZ_HYDiTTrain"] = MZ_HYDiTTrain

hook_HYDiT_main_train_deepspeed.py

@@ -44,6 +44,9 @@ def deepspeed_initialize(args, logger, model, opt, deepspeed_config):
 
     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),
@@ -381,11 +384,17 @@ def Core(args):
             args, model, ema, logger)
 
     if args.training_parts == "lora":
-        loraconfig = LoraConfig(
-            r=args.rank,
-            lora_alpha=args.rank,
-            target_modules=args.target_modules
-        )
+        lora_ckpt = args.lora_ckpt
+        if lora_ckpt is not None:
+            from peft import PeftConfig
+            loraconfig = PeftConfig.from_pretrained(lora_ckpt)
+        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:
@@ -553,7 +562,7 @@ def Core(args):
                 gc.collect()
 
             pbar.step(
-                f"Epoch {epoch}, step {step}, loss {loss.item():.4f}, mean_loss {running_loss / step:.4f}", args.epochs * len(loader), train_steps)
+                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:

hook_HYDiT_run.py

@@ -78,24 +78,40 @@ if __name__ == "__main__":
     margs = hydit.config.get_args()
     margs.model = train_config.get("model", "DiT-g/2")
 
-    margs.task_flag = train_config.get(
-        "task_flag", "lora_porcelain_ema_rank64")
+    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")
@@ -107,16 +123,17 @@ if __name__ == "__main__":
     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")

hook_HYDiT_utils.py

@@ -85,7 +85,7 @@ def easy_sample_images(
     from hydit.modules.models import HUNYUAN_DIT_CONFIG
 
     import traceback
-    with torch.no_grad():
+    with torch.cuda.amp.autocast():
 
         workspace_dir = TRAIN_CONFIG.get("workspace_dir")
         sample_config_file = TRAIN_CONFIG.get("sample_config_file", None)
@@ -132,20 +132,9 @@ def easy_sample_images(
                                            safety_checker=None,
                                            requires_safety_checker=False,
                                            embedder_t5=embedder_t5,
-                                           ).to("cuda")
+                                           )
+        pipeline = pipeline.to("cuda")
         # attr _execution_device is not defined
-        if not hasattr(pipeline, "_execution_device"):
-            time.sleep(1)
-            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,
-                                               ).to("cuda")
 
         style = torch.as_tensor([0, 0] * batch_size, device="cuda")
 
@@ -190,7 +179,7 @@ def easy_sample_images(
                 print(f"Failed to sample images: {e} ")
                 # 打印堆栈信息
                 traceback.print_exc()
-                continue
+                print(f"Failed to sample pipeline: {pipeline} ")
 
             # print("samples:",type(samples),)
             # input("Press Enter to continue...")
@@ -312,7 +301,7 @@ class CustomizeEmbedsModel(nn.Module):
         return self
 
     def forward(self, *args, **kwargs):
-        if self.output_hidden_states or kwargs.get("output_hidden_states", False):
+        if kwargs.get("output_hidden_states", False):
             return {
                 "hidden_states": self.x.to("cuda"),
                 "input_ids": torch.zeros(1, 1),

mz_train_tools_core_HYDiT.py

@@ -345,7 +345,11 @@ def check_required():
 
 
 def MZ_HYDiTTrain_call(args={}):
+
     check_required()
+
+    advanced_config = args.get("advanced_config", {})
+
     args = check_model_auto_download(args)
     # raise Exception(args)
     resolution = args.get("resolution")
@@ -505,6 +509,21 @@ def MZ_HYDiTTrain_call(args={}):
     datetime = time.strftime("%Y%m%d_%H%M%S", time.localtime())
     output_name = f"{workspace_name}_{datetime}"
 
+    lora_ckpt = args.get("base_lora")
+    if lora_ckpt == "empty":
+        lora_ckpt = None
+    elif lora_ckpt == "latest":
+        lora_ckpt = None
+        loras = search_loras(workspace_dir)
+        if len(loras) > 0:
+            lora_ckpt = loras[0]
+    elif not os.path.exists(lora_ckpt):
+        raise Exception(f"未找到指定的lora文件: {lora_ckpt}")
+
+    if lora_ckpt is not None and lora_ckpt.endswith(".safetensors"):
+        # 文件所在目录
+        lora_ckpt = os.path.dirname(lora_ckpt)
+    print(f"使用指定的lora文件: {lora_ckpt}")
     train_config.update({
         "workspace_name": workspace_name,
         "workspace_dir": workspace_dir,
@@ -524,8 +543,23 @@ def MZ_HYDiTTrain_call(args={}):
         "t5_encoder_path": args.get("t5_encoder_path") if args.get("t5_encoder_path") != "none" else None,
         "results_dir": output_dir,
         "task_flag": output_name,
+        "lora_ckpt": lora_ckpt,
     })
 
+    if "target_modules" in train_config:
+        del train_config["target_modules"]
+    for key in advanced_config:
+        if key.startswith("target_modules_"):
+            if advanced_config[key] == "enable":
+                if "target_modules" not in train_config:
+                    train_config["target_modules"] = []
+                train_config["target_modules"].append(
+                    key.replace("target_modules_", ""))
+        elif advanced_config[key] == "enable" or advanced_config[key] == "disable":
+            train_config[key] = advanced_config[key] == "enable"
+        else:
+            train_config[key] = advanced_config[key]
+
     sample_generate = args.get("sample_generate", "enable")
     sample_prompt = args.get("sample_prompt", "")
     sample_config_file = os.path.join(workspace_dir, "sample_config.json")
@@ -645,7 +679,7 @@ def get_sample_images(train_config):
     return None
 
 
-def get_model_from_path(model_path, lora_path, width, height):
+def get_HunYuanDiT_model_from_path(model_path, lora_path, width, height):
 
     from hydit.modules.models import HunYuanDiT, HUNYUAN_DIT_CONFIG
 
@@ -760,7 +794,7 @@ class CustomizeEmbedsModel(nn.Module):
         return self
 
     def forward(self, *args, **kwargs):
-        if self.output_hidden_states or kwargs.get("output_hidden_states", False):
+        if kwargs.get("output_hidden_states", False):
             return {
                 "hidden_states": self.x.to("cuda"),
                 "input_ids": torch.zeros(1, 1),
@@ -821,7 +855,7 @@ def MZ_HYDiTSimpleT2I_call(args={}):
     if unet_data is None or unet_data.get("unet_path") != unet_path or unet_data.get("lora_path") != lora_path:
         clean_unet()
         _lora_path = lora_path if lora_path != "none" else None
-        unet = get_model_from_path(
+        unet = get_HunYuanDiT_model_from_path(
             args.get("unet_path"), _lora_path, width, height)
         Utils.cache_set(unet_cache_key, {
             "model": unet,
@@ -993,3 +1027,19 @@ def MZ_HYDiTSimpleT2I_call(args={}):
                 torch.cuda.empty_cache()
 
             return (Utils.list_tensor2tensor([tensor_image]),)
+
+
+def search_loras(wdirs):
+    loras = []
+    for wdir in wdirs:
+        for root, dirs, files in os.walk(wdir):
+            for file in files:
+                if file.lower().endswith(".safetensors"):
+                    # 并且在同一目录下有相同名字的json文件
+                    json_file = os.path.splitext(file)[0] + ".json"
+                    if os.path.exists(os.path.join(root, json_file)):
+                        loras.append(os.path.join(root, file))
+
+    # 按时间倒序
+    loras = sorted(loras, key=lambda x: os.path.getmtime(x), reverse=True)
+    return loras