train_code/train_csvd.py

#!/usr/bin/env python
'''
    This file is to train Stable Video Diffusion with Conditioning design by my peronal 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
from PIL import Image 
from einops import rearrange, repeat
from pathlib import Path
from omegaconf import OmegaConf
import imageio
import cv2


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 huggingface_hub import create_repo, upload_folder
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,
    UniPCMultistepScheduler,
)
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_controlnet import StableVideoDiffusionControlNetPipeline
from svd.unet_spatio_temporal_condition import UNetSpatioTemporalConditionModel
from svd.temporal_controlnet import ControlNetModel
from utils.img_utils import resize_with_antialiasing
from utils.optical_flow_utils import flow_to_image, filter_uv, bivariate_Gaussian
from data_loader.video_dataset import tokenize_captions
from data_loader.video_this_that_dataset import Video_ThisThat_Dataset, get_thisthat_sam
from train_code.train_svd import import_pretrained_text_encoder

# 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_controlnet.yaml",
        help="Path to pretrained model or model identifier from huggingface.co/models.",
    )

    args = parser.parse_args()
    return args


    
def log_validation(vae, unet, controlnet, 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 = StableVideoDiffusionControlNetPipeline.from_pretrained(
        config["pretrained_model_name_or_path"],        # Still based on regular SVD config
        vae = vae,
        image_encoder = image_encoder,
        unet = 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)      # [0, 255] Range
        ref_image = ref_image.resize((config["width"], config["height"]))


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


        if config["data_loader_type"] == "thisthat":
            condition_img, reflected_motion_bucket_id, controlnet_image_index, coordinate_values = get_thisthat_sam(config, 
                                                                                                                    os.path.join(validation_source_folder, image_name),
                                                                                                                    flip = flip, 
                                                                                                                    store_dir = validation_store_folder,
                                                                                                                    verbose = True)
        else:
            raise NotImplementedError("We don't support such data loader type")



        # Call the pipeline
        with torch.autocast("cuda"):
            frames = pipeline(
                                image = ref_image, 
                                condition_img = condition_img,       # numpy [0,1] range
                                controlnet = accelerator.unwrap_model(controlnet),
                                prompt = tokenized_prompt,
                                use_text = config["use_text"],
                                text_encoder = text_encoder,
                                height = config["height"],
                                width = config["width"],
                                num_frames = config["video_seq_length"], 
                                decode_chunk_size = 8, 
                                motion_bucket_id = reflected_motion_bucket_id,
                                controlnet_image_index = controlnet_image_index,
                                coordinate_values = coordinate_values,
                                num_inference_steps = config["num_inference_steps"],
                                max_guidance_scale = config["inference_max_guidance_scale"],
                                fps = 7,
                                use_instructpix2pix = config["use_instructpix2pix"],
                                noise_aug_strength = config["inference_noise_aug_strength"],
                                controlnet_conditioning_scale = config["outer_conditioning_scale"],
                                inner_conditioning_scale = config["inner_conditioning_scale"],
                                guess_mode = config["inference_guess_mode"],        # False in inference
                                image_guidance_scale = config["image_guidance_scale"],
                              ).frames[0]    

        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, duration=0.05)

        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 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 # In the range [0, 1]
    # TODO: "* (1 - 2e-7) + 1e-7" is not included in previous code, I add it back, don't why whether there is any influence now
    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,
        do_classifier_free_guidance = False,
    ):

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


    # Sanity Check
    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']
    load_unet_path = config['load_unet_path']

    if mixed_precision == 'None':   # For mixed precision use
        mixed_precision = 'no'


    # Default Setting
    revision = None
    variant = "fp16"        # TODO: 这里进行了调整，不知道会有多少区别，现在跟unet training保持一致
    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_controlnet.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 load_unet_path != None:
        print("We will use pretrained UNet path by our, at ", load_unet_path)
        unet = UNetSpatioTemporalConditionModel.from_pretrained(
            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 still use provided UNet path")
        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

    # Init for the Controlnet (check if has pretrained path to load)
    if config["load_controlnet_path"] != None:
        print("We will load pre-trained controlnet from ", config["load_controlnet_path"])
        controlnet = ControlNetModel.from_pretrained(config["load_controlnet_path"], subfolder="controlnet")
    else:
        controlnet = ControlNetModel.from_unet(unet, load_weights_from_unet=True, conditioning_channels=config["conditioning_channels"])


    # Store the config due to the disappearance after accelerator prepare
    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)          # UNet won't be trained in conditioning branch
    controlnet.requires_grad_(False)    # Will turn back to requires grad later on
    if config["use_text"]:
        text_encoder.requires_grad_(False)


    # 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 + unet + image_encoder to gpu and cast to weight_dtype
    vae.to(accelerator.device, dtype=weight_dtype)
    unet.to(accelerator.device, dtype=weight_dtype)     # we don't train UNet anymore, so we cast it here
    image_encoder.to(accelerator.device, dtype=weight_dtype)
    if config["use_text"]:
        text_encoder.to(accelerator.device, dtype=weight_dtype)



    # `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:
                i = len(weights) - 1

                while len(weights) > 0:
                    weights.pop()
                    model = models[i]

                    sub_dir = "controlnet"
                    model.save_pretrained(os.path.join(output_dir, sub_dir))

                    i -= 1

        def load_model_hook(models, input_dir):
            while len(models) > 0:
                # pop models so that they are not loaded again
                model = models.pop()

                # load diffusers style into model
                load_model = ControlNetModel.from_pretrained(input_dir, subfolder="controlnet")
                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:
        controlnet.enable_gradient_checkpointing()


    # Check that all trainable models are in full precision
    low_precision_error_string = (
        " Please make sure to always have all model weights in full float32 precision when starting training - even if"
        " doing mixed precision training, copy of the weights should still be float32."
    )
    if accelerator.unwrap_model(controlnet).dtype != torch.float32:
        raise ValueError(
            f"Controlnet loaded as datatype {accelerator.unwrap_model(controlnet).dtype}. {low_precision_error_string}"
        )
    
    


    ################################ Make Training dataset  ######################################
    if config["data_loader_type"] == "thisthat":    # Only keep thisthat mode now
        train_dataset = Video_ThisThat_Dataset(config, accelerator.device, tokenizer=tokenizer)
    else:
        raise NotImplementedError("We don't support such data loader type")

    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
    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
    
    
    # Make ControlNet require Grad
    controlnet.requires_grad_(True)


    ###################### For partial fine-tune setting #######################
    parameters_list = []
    for name, para in controlnet.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)
            if not name.find("attn") != -1:     # Only block the spatial Transformer
                para.requires_grad = False
            else:
                parameters_list.append(para)
                para.requires_grad = True  
        else:
            parameters_list.append(para)
            para.requires_grad = True

    # Double check the weight that will be trained
    total_params_for_training = 0
    for name, param in controlnet.named_parameters():
        if param.requires_grad:
            total_params_for_training += param.numel()
            print(name + " requires grad update")    
    print("Total parameter that will be trained in controlnet 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`.
    controlnet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        controlnet, 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       TODO: need to check how to use checkpoints from pre-trained weights!!!
    global_step = 0     # Catch the current iteration
    first_epoch = 0
    if resume_from_checkpoint:          # Resume Checkpoints!!!!!
        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:
            accelerator.print(
                f"Checkpoint '{resume_from_checkpoint}' does not exist. Starting a new training run."
            )
            resume_from_checkpoint = None
            initial_global_step = 0
        else:
            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)      # 目前先用text_embeddings 再用encoder_hidden_states感觉好一点

            # 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 test2image in diffusers  TODO: check the data loader conflict
    for epoch in range(first_epoch, num_train_epochs):
        controlnet.train()
        train_loss = 0.0
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(controlnet):
                # batch is a dictionary with video_frames and controlnet_condition
                video_frames = batch["video_frames"].to(weight_dtype).to(accelerator.device, non_blocking=True)        # [-1, 1] range
                condition_img = batch["controlnet_condition"].to(dtype=weight_dtype)        # [0, 1] range
                reflected_motion_bucket_id = batch["reflected_motion_bucket_id"]
                controlnet_image_index = batch["controlnet_image_index"]
                prompt = batch["prompt"]


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

                ##################################### 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 # cond_sigmas
                conditional_latents = vae.encode(conditional_pixel_values).latent_dist.mode()   
                conditional_latents = repeat(conditional_latents, 'b c h w->b f c h w', f=num_frames)       # conditional_latents没有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(weight_dtype).to(latents.device)      # TODO: 我觉得noise这块，sigma算法是最不确定是否正确的地方
                sigmas = sigmas[:, None, None, None, None]
                noisy_latents = latents + torch.randn_like(latents) * sigmas
                inp_noisy_latents = noisy_latents / ((sigmas**2 + 1) ** 0.5)        # multiplied by c_in in paper


                # 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.
                if conditioning_dropout_prob != 0:
                    random_p = torch.rand(bsz, device=latents.device, generator=generator)

                    # Sample masks for the encoder_hidden_states (to replace prompts in InstructPix2Pix). 
                    prompt_mask = random_p < 2 * conditioning_dropout_prob
                    prompt_mask = prompt_mask.reshape(bsz, 1, 1)
                    # Final encoder_hidden_states conditioning.
                    null_conditioning = torch.zeros_like(encoder_hidden_states) # encoder_hidden_states has already been used with .unsqueeze(1)
                    encoder_hidden_states = torch.where(prompt_mask, null_conditioning, encoder_hidden_states) 

                    # Sample masks for the original image latents.
                    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   

                # The Concatenation is move downward with the masking feature


                # 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,       # Note: noise 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)
                ##########################################################################################



                ################################### Get ControlNet Output ###################################

                # Transform controlnet_image_index to the data format we want
                controlnet_image_index = list(controlnet_image_index.cpu().detach().numpy()[0])
                assert condition_img.shape[1] >= len(controlnet_image_index)  

                # Designing the 0/1 mask for Sparse Conditioning
                controlnet_conditioning_mask_shape = list(condition_img.shape)
                controlnet_conditioning_mask_shape[2] = 1       # frame dim
                controlnet_conditioning_mask = torch.zeros(controlnet_conditioning_mask_shape).to(dtype=weight_dtype).to(accelerator.device)
                controlnet_conditioning_mask[:, controlnet_image_index] = 1


                # Add vae latent mask to controlnet noise
                if config["mask_controlnet_vae"]:
                    b, f, c, h, w = conditional_latents.shape

                    # Create a mask: Value less than the threshold is set to be True
                    mask = torch.rand((b, f, 1, h, w), device=accelerator.device) < (1-config["mask_proportion"])      # channel sync
                    # mask[:,0,:,:,:] = 1     # For the first frame, we still keep it

                    # Multiply to the conditional latents, we will just make the mean and variance zero to present those with zero masking
                    masked_conditional_latents = conditional_latents * mask
                    controlnet_inp_noisy_latents = torch.cat([inp_noisy_latents, masked_conditional_latents], dim=2)
                else:
                    controlnet_inp_noisy_latents = torch.cat([inp_noisy_latents, conditional_latents], dim=2)


                # VAE encode
                controlnet_cond = condition_img.flatten(0, 1)
                controlnet_cond = vae.encode(controlnet_cond).latent_dist.mode()


                down_block_res_samples, mid_block_res_sample = controlnet(
                    sample = controlnet_inp_noisy_latents,          
                    timestep = timesteps,
                    encoder_hidden_states = encoder_hidden_states,     
                    added_time_ids = added_time_ids,
                    controlnet_cond = controlnet_cond,
                    return_dict = False,
                    conditioning_scale = config["outer_conditioning_scale"],
                    inner_conditioning_scale = config["inner_conditioning_scale"],
                    guess_mode = False,         # No Guess Mode
                )   

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



                ###################################### Predict Noise ########################################
                # Add vae latent mask to controlnet noise
                if config["mask_unet_vae"]:
                    b, f, c, h, w = conditional_latents.shape

                    # Create a mask
                    mask = torch.rand((b, f, 1, h, w), device=accelerator.device) < (1-config["mask_proportion"])      # channel sync
                    # mask[:,0,:,:,:] = 1     # For the first frame, we still keep it

                    # Multiply to the conditional latents, we will just make the mean and variance zero to present those with zero masking
                    if not config["mask_controlnet_vae"]:  
                        masked_conditional_latents = conditional_latents * mask
                    unet_inp_noisy_latents = torch.cat([inp_noisy_latents, masked_conditional_latents], dim=2)
                else:
                    unet_inp_noisy_latents = torch.cat([inp_noisy_latents, conditional_latents], dim=2)

                # Add with controlnet middle output layers
                model_pred = unet(
                                    unet_inp_noisy_latents,
                                    timesteps, 
                                    encoder_hidden_states, 
                                    added_time_ids = added_time_ids,
                                    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    


                # 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  # What our loss will optimize with
                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.item()/ gradient_accumulation_steps, global_step)        # 我觉得loss的report就用这个avg_loss就行了
                

                # Backpropagate
                accelerator.backward(loss)
                if accelerator.sync_gradients:  # For ControlNet
                    params_to_clip = controlnet.parameters()
                    accelerator.clip_grad_norm_(params_to_clip, max_grad_norm)
                optimizer.step()
                lr_scheduler.step()     # I think constant will take no influence here
                optimizer.zero_grad(set_to_none=True)
                ##########################################################################################


            # 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()}
            progress_bar.set_postfix(**logs)


           ##################################### Validation per XXX iterations #######################################
            if accelerator.is_main_process:
                if global_step > -1 and global_step % validation_step == 0:         # Fixed 100 steps to validate
                    
                    log_validation(
                                    vae,
                                    unet,
                                    controlnet,
                                    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 should be updated together
            global_step += 1

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


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

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