train_code/train_svd.py

#!/usr/bin/env python
'''
    This file is to train stable video diffusion by my personal implementation which is based on diffusers' training example code.
'''

import argparse
import logging
import math
import os, sys
import time
import random
import shutil
import warnings
import cv2
from PIL import Image 
from einops import rearrange, repeat
from pathlib import Path
from omegaconf import OmegaConf
import imageio

import accelerate
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import RandomSampler
import transformers
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import ProjectConfiguration, set_seed
from packaging import version
from PIL import Image
from torchvision import transforms
from tqdm.auto import tqdm
from transformers import AutoTokenizer, PretrainedConfig

import diffusers
from diffusers import (
    AutoencoderKLTemporalDecoder,
    DDPMScheduler,
)
from diffusers.training_utils import EMAModel, compute_snr
from diffusers.optimization import get_scheduler
from diffusers.utils import check_min_version, is_wandb_available, load_image, export_to_video
from diffusers.utils.import_utils import is_xformers_available
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils.torch_utils import randn_tensor
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
if is_wandb_available():
    import wandb


# Import files from the local folder
root_path = os.path.abspath('.')
sys.path.append(root_path)
from svd.pipeline_stable_video_diffusion import StableVideoDiffusionPipeline 
from svd.unet_spatio_temporal_condition import UNetSpatioTemporalConditionModel
from data_loader.video_dataset import Video_Dataset, get_video_frames, tokenize_captions
from utils.img_utils import resize_with_antialiasing


# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
# check_min_version("0.25.0.dev0")
logger = get_logger(__name__)
warnings.filterwarnings('ignore')


###################################################################################################################################################
def parse_args(input_args=None):
    parser = argparse.ArgumentParser(description="Simple example of a ControlNet training script.")
    parser.add_argument(
        "--config_path",
        type=str,
        default="config/train_image2video.yaml",
        help="Path to pretrained model or model identifier from huggingface.co/models.",
    )

    args = parser.parse_args()
    return args


    
def log_validation(vae, unet, image_encoder, text_encoder, tokenizer, config, accelerator, weight_dtype, step, 
                        parent_store_folder = None, force_close_flip = False, use_ambiguous_prompt=False):
    # This function will also be used in other files
    print("Running validation... ")


    # Init
    validation_source_folder = config["validation_img_folder"] 
    

    # Init the pipeline
    pipeline = StableVideoDiffusionPipeline.from_pretrained(
        config["pretrained_model_name_or_path"],
        vae = accelerator.unwrap_model(vae),
        image_encoder = accelerator.unwrap_model(image_encoder),
        unet = accelerator.unwrap_model(unet),
        revision = None,    # Set None directly now
        torch_dtype = weight_dtype,
    )
    pipeline = pipeline.to(accelerator.device)
    pipeline.set_progress_bar_config(disable=True)


    # Process all image in the folder
    frames_collection = []
    for image_name in sorted(os.listdir(validation_source_folder)):
        if accelerator.is_main_process:
            if parent_store_folder is None:
                validation_store_folder = os.path.join(config["validation_store_folder"] + "_" + config["scheduler"], "step_" + str(step), image_name)
            else:
                validation_store_folder = os.path.join(parent_store_folder, image_name)
                
            if os.path.exists(validation_store_folder):
                shutil.rmtree(validation_store_folder)
            os.makedirs(validation_store_folder)

        image_path = os.path.join(validation_source_folder, image_name, 'im_0.jpg')
        ref_image = load_image(image_path)
        ref_image = ref_image.resize((config["width"], config["height"]))
        
        
        # Decide the motion score in SVD (mostly what we use is fix value now)
        if config["motion_bucket_id"] is None:      
            raise NotImplementedError("We need a fixed motion_bucket_id in the config")
        else:
            reflected_motion_bucket_id = config["motion_bucket_id"]
        print("Inference Motion Bucket ID is ", reflected_motion_bucket_id)


        # Prepare text prompt
        if config["use_text"]:
            # Read the file
            file_path = os.path.join(validation_source_folder, image_name, "lang.txt")
            file = open(file_path, 'r')
            prompt = file.readlines()[0]  # Only read the first line
            if use_ambiguous_prompt:
                prompt = prompt.split(" ")[0] + " this to there"
                print("We are creating ambiguous prompt, which is: ", prompt)
        else:
            prompt = ""
        # Use the same tokenize process as the dataset preparation stage
        tokenized_prompt = tokenize_captions(prompt, tokenizer, config, is_train=False).unsqueeze(0).to(accelerator.device)    # Use unsqueeze to expand dim

        # Store the prompt for the sanity check
        f = open(os.path.join(validation_store_folder, "lang_cond.txt"), "a")
        f.write(prompt)
        f.close()


        # Flip the image by chance (it is needed to check whether there is any object position words [left|right] in the prompt text)
        flip = False
        if not force_close_flip:        # force_close_flip is True in testing time; else, we cannot match in the same standard
            if random.random() < config["flip_aug_prob"]:
                if config["use_text"]:
                    if prompt.find("left") == -1 and prompt.find("right") == -1:    # Cannot have position word, like left and right (up and down is ok)
                        flip = True
                else:
                    flip = True
            if flip:
                print("Use flip in validation!")
                ref_image = ref_image.transpose(Image.FLIP_LEFT_RIGHT)


        # Call the model for inference
        with torch.autocast("cuda"):
            frames = pipeline(
                                ref_image, 
                                tokenized_prompt,
                                config["use_text"],
                                text_encoder,
                                height = config["height"],
                                width = config["width"],
                                num_frames = config["video_seq_length"], 
                                num_inference_steps = config["num_inference_steps"],
                                decode_chunk_size = 8, 
                                motion_bucket_id = reflected_motion_bucket_id,
                                fps = 7,
                                noise_aug_strength = config["inference_noise_aug_strength"],
                              ).frames[0]     

        # Store the frames
        # breakpoint()
        for idx, frame in enumerate(frames):
            frame.save(os.path.join(validation_store_folder, str(idx)+".png"))
        imageio.mimsave(os.path.join(validation_store_folder, 'combined.gif'), frames)      # gif storage quality is not high, recommend to check png images

        frames_collection.append(frames)


    # Cleaning process
    del pipeline
    torch.cuda.empty_cache()

    return frames_collection   # Return resuly based on the need


