train.py

import argparse
import os

import torch
import torch.nn.functional as F
import torch.utils.checkpoint

from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from diffusers import AutoencoderKL, DDPMScheduler, DDIMScheduler, UNet2DConditionModel
from diffusers_ import StableDiffusionPipeline
from huggingface_hub import HfFolder, Repository, whoami
from PIL import Image
import numpy as np
from torchvision import transforms
from tqdm.auto import tqdm
from transformers import CLIPTextModel, CLIPTokenizer
from torch import autocast
from utils import inf_save


logger = get_logger(__name__)


def parse_args():
    parser = argparse.ArgumentParser(description="Simple example of a training script.")
    parser.add_argument(
        "--pretrained_model_name",
        type=str,
        default=None,
        required=True,
        help="model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--input_image",
        type=str,
        default=None,
        required=True,
        help="Path to input image to edit.",
    )
    parser.add_argument(
        "--target_text",
        type=str,
        default=None,
        help="The target text describing the output image.",
    )
    parser.add_argument(
        "--source_text",
        type=str,
        default=None,
        help="The source text describing the input image.",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="output",
        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", action="store_true", help="Whether to center crop images before resizing to resolution"
    )
    parser.add_argument(
        "--n_hiper",
        type=int,
        default=5,
        help="Number of hiper embedding",
    )
    parser.add_argument(
        "--emb_train_steps",
        type=int,
        default=1500,
        help="Total number of training steps to perform.",
    )
    parser.add_argument(
        "--emb_learning_rate",
        type=float,
        default=1e-3,
        help="Learning rate for optimizing the embeddings.",
    )
    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(
        "--use_8bit_adam", action="store_false", 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_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
    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."
        ),
    )

    args = parser.parse_args()
    return args


def main():
    args = parse_args()


    # Setting
    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        mixed_precision=args.mixed_precision,
    )

    if args.seed is not None:
        set_seed(args.seed)
        g_cuda = torch.Generator(device='cuda')
        g_cuda.manual_seed(args.seed)
    
    os.makedirs(args.output_dir, exist_ok=True)
    if args.use_8bit_adam:
        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.Adam8bit
    else:
        optimizer_class = torch.optim.Adam

    weight_dtype = torch.float32
    if args.mixed_precision == "fp16":
        weight_dtype = torch.float16
    elif args.mixed_precision == "bf16":
        weight_dtype = torch.bfloat16
        
        
    # Load pretrained models    
    tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_name, subfolder="tokenizer", use_auth_token=True)
    CLIP_text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_name, subfolder="text_encoder", use_auth_token=True)
    vae = AutoencoderKL.from_pretrained(args.pretrained_model_name, subfolder="vae", use_auth_token=True)
    unet = UNet2DConditionModel.from_pretrained(args.pretrained_model_name, subfolder="unet", use_auth_token=True)
    noise_scheduler = DDPMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
    

    # Encode the input image.
    vae.to(accelerator.device, dtype=weight_dtype)
    input_image = Image.open(args.input_image).convert("RGB")
    image_transforms = transforms.Compose(
        [
            transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
            transforms.CenterCrop(args.resolution) if args.center_crop else transforms.RandomCrop(args.resolution),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5]),
        ]
    )
    init_image = image_transforms(input_image)
    init_image = init_image[None].to(device=accelerator.device, dtype=weight_dtype)
    with torch.inference_mode():
        init_latents = vae.encode(init_image).latent_dist.sample()
        init_latents = 0.18215 * init_latents


    # Encode the source and target text.
    CLIP_text_encoder.to(accelerator.device, dtype=weight_dtype)
    text_ids_tgt = tokenizer(args.target_text, padding="max_length", truncation=True, max_length=tokenizer.model_max_length, return_tensors="pt").input_ids
    text_ids_tgt = text_ids_tgt.to(device=accelerator.device)

    text_ids_src = tokenizer(args.source_text,padding="max_length",truncation=True,max_length=tokenizer.model_max_length,return_tensors="pt").input_ids
    text_ids_src = text_ids_src.to(device=accelerator.device)
    with torch.inference_mode():
        target_embeddings = CLIP_text_encoder(text_ids_tgt)[0].float()
        source_embeddings = CLIP_text_encoder(text_ids_src)[0].float()
    
    
    # del vae, CLIP_text_encoder
    del vae, CLIP_text_encoder
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


    # For inference
    ddim_scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
    pipe = StableDiffusionPipeline.from_pretrained(args.pretrained_model_name, scheduler=ddim_scheduler, torch_dtype=torch.float16).to("cuda")
    num_samples = 1 
    guidance_scale = 7.5 
    num_inference_steps = 50
    height = 512
    width = 512
    
    
    # Optimize hiper embedding
    n_hiper = args.n_hiper
    hiper_embeddings = source_embeddings[:,-n_hiper:].clone().detach()
    src_embeddings = source_embeddings[:,:-n_hiper].clone().detach()
    tgt_embeddings = target_embeddings[:,:-n_hiper].clone().detach()
    hiper_embeddings.requires_grad_(True)
    
    
    optimizer = optimizer_class(
        [hiper_embeddings],  
        lr=args.emb_learning_rate,
        betas=(args.adam_beta1, args.adam_beta2),
        eps=args.adam_epsilon,
    )

    unet, optimizer = accelerator.prepare(unet, optimizer)


    def train_loop(pbar, optimizer, hiper_embeddings):
        for step in pbar:
            with accelerator.accumulate(unet):
                
                noise = torch.randn_like(init_latents)
                bsz = init_latents.shape[0]
                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=init_latents.device)
                timesteps = timesteps.long()

                noisy_latents = noise_scheduler.add_noise(init_latents, noise, timesteps)

                source_embeddings = torch.cat([src_embeddings, hiper_embeddings], 1)
                noise_pred = unet(noisy_latents, timesteps, source_embeddings).sample
                loss = F.mse_loss(noise_pred.float(), noise.float(), reduction="mean")

                accelerator.backward(loss)
                optimizer.step()
                optimizer.zero_grad(set_to_none=True)
            
            # Check inference
            if step%100==0:
                inf_emb = []
                inf_emb.append(torch.cat([src_embeddings, hiper_embeddings.clone().detach()], 1))
                inf_emb.append(torch.cat([tgt_embeddings, hiper_embeddings.clone().detach()*0.8], 1))
                inf_emb.append(target_embeddings)
                
                inf_images=[]    
                inf_images.append(input_image)            
                with autocast("cuda"), torch.inference_mode():     
                    for i in range(3):   
                        images = pipe(text_embeddings=inf_emb[i], height=height, width=width, num_images_per_prompt=num_samples,
                            num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=g_cuda).images
                        inf_images.append(images[0])

                save_name = os.path.join(args.output_dir, 'training_results_seed{}_step{}.png'.format(args.seed, step))
                inf_save(inf_images, ['source_img', '[src, hper]', '[tgt, hper]', '[tgt]'], save_name)
                del images
                
                torch.save(hiper_embeddings.cpu(), os.path.join(args.output_dir, "hiper_embeddings_step{}.pt".format(step)))

        accelerator.wait_for_everyone()
    
    progress_bar = tqdm(range(args.emb_train_steps), disable=not accelerator.is_local_main_process)
    progress_bar.set_description("Optimizing embedding")
    
    train_loop(progress_bar, optimizer, hiper_embeddings)


if __name__ == "__main__":
    main()