"""
Stable Diffusion - DreamBooth

Brooke Hamilton
brookehamilton@gmail.com

Tools for training a DreamBooth model. Adapted from https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py
"""
import argparse
import hashlib
import itertools
import math
import os
from pathlib import Path
from typing import Optional

import torch
import torch.nn.functional as F
import torch.utils.checkpoint
from torch.utils.data import Dataset

from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel
from diffusers.optimization import get_scheduler
from huggingface_hub import HfFolder, Repository, whoami
from PIL import Image
from torchvision import transforms
from tqdm.auto import tqdm
from transformers import CLIPTextModel, CLIPTokenizer

from dataclasses import dataclass
import getpass
import datetime

from huggingface_hub import HfFolder, Repository, whoami
from typing import Optional
from pathlib import Path
from huggingface_hub import create_repo



def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: Optional[str] = None):
    if token is None:
        token = HfFolder.get_token()
    if organization is None:
        username = whoami(token)["name"]
        return f"{username}/{model_id}"
    else:
        return f"{organization}/{model_id}"

@dataclass
class DreamBoothConfig():
    """
    This class holds all of the config values for the overall DreamBooth process
    """
    # saved model
    model_out_dir: str      # directory to save the model in
    hub_model_id: str       # huggingface repo to keep in sync with local model_out_dir
    push_to_hub: bool       # whether to push the saved model to the Huggingface Hub
    private_repo: bool      # whether to make the model repo private

    # the pretrained model we want to use
    model_id: str               # e.g. 'runwayml/stable-diffusion-v1-5
    revision: str               # usually None

    # the specific images you want to teach it to replicate
    instance_prompt: str        # e.g. 'a photograph of a vctrnxmas Victorian Christmas card painting'
    instance_data_dir: str

    # prior preservation -- additional images of the generic class
    do_prior_preservation: bool
    class_prompt: str               # e.g. 'a photograph of a Victorian Christmas card painting'
    class_data_dir: str
    num_class_images: int
    sample_batch_size: int

    # image settings
    resolution: int                     # = 512
    center_crop: bool                   # = False

    # training settings
    seed: int                           # = None
    train_text_encoder: bool            # False
    gradient_checkpointing: bool        # False
    learning_rate: int                  # 5e-6,
    use_8bit_adam: bool                 # False,
    adam_beta1: float                   # 0.9
    adam_beta2: float                   # 0.999
    adam_weight_decay: float            # 1e-2
    adam_epsilon: float                 # 1e-08
    max_grad_norm: float                # 1.0
    train_batch_size: int               # 4
    num_train_epochs: int               # 1?
    max_train_steps: int                # 800?
    gradient_accumulation_steps: int    # 1
    lr_scheduler: str                   # "constant"
    lr_warmup_steps: int                # 500
    save_steps: int
    prior_loss_weight: float

