From 464b3b5c675ab4734f5ae769ad0cf85301043e7b Mon Sep 17 00:00:00 2001 From: wailovet Date: Thu, 27 Jun 2024 23:01:02 +0800 Subject: [PATCH] 'update' --- __init__.py | 4 +- hook_HYDiT_main.py | 161 ++++++++++++++++++++++++++++++++--- hook_HYDiT_run.py | 15 ++-- mz_train_tools_core_HYDiT.py | 93 +++++++++++++------- mz_train_tools_utils.py | 2 + 5 files changed, 223 insertions(+), 52 deletions(-) diff --git a/__init__.py b/__init__.py index 189e9df..c8a6dd6 100644 --- a/__init__.py +++ b/__init__.py @@ -58,7 +58,7 @@ class MZ_KohyaSSDatasetConfig: "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}), "force_clear": (["enable", "disable"], {"default": "disable"}), @@ -539,7 +539,7 @@ class MZ_HYDiTTrain: "resolution": ("INT", {"default": 1024}), "batch_size": ("INT", {"default": 1}), "epochs": ("INT", {"default": 50}), - "save_every_n_epochs": ("INT", {"default": 10}), + "ckpt_every": ("INT", {"default": 10}), "base_lora": (["latest", "empty"], {"default": "latest"}), "sample_generate": (["enable", "disable"], {"default": "enable"}), "sample_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}), diff --git a/hook_HYDiT_main.py b/hook_HYDiT_main.py index d33d547..c9eaac4 100644 --- a/hook_HYDiT_main.py +++ b/hook_HYDiT_main.py @@ -83,6 +83,14 @@ T5_ENCODER = { } +global TRAIN_CONFIG + + +def set_train_config(train_config): + global TRAIN_CONFIG + TRAIN_CONFIG = train_config + + def set_t5_encoder_path(path): global T5_ENCODER T5_ENCODER['MT5'] = path @@ -135,12 +143,9 @@ def save_checkpoint(args, rank, logger, model, ema, epoch, train_steps, checkpoi cur_ckpt_save_dir = f"{checkpoint_dir}/{tag}" if rank == 0: if args.use_fp16: - try: - model.module.save_pretrained(cur_ckpt_save_dir) - except Exception as e: - model.save_pretrained(cur_ckpt_save_dir) - else: model.module.save_pretrained(cur_ckpt_save_dir) + else: + model.save_pretrained(cur_ckpt_save_dir) checkpoint_path = "[Not rank 0. Disabled output.]" @@ -183,12 +188,6 @@ def save_checkpoint(args, rank, logger, model, ema, epoch, train_steps, checkpoi except: logger.error(f"Saved failed to {checkpoint_path}") - 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 @@ -441,6 +440,7 @@ def Core(args, LOG): model, ema, start_epoch, start_epoch_step, train_steps = model_resume( args, model, ema, logger) + if args.training_parts == "lora": loraconfig = LoraConfig( r=args.rank, @@ -537,10 +537,15 @@ def Core(args, LOG): optimizer = torch.optim.AdamW( model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + + + easy_sample_images(args, vae, text_encoder, tokenizer, model, embedder_t5, + target_height=768, target_width=1280, train_steps=0) # torch.autograd.set_detect_anomaly(True) # with torch.autograd.detect_anomaly(): with tqdm(total=total_steps, initial=train_steps) as pbar: + # Training loop for epoch in range(start_epoch, args.epochs): logger.info(f" Start random shuffle with seed={seed}") @@ -605,8 +610,9 @@ def Core(args, LOG): train_steps += 1 mean_loss += loss.item() pbar.update(1) - pbar.set_description(f"Epoch {epoch}, step {step}, loss {loss.item():.4f}, mean_loss {mean_loss / step:.4f}") - + pbar.set_description( + f"Epoch {epoch}, step {step}, loss {loss.item():.4f}, mean_loss {mean_loss / step:.4f}") + # =========================================================================== # Log loss values: # =========================================================================== @@ -640,9 +646,11 @@ def Core(args, LOG): gc.collect() 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: + ) and train_steps > 0: logger.info( f" Saving checkpoint at step {train_steps}.") + 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) @@ -708,3 +716,128 @@ def model_resume(args, model, ema, logger): " “If `resume` is True, then either `resume_split` must be true.”") return model, ema, start_epoch, start_epoch_step, train_steps + + +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=5.0, + infer_steps=20, + sampler='ddpm', + train_steps=0, +): + + with torch.no_grad(): + + 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 + + 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) + + # 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, + scheduler=scheduler, + feature_extractor=None, + safety_checker=None, + requires_safety_checker=False, + embedder_t5=embedder_t5, + ) + pipeline.to("cuda") + + 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) + 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, + )[0] + + # 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 diff --git a/hook_HYDiT_run.py b/hook_HYDiT_run.py index 3ced37c..180e725 100644 --- a/hook_HYDiT_run.py +++ b/hook_HYDiT_run.py @@ -5,6 +5,8 @@ import os import sys from types import SimpleNamespace +import torch + # python -m torch.distributed.launch ./hook_HYDi_run.py --sys_path /data/ComfyUI/models/minus_zone_models/train_tools/HunyuanDiT/ @@ -87,11 +89,9 @@ 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.resume_split = train_config.get("resume_split", True) margs.ema_to_module = train_config.get("ema_to_module", True) margs.deepspeed = False @@ -109,8 +109,8 @@ if __name__ == "__main__": 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", "base512") - margs.image_size = train_config.get("image_size", 512) + 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) @@ -122,7 +122,8 @@ if __name__ == "__main__": margs.log_every = train_config.get("log_every", 10) margs.results_dir = train_config.get("results_dir") - margs.mse_loss_weight_type = train_config.get("mse_loss_weight_type", "constant") + margs.mse_loss_weight_type = train_config.get( + "mse_loss_weight_type", "constant") for k, v in train_config.items(): if hasattr(margs, k): @@ -143,4 +144,6 @@ if __name__ == "__main__": hook_HYDiT_main.set_t5_encoder_path( train_config.get("t5_encoder_path")) + hook_HYDiT_main.set_train_config(train_config) + hook_HYDiT_main.Core(margs, LOG) diff --git a/mz_train_tools_core_HYDiT.py b/mz_train_tools_core_HYDiT.py index 7c04555..ebf84e1 100644 --- a/mz_train_tools_core_HYDiT.py +++ b/mz_train_tools_core_HYDiT.py @@ -110,7 +110,7 @@ def MZ_HYDiTDatasetConfig_call(args={}): pil_images = Utils.tensors2pil_list(images) - resolution = args.get("resolution", 512) + resolution = args.get("resolution", 1024) if workspace_name is None or workspace_name == "": raise Exception("训练名称不能为空(workspace_name is required)") @@ -282,13 +282,20 @@ def MZ_HYDiTTrain_call(args={}): # "tokenizer": (["auto"] + folders, {"default": "auto"}), # "t5_encoder": (["auto"] + folders, {"default": "auto"}), + epochs = args.get("epochs", 50) output_dir = os.path.join(workspace_dir, "output") + + datetime = time.strftime("%Y%m%d_%H%M%S", time.localtime()) + output_name = f"{workspace_name}_{datetime}" + train_config.update({ + "workspace_name": workspace_name, + "workspace_dir": workspace_dir, "batch_size": args.get("batch_size", 1), "target_ratios": train_config.get("target_ratios", ["1:1", "3:4", "4:3", "16:9", "9:16"]), "index_file": dataset_mt_json, - "epochs": args.get("epochs", 50), - "ckpt_every": args.get("save_every_n_epochs", 1), + "epochs": epochs, + "ckpt_every": int(args.get("ckpt_every")), "rope_img": f"base{resolution}", "rope_real": args.get("rope_real", "enable") == "enable", "image_size": resolution, @@ -298,11 +305,31 @@ def MZ_HYDiTTrain_call(args={}): "tokenizer_path": args.get("tokenizer_path"), "t5_encoder_path": args.get("t5_encoder_path"), "results_dir": output_dir, - "task_flag": args.get("task_flag", workspace_name), + "task_flag": output_name, }) + sample_generate = args.get("sample_generate", "enable") + sample_prompt = args.get("sample_prompt", "") + sample_config_file = os.path.join(workspace_dir, "sample_config.json") + if sample_generate == "enable": + resolution = args.get("resolution") + sample_config = { + "prompt": sample_prompt, + "negative_prompt": "", + "cfg": 5.0, + "steps": 20, + "width": resolution, + "height": resolution, + } + with open(sample_config_file, "w", encoding="utf-8") as f: + json.dump(sample_config, f, indent=4, ensure_ascii=False) + + train_config["sample_config_file"] = sample_config_file + else: + train_config["sample_config_file"] = None + with open(train_config_path, "w", encoding="utf-8") as f: - json.dump(train_config, f, indent=4) + json.dump(train_config, f, indent=4, ensure_ascii=False) run_hook_HYDiT_pyexec(train_config_path, train_config, HYDiT_tool_dir) return ( @@ -311,7 +338,7 @@ def MZ_HYDiTTrain_call(args={}): def run_hook_HYDiT_pyexec(train_config_path, train_config, HYDiT_tool_dir): - pb = Utils.progress_bar(train_config.get("epochs"), "sdxl") + pb = Utils.progress_bar(0, "sdxl") import traceback import comfy.model_management stop_server = None @@ -365,30 +392,36 @@ def run_hook_HYDiT_pyexec(train_config_path, train_config, HYDiT_tool_dir): def get_sample_images(train_config): - # output_name = train_config.get("output_name") - # sample_images_dir = os.path.join( - # os.path.dirname(train_config.get("dataset_config")), "sample_images" - # ) - # pil_images = [] - # pre_render_texts_x = [] - # if os.path.exists(sample_images_dir): - # image_files = os.listdir(sample_images_dir) - # image_files = list( - # filter(lambda x: x.endswith(".png"), image_files)) - # # 筛选 output_name 前缀 - # image_files = list( - # filter(lambda x: x.startswith(output_name), image_files)) + if train_config.get("sample_config_file") is None: + return None - # image_files = sorted(image_files, key=lambda x: x) + from PIL import Image + output_name = train_config.get("task_flag") - # for image_file in image_files: - # pil_image = Image.open(os.path.join(sample_images_dir, image_file)) - # pil_images.append([pil_image]) - # pre_render_texts_x.append(image_file) - # result = Utils.xy_image( - # pre_render_images=pil_images, - # pre_render_texts_x=pre_render_texts_x, - # pre_render_texts_y=[""], - # ) - # return result + workspace_dir = train_config.get("workspace_dir") + + sample_images_dir = os.path.join(workspace_dir, "sample_images") + + pil_images = [] + pre_render_texts_x = [] + if os.path.exists(sample_images_dir): + image_files = os.listdir(sample_images_dir) + image_files = list( + filter(lambda x: x.endswith(".png"), image_files)) + # 筛选 output_name 前缀 + image_files = list( + filter(lambda x: x.startswith(output_name), image_files)) + + image_files = sorted(image_files, key=lambda x: x) + + for image_file in image_files: + pil_image = Image.open(os.path.join(sample_images_dir, image_file)) + pil_images.append([pil_image]) + pre_render_texts_x.append(image_file) + result = Utils.xy_image( + pre_render_images=pil_images, + pre_render_texts_x=pre_render_texts_x, + pre_render_texts_y=[""], + ) + return result return None diff --git a/mz_train_tools_utils.py b/mz_train_tools_utils.py index c5345eb..b67a9d2 100644 --- a/mz_train_tools_utils.py +++ b/mz_train_tools_utils.py @@ -787,6 +787,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]