def tensor_to_vae_latent(inputs, vae):
    video_length = inputs.shape[1]
    inputs = rearrange(inputs, "b f c h w -> (b f) c h w")
    latents = vae.encode(inputs).latent_dist.mode()
    latents = rearrange(latents, "(b f) c h w -> b f c h w", f=video_length)    # Use f or b to rearrage should have the same effect
    latents = latents * vae.config.scaling_factor

    return latents


def import_pretrained_text_encoder(pretrained_model_name_or_path: str, revision: str):
    ''' Import Text encoder information
    
    '''
    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
    
    else:   # No other cases will be considerred
        raise ValueError(f"{model_class} is not supported.")
    


def rand_log_normal(shape, loc=0., scale=1., device='cpu', dtype=torch.float32):
    """Draws samples from an lognormal distribution."""
    u = torch.rand(shape, dtype=dtype, device=device) * (1 - 2e-7) + 1e-7
    return torch.distributions.Normal(loc, scale).icdf(u).exp()


def get_add_time_ids(
        unet_config,
        expected_add_embed_dim,
        fps,
        motion_bucket_id,
        noise_aug_strength,
        dtype,
        batch_size,
        num_videos_per_prompt,
    ):

    # Construct Basic add_time_ids items
    add_time_ids = [fps, motion_bucket_id, noise_aug_strength]

    passed_add_embed_dim = unet_config.addition_time_embed_dim * len(add_time_ids)
    if expected_add_embed_dim != passed_add_embed_dim:
        raise ValueError(
            f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
        )

    add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
    add_time_ids = add_time_ids.repeat(batch_size * num_videos_per_prompt, 1)

    return add_time_ids


####################################################################################################################################################################