class DreamBoothDataset(Dataset):
    """
    A dataset to prepare the instance and class images with the prompts for fine-tuning the model.
    It pre-processes the images and the tokenizes prompts.
    """

    def __init__(
        self,
        instance_data_dir,
        instance_prompt,
        tokenizer,
        class_data_dir=None,
        class_prompt=None,
        size=512,
        center_crop=False,
    ):
        self.size = size
        self.center_crop = center_crop
        self.tokenizer = tokenizer

        self.instance_data_dir = Path(instance_data_dir)
        if not self.instance_data_dir.exists():
            raise ValueError("Instance images directory doesn't exist.")

        self.instance_images_path = [x for x in Path(self.instance_data_dir).iterdir() if not x.is_dir() and x.suffix in (['.png', '.jpg', '.jpeg'])]
        self.num_instance_images = len(self.instance_images_path)
        self.instance_prompt = instance_prompt
        self._length = self.num_instance_images

        if class_data_dir is not None:
            print('class_data_dir is not None')
            self.class_data_dir = Path(class_data_dir)
            self.class_data_dir.mkdir(parents=True, exist_ok=True)
            self.class_images_path = [x for x in Path(self.class_data_dir).iterdir() if not x.is_dir() and x.suffix in (['.png', '.jpg', '.jpeg'])]
            self.num_class_images = len(self.class_images_path)
            self._length = max(self.num_class_images, self.num_instance_images)
            self.class_prompt = class_prompt
        else:
            self.class_data_dir = None

        self.image_transforms = transforms.Compose(
            [
                transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
                transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size),
                transforms.ToTensor(),
                transforms.Normalize([0.5], [0.5]),
            ]
        )

    def __len__(self):
        return self._length

    def __getitem__(self, index):
        example = {}
        instance_image = Image.open(self.instance_images_path[index % self.num_instance_images])
        if not instance_image.mode == "RGB":
            instance_image = instance_image.convert("RGB")
        example["instance_images"] = self.image_transforms(instance_image)
        example["instance_prompt_ids"] = self.tokenizer(
            self.instance_prompt,
            padding="do_not_pad",
            truncation=True,
            max_length=self.tokenizer.model_max_length,
        ).input_ids

        if self.class_data_dir:
            print('triggered at if self.class_data_dir')
            print('self.num_class_images: ', self.num_class_images)
            class_image = Image.open(self.class_images_path[index % self.num_class_images])
            if not class_image.mode == "RGB":
                class_image = class_image.convert("RGB")
            example["class_images"] = self.image_transforms(class_image)
            example["class_prompt_ids"] = self.tokenizer(
                self.class_prompt,
                padding="do_not_pad",
                truncation=True,
                max_length=self.tokenizer.model_max_length,
            ).input_ids

        return example

'''class DreamBoothDataset(Dataset):
    """
    A dataset to prepare the instance and class images with the prompts for fine-tuning the model.
    It pre-processes the images and then tokenizes prompts.
    """

    def __init__(
        self,
        instance_data_dir,
        instance_prompt,
        tokenizer,
        class_data_dir=None,
        class_prompt=None,
        size=512,
        center_crop=False,
    ):
        self.size = size
        self.center_crop = center_crop
        self.tokenizer = tokenizer

        self.instance_data_dir = Path(instance_data_dir)
        if not self.instance_data_dir.exists():
            raise ValueError("Instance images directory doesn't exist.")

        self.instance_images_path = list(Path(instance_data_dir).iterdir())
        self.num_instance_images = len(self.instance_images_path)
        self.instance_prompt = instance_prompt
        self._length = self.num_instance_images

        if class_data_dir is not None:
            class_data_dir = Path(class_data_dir)
            class_data_dir.mkdir(parents=True, exist_ok=True)
            self.class_images_path = list(class_data_dir.iterdir())
            self.num_class_images = len(self.class_images_path)
            self._length = max(self.num_class_images, self.num_instance_images)
            self.class_prompt = class_prompt
        else:
            self.class_data_dir = None

        self.image_transforms = transforms.Compose(
            [
                transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
                transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size),
                transforms.ToTensor(),
                transforms.Normalize([0.5], [0.5]),
            ]
        )

    def __len__(self):
        return self._length

    def __getitem__(self, index):
        """
        Get a specific example by the index number

        N.B. If the number of instance images and class images is not equal, examples will recycle images as needed
        """
        print(f'getting example {index}')
        example = {}
        instance_image = Image.open(self.instance_images_path[index % self.num_instance_images])
        if not instance_image.mode == "RGB":
            instance_image = instance_image.convert("RGB")
        example["instance_images"] = self.image_transforms(instance_image)
        example["instance_prompt_ids"] = self.tokenizer(
            self.instance_prompt,
            padding="do_not_pad",
            truncation=True,
            max_length=self.tokenizer.model_max_length,
        ).input_ids

        if self.class_data_dir:
            print('self.class_images_path:', self.class_images_path)
            class_image = Image.open(self.class_images_path[index % self.num_class_images])
            if not class_image.mode == "RGB":
                class_image = class_image.convert("RGB")
            example["class_images"] = self.image_transforms(class_image)
            example["class_prompt_ids"] = self.tokenizer(
                self.class_prompt,
                padding="do_not_pad",
                truncation=True,
                max_length=self.tokenizer.model_max_length,
            ).input_ids

        return example'''


