import json from time import time import argparse import logging import os from pathlib import Path import math import numpy as np from PIL import Image from copy import deepcopy import torch import torch.distributed as dist from torch.utils.data import Dataset, DataLoader from torch.utils.data.distributed import DistributedSampler from torchvision import transforms from accelerate import Accelerator from accelerate.utils import ProjectConfiguration, set_seed from diffusers.optimization import get_scheduler from accelerate.utils import DistributedType from peft import LoraConfig, set_peft_model_state_dict, PeftModel, get_peft_model from peft.utils import get_peft_model_state_dict from huggingface_hub import snapshot_download from safetensors.torch import save_file from diffusers.models import AutoencoderKL from OmniGen import OmniGen, OmniGenProcessor from OmniGen.train_helper import DatasetFromJson, TrainDataCollator from OmniGen.train_helper import training_losses from OmniGen.utils import ( create_logger, update_ema, requires_grad, center_crop_arr, crop_arr, vae_encode, vae_encode_list ) def main(args): # Setup accelerator: from accelerate import DistributedDataParallelKwargs as DDPK kwargs = DDPK(find_unused_parameters=False) accelerator = Accelerator( gradient_accumulation_steps=args.gradient_accumulation_steps, mixed_precision=args.mixed_precision, log_with=args.report_to, project_dir=args.results_dir, kwargs_handlers=[kwargs], ) device = accelerator.device accelerator.init_trackers("tensorboard_log", config=args.__dict__) # Setup an experiment folder: checkpoint_dir = f"{args.results_dir}/checkpoints" # Stores saved model checkpoints logger = create_logger(args.results_dir) if accelerator.is_main_process: os.makedirs(checkpoint_dir, exist_ok=True) logger.info(f"Experiment directory created at {args.results_dir}") json.dump(args.__dict__, open(os.path.join(args.results_dir, 'train_args.json'), 'w')) # Create model: if not os.path.exists(args.model_name_or_path): cache_folder = os.getenv('HF_HUB_CACHE') args.model_name_or_path = snapshot_download(repo_id=args.model_name_or_path, cache_dir=cache_folder, ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5']) logger.info(f"Downloaded model to {args.model_name_or_path}") model = OmniGen.from_pretrained(args.model_name_or_path) model.llm.config.use_cache = False model.llm.gradient_checkpointing_enable() model = model.to(device) if args.vae_path is None: print(args.model_name_or_path) vae_path = os.path.join(args.model_name_or_path, "vae") if os.path.exists(vae_path): vae = AutoencoderKL.from_pretrained(vae_path).to(device) else: logger.info("No VAE found in model, downloading stabilityai/sdxl-vae from HF") logger.info("If you have VAE in local folder, please specify the path with --vae_path") vae = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae").to(device) else: vae = AutoencoderKL.from_pretrained(args.vae_path).to(device) weight_dtype = torch.float32 if accelerator.mixed_precision == "fp16": weight_dtype = torch.float16 elif accelerator.mixed_precision == "bf16": weight_dtype = torch.bfloat16 vae.to(dtype=torch.float32) model.to(weight_dtype) processor = OmniGenProcessor.from_pretrained(args.model_name_or_path) requires_grad(vae, False) if args.use_lora: if accelerator.distributed_type == DistributedType.FSDP: raise NotImplementedError("FSDP does not support LoRA") requires_grad(model, False) transformer_lora_config = LoraConfig( r=args.lora_rank, lora_alpha=args.lora_rank, init_lora_weights="gaussian", target_modules=["qkv_proj", "o_proj"], ) model.llm.enable_input_require_grads() model = get_peft_model(model, transformer_lora_config) model.to(weight_dtype) transformer_lora_parameters = list(filter(lambda p: p.requires_grad, model.parameters())) for n,p in model.named_parameters(): print(n, p.requires_grad) opt = torch.optim.AdamW(transformer_lora_parameters, lr=args.lr, weight_decay=args.adam_weight_decay) else: opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.adam_weight_decay) ema = None if args.use_ema: ema = deepcopy(model).to(device) # Create an EMA of the model for use after training requires_grad(ema, False) # Setup data: crop_func = crop_arr if not args.keep_raw_resolution: crop_func = center_crop_arr image_transform = transforms.Compose([ transforms.Lambda(lambda pil_image: crop_func(pil_image, args.max_image_size)), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True) ]) dataset = DatasetFromJson(json_file=args.json_file, image_path=args.image_path, processer=processor, image_transform=image_transform, max_input_length_limit=args.max_input_length_limit, condition_dropout_prob=args.condition_dropout_prob, keep_raw_resolution=args.keep_raw_resolution ) collate_fn = TrainDataCollator(pad_token_id=processor.text_tokenizer.eos_token_id, hidden_size=model.llm.config.hidden_size, keep_raw_resolution=args.keep_raw_resolution) loader = DataLoader( dataset, collate_fn=collate_fn, batch_size=args.batch_size_per_device, shuffle=True, num_workers=args.num_workers, pin_memory=True, drop_last=True, prefetch_factor=2, ) if accelerator.is_main_process: logger.info(f"Dataset contains {len(dataset):,}") num_update_steps_per_epoch = math.ceil(len(loader) / args.gradient_accumulation_steps) max_train_steps = args.epochs * num_update_steps_per_epoch lr_scheduler = get_scheduler( args.lr_scheduler, optimizer=opt, num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps, num_training_steps=max_train_steps * args.gradient_accumulation_steps, ) # Prepare models for training: model.train() # important! This enables embedding dropout for classifier-free guidance if ema is not None: update_ema(ema, model, decay=0) # Ensure EMA is initialized with synced weights ema.eval() # EMA model should always be in eval mode if ema is not None: model, ema = accelerator.prepare(model, ema) else: model = accelerator.prepare(model) opt, loader, lr_scheduler = accelerator.prepare(opt, loader, lr_scheduler) # Variables for monitoring/logging purposes: train_steps, log_steps = 0, 0 running_loss = 0 start_time = time() if accelerator.is_main_process: logger.info(f"Training for {args.epochs} epochs...") for epoch in range(args.epochs): if accelerator.is_main_process: logger.info(f"Beginning epoch {epoch}...") for data in loader: with accelerator.accumulate(model): with torch.no_grad(): output_images = data['output_images'] input_pixel_values = data['input_pixel_values'] if isinstance(output_images, list): output_images = vae_encode_list(vae, output_images, weight_dtype) if input_pixel_values is not None: input_pixel_values = vae_encode_list(vae, input_pixel_values, weight_dtype) else: output_images = vae_encode(vae, output_images, weight_dtype) if input_pixel_values is not None: input_pixel_values = vae_encode(vae, input_pixel_values, weight_dtype) model_kwargs = dict(input_ids=data['input_ids'], input_img_latents=input_pixel_values, input_image_sizes=data['input_image_sizes'], attention_mask=data['attention_mask'], position_ids=data['position_ids'], padding_latent=data['padding_images'], past_key_values=None, return_past_key_values=False) loss_dict = training_losses(model, output_images, model_kwargs) loss = loss_dict["loss"].mean() running_loss += loss.item() accelerator.backward(loss) if args.max_grad_norm is not None and accelerator.sync_gradients: accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm) opt.step() lr_scheduler.step() opt.zero_grad() log_steps += 1 train_steps += 1 accelerator.log({"training_loss": loss.item()}, step=train_steps) if train_steps % args.gradient_accumulation_steps == 0: if accelerator.sync_gradients and ema is not None: update_ema(ema, model) if train_steps % (args.log_every * args.gradient_accumulation_steps) == 0 and train_steps > 0: torch.cuda.synchronize() end_time = time() steps_per_sec = log_steps / args.gradient_accumulation_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() / accelerator.num_processes if accelerator.is_main_process: cur_lr = opt.param_groups[0]["lr"] logger.info(f"(step={int(train_steps/args.gradient_accumulation_steps):07d}) Train Loss: {avg_loss:.4f}, Train Steps/Sec: {steps_per_sec:.2f}, Epoch: {train_steps/len(loader)}, LR: {cur_lr}") # Reset monitoring variables: running_loss = 0 log_steps = 0 start_time = time() if train_steps % (args.ckpt_every * args.gradient_accumulation_steps) == 0 and train_steps > 0: if accelerator.distributed_type == DistributedType.FSDP: state_dict = accelerator.get_state_dict(model) ema_state_dict = accelerator.get_state_dict(ema) if ema is not None else None else: if not args.use_lora: state_dict = model.module.state_dict() ema_state_dict = accelerator.get_state_dict(ema) if ema is not None else None if accelerator.is_main_process: if args.use_lora: checkpoint_path = f"{checkpoint_dir}/{int(train_steps/args.gradient_accumulation_steps):07d}/" os.makedirs(checkpoint_path, exist_ok=True) model.module.save_pretrained(checkpoint_path) else: checkpoint_path = f"{checkpoint_dir}/{int(train_steps/args.gradient_accumulation_steps):07d}/" os.makedirs(checkpoint_path, exist_ok=True) torch.save(state_dict, os.path.join(checkpoint_path, "model.pt")) processor.text_tokenizer.save_pretrained(checkpoint_path) model.llm.config.save_pretrained(checkpoint_path) if ema_state_dict is not None: checkpoint_path = f"{checkpoint_dir}/{int(train_steps/args.gradient_accumulation_steps):07d}_ema" os.makedirs(checkpoint_path, exist_ok=True) torch.save(state_dict, os.path.join(checkpoint_path, "model.pt")) processor.text_tokenizer.save_pretrained(checkpoint_path) model.llm.config.save_pretrained(checkpoint_path) logger.info(f"Saved checkpoint to {checkpoint_path}") dist.barrier() accelerator.end_training() model.eval() if accelerator.is_main_process: logger.info("Done!") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--results_dir", type=str, default="results") parser.add_argument("--model_name_or_path", type=str, default="OmniGen") parser.add_argument("--json_file", type=str) parser.add_argument("--image_path", type=str, default=None) parser.add_argument("--epochs", type=int, default=1400) parser.add_argument("--batch_size_per_device", type=int, default=1) parser.add_argument("--vae_path", type=str, default=None) parser.add_argument("--num_workers", type=int, default=4) parser.add_argument("--log_every", type=int, default=100) parser.add_argument("--ckpt_every", type=int, default=20000) parser.add_argument("--max_grad_norm", type=float, default=1.0) parser.add_argument("--lr", type=float, default=1e-4) parser.add_argument("--max_input_length_limit", type=int, default=1024) parser.add_argument("--condition_dropout_prob", type=float, default=0.1) parser.add_argument("--adam_weight_decay", type=float, default=0.0) parser.add_argument( "--keep_raw_resolution", action="store_true", help="multiple_resolutions", ) parser.add_argument("--max_image_size", type=int, default=1344) parser.add_argument( "--use_lora", action="store_true", ) parser.add_argument( "--lora_rank", type=int, default=8 ) parser.add_argument( "--use_ema", action="store_true", help="Whether or not to use ema.", ) parser.add_argument( "--lr_scheduler", type=str, default="constant", help=( 'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",' ' "constant", "constant_with_warmup"]' ), ) parser.add_argument( "--lr_warmup_steps", type=int, default=1000, help="Number of steps for the warmup in the lr scheduler." ) parser.add_argument( "--report_to", type=str, default="tensorboard", help=( 'The integration to report the results and logs to. Supported platforms are `"tensorboard"`' ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.' ), ) parser.add_argument( "--mixed_precision", type=str, default="bf16", choices=["no", "fp16", "bf16"], help=( "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." ), ) parser.add_argument( "--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.", ) args = parser.parse_args() assert args.max_image_size % 16 == 0, "Image size must be divisible by 16." main(args)