From 661b8f61063c8ded4dbf68a098d6cd6ba8d283c1 Mon Sep 17 00:00:00 2001 From: wailovet Date: Wed, 3 Jul 2024 00:12:57 +0800 Subject: [PATCH] 'update' --- hook_HYDiT_main_train_deepspeed.py | 27 ++++++++++++++++++++------- hook_HYDiT_utils.py | 18 ++++++++++++------ mz_train_tools_core_HYDiT.py | 26 +++++++++++++------------- 3 files changed, 45 insertions(+), 26 deletions(-) diff --git a/hook_HYDiT_main_train_deepspeed.py b/hook_HYDiT_main_train_deepspeed.py index 73d50eb..c25fa9d 100644 --- a/hook_HYDiT_main_train_deepspeed.py +++ b/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() diff --git a/hook_HYDiT_utils.py b/hook_HYDiT_utils.py index 2afcbc2..8d40308 100644 --- a/hook_HYDiT_utils.py +++ b/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 diff --git a/mz_train_tools_core_HYDiT.py b/mz_train_tools_core_HYDiT.py index 5c7fd23..18e4390 100644 --- a/mz_train_tools_core_HYDiT.py +++ b/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)