import argparse import logging import math import os from pathlib import Path import itertools import numpy as np import torch.utils.checkpoint import transformers from accelerate import Accelerator from accelerate.logging import get_logger from accelerate.utils import ProjectConfiguration, set_seed from datasets import load_dataset from PIL import Image from torchvision import transforms from tqdm.auto import tqdm from transformers import AutoTokenizer, PretrainedConfig import diffusers from diffusers import ( AutoencoderKL, UNet2DConditionModel, DDPMScheduler, UniPCMultistepScheduler, ) from diffusers.optimization import get_scheduler from diffusers.utils import check_min_version, is_wandb_available from utils.pipeline import StableHairPipeline from ref_encoder.adapter import * from ref_encoder.reference_control import ReferenceAttentionControl from ref_encoder.reference_unet import ref_unet from ref_encoder.latent_controlnet import ControlNetModel import albumentations as A import cv2 import torch.nn.functional as F # Will error if the minimal version of diffusers is not installed. Remove at your own risks. check_min_version("0.23.0") logger = get_logger(__name__) def concatenate_images(image_files, output_file, type="pil"): if type == "np": image_files = [Image.fromarray(img) for img in image_files] images = image_files # list max_height = max(img.height for img in images) images = [img.resize((img.width, max_height)) for img in images] total_width = sum(img.width for img in images) combined = Image.new('RGB', (total_width, max_height)) x_offset = 0 for img in images: combined.paste(img, (x_offset, 0)) x_offset += img.width combined.save(output_file) def image_grid(imgs, rows, cols): assert len(imgs) == rows * cols w, h = imgs[0].size grid = Image.new("RGB", size=(cols * w, rows * h)) for i, img in enumerate(imgs): grid.paste(img, box=(i % cols * w, i // cols * h)) return grid def log_validation(vae, text_encoder, tokenizer, unet, controlnet, hair_encoder, args, accelerator, weight_dtype, step): logger.info("Running validation... ") controlnet = accelerator.unwrap_model(controlnet) hair_encoder = accelerator.unwrap_model(hair_encoder) pipeline = StableHairPipeline.from_pretrained( args.pretrained_model_name_or_path, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, controlnet=controlnet, safety_checker=None, revision=args.revision, torch_dtype=weight_dtype, ) pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config) pipeline = pipeline.to(accelerator.device) pipeline.set_progress_bar_config(disable=True) validation_ids = args.validation_ids validation_hairs = args.validation_hairs validation_path = os.path.join(args.output_dir, "validation", f"step-{step}") os.makedirs(validation_path, exist_ok=True) _num = 0 for validation_id, validation_hair in zip(validation_ids, validation_hairs): _num += 1 validation_id = np.array(Image.open(validation_id).convert("RGB").resize((512, 512))) validation_hair = np.array(Image.open(validation_hair).convert("RGB").resize((512, 512))) for num in range(args.num_validation_images): with torch.autocast("cuda"): sample = pipeline( prompt="", negative_prompt="", num_inference_steps=30, guidance_scale=2, width=512, height=512, controlnet_condition=validation_id, controlnet_conditioning_scale=1., generator=None, reference_encoder=hair_encoder, ref_image=validation_hair, ).samples concatenate_images([validation_id, validation_hair, (sample * 255.).astype(np.uint8)], output_file=os.path.join(validation_path, str(num)+str(_num)+".jpg"), type="np") def import_model_class_from_model_name_or_path(pretrained_model_name_or_path: str, revision: str): text_encoder_config = PretrainedConfig.from_pretrained( pretrained_model_name_or_path, subfolder="text_encoder", revision=revision, ) model_class = text_encoder_config.architectures[0] if model_class == "CLIPTextModel": from transformers import CLIPTextModel return CLIPTextModel elif model_class == "RobertaSeriesModelWithTransformation": from diffusers.pipelines.alt_diffusion.modeling_roberta_series import RobertaSeriesModelWithTransformation return RobertaSeriesModelWithTransformation else: raise ValueError(f"{model_class} is not supported.") def parse_args(input_args=None): parser = argparse.ArgumentParser(description="Simple example of training script.") parser.add_argument("--noise_offset", type=float, default=0.1, help="The scale of noise offset.") parser.add_argument( "--pretrained_model_name_or_path", type=str, default="", help="Path to pretrained model or model identifier from huggingface.co/models." ) parser.add_argument( "--controlnet_model_name_or_path", type=str, default=None, help="Path to pretrained controlnet model or model identifier from huggingface.co/models." " If not specified controlnet weights are initialized from unet.", ) parser.add_argument( "--train_data_dir", type=str, default="", help=( "A folder containing the training data. Folder contents must follow the structure described in" " https://huggingface.co/docs/datasets/image_dataset#imagefolder. In particular, a `metadata.jsonl` file" " must exist to provide the captions for the images. Ignored if `dataset_name` is specified." ), ) parser.add_argument("--refer_column", type=str, default="reference") parser.add_argument("--source_column", type=str, default="source") parser.add_argument("--target_column", type=str, default="target") parser.add_argument( "--revision", type=str, default=None, required=False, help=( "Revision of pretrained model identifier from huggingface.co/models. Trainable model components should be" " float32 precision." ), ) parser.add_argument( "--output_dir", type=str, default="train_lr1e-5_refunet", 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( "--train_batch_size", type=int, default=1, help="Batch size (per device) for the training dataloader." ) parser.add_argument("--num_train_epochs", type=int, default=1000) 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( "--checkpointing_steps", type=int, default=1000, help=( "Save a checkpoint of the training state every X updates. Checkpoints can be used for resuming training via `--resume_from_checkpoint`. " "In the case that the checkpoint is better than the final trained model, the checkpoint can also be used for inference." "Using a checkpoint for inference requires separate loading of the original pipeline and the individual checkpointed model components." "See https://huggingface.co/docs/diffusers/main/en/training/dreambooth#performing-inference-using-a-saved-checkpoint for step by step" "instructions." ), ) parser.add_argument( "--resume_from_checkpoint", type=str, default=None, help=( "Whether training should be resumed from a previous checkpoint. Use a path saved by" ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.' ), ) 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", 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=1e-5, 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=int, default=500, help="Number of steps for the warmup in the lr scheduler." ) parser.add_argument( "--lr_num_cycles", type=int, default=1, help="Number of hard resets of the lr in cosine_with_restarts scheduler.", ) parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.") parser.add_argument( "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes." ) parser.add_argument( "--dataloader_num_workers", type=int, default=0, help=( "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." ), ) 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( "--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( "--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="fp16", 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( "--enable_xformers_memory_efficient_attention", action="store_true", help="Whether or not to use xformers." ) parser.add_argument( "--max_train_samples", type=int, default=None, help=( "For debugging purposes or quicker training, truncate the number of training examples to this " "value if set." ), ) parser.add_argument( "--proportion_empty_prompts", type=float, default=0, help="Proportion of image prompts to be replaced with empty strings. Defaults to 0 (no prompt replacement).", ) parser.add_argument( "--validation_ids", type=str, default=["/share2/zhangyuxuan/project/stable_hair/test_imgs/ID/girl.jpg", "/share2/zhangyuxuan/project/stable_hair/test_imgs/ID/man.jpg"], nargs="+", help=( "A set of prompts evaluated every `--validation_steps` and logged to `--report_to`." " Provide either a matching number of `--validation_image`s, a single `--validation_image`" " to be used with all prompts, or a single prompt that will be used with all `--validation_image`s." ), ) parser.add_argument( "--validation_hairs", type=str, default=["", ""], nargs="+", help=( "A set of paths to the controlnet conditioning image be evaluated every `--validation_steps`" " and logged to `--report_to`. Provide either a matching number of `--validation_prompt`s, a" " a single `--validation_prompt` to be used with all `--validation_image`s, or a single" " `--validation_image` that will be used with all `--validation_prompt`s." ), ) parser.add_argument( "--num_validation_images", type=int, default=3, help="Number of images to be generated for each `--validation_image`, `--validation_prompt` pair", ) parser.add_argument( "--validation_steps", type=int, default=1000, help=( "Run validation every X steps. Validation consists of running the prompt" " `args.validation_prompt` multiple times: `args.num_validation_images`" " and logging the images." ), ) parser.add_argument( "--tracker_project_name", type=str, default="train", help=( "The `project_name` argument passed to Accelerator.init_trackers for" " more information see https://huggingface.co/docs/accelerate/v0.17.0/en/package_reference/accelerator#accelerate.Accelerator" ), ) if input_args is not None: args = parser.parse_args(input_args) else: args = parser.parse_args() if args.resolution % 8 != 0: raise ValueError( "`--resolution` must be divisible by 8 for consistently sized encoded images between the VAE and the controlnet encoder." ) return args def make_train_dataset(args, tokenizer, accelerator): if args.train_data_dir is not None: dataset = load_dataset('json', data_files=args.train_data_dir) column_names = dataset["train"].column_names # 6. Get the column names for input/target. if args.refer_column is None: refer_column = column_names[0] logger.info(f"image column defaulting to {refer_column}") else: refer_column = args.refer_column if refer_column not in column_names: raise ValueError( f"`--refer_column` value '{args.refer_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}" ) if args.source_column is None: source_column = column_names[1] logger.info(f"source column defaulting to {source_column}") else: source_column = args.source_column if source_column not in column_names: raise ValueError( f"`--source_column` value '{args.source_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}" ) if args.target_column is None: target_column = column_names[1] logger.info(f"target column defaulting to {target_column}") else: target_column = args.target_column if target_column not in column_names: raise ValueError( f"`--target_column` value '{args.target_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}" ) norm = transforms.Normalize([0.5], [0.5]) to_tensor = transforms.ToTensor() prob = 0.7 pixel_transform = A.Compose([ A.SmallestMaxSize(max_size=512), A.CenterCrop(512, 512), A.Affine(scale=(0.5, 1), translate_percent={"x": (-0.1, 0.1), "y": (-0.1, 0.1)}, rotate=(-10, 10), p=0.8), A.OneOf( [ A.PixelDropout(dropout_prob=0.1, p=prob), A.GaussNoise(var_limit=(10.0, 50.0), mean=0, p=prob), A.RandomShadow(shadow_roi=(0.1, 0.1, 0.9, 0.9), p=prob), ] ) ], additional_targets={'image0': 'image', 'image1': 'image'}) hair_transform = A.Compose([ A.SmallestMaxSize(max_size=512), A.CenterCrop(512, 512), A.Affine(scale=(0.9, 1.2), rotate=(-10, 10), p=0.7)] ) def refer_imgaug(image): image = cv2.resize(cv2.cvtColor(image, cv2.COLOR_BGR2RGB), [512, 512]) results = hair_transform(image=image) image = norm(to_tensor(results["image"]/255.)) return image def imgaug(source_image, target_image): source_image = cv2.resize(cv2.cvtColor(source_image, cv2.COLOR_BGR2RGB), [512, 512]) target_image = cv2.resize(cv2.cvtColor(target_image, cv2.COLOR_BGR2RGB), [512, 512]) results = pixel_transform(image=source_image, image0=target_image) source_image, target_image = norm(to_tensor(results["image"]/255.)), norm(to_tensor(results["image0"]/255.)) return source_image, target_image def preprocess_train(examples): source_images = [cv2.imread(image) for image in examples[source_column]] refer_images = [cv2.imread(image) for image in examples[refer_column]] target_images = [cv2.imread(image) for image in examples[target_column]] pair = [imgaug(image1, image2) for image1, image2 in zip(source_images, target_images)] source_images, target_images = zip(*pair) source_images_ls = list(source_images) target_images_ls = list(target_images) refer_images_ls = [refer_imgaug(image) for image in refer_images] examples["source_pixel_values"] = source_images_ls examples["refer_pixel_values"] = refer_images_ls examples["target_pixel_values"] = target_images_ls return examples with accelerator.main_process_first(): train_dataset = dataset["train"].with_transform(preprocess_train) return train_dataset def collate_fn(examples): source_pixel_values = torch.stack([example["source_pixel_values"] for example in examples]) source_pixel_values = source_pixel_values.to(memory_format=torch.contiguous_format).float() refer_pixel_values = torch.stack([example["refer_pixel_values"] for example in examples]) refer_pixel_values = refer_pixel_values.to(memory_format=torch.contiguous_format).float() target_pixel_values = torch.stack([example["target_pixel_values"] for example in examples]) target_pixel_values = target_pixel_values.to(memory_format=torch.contiguous_format).float() return { "source_pixel_values": source_pixel_values, "refer_pixel_values": refer_pixel_values, "target_pixel_values": target_pixel_values, } def main(args): logging_dir = Path(args.output_dir, args.logging_dir) accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) accelerator = Accelerator( gradient_accumulation_steps=args.gradient_accumulation_steps, mixed_precision=args.mixed_precision, log_with=args.report_to, project_config=accelerator_project_config, ) # Make one log on every process with the configuration for debugging. logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO, ) logger.info(accelerator.state, main_process_only=False) if accelerator.is_local_main_process: transformers.utils.logging.set_verbosity_warning() diffusers.utils.logging.set_verbosity_info() else: transformers.utils.logging.set_verbosity_error() diffusers.utils.logging.set_verbosity_error() # If passed along, set the training seed now. if args.seed is not None: set_seed(args.seed) # Handle the repository creation if accelerator.is_main_process: if args.output_dir is not None: os.makedirs(args.output_dir, exist_ok=True) # Load the tokenizer if args.pretrained_model_name_or_path: tokenizer = AutoTokenizer.from_pretrained( args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision, use_fast=False, ) # import correct text encoder class text_encoder_cls = import_model_class_from_model_name_or_path(args.pretrained_model_name_or_path, args.revision) # Load scheduler and models noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler") text_encoder = text_encoder_cls.from_pretrained( args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision ).to(accelerator.device) vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision).to(accelerator.device) unet = UNet2DConditionModel.from_pretrained( args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision ).to(accelerator.device) if args.controlnet_model_name_or_path: logger.info("Loading existing controlnet weights") controlnet = ControlNetModel.from_pretrained(args.controlnet_model_name_or_path).to(accelerator.device) else: logger.info("Initializing controlnet weights from unet") controlnet = ControlNetModel.from_unet(unet).to(accelerator.device) ### load Hair encoder/adapter/reference_control_modules resume = False if resume: Hair_Encoder = ref_unet.from_pretrained( args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision ).to(accelerator.device) pretrained_folder = "" # your checkpoint path _state_dict = torch.load(os.path.join(pretrained_folder, "pytorch_model.bin")) Hair_Encoder.load_state_dict(_state_dict, strict=False) torch.cuda.empty_cache() _state_dict = torch.load(os.path.join(pretrained_folder, "pytorch_model_1.bin")) Hair_Adapter = adapter_injection(unet, dtype=torch.float32, use_resampler=False) Hair_Adapter.load_state_dict(_state_dict, strict=False) torch.cuda.empty_cache() _state_dict = torch.load(os.path.join(pretrained_folder, "pytorch_model_2.bin")) controlnet.load_state_dict(_state_dict, strict=False) torch.cuda.empty_cache() else: Hair_Encoder = ref_unet.from_pretrained( args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision ).to(accelerator.device) Hair_Adapter = adapter_injection(unet, dtype=torch.float32).to(accelerator.device) vae.requires_grad_(False) text_encoder.requires_grad_(False) unet.requires_grad_(False) Hair_Encoder.requires_grad_(True) Hair_Adapter.requires_grad_(True) controlnet.requires_grad_(True) optimizer_class = torch.optim.AdamW # Optimizer creation params_to_optimize = itertools.chain(controlnet.parameters(), Hair_Encoder.parameters(), Hair_Adapter.parameters()) 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, ) train_dataset = make_train_dataset(args, tokenizer, accelerator) train_dataloader = torch.utils.data.DataLoader( train_dataset, shuffle=True, collate_fn=collate_fn, batch_size=args.train_batch_size, num_workers=args.dataloader_num_workers, ) # Scheduler and math around the number of training steps. overrode_max_train_steps = False num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) 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 lr_scheduler = get_scheduler( args.lr_scheduler, optimizer=optimizer, num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, num_training_steps=args.max_train_steps * accelerator.num_processes, num_cycles=args.lr_num_cycles, power=args.lr_power, ) # Prepare everything with our `accelerator`. Hair_Encoder, Hair_Adapter, controlnet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( Hair_Encoder, Hair_Adapter, controlnet, optimizer, train_dataloader, lr_scheduler ) # 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. weight_dtype = torch.float32 if accelerator.mixed_precision == "fp16": weight_dtype = torch.float16 elif accelerator.mixed_precision == "bf16": weight_dtype = torch.bfloat16 # Move vae, unet and text_encoder to device and cast to weight_dtype vae.to(accelerator.device, dtype=weight_dtype) unet.to(accelerator.device, dtype=weight_dtype) text_encoder.to(accelerator.device, dtype=weight_dtype) Hair_Encoder.to(accelerator.device, dtype=torch.float32) Hair_Adapter.to(accelerator.device, dtype=torch.float32) controlnet.to(accelerator.device, dtype=torch.float32) # We need to recalculate our total training steps as the size of the training dataloader may have changed. num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) if overrode_max_train_steps: args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch # Afterwards we recalculate our number of training epochs 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: tracker_config = dict(vars(args)) # tensorboard cannot handle list types for config tracker_config.pop("validation_hairs") tracker_config.pop("validation_ids") accelerator.init_trackers(args.tracker_project_name, config=tracker_config) # 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}") global_step = 0 first_epoch = 0 initial_global_step = 0 progress_bar = tqdm( range(0, args.max_train_steps), initial=initial_global_step, desc="Steps", # Only show the progress bar once on each machine. disable=not accelerator.is_local_main_process, ) null_text_inputs = tokenizer( "", max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt" ).input_ids encoder_hidden_states = text_encoder(null_text_inputs.to(device=accelerator.device))[0] for epoch in range(first_epoch, args.num_train_epochs): for step, batch in enumerate(train_dataloader): with accelerator.accumulate(controlnet): reference_control_writer_train = ReferenceAttentionControl(Hair_Encoder, do_classifier_free_guidance=False, mode='write', fusion_blocks='full') reference_control_reader_train = ReferenceAttentionControl(unet, do_classifier_free_guidance=False, mode='read', fusion_blocks='full') # Convert images to latent space latents = vae.encode(batch["target_pixel_values"].to(dtype=weight_dtype)).latent_dist.sample() latents = latents * vae.config.scaling_factor ref_latents = vae.encode(batch["refer_pixel_values"].to(dtype=weight_dtype)).latent_dist.sample() ref_latents = ref_latents * vae.config.scaling_factor # Sample noise that we'll add to the latents noise = torch.randn_like(latents) if args.noise_offset: # https://www.crosslabs.org//blog/diffusion-with-offset-noise noise += args.noise_offset * torch.randn( (latents.shape[0], latents.shape[1], 1, 1), device=latents.device ) bsz = latents.shape[0] # Sample a random timestep for each image timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (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) # ref_noisy_latents = noise_scheduler.add_noise(ref_latents, noise, timesteps) controlnet_latents = vae.encode(batch["source_pixel_values"].to(dtype=weight_dtype)).latent_dist.sample() controlnet_latents = controlnet_latents * vae.config.scaling_factor # for b in range(bsz): # max_value = torch.max(controlnet_latents[b]) # min_value = torch.min(controlnet_latents[b]) # controlnet_latents[b] = (controlnet_latents[b]-min_value)/(max_value-min_value) down_block_res_samples, mid_block_res_sample = controlnet( noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states.repeat(bsz, 1, 1), controlnet_cond=controlnet_latents, return_dict=False, ) # writer Hair_Encoder( # ref_noisy_latents, ref_latents, timesteps, encoder_hidden_states=encoder_hidden_states.repeat(bsz, 1, 1)) reference_control_reader_train.update(reference_control_writer_train) # Predict the noise residual model_pred = unet( noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states.repeat(bsz, 1, 1).to(dtype=weight_dtype), down_block_additional_residuals=[ sample.to(dtype=weight_dtype) for sample in down_block_res_samples ], mid_block_additional_residual=mid_block_res_sample.to(dtype=weight_dtype), ).sample # clean the reader reference_control_reader_train.clear() # 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}") loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean") accelerator.backward(loss) optimizer.step() lr_scheduler.step() optimizer.zero_grad() # Checks if the accelerator has performed an optimization step behind the scenes if accelerator.sync_gradients: progress_bar.update(1) global_step += 1 if accelerator.is_main_process: if global_step % args.checkpointing_steps == 0: save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}") accelerator.save_state(save_path, safe_serialization=False) logger.info(f"Saved state to {save_path}") if args.validation_ids is not None and global_step % args.validation_steps == 0: log_validation( vae, text_encoder, tokenizer, unet, controlnet, Hair_Encoder, args, accelerator, weight_dtype, global_step, ) logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} progress_bar.set_postfix(**logs) accelerator.log(logs, step=global_step) if global_step >= args.max_train_steps: break # Create the pipeline using using the trained modules and save it. accelerator.wait_for_everyone() accelerator.end_training() if __name__ == "__main__": args = parse_args() main(args)