def main(config):

    # Read Config Setting
    resume_from_checkpoint = config["resume_from_checkpoint"]
    output_dir = config["output_dir"]
    logging_name = config["logging_name"]
    mixed_precision = config["mixed_precision"]
    report_to = config["report_to"]
    pretrained_model_name_or_path = config["pretrained_model_name_or_path"]
    pretrained_tokenizer_name_or_path = config["pretrained_tokenizer_name_or_path"]
    gradient_checkpointing = config["gradient_checkpointing"]
    learning_rate = config["learning_rate"]
    adam_beta1 = config["adam_beta1"]
    adam_beta2 = config["adam_beta2"]
    adam_weight_decay = config["adam_weight_decay"]
    adam_epsilon = config["adam_epsilon"]
    train_batch_size = config["train_batch_size"]
    dataloader_num_workers = config["dataloader_num_workers"]
    gradient_accumulation_steps = config["gradient_accumulation_steps"]
    num_train_iters = config["num_train_iters"]
    lr_warmup_steps = config["lr_warmup_steps"]
    checkpointing_steps = config["checkpointing_steps"]
    process_fps = config["process_fps"]
    train_noise_aug_strength = config["train_noise_aug_strength"]
    use_8bit_adam = config["use_8bit_adam"]
    scale_lr = config["scale_lr"]
    conditioning_dropout_prob = config["conditioning_dropout_prob"]
    checkpoints_total_limit = config["checkpoints_total_limit"]
    validation_step = config["validation_step"]
    partial_finetune = config['partial_finetune']


    # Default Setting
    revision = None
    variant = "fp16"
    lr_scheduler = "constant"
    max_grad_norm = 1.0
    tracker_project_name = "img2video"
    num_videos_per_prompt = 1
    seed = 42
    # No CFG in training now



    # Define the accelerator
    logging_dir = Path(output_dir, logging_name)
    accelerator_project_config = ProjectConfiguration(project_dir=output_dir, logging_dir=logging_dir)
    accelerator = Accelerator(
        gradient_accumulation_steps = gradient_accumulation_steps,
        mixed_precision = mixed_precision,
        log_with = report_to,
        project_config = accelerator_project_config,
    )
    generator = torch.Generator(device=accelerator.device).manual_seed(seed)


    # 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()


    # Handle the repository creation
    if accelerator.is_main_process and resume_from_checkpoint != "latest":      # For the latest checkpoint version, we don't need to delete our folders
        # Validation file
        validation_store_folder = config["validation_store_folder"] + "_" + config["scheduler"]
        print("We will remove ", validation_store_folder)
        if os.path.exists(validation_store_folder):
            archive_name = validation_store_folder + "_archive"
            if os.path.exists(archive_name):
                shutil.rmtree(archive_name)
            print("We will move to archive ", archive_name)
            os.rename(validation_store_folder, archive_name)
        os.makedirs(validation_store_folder)

        # Output Dir
        if os.path.exists(output_dir):
            shutil.rmtree(output_dir)
            # os.makedirs(output_dir, exist_ok=True)

        # Log
        if os.path.exists("runs"):
            shutil.rmtree("runs")

        # Copy the config to here
        os.system(" cp config/train_image2video.yaml " + validation_store_folder + "/")
            

    # Load All Module Needed
    feature_extractor = CLIPImageProcessor.from_pretrained(
        pretrained_model_name_or_path, subfolder="feature_extractor", revision=revision
    )   # This instance has now weight, they are just seeting file
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        pretrained_model_name_or_path, subfolder="image_encoder", revision=revision, variant=variant
    )
    vae = AutoencoderKLTemporalDecoder.from_pretrained(
        pretrained_model_name_or_path, subfolder="vae", revision=revision, variant=variant
    )
    if config["load_unet_path"] != None:
        print("We will load UNet from ", config["load_unet_path"])
        unet = UNetSpatioTemporalConditionModel.from_pretrained(
            config["load_unet_path"], 
            subfolder = "unet", 
            low_cpu_mem_usage = True,       
        )   # For the variant, we don't have fp16 version, so we will read from fp32
    else:
        print("We will only use SVD pretrained UNet")
        unet = UNetSpatioTemporalConditionModel.from_pretrained(
            pretrained_model_name_or_path, 
            subfolder = "unet", 
            low_cpu_mem_usage = True,
            variant = variant,
        )
    
    # Prepare for the tokenizer if use text
    tokenizer = AutoTokenizer.from_pretrained(
        pretrained_tokenizer_name_or_path,
        subfolder = "tokenizer",
        revision = revision,
        use_fast = False,
    )
    
    if config["use_text"]:
        # Clip Text Encoder
        text_encoder_cls = import_pretrained_text_encoder(pretrained_tokenizer_name_or_path, revision)
        text_encoder = text_encoder_cls.from_pretrained(
            pretrained_tokenizer_name_or_path, subfolder = "text_encoder", revision = revision, variant = variant
        )
    else:
        text_encoder = None


    # Store the config due to the disappearance after accelerator prepare (This is written to handle some unknown phenomenon)
    unet_config = unet.config
    expected_add_embed_dim = unet.add_embedding.linear_1.in_features
    

    # Freeze vae + feature_extractor + image_encoder, but set unet to trainable
    vae.requires_grad_(False)       
    image_encoder.requires_grad_(False)
    unet.requires_grad_(False)  # Will switch back to train mode later on
    if config["use_text"]:
        text_encoder.requires_grad_(False)      # All set with no grad needed (like VAE) follow other T2I papers



    # For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora unet) to half-precision
    # as these weights 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 + image_encoder to gpu and cast to weight_dtype
    vae.to(accelerator.device, dtype=weight_dtype)
    image_encoder.to(accelerator.device, dtype=weight_dtype)
    # unet.to(accelerator.device, dtype=weight_dtype)
    if config["use_text"]:
        text_encoder.to(accelerator.device, dtype=weight_dtype)



    # Acceleration: `accelerate` 0.16.0 will have better support for customized saving
    if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
        # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
        def save_model_hook(models, weights, output_dir):
            if accelerator.is_main_process:

                for i, model in enumerate(models):
                    model.save_pretrained(os.path.join(output_dir, "unet"))

                    # make sure to pop weight so that corresponding model is not saved again
                    weights.pop()

        def load_model_hook(models, input_dir):

            for i in range(len(models)):
                # pop models so that they are not loaded again
                model = models.pop()

                # load diffusers style into model
                load_model = UNetSpatioTemporalConditionModel.from_pretrained(input_dir, subfolder="unet")
                model.register_to_config(**load_model.config)

                model.load_state_dict(load_model.state_dict())
                del load_model

        accelerator.register_save_state_pre_hook(save_model_hook)
        accelerator.register_load_state_pre_hook(load_model_hook)
        
    if gradient_checkpointing:
        unet.enable_gradient_checkpointing()



    ################################ Make Training dataset  ###############################
    train_dataset = Video_Dataset(config, device = accelerator.device, tokenizer=tokenizer)
    sampler = RandomSampler(train_dataset)
    train_dataloader = torch.utils.data.DataLoader(
        train_dataset,
        sampler = sampler,
        batch_size = train_batch_size,
        num_workers = dataloader_num_workers * accelerator.num_processes,
    )       
    #######################################################################################


    ####################################### Optimizer  Setting #####################################################################
    if scale_lr:
        learning_rate = (
            learning_rate * gradient_accumulation_steps * train_batch_size * accelerator.num_processes
        )

    # 8bit adam to save more memory (Usally we need this to save the memory)
    if use_8bit_adam:       
        try:
            import bitsandbytes as bnb
        except ImportError:
            raise ImportError(
                "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
            )

        optimizer_cls = bnb.optim.AdamW8bit
    else:
        optimizer_cls = torch.optim.AdamW


    # Switch back to unet in training mode
    unet.requires_grad_(True) 


    ############################## For partial fine-tune setting ##############################
    parameters_list = []
    for name, param in unet.named_parameters():
        if partial_finetune:    # The partial finetune we use is to only train attn layers, which will be ~190M params (TODO:needs to check later for exact value)
            # Full Spatial: .transformer_blocks. && spatial_
            # Attn + All emb: attn && emb
            if name.find("attn") != -1 or name.find("emb") != -1:     # Only block the spatial Transformer
                parameters_list.append(param)
                param.requires_grad = True  
            else:
                param.requires_grad = False 
        else:
            parameters_list.append(param)
            param.requires_grad = True

    # Double check what will be trained
    total_params_for_training = 0
    # if os.path.exists("param_lists.txt"):
    #     os.remove("param_lists.txt")
    # file1 = open("param_lists.txt","a")
    for name, param in unet.named_parameters():
        # file1.write(name + "\n")
        if param.requires_grad:
            total_params_for_training += param.numel()
            print(name + " requires grad update")     
    print("Total parameter that will be trained has ", total_params_for_training)
    ##########################################################################################

    # Optimizer creation
    optimizer = optimizer_cls(
        parameters_list,
        lr = learning_rate,
        betas = (adam_beta1, adam_beta2),
        weight_decay = adam_weight_decay,
        eps = adam_epsilon,
    )


    # Scheduler and Training steps
    dataset_length = len(train_dataset)
    print("Dataset length read from the train side is ", dataset_length)
    num_update_steps_per_epoch = math.ceil(dataset_length / gradient_accumulation_steps)
    max_train_steps = num_train_iters * train_batch_size

    # Learning Rate Scheduler  (we all use constant)
    lr_scheduler = get_scheduler(
        "constant",
        optimizer = optimizer,
        num_warmup_steps = lr_warmup_steps * accelerator.num_processes,
        num_training_steps = max_train_steps *  accelerator.num_processes,
        num_cycles = 1,
        power = 1.0,
    )
    #####################################################################################################################################
    


    # Prepare everything with our `accelerator`.
    unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        unet, optimizer, train_dataloader, lr_scheduler
    )   


    # We need to RECALCULATE our total training steps as the size of the training dataloader may have changed.
    print("accelerator.num_processes is ", accelerator.num_processes)
    print("num_train_iters is ", num_train_iters)
    num_train_epochs = math.ceil(num_train_iters * accelerator.num_processes * gradient_accumulation_steps / dataset_length) 
    print("num_train_epochs is ", num_train_epochs)

    # We need to initialize the trackers we use, and also store our configuration.
    if accelerator.is_main_process: # Only on the main process!
        tracker_config = dict(vars(args))
        accelerator.init_trackers(tracker_project_name, tracker_config)



    # Train!
    logger.info("***** Running training *****")
    logger.info(f"  Dataset Length = {dataset_length}")
    logger.info(f"  Num Epochs = {num_train_epochs}")
    logger.info(f"  Instantaneous batch size per device = {train_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {max_train_steps}")



    # Load the Closest / Best weight   
    global_step = 0     # Catch the current iteration
    first_epoch = 0
    if resume_from_checkpoint:
        if resume_from_checkpoint != "latest":
            path = os.path.basename(resume_from_checkpoint)
        else:
            # Get the most recent checkpoint
            dirs = os.listdir(output_dir)
            dirs = [d for d in dirs if d.startswith("checkpoint")]
            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
            path = dirs[-1] if len(dirs) > 0 else None
            print("We will resume the latest weight ", path)

        if path is None:     # Don't resume
            accelerator.print(
                f"Checkpoint '{resume_from_checkpoint}' does not exist. Starting a new training run."
            )
            resume_from_checkpoint = None
            initial_global_step = 0
        else:       # Resume from the closest checkpoint
            accelerator.print(f"Resuming from checkpoint {path}")
            accelerator.load_state(os.path.join(output_dir, path))
            global_step = int(path.split("-")[1])

            initial_global_step = global_step
            first_epoch = global_step // num_update_steps_per_epoch
    else:
        initial_global_step = 0


    if accelerator.is_main_process:
        print("Initial Learning rate is ", optimizer.param_groups[0]['lr'])
        print("global_step will start from ", global_step)

    progress_bar = tqdm(
                            range(initial_global_step, 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,
                        )

    

    # Prepare tensorboard log
    writer = SummaryWriter() 


    ######################################################### Auxiliary Function #################################################################
    def encode_clip(pixel_values, prompt):
        ''' Encoder hidden states input source
            pixel_values:   first frame pixel information
            prompt:         language prompt with takenized
        '''

        ########################################## Prepare the Text Embedding #####################################################
        # pixel_values is in the range [-1, 1]
        pixel_values = resize_with_antialiasing(pixel_values, (224, 224))  
        pixel_values = (pixel_values + 1.0) / 2.0   # [-1, 1] -> [0, 1]

        # Normalize the image with for CLIP input
        pixel_values = feature_extractor(
            images=pixel_values,
            do_normalize=True,
            do_center_crop=False,
            do_resize=False,
            do_rescale=False,
            return_tensors="pt",
        ).pixel_values

        # The following is the same as _encode_image in SVD pipeline
        pixel_values = pixel_values.to(device=accelerator.device, dtype=weight_dtype)
        image_embeddings = image_encoder(pixel_values).image_embeds
        image_embeddings = image_embeddings.unsqueeze(1)

        # duplicate image embeddings for each generation per prompt, using mps friendly method
        bs_embed, seq_len, _ = image_embeddings.shape
        image_embeddings = image_embeddings.repeat(1, num_videos_per_prompt, 1)
        encoder_hidden_states = image_embeddings.view(bs_embed * num_videos_per_prompt, seq_len, -1)
        #############################################################################################################################



        ########################################## Prepare the Text embedding if needed #############################################
        if config["use_text"]:
            text_embeddings = text_encoder(prompt)[0]
            
            # Concat two embeddings together on dim 1
            encoder_hidden_states = torch.cat((text_embeddings, encoder_hidden_states), dim=1)
        
            # Layer norm on the last dim
            layer_norm = nn.LayerNorm((78, 1024)).to(device=accelerator.device, dtype=weight_dtype)
            encoder_hidden_states_norm = layer_norm(encoder_hidden_states)

            # Return
            return encoder_hidden_states_norm

        else:   # Just return back default on
            return encoder_hidden_states
        #############################################################################################################################

    ####################################################################################################################################################


    ############################################################################################################################
    # For the training, we mimic the code from T2I in diffusers
    for epoch in range(first_epoch, num_train_epochs):
        unet.train()
        train_loss = 0.0
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(unet):
                # batch is a torch tensor with range of [-1, 1] but no other pre-porcessing
                video_frames = batch["video_frames"].to(weight_dtype).to(accelerator.device, non_blocking=True)
                reflected_motion_bucket_id = batch["reflected_motion_bucket_id"]
                prompt = batch["prompt"]
                

                # Images to VAE latent space
                latents = tensor_to_vae_latent(video_frames, vae)


                ##################################### Add Noise ########################################
                bsz, num_frames = latents.shape[:2]
            
                # Encode the first frame
                conditional_pixel_values = video_frames[:, 0, :, :, :]      # First frame
                # Following AnimateSomething, we use constant to repace cond_sigmas
                conditional_pixel_values = conditional_pixel_values + torch.randn_like(conditional_pixel_values) * train_noise_aug_strength        
                conditional_latents = vae.encode(conditional_pixel_values).latent_dist.mode()           # mode() returns mean value no std influence
                conditional_latents = repeat(conditional_latents, 'b c h w->b f c h w', f=num_frames)   # copied across the frame axis to be the same shape as noise


                # Add noise to the latents according to the noise magnitude at each timestep
                # This is the forward diffusion process
                sigmas = rand_log_normal(shape=[bsz,], loc=config["noise_mean"], scale=config["noise_std"]).to(latents.device)
                sigmas = sigmas[:, None, None, None, None]
                noisy_latents = latents + torch.randn_like(latents) * sigmas
                inp_noisy_latents = noisy_latents / ((sigmas**2 + 1) ** 0.5)


                # For the encoder hidden states based on the first frame and prompt
                encoder_hidden_states = encode_clip(video_frames[:, 0, :, :, :].float(), prompt)     # First Frame + Text Prompt


                # Conditioning dropout to support classifier-free guidance during inference. For more details
                # check out the section 3.2.1 of the original paper https://arxiv.org/abs/2211.09800 (InstructPix2Pix).
                if conditioning_dropout_prob != 0:
                    random_p = torch.rand(bsz, device=latents.device, generator=generator)
                    
                    # Sample masks for the edit prompts.
                    prompt_mask = random_p < 2 * conditioning_dropout_prob
                    prompt_mask = prompt_mask.reshape(bsz, 1, 1)
                    # Final text conditioning.
                    null_conditioning = torch.zeros_like(encoder_hidden_states)
                    encoder_hidden_states = torch.where(prompt_mask, null_conditioning, encoder_hidden_states)

                    # Sample masks for the original images.
                    image_mask_dtype = conditional_latents.dtype
                    image_mask = 1 - ((random_p >= conditioning_dropout_prob).to(image_mask_dtype) * (random_p < 3 * conditioning_dropout_prob).to(image_mask_dtype))
                    image_mask = image_mask.reshape(bsz, 1, 1, 1)

                    # Final image conditioning.
                    conditional_latents = image_mask * conditional_latents


                # Concatenate the `conditional_latents` with the `noisy_latents`.
                inp_noisy_latents = torch.cat([inp_noisy_latents, conditional_latents], dim=2)


                # GT noise
                target = latents
                ##########################################################################################


                ################################ Other Embedding and Conditioning ###################################
                reflected_motion_bucket_id = torch.sum(reflected_motion_bucket_id)/len(reflected_motion_bucket_id)
                reflected_motion_bucket_id = int(reflected_motion_bucket_id.cpu().detach().numpy())
                # print("Training reflected_motion_bucket_id is ", reflected_motion_bucket_id)
                
                added_time_ids = get_add_time_ids(
                                    unet_config,
                                    expected_add_embed_dim,
                                    process_fps,
                                    reflected_motion_bucket_id,
                                    train_noise_aug_strength,       
                                    weight_dtype,
                                    train_batch_size,
                                    num_videos_per_prompt,
                                )       # The same as SVD pipeline's _get_add_time_ids
                added_time_ids = added_time_ids.to(accelerator.device)

                timesteps = torch.Tensor([0.25 * sigma.log() for sigma in sigmas]).to(accelerator.device)
                #####################################################################################################



                ###################################### Predict Noise ######################################
                model_pred = unet(
                                    inp_noisy_latents,
                                    timesteps, 
                                    encoder_hidden_states, 
                                    added_time_ids = added_time_ids
                                  ).sample      

                # Denoise the latents
                c_out = -sigmas / ((sigmas**2 + 1)**0.5)
                c_skip = 1 / (sigmas**2 + 1)
                denoised_latents = model_pred * c_out + c_skip * noisy_latents
                weighing = (1 + sigmas ** 2) * (sigmas**-2.0)
                ##########################################################################################


                ############################### Calculate Loss and Update Optimizer #######################
                # MSE loss
                loss = torch.mean(
                    (  weighing.float() * (denoised_latents.float() - target.float())**2 ).reshape(target.shape[0], -1),
                    dim=1,
                )
                loss = loss.mean()

                # Gather the losses across all processes for logging (if we use distributed training).
                avg_loss = accelerator.gather(loss.repeat(train_batch_size)).mean()
                train_loss += avg_loss.item() / gradient_accumulation_steps

                # Update Tensorboard
                writer.add_scalar('Loss/train-Loss-Step', avg_loss, global_step)
                

                # Backpropagate
                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)
                accelerator.log({"train_loss": train_loss}, step=global_step)
                train_loss = 0.0

                ########################################## Checkpoints #########################################
                if global_step != 0 and global_step % checkpointing_steps == 0:
                    if accelerator.is_main_process:
                        start = time.time()
                        # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
                        if checkpoints_total_limit is not None:
                            checkpoints = os.listdir(output_dir)
                            checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
                            checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))

                            # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
                            if len(checkpoints) >= checkpoints_total_limit:
                                num_to_remove = len(checkpoints) - checkpoints_total_limit + 1
                                removing_checkpoints = checkpoints[0:num_to_remove]

                                logger.info(
                                    f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
                                )
                                logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")

                                for removing_checkpoint in removing_checkpoints:
                                    removing_checkpoint = os.path.join(output_dir, removing_checkpoint)
                                    shutil.rmtree(removing_checkpoint)

                        save_path = os.path.join(output_dir, f"checkpoint-{global_step}")
                        accelerator.save_state(save_path)
                        logger.info(f"Saved state to {save_path}")
                        print("Save time use " + str(time.time() - start) + " s")
                ########################################################################################################


            # Update Log
            logs = {"step_loss": loss.detach().item(), "lr": optimizer.param_groups[0]['lr']}
            progress_bar.set_postfix(**logs)


           ##################################### Validation per XXX iterations #######################################
            if accelerator.is_main_process:
                if global_step % validation_step == 0:         # Fixed 100 steps to validate
                    
                    if config["validation_img_folder"] is not None:
                        log_validation(
                                        vae,
                                        unet,
                                        image_encoder,
                                        text_encoder,
                                        tokenizer,
                                        config,
                                        accelerator,
                                        weight_dtype,
                                        global_step,
                                        use_ambiguous_prompt = config["mix_ambiguous"],
                                    )

            ###############################################################################################################

            # Update Steps and Break if needed
            global_step += 1

            if global_step >= max_train_steps:
                break
    
    ############################################################################################################################


if __name__ == "__main__":
    args = parse_args()

    config = OmegaConf.load(args.config_path)
    main(config)