class PromptDataset(Dataset):
    "A simple dataset to prepare the prompts to generate class images on multiple GPUs."

    def __init__(self, prompt, num_samples):
        self.prompt = prompt
        self.num_samples = num_samples

    def __len__(self):
        return self.num_samples

    def __getitem__(self, index):
        example = {}
        example["prompt"] = self.prompt
        example["index"] = index
        return example



class Collater():
    def __init__(self, do_prior_preservation, tokenizer):
        self.do_prior_preservation = do_prior_preservation
        self.tokenizer = tokenizer
    def __call__(self, examples):
        input_ids = [example["instance_prompt_ids"] for example in examples]
        pixel_values = [example["instance_images"] for example in examples]

        # Concat class and instance examples for prior preservation.
        # We do this to avoid doing two forward passes.
        if self.do_prior_preservation:
            input_ids += [example["class_prompt_ids"] for example in examples]
            pixel_values += [example["class_images"] for example in examples]

        pixel_values = torch.stack(pixel_values)
        pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()

        input_ids = self.tokenizer.pad(
            {"input_ids": input_ids},
            padding="max_length",
            max_length=self.tokenizer.model_max_length,
            return_tensors="pt",
        ).input_ids

        batch = {
            "input_ids": input_ids,
            "pixel_values": pixel_values,
        }
        return batch


