From d87d098bc83f1ee28651c69944c06e25c4525957 Mon Sep 17 00:00:00 2001 From: wailovet Date: Sun, 30 Jun 2024 01:31:22 +0800 Subject: [PATCH] 'update' --- __init__.py | 69 +++++++++++++++++++++++------- hook_HYDiT_main_train_deepspeed.py | 21 ++++++--- hook_HYDiT_run.py | 25 +++++++++-- hook_HYDiT_utils.py | 21 +++------ mz_train_tools_core_HYDiT.py | 56 ++++++++++++++++++++++-- 5 files changed, 147 insertions(+), 45 deletions(-) diff --git a/__init__.py b/__init__.py index 9acde5a..19f0b59 100644 --- a/__init__.py +++ b/__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 diff --git a/hook_HYDiT_main_train_deepspeed.py b/hook_HYDiT_main_train_deepspeed.py index 521d3fc..73d50eb 100644 --- a/hook_HYDiT_main_train_deepspeed.py +++ b/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: diff --git a/hook_HYDiT_run.py b/hook_HYDiT_run.py index e5300e3..8667e6f 100644 --- a/hook_HYDiT_run.py +++ b/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") diff --git a/hook_HYDiT_utils.py b/hook_HYDiT_utils.py index eadbe30..2afcbc2 100644 --- a/hook_HYDiT_utils.py +++ b/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), diff --git a/mz_train_tools_core_HYDiT.py b/mz_train_tools_core_HYDiT.py index dd34ae2..a8b8973 100644 --- a/mz_train_tools_core_HYDiT.py +++ b/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