""" Copyright 2022 HuggingFace, ShivamShrirao Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. """ import keyboard import gradio as gr import argparse import random import hashlib import itertools import json import math import os import copy from contextlib import nullcontext from pathlib import Path import shutil import torch import torch.nn.functional as F import torch.utils.checkpoint import numpy as np from accelerate import Accelerator from accelerate.logging import get_logger from accelerate.utils import set_seed from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, DiffusionPipeline, UNet2DConditionModel,DiffusionPipeline, DPMSolverMultistepScheduler,EulerDiscreteScheduler from diffusers.optimization import get_scheduler from torchvision.transforms import functional from tqdm.auto import tqdm from transformers import CLIPTextModel, CLIPTokenizer from typing import Dict, List, Generator, Tuple from PIL import Image, ImageFile from diffusers.utils.import_utils import is_xformers_available from trainer_util import * from dataloaders_util import * from discriminator import Discriminator2D from lion_pytorch import Lion logger = get_logger(__name__) def parse_args(): parser = argparse.ArgumentParser(description="Simple example of a training script.") parser.add_argument( "--revision", type=str, default=None, required=False, help="Revision of pretrained model identifier from huggingface.co/models.", ) parser.add_argument( "--attention", type=str, choices=["xformers", "flash_attention"], default="xformers", help="Type of attention to use." ) parser.add_argument( "--model_variant", type=str, default='base', required=False, help="Train Base/Inpaint/Depth2Img", ) parser.add_argument( "--aspect_mode", type=str, default='dynamic', required=False, help="Path to pretrained model or model identifier from huggingface.co/models.", ) parser.add_argument( "--aspect_mode_action_preference", type=str, default='add', required=False, help="Path to pretrained model or model identifier from huggingface.co/models.", ) parser.add_argument('--use_lion',default=False,action="store_true", help='Use the new LION optimizer') parser.add_argument('--use_ema',default=False,action="store_true", help='Use EMA for finetuning') parser.add_argument('--clip_penultimate',default=False,action="store_true", help='Use penultimate CLIP layer for text embedding') parser.add_argument("--conditional_dropout", type=float, default=None,required=False, help="Conditional dropout probability") parser.add_argument('--disable_cudnn_benchmark', default=False, action="store_true") parser.add_argument('--use_text_files_as_captions', default=False, action="store_true") parser.add_argument( "--sample_from_batch", type=int, default=0, help=("Number of prompts to sample from the batch for inference"), ) parser.add_argument( "--flatten_sample_folder", default=True, action="store_true", help="Will save samples in one folder instead of per-epoch", ) parser.add_argument( "--stop_text_encoder_training", type=int, default=999999999999999, help=("The epoch at which the text_encoder is no longer trained"), ) parser.add_argument( "--use_bucketing", default=False, action="store_true", help="Will save and generate samples before training", ) parser.add_argument( "--regenerate_latent_cache", default=False, action="store_true", help="Will save and generate samples before training", ) parser.add_argument( "--sample_on_training_start", default=False, action="store_true", help="Will save and generate samples before training", ) parser.add_argument( "--add_class_images_to_dataset", default=False, action="store_true", help="will generate and add class images to the dataset without using prior reservation in training", ) parser.add_argument( "--auto_balance_concept_datasets", default=False, action="store_true", help="will balance the number of images in each concept dataset to match the minimum number of images in any concept dataset", ) parser.add_argument( "--sample_aspect_ratios", default=False, action="store_true", help="sample different aspect ratios for each image", ) parser.add_argument( "--dataset_repeats", type=int, default=1, help="repeat the dataset this many times", ) parser.add_argument( "--save_every_n_epoch", type=int, default=1, help="save on epoch finished", ) parser.add_argument( "--pretrained_model_name_or_path", type=str, default=None, required=True, help="Path to pretrained model or model identifier from huggingface.co/models.", ) parser.add_argument( "--pretrained_vae_name_or_path", type=str, default=None, help="Path to pretrained vae or vae identifier from huggingface.co/models.", ) parser.add_argument( "--tokenizer_name", type=str, default=None, help="Pretrained tokenizer name or path if not the same as model_name", ) parser.add_argument( "--instance_data_dir", type=str, default=None, help="A folder containing the training data of instance images.", ) parser.add_argument( "--class_data_dir", type=str, default=None, help="A folder containing the training data of class images.", ) parser.add_argument( "--instance_prompt", type=str, default=None, help="The prompt with identifier specifying the instance", ) parser.add_argument( "--class_prompt", type=str, default=None, help="The prompt to specify images in the same class as provided instance images.", ) parser.add_argument( "--save_sample_prompt", type=str, default=None, help="The prompt used to generate sample outputs to save.", ) parser.add_argument( "--n_save_sample", type=int, default=4, help="The number of samples to save.", ) parser.add_argument( "--sample_height", type=int, default=512, help="The number of samples to save.", ) parser.add_argument( "--sample_width", type=int, default=512, help="The number of samples to save.", ) parser.add_argument( "--save_guidance_scale", type=float, default=7.5, help="CFG for save sample.", ) parser.add_argument( "--save_infer_steps", type=int, default=30, help="The number of inference steps for save sample.", ) parser.add_argument( "--with_prior_preservation", default=False, action="store_true", help="Flag to add prior preservation loss.", ) parser.add_argument("--prior_loss_weight", type=float, default=1.0, help="The weight of prior preservation loss.") parser.add_argument( "--with_offset_noise", default=False, action="store_true", help="Flag to offset noise applied to latents.", ) parser.add_argument("--offset_noise_weight", type=float, default=0.1, help="The weight of offset noise applied during training.") parser.add_argument( "--num_class_images", type=int, default=100, help=( "Minimal class images for prior preservation loss. If not have enough images, additional images will be" " sampled with class_prompt." ), ) parser.add_argument( "--output_dir", type=str, default="text-inversion-model", help="The output directory where the model predictions and checkpoints will be written.", ) parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.") parser.add_argument( "--resolution", type=int, default=512, help=( "The resolution for input images, all the images in the train/validation dataset will be resized to this" " resolution" ), ) parser.add_argument( "--center_crop", default=False, action="store_true", help="Whether to center crop images before resizing to resolution" ) parser.add_argument("--train_text_encoder", default=False, action="store_true", help="Whether to train the text encoder") parser.add_argument( "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader." ) parser.add_argument( "--sample_batch_size", type=int, default=4, help="Batch size (per device) for sampling images." ) parser.add_argument("--num_train_epochs", type=int, default=1) parser.add_argument( "--max_train_steps", type=int, default=None, help="Total number of training steps to perform. If provided, overrides num_train_epochs.", ) parser.add_argument( "--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.", ) parser.add_argument( "--gradient_checkpointing", default=False, action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.", ) parser.add_argument( "--learning_rate", type=float, default=5e-6, help="Initial learning rate (after the potential warmup period) to use.", ) parser.add_argument( "--scale_lr", action="store_true", default=False, help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.", ) 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=float, default=500, help="Number of steps for the warmup in the lr scheduler." ) parser.add_argument( "--use_8bit_adam", default=False, action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes." ) parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.") parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.") parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.") parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer") parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") parser.add_argument("--push_to_hub", default=False, action="store_true", help="Whether or not to push the model to the Hub.") parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") parser.add_argument( "--hub_model_id", type=str, default=None, help="The name of the repository to keep in sync with the local `output_dir`.", ) parser.add_argument( "--logging_dir", type=str, default="logs", help=( "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to" " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***." ), ) parser.add_argument("--log_interval", type=int, default=10, help="Log every N steps.") parser.add_argument("--sample_step_interval", type=int, default=100000000000000, help="Sample images every N steps.") parser.add_argument( "--mixed_precision", type=str, default="no", choices=["no", "fp16", "bf16","tf32"], help=( "Whether to use mixed precision. Choose" "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10." "and an Nvidia Ampere GPU." ), ) parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank") parser.add_argument( "--concepts_list", type=str, default=None, help="Path to json containing multiple concepts, will overwrite parameters like instance_prompt, class_prompt, etc.", ) parser.add_argument("--save_sample_controlled_seed", type=int, action='append', help="Set a seed for an extra sample image to be constantly saved.") parser.add_argument("--detect_full_drive", default=True, action="store_true", help="Delete checkpoints when the drive is full.") parser.add_argument("--send_telegram_updates", default=False, action="store_true", help="Send Telegram updates.") parser.add_argument("--telegram_chat_id", type=str, default="0", help="Telegram chat ID.") parser.add_argument("--telegram_token", type=str, default="0", help="Telegram token.") parser.add_argument("--use_deepspeed_adam", default=False, action="store_true", help="Use experimental DeepSpeed Adam 8.") parser.add_argument('--append_sample_controlled_seed_action', action='append') parser.add_argument('--add_sample_prompt', type=str, action='append') parser.add_argument('--use_image_names_as_captions', default=False, action="store_true") parser.add_argument('--shuffle_captions', default=False, action="store_true") parser.add_argument("--masked_training", default=False, required=False, action='store_true', help="Whether to mask parts of the image during training") parser.add_argument("--normalize_masked_area_loss", default=False, required=False, action='store_true', help="Normalize the loss, to make it independent of the size of the masked area") parser.add_argument("--unmasked_probability", type=float, default=1, required=False, help="Probability of training a step without a mask") parser.add_argument("--max_denoising_strength", type=float, default=1, required=False, help="Max denoising steps to train on") parser.add_argument('--add_mask_prompt', type=str, default=None, action="append", dest="mask_prompts", help="Prompt for automatic mask creation") parser.add_argument('--with_gan', default=False, action="store_true", help="Use GAN (experimental)") parser.add_argument("--gan_weight", type=float, default=0.2, required=False, help="Strength of effect GAN has on training") parser.add_argument("--gan_warmup", type=float, default=0, required=False, help="Slowly increases GAN weight from zero over this many steps, useful when initializing a GAN discriminator from scratch") parser.add_argument('--discriminator_config', default="configs/discriminator_large.json", help="Location of config file to use when initializing a new GAN discriminator") parser.add_argument('--sample_from_ema', default=True, action="store_true", help="Generate sample images using the EMA model") parser.add_argument('--run_name', type=str, default=None, help="Adds a custom identifier to the sample and checkpoint directories") args = parser.parse_args() env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) if env_local_rank != -1 and env_local_rank != args.local_rank: args.local_rank = env_local_rank return args def main(): print(f" {bcolors.OKBLUE}Booting Up StableTuner{bcolors.ENDC}") print(f" {bcolors.OKBLUE}Please wait a moment as we load up some stuff...{bcolors.ENDC}") #torch.cuda.set_per_process_memory_fraction(0.5) args = parse_args() #temp arg args.batch_tokens = None if args.disable_cudnn_benchmark: torch.backends.cudnn.benchmark = False else: torch.backends.cudnn.benchmark = True if args.send_telegram_updates: send_telegram_message(f"Booting up StableTuner!\n", args.telegram_chat_id, args.telegram_token) logging_dir = Path(args.output_dir, "logs", args.logging_dir) if args.run_name: main_sample_dir = os.path.join(args.output_dir, f"samples_{args.run_name}") else: main_sample_dir = os.path.join(args.output_dir, "samples") if os.path.exists(main_sample_dir): shutil.rmtree(main_sample_dir) os.makedirs(main_sample_dir) #create logging directory if not logging_dir.exists(): logging_dir.mkdir(parents=True) #create output directory if not Path(args.output_dir).exists(): Path(args.output_dir).mkdir(parents=True) accelerator = Accelerator( gradient_accumulation_steps=args.gradient_accumulation_steps, mixed_precision=args.mixed_precision if args.mixed_precision != 'tf32' else 'no', log_with="tensorboard", logging_dir=logging_dir, ) # Currently, it's not possible to do gradient accumulation when training two models with accelerate.accumulate # This will be enabled soon in accelerate. For now, we don't allow gradient accumulation when training two models. # TODO (patil-suraj): Remove this check when gradient accumulation with two models is enabled in accelerate. if args.train_text_encoder and args.gradient_accumulation_steps > 1 and accelerator.num_processes > 1: raise ValueError( "Gradient accumulation is not supported when training the text encoder in distributed training. " "Please set gradient_accumulation_steps to 1. This feature will be supported in the future." ) if args.seed is not None: set_seed(args.seed) if args.concepts_list is None: args.concepts_list = [ { "instance_prompt": args.instance_prompt, "class_prompt": args.class_prompt, "instance_data_dir": args.instance_data_dir, "class_data_dir": args.class_data_dir } ] else: with open(args.concepts_list, "r") as f: args.concepts_list = json.load(f) if args.with_prior_preservation or args.add_class_images_to_dataset: pipeline = None for concept in args.concepts_list: class_images_dir = Path(concept["class_data_dir"]) class_images_dir.mkdir(parents=True, exist_ok=True) cur_class_images = len(list(class_images_dir.iterdir())) if cur_class_images < args.num_class_images: torch_dtype = torch.float16 if accelerator.device.type == "cuda" else torch.float32 if pipeline is None: pipeline = DiffusionPipeline.from_pretrained( args.pretrained_model_name_or_path, safety_checker=None, vae=AutoencoderKL.from_pretrained(args.pretrained_vae_name_or_path or args.pretrained_model_name_or_path,subfolder=None if args.pretrained_vae_name_or_path else "vae" ,safe_serialization=True), torch_dtype=torch_dtype, requires_safety_checker=False, ) pipeline.set_progress_bar_config(disable=True) pipeline.to(accelerator.device) #if args.use_bucketing == False: num_new_images = args.num_class_images - cur_class_images logger.info(f"Number of class images to sample: {num_new_images}.") sample_dataset = PromptDataset(concept["class_prompt"], num_new_images) sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size) sample_dataloader = accelerator.prepare(sample_dataloader) #else: #create class images that match up to the concept target buckets # instance_images_dir = Path(concept["instance_data_dir"]) # cur_instance_images = len(list(instance_images_dir.iterdir())) #target_wh = min(self.aspects, key=lambda aspects:abs(aspects[0]/aspects[1] - image_aspect)) # num_new_images = cur_instance_images - cur_class_images with torch.autocast("cuda"): for example in tqdm( sample_dataloader, desc="Generating class images", disable=not accelerator.is_local_main_process ): with torch.autocast("cuda"): images = pipeline(example["prompt"],height=args.resolution,width=args.resolution).images for i, image in enumerate(images): hash_image = hashlib.sha1(image.tobytes()).hexdigest() image_filename = class_images_dir / f"{example['index'][i] + cur_class_images}-{hash_image}.jpg" image.save(image_filename) del pipeline if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() # Load the tokenizer if args.tokenizer_name: tokenizer = CLIPTokenizer.from_pretrained(args.tokenizer_name ) elif args.pretrained_model_name_or_path: #print(os.getcwd()) tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder="tokenizer" ) # Load models and create wrapper for stable diffusion #text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder" ) text_encoder = CLIPTextModel.from_pretrained( args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision, ) vae = AutoencoderKL.from_pretrained( args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision, ) unet = UNet2DConditionModel.from_pretrained( args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision, torch_dtype=torch.float32 ) if args.with_gan: if os.path.isdir(os.path.join(args.pretrained_model_name_or_path, "discriminator")): discriminator = Discriminator2D.from_pretrained( args.pretrained_model_name_or_path, subfolder="discriminator", revision=args.revision, ) else: print(f" {bcolors.WARNING}Discriminator network (GAN) not found. Initializing a new network. It may take a very large number of steps to train.{bcolors.ENDC}") if not args.gan_warmup: print(f" {bcolors.WARNING}Consider using --gan_warmup to stabilize the model while the discriminator is being trained.{bcolors.ENDC}") with open(args.discriminator_config, "r") as f: discriminator_config = json.load(f) discriminator = Discriminator2D.from_config(discriminator_config) if is_xformers_available() and args.attention=='xformers': try: vae.enable_xformers_memory_efficient_attention() unet.enable_xformers_memory_efficient_attention() if args.with_gan: discriminator.enable_xformers_memory_efficient_attention() except Exception as e: logger.warning( "Could not enable memory efficient attention. Make sure xformers is installed" f" correctly and a GPU is available: {e}" ) elif args.attention=='flash_attention': replace_unet_cross_attn_to_flash_attention() if args.use_ema == True: if os.path.isdir(os.path.join(args.pretrained_model_name_or_path, "unet_ema")): ema_unet = UNet2DConditionModel.from_pretrained( args.pretrained_model_name_or_path, subfolder="unet_ema", revision=args.revision, torch_dtype=torch.float32 ) else: ema_unet = copy.deepcopy(unet) ema_unet.config["step"] = 0 for param in ema_unet.parameters(): param.requires_grad = False if args.model_variant == "depth2img": d2i = Depth2Img(unet,text_encoder,args.mixed_precision,args.pretrained_model_name_or_path,accelerator) vae.requires_grad_(False) vae.enable_slicing() if not args.train_text_encoder: text_encoder.requires_grad_(False) if args.gradient_checkpointing: unet.enable_gradient_checkpointing() if args.train_text_encoder: text_encoder.gradient_checkpointing_enable() if args.with_gan: discriminator.enable_gradient_checkpointing() if args.scale_lr: args.learning_rate = ( args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes ) # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs if args.use_8bit_adam and args.use_deepspeed_adam==False and args.use_lion==False: try: import bitsandbytes as bnb except ImportError: raise ImportError( "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." ) optimizer_class = bnb.optim.AdamW8bit print("Using 8-bit Adam") elif args.use_8bit_adam and args.use_deepspeed_adam==True: try: from deepspeed.ops.adam import DeepSpeedCPUAdam except ImportError: raise ImportError( "To use 8-bit DeepSpeed Adam, try updating your cuda and deepspeed integrations." ) optimizer_class = DeepSpeedCPUAdam elif args.use_lion == True: print("Using LION optimizer") optimizer_class = Lion elif args.use_deepspeed_adam==False and args.use_lion==False and args.use_8bit_adam==False: optimizer_class = torch.optim.AdamW params_to_optimize = ( itertools.chain(unet.parameters(), text_encoder.parameters()) if args.train_text_encoder else unet.parameters() ) if args.use_lion == False: optimizer = optimizer_class( params_to_optimize, lr=args.learning_rate, betas=(args.adam_beta1, args.adam_beta2), weight_decay=args.adam_weight_decay, eps=args.adam_epsilon, ) if args.with_gan: optimizer_discriminator = optimizer_class( discriminator.parameters(), lr=args.learning_rate, betas=(args.adam_beta1, args.adam_beta2), weight_decay=args.adam_weight_decay, eps=args.adam_epsilon, ) else: optimizer = optimizer_class( params_to_optimize, lr=args.learning_rate, betas=(args.adam_beta1, args.adam_beta2), weight_decay=args.adam_weight_decay, #eps=args.adam_epsilon, ) if args.with_gan: optimizer_discriminator = optimizer_class( discriminator.parameters(), lr=args.learning_rate, betas=(args.adam_beta1, args.adam_beta2), weight_decay=args.adam_weight_decay, #eps=args.adam_epsilon, ) noise_scheduler = DDPMScheduler.from_config(args.pretrained_model_name_or_path, subfolder="scheduler") if args.use_bucketing: train_dataset = AutoBucketing( concepts_list=args.concepts_list, use_image_names_as_captions=args.use_image_names_as_captions, shuffle_captions=args.shuffle_captions, batch_size=args.train_batch_size, tokenizer=tokenizer, add_class_images_to_dataset=args.add_class_images_to_dataset, balance_datasets=args.auto_balance_concept_datasets, resolution=args.resolution, with_prior_loss=False,#args.with_prior_preservation, repeats=args.dataset_repeats, use_text_files_as_captions=args.use_text_files_as_captions, aspect_mode=args.aspect_mode, action_preference=args.aspect_mode_action_preference, seed=args.seed, model_variant=args.model_variant, extra_module=None if args.model_variant != "depth2img" else d2i, mask_prompts=args.mask_prompts, load_mask=args.masked_training, ) else: train_dataset = NormalDataset( concepts_list=args.concepts_list, tokenizer=tokenizer, with_prior_preservation=args.with_prior_preservation, size=args.resolution, center_crop=args.center_crop, num_class_images=args.num_class_images, use_image_names_as_captions=args.use_image_names_as_captions, shuffle_captions=args.shuffle_captions, repeats=args.dataset_repeats, use_text_files_as_captions=args.use_text_files_as_captions, seed = args.seed, model_variant=args.model_variant, extra_module=None if args.model_variant != "depth2img" else d2i, mask_prompts=args.mask_prompts, load_mask=args.masked_training, ) def collate_fn(examples): #print(examples) #print('test') input_ids = [example["instance_prompt_ids"] for example in examples] tokens = input_ids pixel_values = [example["instance_images"] for example in examples] mask = None if "mask" in examples[0]: mask = [example["mask"] for example in examples] if args.model_variant == 'depth2img': depth = [example["instance_depth_images"] for example in examples] #print('test') # Concat class and instance examples for prior preservation. # We do this to avoid doing two forward passes. if args.with_prior_preservation: input_ids += [example["class_prompt_ids"] for example in examples] pixel_values += [example["class_images"] for example in examples] if "mask" in examples[0]: mask += [example["class_mask"] for example in examples] if args.model_variant == 'depth2img': depth = [example["class_depth_images"] for example in examples] mask_values = None if mask is not None: mask_values = torch.stack(mask) mask_values = mask_values.to(memory_format=torch.contiguous_format).float() if args.model_variant == 'depth2img': depth_values = torch.stack(depth) depth_values = depth_values.to(memory_format=torch.contiguous_format).float() ### no need to do it now when it's loaded by the multiAspectsDataset #if args.with_prior_preservation: # input_ids += [example["class_prompt_ids"] for example in examples] # pixel_values += [example["class_images"] for example in examples] #print(pixel_values) #unpack the pixel_values from tensor to list pixel_values = torch.stack(pixel_values) pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float() input_ids = tokenizer.pad( {"input_ids": input_ids}, padding="max_length", max_length=tokenizer.model_max_length, return_tensors="pt",\ ).input_ids extra_values = None if args.model_variant == 'depth2img': extra_values = depth_values return { "input_ids": input_ids, "pixel_values": pixel_values, "extra_values": extra_values, "mask_values": mask_values, "tokens": tokens } train_dataloader = torch.utils.data.DataLoader( train_dataset, batch_size=args.train_batch_size, shuffle=False, collate_fn=collate_fn, pin_memory=True ) #get the length of the dataset train_dataset_length = len(train_dataset) #code to check if latent cache needs to be resaved #check if last_run.json file exists in logging_dir if os.path.exists(logging_dir / "last_run.json"): #if it exists, load it with open(logging_dir / "last_run.json", "r") as f: last_run = json.load(f) last_run_batch_size = last_run["batch_size"] last_run_dataset_length = last_run["dataset_length"] if last_run_batch_size != args.train_batch_size: print(f" {bcolors.WARNING}The batch_size has changed since the last run. Regenerating Latent Cache.{bcolors.ENDC}") args.regenerate_latent_cache = True #save the new batch_size and dataset_length to last_run.json if last_run_dataset_length != train_dataset_length: print(f" {bcolors.WARNING}The dataset length has changed since the last run. Regenerating Latent Cache.{bcolors.ENDC}") args.regenerate_latent_cache = True #save the new batch_size and dataset_length to last_run.json with open(logging_dir / "last_run.json", "w") as f: json.dump({"batch_size": args.train_batch_size, "dataset_length": train_dataset_length}, f) else: #if it doesn't exist, create it last_run = {"batch_size": args.train_batch_size, "dataset_length": train_dataset_length} #create the file with open(logging_dir / "last_run.json", "w") as f: json.dump(last_run, f) weight_dtype = torch.float32 if accelerator.mixed_precision == "fp16": print("Using fp16") weight_dtype = torch.float16 elif accelerator.mixed_precision == "bf16": print("Using bf16") weight_dtype = torch.bfloat16 elif args.mixed_precision == "tf32": torch.backends.cuda.matmul.allow_tf32 = True #torch.set_float32_matmul_precision("medium") # Move text_encode and vae to gpu. # For mixed precision training we cast the text_encoder and vae weights to half-precision # as these models are only used for inference, keeping weights in full precision is not required. vae.to(accelerator.device, dtype=weight_dtype) if args.use_ema == True: ema_unet.to(accelerator.device) if not args.train_text_encoder: text_encoder.to(accelerator.device, dtype=weight_dtype) if args.use_bucketing: wh = set([tuple(x.target_wh) for x in train_dataset.image_train_items]) else: wh = set([tuple([args.resolution, args.resolution]) for x in train_dataset.image_paths]) full_mask_by_aspect = {shape: vae.encode(torch.zeros(1, 3, shape[1], shape[0]).to(accelerator.device, dtype=weight_dtype)).latent_dist.mean * 0.18215 for shape in wh} cached_dataset = CachedLatentsDataset(batch_size=args.train_batch_size, text_encoder=text_encoder, tokenizer=tokenizer, dtype=weight_dtype, model_variant=args.model_variant, shuffle_per_epoch="False", args = args,) gen_cache = False data_len = len(train_dataloader) latent_cache_dir = Path(args.output_dir, "logs", "latent_cache") #check if latents_cache.pt exists in the output_dir if not os.path.exists(latent_cache_dir): os.makedirs(latent_cache_dir) for i in range(0,data_len-1): if not os.path.exists(os.path.join(latent_cache_dir, f"latents_cache_{i}.pt")): gen_cache = True break if args.regenerate_latent_cache == True: files = os.listdir(latent_cache_dir) gen_cache = True for file in files: os.remove(os.path.join(latent_cache_dir,file)) if gen_cache == False : print(f" {bcolors.OKGREEN}Loading Latent Cache from {latent_cache_dir}{bcolors.ENDC}") del vae if not args.train_text_encoder: del text_encoder if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() #load all the cached latents into a single dataset for i in range(0,data_len-1): cached_dataset.add_pt_cache(os.path.join(latent_cache_dir,f"latents_cache_{i}.pt")) if gen_cache == True: #delete all the cached latents if they exist to avoid problems print(f" {bcolors.WARNING}Generating latents cache...{bcolors.ENDC}") train_dataset = LatentsDataset([], [], [], [], [], []) counter = 0 ImageFile.LOAD_TRUNCATED_IMAGES = True with torch.no_grad(): for batch in tqdm(train_dataloader, desc="Caching latents", bar_format='%s{l_bar}%s%s{bar}%s%s{r_bar}%s'%(bcolors.OKBLUE,bcolors.ENDC, bcolors.OKBLUE, bcolors.ENDC,bcolors.OKBLUE,bcolors.ENDC,)): cached_extra = None cached_mask = None batch["pixel_values"] = batch["pixel_values"].to(accelerator.device, non_blocking=True, dtype=weight_dtype) batch["input_ids"] = batch["input_ids"].to(accelerator.device, non_blocking=True) cached_latent = vae.encode(batch["pixel_values"]).latent_dist if batch["mask_values"] is not None: cached_mask = functional.resize(batch["mask_values"], size=cached_latent.mean.shape[2:]) if batch["mask_values"] is not None and args.model_variant == "inpainting": batch["mask_values"] = batch["mask_values"].to(accelerator.device, non_blocking=True, dtype=weight_dtype) cached_extra = vae.encode(batch["pixel_values"] * (1 - batch["mask_values"])).latent_dist if args.model_variant == "depth2img": batch["extra_values"] = batch["extra_values"].to(accelerator.device, non_blocking=True, dtype=weight_dtype) cached_extra = functional.resize(batch["extra_values"], size=cached_latent.mean.shape[2:]) if args.train_text_encoder: cached_text_enc = batch["input_ids"] else: cached_text_enc = text_encoder(batch["input_ids"])[0] train_dataset.add_latent(cached_latent, cached_text_enc, cached_mask, cached_extra, batch["tokens"]) del batch del cached_latent del cached_text_enc del cached_mask del cached_extra torch.save(train_dataset, os.path.join(latent_cache_dir,f"latents_cache_{counter}.pt")) cached_dataset.add_pt_cache(os.path.join(latent_cache_dir,f"latents_cache_{counter}.pt")) counter += 1 train_dataset = LatentsDataset([], [], [], [], [], []) #if counter % 300 == 0: #train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=1, collate_fn=lambda x: x, shuffle=False) # gc.collect() # torch.cuda.empty_cache() # accelerator.free_memory() #clear vram after caching latents del vae if not args.train_text_encoder: del text_encoder if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() #load all the cached latents into a single dataset train_dataloader = torch.utils.data.DataLoader(cached_dataset, batch_size=1, collate_fn=lambda x: x, shuffle=False) print(f" {bcolors.OKGREEN}Latents are ready.{bcolors.ENDC}") # Scheduler and math around the number of training steps. overrode_max_train_steps = False num_update_steps_per_epoch = len(train_dataloader) if args.max_train_steps is None: args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch overrode_max_train_steps = True if args.lr_warmup_steps < 1: args.lr_warmup_steps = math.floor(args.lr_warmup_steps * args.max_train_steps / args.gradient_accumulation_steps) lr_scheduler = get_scheduler( args.lr_scheduler, optimizer=optimizer, num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps, num_training_steps=args.max_train_steps, ) if args.train_text_encoder and not args.use_ema: unet, text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( unet, text_encoder, optimizer, train_dataloader, lr_scheduler ) elif args.train_text_encoder and args.use_ema: unet, text_encoder, ema_unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( unet, text_encoder, ema_unet, optimizer, train_dataloader, lr_scheduler ) elif not args.train_text_encoder and args.use_ema: unet, ema_unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( unet, ema_unet, optimizer, train_dataloader, lr_scheduler ) elif not args.train_text_encoder and not args.use_ema: unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( unet, optimizer, train_dataloader, lr_scheduler ) if args.with_gan: lr_scheduler_discriminator = get_scheduler( args.lr_scheduler, optimizer=optimizer_discriminator, num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps, num_training_steps=args.max_train_steps, ) discriminator, optimizer_discriminator, lr_scheduler_discriminator = accelerator.prepare(discriminator, optimizer_discriminator, lr_scheduler_discriminator) # We need to recalculate our total training steps as the size of the training dataloader may have changed. num_update_steps_per_epoch = len(train_dataloader) if overrode_max_train_steps: args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch #print(args.max_train_steps, num_update_steps_per_epoch) # Afterwards we recalculate our number of training epochs #print(args.max_train_steps, num_update_steps_per_epoch) args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) # We need to initialize the trackers we use, and also store our configuration. # The trackers initializes automatically on the main process. if accelerator.is_main_process: accelerator.init_trackers("dreambooth") # Train! total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps logger.info("***** Running training *****") logger.info(f" Num examples = {len(train_dataset)}") logger.info(f" Num batches each epoch = {len(train_dataloader)}") logger.info(f" Num Epochs = {args.num_train_epochs}") logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") logger.info(f" Total optimization steps = {args.max_train_steps}") def mid_train_playground(step): tqdm.write(f"{bcolors.WARNING} Booting up GUI{bcolors.ENDC}") epoch = step // num_update_steps_per_epoch if args.train_text_encoder and args.stop_text_encoder_training == True: text_enc_model = accelerator.unwrap_model(text_encoder,True) elif args.train_text_encoder and args.stop_text_encoder_training > epoch: text_enc_model = accelerator.unwrap_model(text_encoder,True) elif args.train_text_encoder == False: text_enc_model = CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder" ) elif args.train_text_encoder and args.stop_text_encoder_training <= epoch: if 'frozen_directory' in locals(): text_enc_model = CLIPTextModel.from_pretrained(frozen_directory, subfolder="text_encoder") else: text_enc_model = accelerator.unwrap_model(text_encoder,True) scheduler = DPMSolverMultistepScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler") unwrapped_unet = accelerator.unwrap_model(ema_unet if args.use_ema else unet,True) pipeline = DiffusionPipeline.from_pretrained( args.pretrained_model_name_or_path, unet=unwrapped_unet, text_encoder=text_enc_model, vae=AutoencoderKL.from_pretrained(args.pretrained_vae_name_or_path or args.pretrained_model_name_or_path,subfolder=None if args.pretrained_vae_name_or_path else "vae", safe_serialization=True), safety_checker=None, torch_dtype=weight_dtype, local_files_only=False, requires_safety_checker=False, ) pipeline.scheduler = scheduler if is_xformers_available() and args.attention=='xformers': try: vae.enable_xformers_memory_efficient_attention() unet.enable_xformers_memory_efficient_attention() except Exception as e: logger.warning( "Could not enable memory efficient attention. Make sure xformers is installed" f" correctly and a GPU is available: {e}" ) elif args.attention=='flash_attention': replace_unet_cross_attn_to_flash_attention() pipeline = pipeline.to(accelerator.device) def inference(prompt, negative_prompt, num_samples, height=512, width=512, num_inference_steps=50,seed=-1,guidance_scale=7.5): with torch.autocast("cuda"), torch.inference_mode(): if seed != -1: if g_cuda is None: g_cuda = torch.Generator(device='cuda') else: g_cuda.manual_seed(int(seed)) else: seed = random.randint(0, 100000) g_cuda = torch.Generator(device='cuda') g_cuda.manual_seed(seed) return pipeline( prompt, height=int(height), width=int(width), negative_prompt=negative_prompt, num_images_per_prompt=int(num_samples), num_inference_steps=int(num_inference_steps), guidance_scale=guidance_scale, generator=g_cuda).images, seed with gr.Blocks() as demo: with gr.Row(): with gr.Column(): prompt = gr.Textbox(label="Prompt", value="photo of zwx dog in a bucket") negative_prompt = gr.Textbox(label="Negative Prompt", value="") run = gr.Button(value="Generate") with gr.Row(): num_samples = gr.Number(label="Number of Samples", value=4) guidance_scale = gr.Number(label="Guidance Scale", value=7.5) with gr.Row(): height = gr.Number(label="Height", value=512) width = gr.Number(label="Width", value=512) with gr.Row(): num_inference_steps = gr.Slider(label="Steps", value=25) seed = gr.Number(label="Seed", value=-1) with gr.Column(): gallery = gr.Gallery() seedDisplay = gr.Number(label="Used Seed:", value=0) run.click(inference, inputs=[prompt, negative_prompt, num_samples, height, width, num_inference_steps,seed, guidance_scale], outputs=[gallery,seedDisplay]) demo.launch(share=True,prevent_thread_lock=True) tqdm.write(f"{bcolors.WARNING}Gradio Session is active, Press 'F12' to resume training{bcolors.ENDC}") keyboard.wait('f12') demo.close() del demo del text_enc_model del unwrapped_unet del pipeline return def save_and_sample_weights(step,context='checkpoint',save_model=True): try: #check how many folders are in the output dir #if there are more than 5, delete the oldest one #save the model #save the optimizer #save the lr_scheduler #save the args height = args.sample_height width = args.sample_width batch_prompts = [] if args.sample_from_batch > 0: num_samples = args.sample_from_batch if args.sample_from_batch < args.train_batch_size else args.train_batch_size batch_prompts = [] tokens = args.batch_tokens if tokens != None: allPrompts = list(set([tokenizer.decode(p).replace('<|endoftext|>','').replace('<|startoftext|>', '') for p in tokens])) if len(allPrompts) < num_samples: num_samples = len(allPrompts) batch_prompts = random.sample(allPrompts, num_samples) if args.sample_aspect_ratios: #choose random aspect ratio from ASPECTS aspect_ratio = random.choice(ASPECTS) height = aspect_ratio[0] width = aspect_ratio[1] if os.path.exists(args.output_dir): if args.detect_full_drive==True: folders = os.listdir(args.output_dir) #check how much space is left on the drive total, used, free = shutil.disk_usage("/") if (free // (2**30)) < 4: #folders.remove("0") #get the folder with the lowest number #oldest_folder = min(folder for folder in folders if folder.isdigit()) tqdm.write(f"{bcolors.FAIL}Drive is almost full, Please make some space to continue training.{bcolors.ENDC}") if args.send_telegram_updates: try: send_telegram_message(f"Drive is almost full, Please make some space to continue training.", args.telegram_chat_id, args.telegram_token) except: pass #count time import time start_time = time.time() import platform while input("Press Enter to continue... if you're on linux we'll wait 5 minutes for you to make space and continue"): #check if five minutes have passed #check if os is linux if 'Linux' in platform.platform(): if time.time() - start_time > 300: break #oldest_folder_path = os.path.join(args.output_dir, oldest_folder) #shutil.rmtree(oldest_folder_path) # Create the pipeline using using the trained modules and save it. if accelerator.is_main_process: if 'step' in context: #what is the current epoch epoch = step // num_update_steps_per_epoch else: epoch = step if args.train_text_encoder and args.stop_text_encoder_training == True: text_enc_model = accelerator.unwrap_model(text_encoder,True) elif args.train_text_encoder and args.stop_text_encoder_training > epoch: text_enc_model = accelerator.unwrap_model(text_encoder,True) elif args.train_text_encoder == False: text_enc_model = CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder" ) elif args.train_text_encoder and args.stop_text_encoder_training <= epoch: if 'frozen_directory' in locals(): text_enc_model = CLIPTextModel.from_pretrained(frozen_directory, subfolder="text_encoder") else: text_enc_model = accelerator.unwrap_model(text_encoder,True) #scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False) #scheduler = EulerDiscreteScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler", prediction_type="v_prediction") scheduler = DPMSolverMultistepScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler") unwrapped_unet = accelerator.unwrap_model(unet,True) pipeline = DiffusionPipeline.from_pretrained( args.pretrained_model_name_or_path, unet=unwrapped_unet, text_encoder=text_enc_model, vae=AutoencoderKL.from_pretrained(args.pretrained_vae_name_or_path or args.pretrained_model_name_or_path,subfolder=None if args.pretrained_vae_name_or_path else "vae",), safety_checker=None, torch_dtype=weight_dtype, local_files_only=False, requires_safety_checker=False, ) pipeline.scheduler = scheduler if is_xformers_available() and args.attention=='xformers': try: unet.enable_xformers_memory_efficient_attention() except Exception as e: logger.warning( "Could not enable memory efficient attention. Make sure xformers is installed" f" correctly and a GPU is available: {e}" ) elif args.attention=='flash_attention': replace_unet_cross_attn_to_flash_attention() if args.run_name: save_dir = os.path.join(args.output_dir, f"{context}_{step}_{args.run_name}") else: save_dir = os.path.join(args.output_dir, f"{context}_{step}") if args.flatten_sample_folder: sample_dir = main_sample_dir else: sample_dir = os.path.join(main_sample_dir, f"{context}_{step}") #if sample dir path does not exist, create it if args.stop_text_encoder_training == True: save_dir = frozen_directory if save_model: pipeline.save_pretrained(save_dir,safe_serialization=True) if args.with_gan: discriminator.save_pretrained(os.path.join(save_dir, "discriminator"), safe_serialization=True) if args.use_ema: ema_unet.save_pretrained(os.path.join(save_dir, "unet_ema"), safe_serialization=True) with open(os.path.join(save_dir, "args.json"), "w") as f: json.dump(args.__dict__, f, indent=2) if args.stop_text_encoder_training == True: #delete every folder in frozen_directory but the text encoder for folder in os.listdir(save_dir): if folder != "text_encoder" and os.path.isdir(os.path.join(save_dir, folder)): shutil.rmtree(os.path.join(save_dir, folder)) imgs = [] if args.use_ema and args.sample_from_ema: pipeline.unet = ema_unet for param in unet.parameters(): param.requires_grad = False if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() if args.add_sample_prompt is not None or batch_prompts != [] and args.stop_text_encoder_training != True: prompts = [] if args.add_sample_prompt is not None: for prompt in args.add_sample_prompt: prompts.append(prompt) if batch_prompts != []: for prompt in batch_prompts: prompts.append(prompt) pipeline = pipeline.to(accelerator.device) pipeline.set_progress_bar_config(disable=True) #sample_dir = os.path.join(save_dir, "samples") #if sample_dir exists, delete it if os.path.exists(sample_dir): if not args.flatten_sample_folder: shutil.rmtree(sample_dir) os.makedirs(sample_dir, exist_ok=True) with torch.autocast("cuda"), torch.inference_mode(): if args.send_telegram_updates: try: send_telegram_message(f"Generating samples for {step} {context}", args.telegram_chat_id, args.telegram_token) except: pass n_sample = args.n_save_sample if args.save_sample_controlled_seed: n_sample += len(args.save_sample_controlled_seed) progress_bar_sample = tqdm(total=len(prompts)*n_sample,desc="Generating samples") for samplePrompt in prompts: sampleIndex = prompts.index(samplePrompt) #convert sampleIndex to number in words # Data to be written sampleProperties = { "samplePrompt" : samplePrompt } # Serializing json json_object = json.dumps(sampleProperties, indent=4) if args.flatten_sample_folder: sampleName = f"{context}_{step}_prompt_{sampleIndex+1}" else: sampleName = f"prompt_{sampleIndex+1}" if not args.flatten_sample_folder: os.makedirs(os.path.join(sample_dir,sampleName), exist_ok=True) if args.model_variant == 'inpainting': conditioning_image = torch.zeros(1, 3, height, width) mask = torch.ones(1, 1, height, width) if args.model_variant == 'depth2img': #pil new white image test_image = Image.new('RGB', (width, height), (255, 255, 255)) depth_image = Image.new('RGB', (width, height), (255, 255, 255)) depth = np.array(depth_image.convert("L")) depth = depth.astype(np.float32) / 255.0 depth = depth[None, None] depth = torch.from_numpy(depth) for i in range(n_sample): #check if the sample is controlled by a seed if i < args.n_save_sample: if args.model_variant == 'inpainting': images = pipeline(samplePrompt, conditioning_image, mask, height=height,width=width, guidance_scale=args.save_guidance_scale, num_inference_steps=args.save_infer_steps).images if args.model_variant == 'depth2img': images = pipeline(samplePrompt,image=test_image, guidance_scale=args.save_guidance_scale, num_inference_steps=args.save_infer_steps,strength=1.0).images elif args.model_variant == 'base': images = pipeline(samplePrompt,height=height,width=width, guidance_scale=args.save_guidance_scale, num_inference_steps=args.save_infer_steps).images if not args.flatten_sample_folder: images[0].save(os.path.join(sample_dir,sampleName, f"{sampleName}_{i}.png")) else: images[0].save(os.path.join(sample_dir, f"{sampleName}_{i}.png")) else: seed = args.save_sample_controlled_seed[i - args.n_save_sample] generator = torch.Generator("cuda").manual_seed(seed) if args.model_variant == 'inpainting': images = pipeline(samplePrompt,conditioning_image, mask,height=height,width=width, guidance_scale=args.save_guidance_scale, num_inference_steps=args.save_infer_steps, generator=generator).images if args.model_variant == 'depth2img': images = pipeline(samplePrompt,image=test_image, guidance_scale=args.save_guidance_scale, num_inference_steps=args.save_infer_steps,generator=generator,strength=1.0).images elif args.model_variant == 'base': images = pipeline(samplePrompt,height=height,width=width, guidance_scale=args.save_guidance_scale, num_inference_steps=args.save_infer_steps, generator=generator).images if not args.flatten_sample_folder: images[0].save(os.path.join(sample_dir,sampleName, f"{sampleName}_controlled_seed_{str(seed)}.png")) else: images[0].save(os.path.join(sample_dir, f"{sampleName}_controlled_seed_{str(seed)}.png")) progress_bar_sample.update(1) if args.send_telegram_updates: imgs = [] #get all the images from the sample folder if not args.flatten_sample_folder: dir = os.listdir(os.path.join(sample_dir,sampleName)) else: dir = sample_dir for file in dir: if file.endswith(".png"): #open the image with pil img = Image.open(os.path.join(sample_dir,sampleName,file)) imgs.append(img) try: send_media_group(args.telegram_chat_id,args.telegram_token,imgs, caption=f"Samples for the {step} {context} using the prompt:\n\n{samplePrompt}") except: pass del pipeline del unwrapped_unet for param in unet.parameters(): param.requires_grad = True if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() if save_model == True: tqdm.write(f"{bcolors.OKGREEN}Weights saved to {save_dir}{bcolors.ENDC}") elif save_model == False and len(imgs) > 0: del imgs tqdm.write(f"{bcolors.OKGREEN}Samples saved to {sample_dir}{bcolors.ENDC}") except Exception as e: tqdm.write(e) tqdm.write(f"{bcolors.FAIL} Error occured during sampling, skipping.{bcolors.ENDC}") pass @torch.no_grad() def update_ema(ema_model, model): ema_step = ema_model.config["step"] decay = min((ema_step + 1) / (ema_step + 10), 0.9999) ema_model.config["step"] += 1 for (s_param, param) in zip(ema_model.parameters(), model.parameters()): if param.requires_grad: s_param.add_((1 - decay) * (param - s_param)) else: s_param.copy_(param) # Only show the progress bar once on each machine. progress_bar = tqdm(range(args.max_train_steps),bar_format='%s{l_bar}%s%s{bar}%s%s{r_bar}%s'%(bcolors.OKBLUE,bcolors.ENDC, bcolors.OKBLUE, bcolors.ENDC,bcolors.OKBLUE,bcolors.ENDC,), disable=not accelerator.is_local_main_process) progress_bar_inter_epoch = tqdm(range(num_update_steps_per_epoch),bar_format='%s{l_bar}%s%s{bar}%s%s{r_bar}%s'%(bcolors.OKBLUE,bcolors.ENDC, bcolors.OKGREEN, bcolors.ENDC,bcolors.OKBLUE,bcolors.ENDC,), disable=not accelerator.is_local_main_process) progress_bar_e = tqdm(range(args.num_train_epochs),bar_format='%s{l_bar}%s%s{bar}%s%s{r_bar}%s'%(bcolors.OKBLUE,bcolors.ENDC, bcolors.OKGREEN, bcolors.ENDC,bcolors.OKBLUE,bcolors.ENDC,), disable=not accelerator.is_local_main_process) progress_bar.set_description("Overall Steps") progress_bar_inter_epoch.set_description("Steps To Epoch") progress_bar_e.set_description("Overall Epochs") global_step = 0 loss_avg = AverageMeter("loss_avg", max_eta=0.999) gan_loss_avg = AverageMeter("gan_loss_avg", max_eta=0.999) text_enc_context = nullcontext() if args.train_text_encoder else torch.no_grad() if args.send_telegram_updates: try: send_telegram_message(f"Starting training with the following settings:\n\n{format_dict(args.__dict__)}", args.telegram_chat_id, args.telegram_token) except: pass try: tqdm.write(f"{bcolors.OKBLUE}Starting Training!{bcolors.ENDC}") try: def toggle_gui(event=None): if keyboard.is_pressed("ctrl") and keyboard.is_pressed("shift") and keyboard.is_pressed("g"): tqdm.write(f"{bcolors.WARNING}GUI will boot as soon as the current step is done.{bcolors.ENDC}") nonlocal mid_generation if mid_generation == True: mid_generation = False tqdm.write(f"{bcolors.WARNING}Cancelled GUI.{bcolors.ENDC}") else: mid_generation = True def toggle_checkpoint(event=None): if keyboard.is_pressed("ctrl") and keyboard.is_pressed("shift") and keyboard.is_pressed("s") and not keyboard.is_pressed("alt"): tqdm.write(f"{bcolors.WARNING}Saving the model as soon as this epoch is done.{bcolors.ENDC}") nonlocal mid_checkpoint if mid_checkpoint == True: mid_checkpoint = False tqdm.write(f"{bcolors.WARNING}Cancelled Checkpointing.{bcolors.ENDC}") else: mid_checkpoint = True def toggle_sample(event=None): if keyboard.is_pressed("ctrl") and keyboard.is_pressed("shift") and keyboard.is_pressed("p") and not keyboard.is_pressed("alt"): tqdm.write(f"{bcolors.WARNING}Sampling will begin as soon as this epoch is done.{bcolors.ENDC}") nonlocal mid_sample if mid_sample == True: mid_sample = False tqdm.write(f"{bcolors.WARNING}Cancelled Sampling.{bcolors.ENDC}") else: mid_sample = True def toggle_checkpoint_step(event=None): if keyboard.is_pressed("ctrl") and keyboard.is_pressed("shift") and keyboard.is_pressed("alt") and keyboard.is_pressed("s"): tqdm.write(f"{bcolors.WARNING}Saving the model as soon as this step is done.{bcolors.ENDC}") nonlocal mid_checkpoint_step if mid_checkpoint_step == True: mid_checkpoint_step = False tqdm.write(f"{bcolors.WARNING}Cancelled Checkpointing.{bcolors.ENDC}") else: mid_checkpoint_step = True def toggle_sample_step(event=None): if keyboard.is_pressed("ctrl") and keyboard.is_pressed("shift") and keyboard.is_pressed("alt") and keyboard.is_pressed("p"): tqdm.write(f"{bcolors.WARNING}Sampling will begin as soon as this step is done.{bcolors.ENDC}") nonlocal mid_sample_step if mid_sample_step == True: mid_sample_step = False tqdm.write(f"{bcolors.WARNING}Cancelled Sampling.{bcolors.ENDC}") else: mid_sample_step = True def toggle_quit_and_save_epoch(event=None): if keyboard.is_pressed("ctrl") and keyboard.is_pressed("shift") and keyboard.is_pressed("q") and not keyboard.is_pressed("alt"): tqdm.write(f"{bcolors.WARNING}Quitting and saving the model as soon as this epoch is done.{bcolors.ENDC}") nonlocal mid_quit if mid_quit == True: mid_quit = False tqdm.write(f"{bcolors.WARNING}Cancelled Quitting.{bcolors.ENDC}") else: mid_quit = True def toggle_quit_and_save_step(event=None): if keyboard.is_pressed("ctrl") and keyboard.is_pressed("shift") and keyboard.is_pressed("alt") and keyboard.is_pressed("q"): tqdm.write(f"{bcolors.WARNING}Quitting and saving the model as soon as this step is done.{bcolors.ENDC}") nonlocal mid_quit_step if mid_quit_step == True: mid_quit_step = False tqdm.write(f"{bcolors.WARNING}Cancelled Quitting.{bcolors.ENDC}") else: mid_quit_step = True def help(event=None): if keyboard.is_pressed("ctrl") and keyboard.is_pressed("h"): print_instructions() keyboard.on_press_key("g", toggle_gui) keyboard.on_press_key("s", toggle_checkpoint) keyboard.on_press_key("p", toggle_sample) keyboard.on_press_key("s", toggle_checkpoint_step) keyboard.on_press_key("p", toggle_sample_step) keyboard.on_press_key("q", toggle_quit_and_save_epoch) keyboard.on_press_key("q", toggle_quit_and_save_step) keyboard.on_press_key("h", help) print_instructions() except Exception as e: pass mid_generation = False mid_checkpoint = False mid_sample = False mid_checkpoint_step = False mid_sample_step = False mid_quit = False mid_quit_step = False #lambda set mid_generation to true if args.run_name: frozen_directory = os.path.join(args.output_dir, f"frozen_text_encoder_{args.run_name}") else: frozen_directory = os.path.join(args.output_dir, "frozen_text_encoder") unet_stats = {} discriminator_stats = {} os.makedirs(main_sample_dir, exist_ok=True) with open(os.path.join(main_sample_dir, "args.json"), "w") as f: json.dump(args.__dict__, f, indent=2) if args.with_gan: with open(os.path.join(main_sample_dir, "discriminator_config.json"), "w") as f: json.dump(discriminator.config, f, indent=2) for epoch in range(args.num_train_epochs): #every 10 epochs print instructions unet.train() if args.train_text_encoder: text_encoder.train() #save initial weights if args.sample_on_training_start==True and epoch==0: save_and_sample_weights(epoch,'start',save_model=False) if args.train_text_encoder and args.stop_text_encoder_training == epoch: args.stop_text_encoder_training = True if accelerator.is_main_process: tqdm.write(f"{bcolors.WARNING} Stopping text encoder training{bcolors.ENDC}") current_percentage = (epoch/args.num_train_epochs)*100 #round to the nearest whole number current_percentage = round(current_percentage,0) try: send_telegram_message(f"Text encoder training stopped at epoch {epoch} which is {current_percentage}% of training. Freezing weights and saving.", args.telegram_chat_id, args.telegram_token) except: pass if os.path.exists(frozen_directory): #delete the folder if it already exists shutil.rmtree(frozen_directory) os.mkdir(frozen_directory) save_and_sample_weights(epoch,'epoch') args.stop_text_encoder_training = epoch progress_bar_inter_epoch.reset(total=num_update_steps_per_epoch) for step, batch in enumerate(train_dataloader): with accelerator.accumulate(unet): # Convert images to latent space with torch.no_grad(): latent_dist = batch[0][0] latents = latent_dist.sample() * 0.18215 if args.model_variant == 'inpainting': mask = batch[0][2] mask_mean = batch[0][3] conditioning_latent_dist = batch[0][4] conditioning_latents = conditioning_latent_dist.sample() * 0.18215 if args.model_variant == 'depth2img': depth = batch[0][4] if args.sample_from_batch > 0: args.batch_tokens = batch[0][5] # Sample noise that we'll add to the latents # and some extra bits to make it so that the model learns to change the zero-frequency of the component freely # https://www.crosslabs.org/blog/diffusion-with-offset-noise if (args.with_offset_noise == True): noise = torch.randn_like(latents) + (args.offset_noise_weight * torch.randn(latents.shape[0], latents.shape[1], 1, 1).to(accelerator.device)) else: noise = torch.randn_like(latents) bsz = latents.shape[0] # Sample a random timestep for each image timesteps = torch.randint(0, int(noise_scheduler.config.num_train_timesteps * args.max_denoising_strength), (bsz,), device=latents.device) timesteps = timesteps.long() # Add noise to the latents according to the noise magnitude at each timestep # (this is the forward diffusion process) noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps) # Get the text embedding for conditioning with text_enc_context: if args.train_text_encoder: if args.clip_penultimate == True: encoder_hidden_states = text_encoder(batch[0][1],output_hidden_states=True) encoder_hidden_states = text_encoder.text_model.final_layer_norm(encoder_hidden_states['hidden_states'][-2]) else: encoder_hidden_states = text_encoder(batch[0][1])[0] else: encoder_hidden_states = batch[0][1] # Predict the noise residual mask=None if args.model_variant == 'inpainting': if mask is not None and random.uniform(0, 1) < args.unmasked_probability: # for some steps, predict the unmasked image conditioning_latents = torch.stack([full_mask_by_aspect[tuple([latents.shape[3]*8, latents.shape[2]*8])].squeeze()] * bsz) mask = torch.ones(bsz, 1, latents.shape[2], latents.shape[3]).to(accelerator.device, dtype=weight_dtype) noisy_inpaint_latents = torch.concat([noisy_latents, mask, conditioning_latents], 1) model_pred = unet(noisy_inpaint_latents, timesteps, encoder_hidden_states).sample elif args.model_variant == 'depth2img': noisy_depth_latents = torch.cat([noisy_latents, depth], dim=1) model_pred = unet(noisy_depth_latents, timesteps, encoder_hidden_states, depth).sample elif args.model_variant == "base": model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample # Get the target for loss depending on the prediction type if noise_scheduler.config.prediction_type == "epsilon": target = noise elif noise_scheduler.config.prediction_type == "v_prediction": target = noise_scheduler.get_velocity(latents, noise, timesteps) else: raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}") # GAN stuff # Input: noisy_latents # True output: target # Fake output: model_pred if args.with_gan: # Turn on learning for the discriminator, and do an optimization step for param in discriminator.parameters(): param.requires_grad = True pred_fake = discriminator(torch.cat((noisy_latents, model_pred), 1).detach(), encoder_hidden_states) pred_real = discriminator(torch.cat((noisy_latents, target), 1), encoder_hidden_states) discriminator_loss = F.mse_loss(pred_fake, torch.zeros_like(pred_fake), reduction="mean") + F.mse_loss(pred_real, torch.ones_like(pred_real), reduction="mean") if discriminator_loss.isnan(): tqdm.write(f"{bcolors.WARNING}Discriminator loss is NAN, skipping GAN update.{bcolors.ENDC}") else: accelerator.backward(discriminator_loss) if accelerator.sync_gradients: accelerator.clip_grad_norm_(discriminator.parameters(), args.max_grad_norm) optimizer_discriminator.step() lr_scheduler_discriminator.step() # Hack to fix NaNs caused by GAN training for name, p in discriminator.named_parameters(): if p.isnan().any(): fix_nans_(p, name, discriminator_stats[name]) else: (std, mean) = torch.std_mean(p) discriminator_stats[name] = (std.item(), mean.item()) del std, mean optimizer_discriminator.zero_grad() del pred_real, pred_fake, discriminator_loss # Turn off learning for the discriminator for the generator optimization step for param in discriminator.parameters(): param.requires_grad = False if args.with_prior_preservation: # Chunk the noise and noise_pred into two parts and compute the loss on each part separately. """ noise_pred, noise_pred_prior = torch.chunk(noise_pred, 2, dim=0) noise, noise_prior = torch.chunk(noise, 2, dim=0) # Compute instance loss loss = F.mse_loss(noise_pred.float(), noise.float(), reduction="none").mean([1, 2, 3]).mean() # Compute prior loss prior_loss = F.mse_loss(noise_pred_prior.float(), noise_prior.float(), reduction="mean") # Add the prior loss to the instance loss. loss = loss + args.prior_loss_weight * prior_loss """ # Chunk the noise and model_pred into two parts and compute the loss on each part separately. model_pred, model_pred_prior = torch.chunk(model_pred, 2, dim=0) target, target_prior = torch.chunk(target, 2, dim=0) if mask is not None and args.model_variant != "inpainting": loss = masked_mse_loss(model_pred.float(), target.float(), mask, reduction="none").mean([1, 2, 3]).mean() prior_loss = masked_mse_loss(model_pred_prior.float(), target_prior.float(), mask, reduction="mean") else: loss = F.mse_loss(model_pred.float(), target.float(), reduction="none").mean([1, 2, 3]).mean() prior_loss = F.mse_loss(model_pred_prior.float(), target_prior.float(), reduction="mean") # Add the prior loss to the instance loss. loss = loss + args.prior_loss_weight * prior_loss if mask is not None and args.normalize_masked_area_loss: loss = loss / mask_mean else: if mask is not None and args.model_variant != "inpainting": loss = masked_mse_loss(model_pred.float(), target.float(), mask, reduction="none").mean([1, 2, 3]) loss = loss.mean() else: loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean") if mask is not None and args.normalize_masked_area_loss: loss = loss / mask_mean base_loss = loss if args.with_gan: # Add loss from the GAN pred_fake = discriminator(torch.cat((noisy_latents, model_pred), 1), encoder_hidden_states) gan_loss = F.mse_loss(pred_fake, torch.ones_like(pred_fake), reduction="mean") if gan_loss.isnan(): tqdm.write(f"{bcolors.WARNING}GAN loss is NAN, skipping GAN loss.{bcolors.ENDC}") else: gan_weight = args.gan_weight if args.gan_warmup and global_step < args.gan_warmup: gan_weight *= global_step / args.gan_warmup loss += gan_weight * gan_loss del pred_fake accelerator.backward(loss) if accelerator.sync_gradients: params_to_clip = ( itertools.chain(unet.parameters(), text_encoder.parameters()) if args.train_text_encoder else unet.parameters() ) accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm) optimizer.step() lr_scheduler.step() # Hack to fix NaNs caused by GAN training for name, p in unet.named_parameters(): if p.isnan().any(): fix_nans_(p, name, unet_stats[name]) else: (std, mean) = torch.std_mean(p) unet_stats[name] = (std.item(), mean.item()) del std, mean optimizer.zero_grad() loss_avg.update(base_loss.detach_()) if args.with_gan and not gan_loss.isnan(): gan_loss_avg.update(gan_loss.detach_()) if args.use_ema == True: update_ema(ema_unet, unet) del loss, model_pred if args.with_prior_preservation: del model_pred_prior logs = {"loss": loss_avg.avg.item(), "lr": lr_scheduler.get_last_lr()[0]} if args.with_gan: logs["gan_loss"] = gan_loss_avg.avg.item() progress_bar.set_postfix(**logs) if not global_step % args.log_interval: accelerator.log(logs, step=global_step) if global_step > 0 and not global_step % args.sample_step_interval: save_and_sample_weights(global_step,'step',save_model=False) progress_bar.update(1) progress_bar_inter_epoch.update(1) progress_bar_e.refresh() global_step += 1 if mid_quit_step==True: accelerator.wait_for_everyone() save_and_sample_weights(global_step,'quit_step') quit() if mid_generation==True: mid_train_playground(global_step) mid_generation=False if mid_checkpoint_step == True: save_and_sample_weights(global_step,'step',save_model=True) mid_checkpoint_step=False mid_sample_step=False elif mid_sample_step == True: save_and_sample_weights(global_step,'step',save_model=False) mid_sample_step=False if global_step >= args.max_train_steps: break progress_bar_e.update(1) if mid_quit==True: accelerator.wait_for_everyone() save_and_sample_weights(epoch,'quit_epoch') quit() if epoch == args.num_train_epochs - 1: save_and_sample_weights(epoch,'epoch',True) elif args.save_every_n_epoch and (epoch + 1) % args.save_every_n_epoch == 0: save_and_sample_weights(epoch,'epoch',True) elif mid_checkpoint==True: save_and_sample_weights(epoch,'epoch',True) mid_checkpoint=False mid_sample=False elif mid_sample==True: save_and_sample_weights(epoch,'epoch',False) mid_sample=False accelerator.wait_for_everyone() except Exception: try: send_telegram_message("Something went wrong while training! :(", args.telegram_chat_id, args.telegram_token) #save_and_sample_weights(global_step,'checkpoint') send_telegram_message(f"Saved checkpoint {global_step} on exit", args.telegram_chat_id, args.telegram_token) except Exception: pass raise except KeyboardInterrupt: send_telegram_message("Training stopped", args.telegram_chat_id, args.telegram_token) try: send_telegram_message("Training finished!", args.telegram_chat_id, args.telegram_token) except: pass accelerator.end_training() if __name__ == "__main__": main()