'update'

hook_HYDiT_main_train_deepspeed.py

@@ -8,6 +8,7 @@ from functools import partial
 from glob import glob
 from pathlib import Path
 import numpy as np
+import safetensors.torch
 
 import deepspeed
 import torch
@@ -382,12 +383,19 @@ def Core(args):
     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:
-            from peft import PeftConfig
-            loraconfig = PeftConfig.from_pretrained(lora_ckpt)
+            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,
@@ -395,10 +403,10 @@ def Core(args):
                 target_modules=args.target_modules
             )
 
-        if args.use_fp16:
-            model.module = get_peft_model(model.module, loraconfig)
-        else:
-            model = get_peft_model(model, loraconfig)
+            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}")
 
@@ -409,6 +417,11 @@ def Core(args):
     # 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()

hook_HYDiT_utils.py

@@ -291,8 +291,8 @@ from torch import nn
 
 class CustomizeEmbedsModel(nn.Module):
     dtype = torch.float16
-    x = torch.zeros(1, 1, 256, 2048)
-
+    # x = torch.zeros(1, 1, 256, 2048)
+    x = None
     def __init__(self, *args, **kwargs):
         super().__init__()
 
@@ -300,7 +300,13 @@ class CustomizeEmbedsModel(nn.Module):
         self.dtype = torch.float16
         return self
 
-    def forward(self, *args, **kwargs):
+    def forward(self, *args, **kwargs): 
+        input_ids = kwargs.get("input_ids", None)
+        if self.x is None:
+            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"),
@@ -315,15 +321,15 @@ class CustomizeTokenizer(dict):
     input_ids = torch.zeros(1, 256)
     attention_mask = torch.zeros(1, 256)
 
-    def __init__(self):
+    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):
+    def tokenize(self, text):  
         return text
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args, **kwargs): 
         return self
 
 

mz_train_tools_core_HYDiT.py

@@ -72,7 +72,8 @@ def MZ_HYDiTInitClone_call(args={}):
         with open(os.path.join(hunyuan_lora_dir, "hydit/diffusion/pipeline.py"), "r", encoding="utf-8") as f:
             pre_replace = "device = self._execution_device"
             content = f.read()
-            content = content.replace(pre_replace, "device = torch.device('cuda')")
+            content = content.replace(
+                pre_replace, "device = torch.device('cuda')")
         with open(os.path.join(hunyuan_lora_dir, "hydit/diffusion/pipeline.py"), "w", encoding="utf-8") as f:
             f.write(content)
 
@@ -442,6 +443,9 @@ def MZ_HYDiTTrain_call(args={}):
 
     # python ./hydit/data_loader/csv2arrow.py ./dataset/porcelain/csvfile/image_text.csv ./dataset/porcelain/arrows
     arrows_dir = os.path.join(workspace_dir, "arrows")
+    if os.path.exists(arrows_dir):
+        shutil.rmtree(arrows_dir)
+
     os.makedirs(arrows_dir, exist_ok=True)
     csv2arrow_exec = os.path.join(
         HYDiT_tool_dir, "hydit", "data_loader", "csv2arrow.py")
@@ -534,15 +538,12 @@ def MZ_HYDiTTrain_call(args={}):
         lora_ckpt = None
     elif lora_ckpt == "latest":
         lora_ckpt = None
-        loras = search_loras(workspace_dir)
+        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,
@@ -564,6 +565,7 @@ def MZ_HYDiTTrain_call(args={}):
         "results_dir": output_dir,
         "task_flag": output_name,
         "lora_ckpt": lora_ckpt,
+        "rank": int(args.get("rank")),
     })
 
     if "target_modules" in train_config:
@@ -705,10 +707,9 @@ def get_HunYuanDiT_model_from_path(model_path, lora_path, width, height):
 
     from hydit.modules.models import HUNYUAN_DIT_CONFIG
     import importlib
-    import hydit.modules.models 
+    import hydit.modules.models
     importlib.reload(hydit.modules.models)
     HunYuanDiT = hydit.modules.models.HunYuanDiT
-    
 
     from types import SimpleNamespace
 
@@ -1091,13 +1092,12 @@ def search_loras(wdirs):
         for root, dirs, files in os.walk(wdir):
             # 排除隐藏文件夹
             dirs[:] = [d for d in dirs if not d.startswith('.')]
-
+            # 存在adapter_config.json文件并且存在adapter_model.safetensors文件
             for file in files:
-                if file.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))
+                if file == "adapter_config.json":
+                    lora_path = os.path.join(root, "adapter_model.safetensors")
+                    if os.path.exists(lora_path):
+                        loras.append(root)
 
     # 按时间倒序
     loras = sorted(loras, key=lambda x: os.path.getmtime(x), reverse=True)