class DreamBoothRunner():
    """
    This class coordinates a dream booth training job
    """
    def __init__(self, DreamBoothConfig: DreamBoothConfig):
        print('='*50)
        print('Initialized DreamBooth model training job')
        print('='*50)
        self.config = DreamBoothConfig

        # Print some basic info about what we're doing as a sanity check
        print(f'Training DreamBooth ')

        if self.config.seed is not None:
            print(f'Setting seed to {self.config.seed}')
            set_seed(self.config.seed)

        # Make sure model_out_dir exists
        os.makedirs(self.config.model_out_dir, exist_ok=True)

        # Log in to Hugging Face
        print('Authenticating to HuggingFace')
        if "HUGGINGFACE_ACCESS_TOKEN" in os.environ:
            self.auth_token = os.environ['HUGGINGFACE_ACCESS_TOKEN']
        else:
            self.auth_token = getpass.getpass(prompt='HuggingFace access token:')


        print('Starting Accelerator')
        self.accelerator = Accelerator(gradient_accumulation_steps=1,
            mixed_precision="fp16" if torch.cuda.is_available() else "no", #["no", "fp16", "bf16"]
            log_with="tensorboard",
            logging_dir='logs')


        # Handle the huggingface repository creation
        if self.accelerator.is_main_process:
            if self.config.push_to_hub:
                if self.config.hub_model_id is None:
                    repo_name = get_full_repo_name(Path(self.config.model_out_dir).name, token=self.auth_token)
                else:
                    repo_name = self.config.hub_model_id

                print('Repo name will be: ', repo_name)
                self.repo = Repository(self.config.model_out_dir, clone_from=repo_name, use_auth_token=self.auth_token)
                print('Created repo at Run.repo')

                with open(os.path.join(self.config.model_out_dir, ".gitignore"), "w+") as gitignore:
                    if "step_*" not in gitignore:
                        gitignore.write("step_*\n")
                    if "epoch_*" not in gitignore:
                        gitignore.write("epoch_*\n")
            elif self.config.model_out_dir is not None:
                print('hit condition line 355')
                os.makedirs(self.config.model_out_dir, exist_ok=True)




        self.tokenizer = CLIPTokenizer.from_pretrained(
            self.config.model_id,
            subfolder="tokenizer",
            revision=self.config.revision,
            use_auth_token=self.auth_token
        )

        # Load models and create wrapper for stable diffusion
        print('Getting CLIPTextModel text encoder')
        self.text_encoder = CLIPTextModel.from_pretrained(
            self.config.model_id,
            subfolder="text_encoder",
            revision=self.config.revision,
            use_auth_token=self.auth_token
        )
        if not self.config.train_text_encoder:
            self.text_encoder.requires_grad_(False)
        if self.config.train_text_encoder:
            if self.config.gradient_checkpointing:
                self.text_encoder.gradient_checkpointing_enable()

        print('Getting AutoencoderKL vae')
        self.vae = AutoencoderKL.from_pretrained(
            self.config.model_id,
            subfolder="vae",
            revision=self.config.revision,
            use_auth_token=self.auth_token
        )
        self.vae.requires_grad_(False)

        print('Gettting UNet2DConditionModel unet')
        self.unet = UNet2DConditionModel.from_pretrained(
            self.config.model_id,
            subfolder="unet",
            revision=self.config.revision,
            use_auth_token=self.auth_token
        )
        if self.config.gradient_checkpointing:
            self.unet.enable_gradient_checkpointing()

        print('Getting DDPMScheduler noise scheduler')
        self.noise_scheduler = DDPMScheduler.from_config(
            self.config.model_id,
            subfolder="scheduler",
            use_auth_token=self.auth_token)

        # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs
        print('Setting optimizer class')
        if self.config.use_8bit_adam:
            print('hit condition 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`."
                )

            self.optimizer_class = bnb.optim.AdamW8bit
        else:
            print('not using 8bit adam')
            self.optimizer_class = torch.optim.AdamW

        print('Getting optimizer')
        params_to_optimize = (
            itertools.chain(self.unet.parameters(),
            self.text_encoder.parameters()) if self.config.train_text_encoder else self.unet.parameters()
            )
        self.optimizer = self.optimizer_class(
            params_to_optimize,
            lr=self.config.learning_rate,
            betas=(self.config.adam_beta1, self.config.adam_beta2),
            weight_decay=self.config.adam_weight_decay,
            eps=self.config.adam_epsilon,
        )

        print('within __init__(), triggering self.generate_class_images() prior to making the dataloader')
        if self.config.do_prior_preservation:
            self.generate_class_images()

        print('Loading image dataset')
        self.train_dataset = DreamBoothDataset(
            instance_data_dir=self.config.instance_data_dir,
            instance_prompt=self.config.instance_prompt,
            class_data_dir=self.config.class_data_dir if self.config.do_prior_preservation else None,
            class_prompt=self.config.class_prompt,
            tokenizer=self.tokenizer,
            size=self.config.resolution,
            center_crop=self.config.center_crop,
            )



        collater = Collater(do_prior_preservation=self.config.do_prior_preservation, tokenizer=self.tokenizer)


        '''def collate_fn(examples):
            """
            Custom collate function to pass to the DataLoader. This concats class and instance examples
            for prior preservation to avoid doing two forward passes
            """
            input_ids = [example["instance_prompt_ids"] for example in examples]
            pixel_values = [example["instance_images"] for example in examples]

            # Concat class and instance examples for prior preservation.
            # We do this to avoid doing two forward passes.
            if self.config.do_prior_preservation:
                input_ids += [example["class_prompt_ids"] for example in examples]
                pixel_values += [example["class_images"] for example in examples]

            pixel_values = torch.stack(pixel_values)
            pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()

            input_ids = self.tokenizer.pad(
                {"input_ids": input_ids},
                padding="max_length",
                max_length=self.tokenizer.model_max_length,
                return_tensors="pt",
            ).input_ids

            batch = {
                "input_ids": input_ids,
                "pixel_values": pixel_values,
            }
            return batch'''




        print('Setting up the DataLoader')
        self.train_dataloader = torch.utils.data.DataLoader(
            self.train_dataset, batch_size=self.config.train_batch_size, shuffle=True, collate_fn=collater, num_workers=1)

        print('Determining number of training steps')
        # Scheduler and math around the number of training steps.
        overrode_max_train_steps = False
        print('len(self.train_dataloader):', len(self.train_dataloader))
        print('self.config.gradient_accumulation_steps:', self.config.gradient_accumulation_steps)
        num_update_steps_per_epoch = math.ceil(len(self.train_dataloader) / self.config.gradient_accumulation_steps)
        print('num_update_steps_per_epoch:', num_update_steps_per_epoch)
        if self.config.max_train_steps is None:
            self.config.max_train_steps = self.config.num_train_epochs * num_update_steps_per_epoch
            overrode_max_train_steps = True
        print('max_train_steps:', self.config.max_train_steps)
        print('overrode_max_train_steps:', overrode_max_train_steps)

        print('Getting lr_scheduler')
        self.lr_scheduler = get_scheduler(
            self.config.lr_scheduler,
            optimizer=self.optimizer,
            num_warmup_steps=self.config.lr_warmup_steps * self.config.gradient_accumulation_steps,
            num_training_steps=self.config.max_train_steps * self.config.gradient_accumulation_steps,
            )

        print('Preparing unet, optimizer, train_dataloader, and lr_scheduler with accelerator')
        if self.config.train_text_encoder:
            self.unet, self.text_encoder, self.optimizer, self.train_dataloader, self.lr_scheduler = self.accelerator.prepare(
                self.unet, self.text_encoder, self.optimizer, self.train_dataloader, self.lr_scheduler
            )
        else:
            self.unet, self.optimizer, self.train_dataloader, self.lr_scheduler = self.accelerator.prepare(
                self.unet, self.optimizer, self.train_dataloader, self.lr_scheduler
            )

        print('Setting up weight dtype')
        self.weight_dtype = torch.float32
        if self.accelerator.mixed_precision == "fp16":
            self.weight_dtype = torch.float16
        elif self.accelerator.mixed_precision == "bf16":
            self.weight_dtype = torch.bfloat16

        print('Move text_encode and vae to correct accelerator device')
        # 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.
        self.vae.to(self.accelerator.device, dtype=self.weight_dtype)
        if not self.config.train_text_encoder:
            self.text_encoder.to(self.accelerator.device, dtype=self.weight_dtype)

        print('Recalculating total training steps and training epochs')
        # 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(self.train_dataloader) / self.config.gradient_accumulation_steps)
        if overrode_max_train_steps:
            self.config.max_train_steps = self.config.num_train_epochs * num_update_steps_per_epoch
        # Afterwards we recalculate our number of training epochs
        self.config.num_train_epochs = math.ceil(self.config.max_train_steps / num_update_steps_per_epoch)

        print('Initializing trackers and storing configuration')
        # We need to initialize the trackers we use, and also store our configuration.
        # The trackers initializes automatically on the main process.
        if self.accelerator.is_main_process:
            self.accelerator.init_trackers("dreambooth", config=vars(self.config))

    def generate_class_images(self):
        print('='*50)
        print('Generating new class images for prior preservation')
        print('='*50)
        #if args.with_prior_preservation:
        class_images_dir = Path(self.config.class_data_dir)
        if not class_images_dir.exists():
            class_images_dir.mkdir(parents=True)
        cur_class_images = len(list(class_images_dir.iterdir()))
        print(f'We want to create {self.config.num_class_images} total class images')
        print(f'There are currently {cur_class_images} in the class_images_dir {self.config.class_data_dir}')

        if cur_class_images < self.config.num_class_images:
            num_new_images = self.config.num_class_images - cur_class_images
            print(f"Number of class images to sample: {num_new_images}")

            print(f'Generating {num_new_images} new images')
            torch_dtype = torch.float16 if self.accelerator.device.type == "cuda" else torch.float32
            pipeline = StableDiffusionPipeline.from_pretrained(   # was previously DiffusionPipeline
                self.config.model_id,
                torch_dtype=torch_dtype,
                safety_checker=None,
                revision=self.config.revision,
                use_auth_token=self.auth_token
            )
            pipeline.set_progress_bar_config(disable=True)
            print(f'Setting up prompt data set with class prompt: "{self.config.class_prompt}"')
            sample_dataset = PromptDataset(self.config.class_prompt, num_new_images)
            print('Setting up sample dataloader')
            sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=self.config.sample_batch_size)
            print('Preparing sample dataloader with accelerator')
            sample_dataloader = self.accelerator.prepare(sample_dataloader)
            pipeline.to(self.accelerator.device)
            print(f'Accelerator device: {self.accelerator.device}')

            print('Now generating images')
            for example in tqdm(
                sample_dataloader, desc="Generating class images", disable=not self.accelerator.is_local_main_process
            ):
                images = pipeline(example["prompt"]).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)
                    print(f'Generated image at {image_filename}')

            del pipeline
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

        else:
            print('Not generating images because there are already enough images in the class images directory')

            '''# Update these attributes now that we have class images in the directory
            print(f'Before, self.class_images_path: {self.class_images_path}')
            self.class_images_path = list(self.class_data_dir.iterdir())
            print(f'After, self.class_images_path: {self.class_images_path}')

            print(f'Before, self.num_class_images: {self.num_class_images}')
            self.num_class_images = len(self.class_images_path)
            print(f'After, self.num_class_images: {self.num_class_images}')

            print(f'Before, self._length: {self._length}')
            self._length = max(self.num_class_images, self.num_instance_images)
            print(f'After, self._length: {self._length}')'''

    # Train!
    def train(self):
        total_batch_size = self.config.train_batch_size * self.accelerator.num_processes * self.config.gradient_accumulation_steps

        print("***** Running training *****")
        print(f"  Num examples = {len(self.train_dataset)}")
        print(f"  Num batches each epoch = {len(self.train_dataloader)}")
        print(f"  Num Epochs = {self.config.num_train_epochs}")
        print(f"  Instantaneous batch size per device = {self.config.train_batch_size}")
        print(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
        print(f"  Gradient Accumulation steps = {self.config.gradient_accumulation_steps}")
        print(f"  Total optimization steps = {self.config.max_train_steps}")
        # Only show the progress bar once on each machine.
        progress_bar = tqdm(range(self.config.max_train_steps), disable=not self.accelerator.is_local_main_process)
        progress_bar.set_description("Steps")
        global_step = 0

        for epoch in range(self.config.num_train_epochs):
            #print(f'Doing global step: {global_step} at time {datetime.datetime.now()}')
            self.unet.train()
            if self.config.train_text_encoder:
                #print(f'Training text_encoder for global step: {global_step}')
                self.text_encoder.train()
            for step, batch in enumerate(self.train_dataloader):
                #print(f'Step: {step}')
                with self.accelerator.accumulate(self.unet):
                    # Convert images to latent space
                    #print('converting images to latent space')
                    latents = self.vae.encode(batch["pixel_values"].to(dtype=self.weight_dtype)).latent_dist.sample()
                    latents = latents * 0.18215

                    # Sample noise that we'll add to the latents
                    #print('sampling noise to add to latents')
                    noise = torch.randn_like(latents)
                    bsz = latents.shape[0]
                    # Sample a random timestep for each image
                    timesteps = torch.randint(0, self.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)
                    #print('Adding noise to the latents')
                    noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)

                    # Get the text embedding for conditioning
                    #print('Getting the text embedding for conditioning')
                    encoder_hidden_states = self.text_encoder(batch["input_ids"])[0]

                    # Predict the noise residual
                    #print('predicting the noise residual')
                    model_pred = self.unet(noisy_latents, timesteps, encoder_hidden_states).sample

                    # Get the target for loss depending on the prediction type
                    #print('Getting the target for loss depending on the prediction type')
                    #if self.noise_scheduler.predict_epsilon:
                    if self.noise_scheduler.config.prediction_type == "epsilon":
                        target = noise
                    elif self.noise_scheduler.config.prediction_type == "v_prediction":
                    #else:
                        target = self.noise_scheduler.get_velocity(latents, noise, timesteps)
                    else:
                        raise ValueError(f"Unknown prediction type {self.noise_scheduler.config.prediction_type}")

                    if self.config.do_prior_preservation:
                        # 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)

                        # Compute instance loss
                        loss = F.mse_loss(model_pred.float(), target.float(), reduction="none").mean([1, 2, 3]).mean()

                        # Compute prior loss
                        prior_loss = F.mse_loss(model_pred_prior.float(), target_prior.float(), reduction="mean")

                        # Add the prior loss to the instance loss.
                        loss = loss + self.config.prior_loss_weight * prior_loss
                    else:
                        loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")

                    #print('Doing backward')
                    self.accelerator.backward(loss)
                    #print('Syncing gradients')
                    if self.accelerator.sync_gradients:
                        params_to_clip = (
                            itertools.chain(self.unet.parameters(), self.text_encoder.parameters())
                            if self.config.train_text_encoder
                            else self.unet.parameters()
                        )
                        self.accelerator.clip_grad_norm_(params_to_clip, self.config.max_grad_norm)
                    #print('Doing optimizer step')
                    self.optimizer.step()
                    #print('Doing lr_scheduler step')
                    self.lr_scheduler.step()
                    #print('Optimizer zero_grad')
                    self.optimizer.zero_grad()

                # Checks if the accelerator has performed an optimization step behind the scenes
                if self.accelerator.sync_gradients:
                    #print('Updating progress bar')
                    progress_bar.update(1)
                    global_step += 1

                    if global_step % self.config.save_steps == 0:
                        if self.accelerator.is_main_process:
                            pipeline = StableDiffusionPipeline.from_pretrained(
                                self.config.model_id,
                                unet=self.accelerator.unwrap_model(self.unet),
                                text_encoder=self.accelerator.unwrap_model(self.text_encoder),
                                revision=self.config.revision,
                                use_auth_token=self.auth_token
                            )
                            save_path = os.path.join(self.config.model_out_dir, f"checkpoint-{global_step}")
                            print(f'Saving checkpoint at {save_path}')
                            pipeline.save_pretrained(save_path)

                logs = {"loss": loss.detach().item(), "lr": self.lr_scheduler.get_last_lr()[0]}
                progress_bar.set_postfix(**logs)
                self.accelerator.log(logs, step=global_step)

                if global_step >= self.config.max_train_steps:
                    break

            self.accelerator.wait_for_everyone()

        # Create the pipeline using using the trained modules and save it.
        if self.accelerator.is_main_process:
            pipeline = StableDiffusionPipeline.from_pretrained(
                self.config.model_id,
                unet=self.accelerator.unwrap_model(self.unet),
                text_encoder=self.accelerator.unwrap_model(self.text_encoder),
                revision=self.config.revision,
                use_auth_token=self.auth_token
            )
            print(f'Saving model to {self.config.model_out_dir}')
            pipeline.save_pretrained(self.config.model_out_dir)

            if self.config.push_to_hub:
                print('Pushing to hub...')
                self.repo.push_to_hub(commit_message="End of training", blocking=False, auto_lfs_prune=True)
                print('Pushed to hub')

        self.accelerator.end_training()
        print('Training done')

