diff --git a/hook_HYDiT_config.py b/hook_HYDiT_config.py deleted file mode 100644 index a98fcfa..0000000 --- a/hook_HYDiT_config.py +++ /dev/null @@ -1,235 +0,0 @@ - - -import argparse - -from hydit.constants import * -from hydit.modules.models import HUNYUAN_DIT_CONFIG, HUNYUAN_DIT_MODELS -from hydit.diffusion.gaussian_diffusion import ModelVarType - - -def model_var_type(value): - try: - return ModelVarType[value] - except KeyError: - raise ValueError( - f"Invalid choice '{value}', valid choices are {[v.name for v in ModelVarType]}") - - -def get_args(default_args=None): - parser = argparse.ArgumentParser() - parser.add_argument("--task-flag", type=str) - - # General Setting - parser.add_argument("--batch-size", type=int, - default=1, help="Per-GPU batch size") - parser.add_argument('--seed', type=int, default=42, - help="A seed for all the prompts.") - parser.add_argument("--use-fp16", action="store_true", - help="Use FP16 precision.") - parser.add_argument("--no-fp16", dest="use_fp16", action="store_false") - parser.set_defaults(use_fp16=True) - parser.add_argument("--extra-fp16", action="store_true", - help="Use extra fp16 for vae and text_encoder.") - - # HunYuan-DiT - parser.add_argument( - "--model", type=str, choices=list(HUNYUAN_DIT_CONFIG.keys()), default='DiT-g/2') - parser.add_argument("--image-size", type=int, nargs='+', default=[1024, 1024], - help='Image size (h, w). If a single value is provided, the image will be treated to ' - '(value, value).') - parser.add_argument("--qk-norm", action="store_true", - help="Query Key normalization. See http://arxiv.org/abs/2302.05442 for details.") - parser.set_defaults(qk_norm=True) - parser.add_argument( - "--norm", type=str, choices=["rms", "laryer"], default="layer", help="Normalization layer type") - parser.add_argument("--text-states-dim", type=int, - default=1024, help="Hidden size of CLIP text encoder.") - parser.add_argument("--text-len", type=int, default=77, - help="Token length of CLIP text encoder output.") - parser.add_argument("--text-states-dim-t5", type=int, - default=2048, help="Hidden size of CLIP text encoder.") - parser.add_argument("--text-len-t5", type=int, default=256, - help="Token length of T5 text encoder output.") - - # LoRA config - parser.add_argument("--training-parts", type=str, default='all', - choices=['all', 'lora'], help="Training parts") - parser.add_argument("--rank", type=int, default=64, help="Rank of LoRA") - parser.add_argument("--lora-ckpt", type=str, - default=None, help="LoRA checkpoint") - parser.add_argument('--target-modules', type=str, nargs='+', default=['Wqkv', 'q_proj', 'kv_proj', 'out_proj'], - help="Target modules for LoRA fine tune") - parser.add_argument("--output-merge-path", type=str, - default=None, help="Output path for merged model") - - # controlnet config - parser.add_argument("--control-type", type=str, default='canny', - choices=['canny', 'depth', 'pose'], help="Controlnet condition type") - parser.add_argument("--control-weight", type=float, - default=1.0, help="Controlnet weight") - parser.add_argument("--condition-image-path", type=str, - default=None, help="Inference condition image path") - - # Diffusion - parser.add_argument("--learn-sigma", action="store_true", - help="Learn extra channels for sigma.") - parser.add_argument("--no-learn-sigma", - dest="learn_sigma", action="store_false") - parser.set_defaults(learn_sigma=True) - parser.add_argument("--predict-type", type=str, default="v_prediction", - help="Diffusion predict type") - parser.add_argument("--noise-schedule", type=str, default="scaled_linear", - help="Noise schedule") - parser.add_argument("--beta-start", type=float, - default=0.00085, help="Beta start value") - parser.add_argument("--beta-end", type=float, - default=0.03, help="Beta end value") - parser.add_argument("--sigma-small", action="store_true") - parser.add_argument("--mse-loss-weight-type", type=str, default="constant", - help="Min-SNR-gamma. Can be constant or min_snr_ where gamma is a integer. 5 is recommended in the paper.") - parser.add_argument("--model-var-type", type=model_var_type, - default=None, help="Specify the model variable type.") - parser.add_argument("--noise-offset", type=float, default=0.0, - help="Add extra noise to the input image.") - - # ======================================================================================================== - # Inference - # ======================================================================================================== - - # Basic Setting - parser.add_argument("--prompt", type=str, default="一只小猫", - help="The prompt for generating images.") - parser.add_argument("--model-root", type=str, - default="ckpts", help="Model root path.") - - # Model setting - parser.add_argument("--load-key", type=str, choices=[ - "ema", "module", "distill", 'merge'], default="ema", help="Load model key for HunYuanDiT checkpoint.") - parser.add_argument('--size-cond', type=int, nargs='+', default=[1024, 1024], - help="Size condition used in sampling. 2 values are required for height and width. " - "If a single value is provided, the image will be treated to (value, value).") - parser.add_argument('--target-ratios', type=str, nargs='+', default=None, - help="Target ratios for multi-resolution training.") - parser.add_argument("--cfg-scale", type=float, default=6.0, - help="Guidance scale for classifier-free.") - parser.add_argument("--negative", type=str, - default=None, help="Negative prompt.") - - # Acceleration - parser.add_argument("--infer-mode", type=str, - choices=["fa", "torch", "trt"], default="torch", help="Inference mode") - parser.add_argument("--onnx-workdir", type=str, - default="onnx_model", help="Path to save ONNX model") - - # Sampling - parser.add_argument("--sampler", type=str, - default="ddpm", help="Diffusion sampler") - parser.add_argument("--infer-steps", type=int, - default=100, help="Inference steps") - - # Prompt enhancement - parser.add_argument("--enhance", action="store_true", - help="Enhance prompt with dialoggen.") - parser.add_argument("--no-enhance", dest="enhance", action="store_false") - parser.add_argument( - "--load-4bit", help="load DialogGen model with 4bit quantization.", action="store_true") - parser.set_defaults(enhance=True) - - # App - parser.add_argument("--lang", type=str, default="zh", - choices=["zh", "en"], help="Language") - - # ======================================================================================================== - # Training - # ======================================================================================================== - - # Basic Setting - parser.add_argument("--lr", type=float, default=1e-4) - parser.add_argument("--epochs", type=int, default=1400) - parser.add_argument("--max-training-steps", type=int, default=10_000_000) - parser.add_argument("--gc-interval", type=int, default=40, - help='To address the memory bottleneck encountered during the preprocessing of the dataset,' - ' memory fragments are reclaimed here by invoking the gc.collect() function.') - parser.add_argument("--log-every", type=int, default=100) - parser.add_argument("--ckpt-every", type=int, default=100_000, - help="Create a ckpt every a few steps.") - parser.add_argument("--ckpt-latest-every", type=int, default=10_000, - help="Create a ckpt named `latest.pt` every a few steps.") - parser.add_argument("--num-workers", type=int, default=4) - parser.add_argument("--global-seed", type=int, default=1234) - parser.add_argument("--warmup-min-lr", type=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") - parser.add_argument("--rope-img", type=str, default=None, choices=['extend', 'base512', 'base1024'], - help="Extend or interpolate the positional embedding of the image.") - parser.add_argument("--rope-real", action="store_true", - help="Use real part and imaginary part separately for RoPE.") - - # Classifier-free - 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") - - # Directory - parser.add_argument("--results-dir", type=str, default="results") - parser.add_argument("--resume", type=str, default=None, - help="Resume experiment from a checkpoint") - parser.add_argument("--no-strict", dest="strict", - action="store_false", help="Strict loading of checkpoint") - parser.set_defaults(strict=True) - parser.add_argument("--resume-deepspeed", action="store_true", - help="Resume model and ema states from a checkpoint saved by Deepspeed version of DIT.") - parser.add_argument("--resume-split", action="store_true", - help="Resume model and ema states from two checkpoint separated from DeepSpeed ckpt.") - parser.add_argument("--ema-to-module", action="store_true", - help="If true, initialize the module with EMA weights.") - parser.add_argument("--module-to-ema", action="store_true", - help="if true, initialize the ema with Module weights.") - - # Dataset - parser.add_argument("--index-file", type=str, nargs='+', - help="During training, provide a JSON file with data indices.") - parser.add_argument("--random-flip", action="store_true", - help="Random flip image") - parser.add_argument("--reset-loader", action="store_true", - help="Reset the data loader. It is useful when resuming from a checkpoint but switch to a new dataset.") - parser.add_argument("--multireso", action="store_true", - help="Use multi-resolution training.") - parser.add_argument("--reso-step", type=int, default=None, - help="Step size for multi-resolution training.") - - # Additional condition - parser.add_argument("--random-shrink-size-cond", action="store_true", - help="Randomly shrink the original size condition.") - parser.add_argument("--merge-src-cond", action="store_true", - help="Merge the source condition into a single value.") - - # EMA Model - parser.add_argument("--use-ema", action="store_true", help="Use EMA model") - parser.add_argument("--ema-dtype", type=str, choices=['fp16', 'fp32', 'none'], default="none", - help="EMA data type. If none, use the same data type as the model.") - parser.add_argument("--ema-decay", type=float, default=None, - help="EMA decay rate. If None, use the default value of the model.") - parser.add_argument("--ema-warmup", action="store_true", - help="EMA warmup. If True, perform ema_decay warmup from 0 to ema_decay.") - parser.add_argument("--ema-warmup-power", type=float, default=None, - help="EMA power. If None, use the default value of the model.") - parser.add_argument("--ema-reset-decay", action="store_true", - help="Reset EMA decay to 0 and restart increasing the EMA decay." - "Only works when --ema-warmup is enabled.") - # Acceleration - 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.") - - - args = parser.parse_args(default_args) - - return args diff --git a/hook_HYDiT_main.py b/hook_HYDiT_main.py deleted file mode 100644 index 1d7ff01..0000000 --- a/hook_HYDiT_main.py +++ /dev/null @@ -1,820 +0,0 @@ - -import argparse -import json -import logging -import os -import sys -from types import SimpleNamespace -import gc -import json -import os -import random -import sys -import time -from functools import partial -from glob import glob -from pathlib import Path -import numpy as np - -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.utils.data import DataLoader -from torchvision.transforms import functional as TF -from diffusers.models import AutoencoderKL -from transformers import BertModel, BertTokenizer, logging as tf_logging - -# from hydit.constants import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER -from hydit.lr_scheduler import WarmupLR -from hydit.data_loader.arrow_load_stream import TextImageArrowStream -from hydit.diffusion import create_diffusion -from hydit.modules.fp16_layers import Float16Module -from hydit.modules.models import HUNYUAN_DIT_MODELS -from hydit.modules.posemb_layers import init_image_posemb -from hydit.utils.tools import set_seeds, get_trainable_params -from IndexKits.index_kits import ResolutionGroup -from IndexKits.index_kits.sampler import DistributedSamplerWithStartIndex, BlockDistributedSampler -from peft import LoraConfig, get_peft_model -from tqdm import tqdm - - -UNET_PATH = "ckpts/t2i/model/pytorch_model_ema.pt" -def set_unet_path(path): - global UNET_PATH - UNET_PATH = path -VAE_EMA_PATH = "ckpts/t2i/sdxl-vae-fp16-fix" -def set_vae_ema_path(path): - global VAE_EMA_PATH - VAE_EMA_PATH = path -TOKENIZER = "ckpts/t2i/tokenizer" -def set_tokenizer_path(path): - global TOKENIZER - TOKENIZER = path -TEXT_ENCODER = 'ckpts/t2i/clip_text_encoder' -def set_text_encoder_path(path): - global TEXT_ENCODER - TEXT_ENCODER = path -T5_ENCODER = { - 'MT5': 'ckpts/t2i/mt5', - 'attention_mask': True, - 'layer_index': -1, - 'attention_pool': True, - 'torch_dtype': torch.float16, - 'learnable_replace': True -} -def set_t5_encoder_path(path): - global T5_ENCODER - T5_ENCODER['MT5'] = path -global TRAIN_CONFIG -def set_train_config(train_config): - global TRAIN_CONFIG - TRAIN_CONFIG = train_config - -def create_logger(log_dir=None): - logger = logging.getLogger(__name__) - logger.addHandler(logging.NullHandler()) - return logger - - -def create_exp_folder(args, rank): - if rank == 0: - os.makedirs(args.results_dir, exist_ok=True) - - # print('args.results_dir', args.results_dir) - _existed_experiments = os.listdir(args.results_dir) - # print('_existed_experiments', _existed_experiments) - existed_experiments = [] - for x in _existed_experiments: - if os.path.isdir(os.path.join(args.results_dir, x)): - if x.split('-')[0].isdigit(): - existed_experiments.append(x) - - if len(existed_experiments) == 0: - experiment_index = 1 - else: - existed_experiments.sort() - # print('existed_experiments', existed_experiments) - experiment_index = max([int(x.split('-')[0]) - for x in existed_experiments]) + 1 - model_string_name = args.task_flag if args.task_flag else args.model.replace( - "/", "-") - # Create an experiment folder - experiment_dir = f"{args.results_dir}/{experiment_index:03d}-{model_string_name}" - # Stores saved model checkpoints - checkpoint_dir = f"{experiment_dir}/checkpoints" - if rank == 0: - os.makedirs(checkpoint_dir, exist_ok=True) - logger = create_logger(experiment_dir) - logger.info(f"Experiment directory created at {experiment_dir}") - else: - logger = create_logger() - experiment_dir = "" - logger.info = print - return experiment_dir, checkpoint_dir, logger - - -def save_checkpoint(args, rank, logger, model, ema, epoch, train_steps, checkpoint_dir): - def save_lora_weight(checkpoint_dir, client_state, tag=f"{train_steps:07d}.pt"): - cur_ckpt_save_dir = f"{checkpoint_dir}/{tag}" - if rank == 0: - if args.use_fp16: - model.module.save_pretrained(cur_ckpt_save_dir) - else: - model.save_pretrained(cur_ckpt_save_dir) - - checkpoint_path = "[Not rank 0. Disabled output.]" - - client_state = { - "steps": train_steps, - "epoch": epoch, - "args": args - } - if ema is not None: - client_state['ema'] = ema.state_dict() - - dst_paths = [] - - if train_steps % args.ckpt_every == 0: - checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt" - try: - if args.training_parts == "lora": - save_lora_weight(checkpoint_dir, client_state, - tag=f"{train_steps:07d}.pt") - else: - model.save_checkpoint( - checkpoint_dir, client_state=client_state, tag=f"{train_steps:07d}.pt") - dst_paths.append(checkpoint_path) - logger.info(f"Saved checkpoint to {checkpoint_path}") - except: - logger.error(f"Saved failed to {checkpoint_path}") - - if train_steps % args.ckpt_latest_every == 0 or train_steps == args.max_training_steps: - save_name = "latest.pt" - checkpoint_path = f"{checkpoint_dir}/{save_name}" - try: - if args.training_parts == "lora": - save_lora_weight(checkpoint_dir, client_state, - tag=f"{save_name}") - else: - model.save_checkpoint( - checkpoint_dir, client_state=client_state, tag=f"{save_name}") - dst_paths.append(checkpoint_path) - logger.info(f"Saved checkpoint to {checkpoint_path}") - except: - logger.error(f"Saved failed to {checkpoint_path}") - - return checkpoint_path - - -@torch.no_grad() -def prepare_model_inputs(args, batch, device, vae, text_encoder, text_encoder_t5, freqs_cis_img): - image, text_embedding, text_embedding_mask, text_embedding_t5, text_embedding_mask_t5, kwargs = batch - - # clip & mT5 text embedding - text_embedding = text_embedding.to(device) - text_embedding_mask = text_embedding_mask.to(device) - encoder_hidden_states = text_encoder( - text_embedding.to(device), - attention_mask=text_embedding_mask.to(device), - )[0] - text_embedding_t5 = text_embedding_t5.to(device).squeeze(1) - text_embedding_mask_t5 = text_embedding_mask_t5.to(device).squeeze(1) - with torch.no_grad(): - output_t5 = text_encoder_t5( - input_ids=text_embedding_t5, - attention_mask=text_embedding_mask_t5 if T5_ENCODER['attention_mask'] else None, - output_hidden_states=True - ) - encoder_hidden_states_t5 = output_t5['hidden_states'][T5_ENCODER['layer_index']].detach( - ) - - # additional condition - image_meta_size = kwargs['image_meta_size'].to(device) - style = kwargs['style'].to(device) - - if args.extra_fp16: - image = image.half() - image_meta_size = image_meta_size.half() if image_meta_size is not None else None - - # Map input images to latent space + normalize latents: - image = image.to(device) - vae_scaling_factor = vae.config.scaling_factor - latents = vae.encode(image).latent_dist.sample().mul_( - vae_scaling_factor).to(device) - - # positional embedding - _, _, height, width = image.shape - reso = f"{height}x{width}" - - # print("freqs_cis_img:", freqs_cis_img.keys()) - # print("freqs_cis_img[reso]:", freqs_cis_img[reso]) - - cos_cis_img, sin_cis_img = freqs_cis_img[reso] - - # Model conditions - model_kwargs = dict( - encoder_hidden_states=encoder_hidden_states, - text_embedding_mask=text_embedding_mask, - encoder_hidden_states_t5=encoder_hidden_states_t5, - text_embedding_mask_t5=text_embedding_mask_t5, - image_meta_size=image_meta_size, - style=style, - cos_cis_img=cos_cis_img, - sin_cis_img=sin_cis_img, - ) - - return latents, model_kwargs - - -def Core(args, LOG): - if args.training_parts == "lora": - args.use_ema = False - - assert torch.cuda.is_available(), "Training currently requires at least one GPU." - - batch_size = args.batch_size - grad_accu_steps = args.grad_accu_steps - global_batch_size = batch_size * grad_accu_steps - - seed = args.global_seed - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - print(f"Starting ,seed={seed}.") - rank = 0 - # Setup an experiment folder - experiment_dir, checkpoint_dir, logger = create_exp_folder(args, rank) - - # Log all the arguments - logger.info(sys.argv) - logger.info(str(args)) - # Save to a json file - args_dict = vars(args) - with open(f"{experiment_dir}/args.json", 'w') as f: - json.dump(args_dict, f, indent=4) - - # Disable the message "Some weights of the model checkpoint at ... were not used when initializing BertModel." - # If needed, just comment the following line. - tf_logging.set_verbosity_error() - - # =========================================================================== - # Building HYDIT - # =========================================================================== - - logger.info("Building HYDIT Model.") - - # --------------------------------------------------------------------------- - # Training sample base size, such as 256/512/1024. Notice that this size is - # just a base size, not necessary the actual size of training samples. Actual - # size of the training samples are correlated with `resolutions` when enabling - # multi-resolution training. - # --------------------------------------------------------------------------- - image_size = args.image_size - if type(image_size) == int: - image_size = [image_size, image_size] - if len(image_size) == 1: - image_size = [image_size[0], image_size[0]] - if len(image_size) != 2: - raise ValueError(f"Invalid image size: {args.image_size}") - assert image_size[0] % 8 == 0 and image_size[1] % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder). " \ - f"got {image_size}" - latent_size = [image_size[0] // 8, image_size[1] // 8] - - # initialize model by deepspeed - model = HUNYUAN_DIT_MODELS[args.model](args, - input_size=latent_size, - log_fn=logger.info, - ) - # Multi-resolution / Single-resolution training. - if args.multireso: - resolutions = ResolutionGroup(image_size[0], - align=16, - step=args.reso_step, - target_ratios=args.target_ratios).data - else: - resolutions = ResolutionGroup(image_size[0], - align=16, - target_ratios=['1:1']).data - - freqs_cis_img = init_image_posemb(args.rope_img, - resolutions=resolutions, - patch_size=model.patch_size, - hidden_size=model.hidden_size, - num_heads=model.num_heads, - log_fn=logger.info, - rope_real=args.rope_real, - ) - - # Create EMA model and convert to fp16 if needed. - ema = None - device = "cuda" - if args.use_ema: - raise ValueError("Not support EMA model.") - - # Setup FP16 main model: - if args.use_fp16: - model = Float16Module(model, args) - logger.info( - f" Using main model with data type {'fp16' if args.use_fp16 else 'fp32'}") - - diffusion = create_diffusion( - noise_schedule=args.noise_schedule, - predict_type=args.predict_type, - learn_sigma=args.learn_sigma, - mse_loss_weight_type=args.mse_loss_weight_type, - beta_start=args.beta_start, - beta_end=args.beta_end, - noise_offset=args.noise_offset, - ) - - # Setup VAE - logger.info(f" Loading vae from {VAE_EMA_PATH}") - - if os.path.isdir(VAE_EMA_PATH): - vae = AutoencoderKL.from_pretrained(VAE_EMA_PATH) - else: - vae = AutoencoderKL.from_single_file(VAE_EMA_PATH) - # Setup BERT text encoder - logger.info(f" Loading Bert text encoder from {TEXT_ENCODER}") - text_encoder = BertModel.from_pretrained( - TEXT_ENCODER, False, revision=None) - # Setup BERT tokenizer: - logger.info(f" Loading Bert tokenizer from {TOKENIZER}") - tokenizer = BertTokenizer.from_pretrained(TOKENIZER) - # Setup T5 text encoder - from hydit.modules.text_encoder import MT5Embedder - mt5_path = T5_ENCODER['MT5'] - embedder_t5 = MT5Embedder( - mt5_path, torch_dtype=T5_ENCODER['torch_dtype'], max_length=args.text_len_t5) - tokenizer_t5 = embedder_t5.tokenizer - text_encoder_t5 = embedder_t5.model - - if args.extra_fp16: - logger.info(f" Using fp16 for extra modules: vae, text_encoder") - vae = vae.half().to(device) - text_encoder = text_encoder.half().to(device) - text_encoder_t5 = text_encoder_t5.half().to(device) - else: - vae = vae.to(device) - text_encoder = text_encoder.to(device) - text_encoder_t5 = text_encoder_t5.to(device) - - logger.info( - f" Optimizer parameters: lr={args.lr}, weight_decay={args.weight_decay}") - - # =========================================================================== - # Building Dataset - # =========================================================================== - - logger.info(f"Building Streaming Dataset.") - logger.info(f" Loading index file {args.index_file} (v2)") - world_size = 1 - - dataset = TextImageArrowStream(args=args, - resolution=image_size[0], - random_flip=args.random_flip, - log_fn=logger.info, - index_file=args.index_file, - multireso=args.multireso, - batch_size=batch_size, - world_size=world_size, - random_shrink_size_cond=args.random_shrink_size_cond, - merge_src_cond=args.merge_src_cond, - uncond_p=args.uncond_p, - text_ctx_len=args.text_len, - tokenizer=tokenizer, - uncond_p_t5=args.uncond_p_t5, - text_ctx_len_t5=args.text_len_t5, - tokenizer_t5=tokenizer_t5, - ) - if args.multireso: - sampler = BlockDistributedSampler(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed, - shuffle=False, drop_last=True, batch_size=batch_size) - else: - sampler = DistributedSamplerWithStartIndex(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed, - shuffle=False, drop_last=True) - loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, sampler=sampler, - num_workers=args.num_workers, pin_memory=True, drop_last=True) - logger.info(f" Dataset contains {len(dataset):,} images.") - logger.info(f" Index file: {args.index_file}.") - if args.multireso: - logger.info(f' Using MultiResolutionBucketIndexV2 with step {dataset.index_manager.step} ' - f'and base size {dataset.index_manager.base_size}') - logger.info(f'\n {dataset.index_manager.resolutions}') - - # =========================================================================== - # Loading parameter - # =========================================================================== - - logger.info(f"Loading parameter") - start_epoch = 0 - start_epoch_step = 0 - train_steps = 0 - # Resume checkpoint if needed - 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, - lora_alpha=args.rank, - target_modules=args.target_modules - ) - 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}") - - # =========================================================================== - # Training - # =========================================================================== - - model.to(device) - model.train() - if args.use_ema: - ema.eval() - - print(f" Worker {rank} ready.") - - iters_per_epoch = len(loader) - logger.info( - " ****************************** Running training ******************************") - logger.info(f" Number GPUs: {world_size}") - logger.info(f" Number training samples: {len(dataset):,}") - logger.info( - f" Number parameters: {sum(p.numel() for p in model.parameters()):,}") - logger.info( - f" Number trainable params: {sum(p.numel() for p in get_trainable_params(model)):,}") - logger.info( - " ------------------------------------------------------------------------------") - logger.info(f" Iters per epoch: {iters_per_epoch:,}") - logger.info(f" Batch size per device: {batch_size}") - logger.info( - f" Batch size all device: {batch_size * world_size * grad_accu_steps:,} (world_size * batch_size * grad_accu_steps)") - logger.info(f" Gradient Accu steps: {args.grad_accu_steps}") - logger.info( - f" Total optimization steps: {args.epochs * iters_per_epoch // grad_accu_steps:,}") - - logger.info( - f" Training epochs: {start_epoch}/{args.epochs}") - logger.info( - f" Training epoch steps: {start_epoch_step:,}/{iters_per_epoch:,}") - logger.info( - f" Training total steps: {train_steps:,}/{min(args.max_training_steps, args.epochs * iters_per_epoch):,}") - logger.info( - " ------------------------------------------------------------------------------") - logger.info(f" Noise schedule: {args.noise_schedule}") - logger.info( - f" Beta limits: ({args.beta_start}, {args.beta_end})") - logger.info(f" Learn sigma: {args.learn_sigma}") - logger.info(f" Prediction type: {args.predict_type}") - logger.info(f" Noise offset: {args.noise_offset}") - - logger.info( - " ------------------------------------------------------------------------------") - logger.info( - f" Using EMA model: {args.use_ema} ({args.ema_dtype})") - if args.use_ema: - logger.info( - f" Using EMA decay: {ema.max_value if args.use_ema else None}") - logger.info( - f" Using EMA warmup power: {ema.power if args.use_ema else None}") - logger.info(f" Using main model fp16: {args.use_fp16}") - logger.info(f" Using extra modules fp16: {args.extra_fp16}") - logger.info( - " ------------------------------------------------------------------------------") - logger.info(f" Experiment directory: {experiment_dir}") - logger.info( - " *******************************************************************************") - - if args.gc_interval > 0: - gc.disable() - gc.collect() - - # Variables for monitoring/logging purposes: - log_steps = 0 - running_loss = 0 - start_time = time.time() - - if args.async_ema: - ema_stream = torch.cuda.Stream() - - total_steps = args.epochs * len(loader) - - if args.use_fp16: - optimizer = torch.optim.AdamW( - model.parameters(), lr=args.lr, weight_decay=args.weight_decay, eps=TRAIN_CONFIG.get("adam_epsilon", 1e-7)) - else: - 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}") - # Makesure all processors use the same seed to shuffle dataset. - dataset.shuffle(seed=args.global_seed + epoch, fast=True) - logger.info(f" End of random shuffle") - - # Move sampler to start_index - if not args.multireso: - start_index = start_epoch_step * world_size * batch_size - if start_index != sampler.start_index: - sampler.start_index = start_index - # Reset start_epoch_step to zero, to ensure next epoch will start from the beginning. - start_epoch_step = 0 - logger.info( - f" Iters left this epoch: {len(loader):,}") - - logger.info(f" Beginning epoch {epoch}...") - step = 0 - mean_loss = 0 - for batch in loader: - step += 1 - - # Zero gradients, backward pass, and optimizer step: - optimizer.zero_grad() - latents, model_kwargs = prepare_model_inputs( - args, batch, device, vae, text_encoder, text_encoder_t5, freqs_cis_img) - # with torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True): - - if True: - # training model by deepspeed while use fp16 - if args.use_fp16: - if args.use_ema and args.async_ema: - with torch.cuda.stream(ema_stream): - ema.update(model.module.module, step=step) - torch.cuda.current_stream().wait_stream(ema_stream) - - with torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True): - loss_dict = diffusion.training_losses( - model=model, x_start=latents, model_kwargs=model_kwargs) - - loss = loss_dict["loss"].mean() - - # print(f"step={step}, loss={loss.item()}") - - # model.backward(loss) - - loss.backward() - - # last_batch_iteration = ( - # train_steps + 1) // (global_batch_size // (batch_size * world_size)) - # model.step( - # lr_kwargs={'last_batch_iteration': last_batch_iteration}) - - optimizer.step() - - if args.use_ema and not args.async_ema or (args.async_ema and step == len(loader) - 1): - if args.use_fp16: - ema.update(model.module.module, step=step) - else: - ema.update(model.module, step=step) - - 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}") - - # =========================================================================== - # Log loss values: - # =========================================================================== - # running_loss += loss.item() - # log_steps += 1 - # if train_steps % args.log_every == 0: - # # Measure training speed: - # torch.cuda.synchronize() - # end_time = time.time() - # steps_per_sec = log_steps / (end_time - start_time) - # # Reduce loss history over all processes: - # avg_loss = torch.tensor( - # running_loss / log_steps, device=device) - # # dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM) - - # # avg_loss = avg_loss.item() / world_size - # # get lr from deepspeed fused optimizer - # logger.info(f"(step={train_steps:07d}) " + - # (f"(update_step={train_steps // args.grad_accu_steps:07d}) " if args.grad_accu_steps > 1 else "") + - # f"Train Loss: {avg_loss:.4f}, " - # # f"Lr: {opt.param_groups[0]['lr']:.6g}, " - # f"Steps/Sec: {steps_per_sec:.2f}, " - # f"Samples/Sec: {int(steps_per_sec * batch_size * world_size):d}") - # # Reset monitoring variables: - # running_loss = 0 - # log_steps = 0 - # start_time = time.time() - - # collect gc: - if args.gc_interval > 0 and (step % args.gc_interval == 0): - 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: - 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) - - if train_steps >= args.max_training_steps: - logger.info(f"Breaking step loop at {train_steps}.") - break - - LOG({ - "type": "sample_images", - "global_step": train_steps, - "total_steps": total_steps, - # "latent": noise_pred_latent_path, - }) - - if train_steps >= args.max_training_steps: - logger.info(f"Breaking epoch loop at {epoch}.") - break - - -def model_resume(args, model, ema, logger): - """ - Load pretrained weights. - """ - start_epoch = 0 - start_epoch_step = 0 - train_steps = 0 - resume_path = UNET_PATH - if not Path(resume_path).exists(): - raise FileNotFoundError( - f" Cannot find checkpoint from {resume_path}") - - logger.info(f"Resume from checkpoint {resume_path}") - - if args.resume_split: - # Resume main model - - resume_ckpt_module = torch.load( - resume_path, map_location=lambda storage, loc: storage) - model.load_state_dict(resume_ckpt_module, strict=False) - - # Resume ema model - if args.use_ema: - if args.module_to_ema: - if "resume_ckpt_module" in locals(): - logger.info(f" Resume ema model from main states.") - ema.load_state_dict(resume_ckpt_module, strict=args.strict) - else: - logger.info(f" Resume ema model from module states.") - resume_ckpt_module = torch.load( - resume_path, map_location=lambda storage, loc: storage) - ema.load_state_dict(resume_ckpt_module, strict=args.strict) - else: - if "resume_ckpt_ema" in locals(): - logger.info(f" Resume ema model from EMA states.") - ema.load_state_dict(resume_ckpt_ema, strict=args.strict) - else: - logger.info(f" Resume ema model from EMA states.") - resume_ckpt_ema = torch.load(resume_path, - map_location=lambda storage, loc: storage) - ema.load_state_dict(resume_ckpt_ema, strict=args.strict) - else: - raise ValueError( - " “If `resume` is True, then either `resume_split` must be true.”") - - return model, ema, start_epoch, start_epoch_step, train_steps - - -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 - - -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(): - - 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, - ).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) - - try: - 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] - except Exception as e: - print(f"Failed to sample images: {e}") - continue - - # 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 11ed62c..20f9422 100644 --- a/hook_HYDiT_run.py +++ b/hook_HYDiT_run.py @@ -49,10 +49,7 @@ if __name__ == "__main__": print("HYDi run hook") - try: - from . import hook_HYDiT_main - except Exception as e: - import hook_HYDiT_main + try: from . import hook_HYDiT_main_train_deepspeed except Exception as e: