import os
from typing import Optional, Union, Tuple, List, Callable, Dict
from tqdm.notebook import tqdm
import torch
from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL
import torch.nn.functional as nnf
import numpy as np
import abc
import ptp_utils
import seq_aligner
import shutil
from torch.optim.adam import Adam
from PIL import Image
from transformers import AutoTokenizer, CLIPTextModel, CLIPTokenizer
from einops import rearrange

from tuneavideo.models.unet import UNet3DConditionModel
from tuneavideo.pipelines.pipeline_tuneavideo import TuneAVideoPipeline

import cv2
import argparse
from omegaconf import OmegaConf

scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
MY_TOKEN = ''
LOW_RESOURCE = False
NUM_DDIM_STEPS = 50
GUIDANCE_SCALE = 7.5
MAX_NUM_WORDS = 77
device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
IRC = True

# need to adjust
cross_replace_steps = {'default_': .2,}
self_replace_steps = .5
mask_th = (.3, .3)
video_len = 8

def main(
    pretrained_model_path: str,
    image_path: str,
    prompt: str,
    prompts: Tuple[str],
    blend_word: Tuple[str],
    eq_params: Dict,
    gif_folder: str,
    gif_name_1: str,
    gif_name_2: str,
    IRC: bool,
):
    blend_word = (((blend_word[0],), (blend_word[1],)))
    eq_params["words"] = (eq_params["words"],)
    eq_params["values"] = (eq_params["values"],)
    eq_params = dict(eq_params)
    prompts = list(prompts)
    if not os.path.exists(gif_folder):
        os.makedirs(gif_folder)

    # Load the tokenizer
    tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer")
    # Load models and create wrapper for stable diffusion
    text_encoder = CLIPTextModel.from_pretrained(
        pretrained_model_path,
        subfolder="text_encoder",
    )
    vae = AutoencoderKL.from_pretrained(
        pretrained_model_path,
        subfolder="vae",
    )
    unet = UNet3DConditionModel.from_pretrained(pretrained_model_path, subfolder="unet")
    ldm_stable = TuneAVideoPipeline(
        vae=vae,
        text_encoder=text_encoder,
        tokenizer=tokenizer,
        unet=unet,
        scheduler=scheduler,
    ).to(device)

    try:
        ldm_stable.disable_xformers_memory_efficient_attention()
    except AttributeError:
        print("Attribute disable_xformers_memory_efficient_attention() is missing")
    tokenizer = ldm_stable.tokenizer # Tokenizer of class: [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer)
    # A tokenizer breaks a stream of text into tokens, usually by looking for whitespace (tabs, spaces, new lines).

    class LocalBlend:
        
        def get_mask(self, maps, alpha, use_pool): # alpha is a word map
            k = 1
            maps = (maps * alpha).sum(-1).mean(2) # [2, 80, 1, 16, 16, 77], [2, 1, 1, 1, 1, 77]
            if use_pool:
                maps = nnf.max_pool2d(maps, (k * 2 + 1, k * 2 +1), (1, 1), padding=(k, k))
            mask = nnf.interpolate(maps, size=(x_t.shape[3:]))
            mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0]
            mask = mask.gt(self.th[1-int(use_pool)])
            mask = mask[:1] + mask
            return mask
        
        def __call__(self, x_t, attention_store, step):
            self.counter += 1
            if self.counter > self.start_blend:
                # attention_store["down_cross"]: 4, attention_store["up_cross"]:6, attention_store["down_cross"][0]: torch.Size([32, 1024, 77])
                maps = attention_store["down_cross"][2:4] + attention_store["up_cross"][:3]
                # maps = [item.reshape(self.alpha_layers.shape[0], -1, 1, 16, 16, MAX_NUM_WORDS) for item in maps]
                maps = [item.reshape(self.alpha_layers.shape[0], -1, 8, 16, 16, MAX_NUM_WORDS) for item in maps]
                maps = torch.cat(maps, dim=2)
                # self.alpha_layers: torch.Size([2, 1, 1, 1, 1, 77])
                mask = self.get_mask(maps, self.alpha_layers, True)
                if self.substruct_layers is not None:
                    maps_sub = ~self.get_mask(maps, self.substruct_layers, False)
                    mask = mask * maps_sub
                mask = mask.float()
                mask = mask.reshape(-1, 1, mask.shape[-3], mask.shape[-2], mask.shape[-1])

                x_t = x_t[:1] + mask * (x_t - x_t[:1]) # line13 algorithm
            return x_t
        
        def __init__(self, prompts: List[str], words: [List[List[str]]], substruct_words=None, start_blend=0.2, th=(.3, .3)):
            alpha_layers = torch.zeros(len(prompts),  1, 1, 1, 1, MAX_NUM_WORDS)
            for i, (prompt, words_) in enumerate(zip(prompts, words)):
                if type(words_) is str:
                    words_ = [words_]
                for word in words_:
                    ind = ptp_utils.get_word_inds(prompt, word, tokenizer)
                    alpha_layers[i, :, :, :, :, ind] = 1
            
            if substruct_words is not None:
                substruct_layers = torch.zeros(len(prompts),  1, 1, 1, 1, MAX_NUM_WORDS)
                for i, (prompt, words_) in enumerate(zip(prompts, substruct_words)):
                    if type(words_) is str:
                        words_ = [words_]
                    for word in words_:
                        ind = ptp_utils.get_word_inds(prompt, word, tokenizer)
                        substruct_layers[i, :, :, :, :, ind] = 1
                self.substruct_layers = substruct_layers.to(device)
            else:
                self.substruct_layers = None
            self.alpha_layers = alpha_layers.to(device)
            self.start_blend = int(start_blend * NUM_DDIM_STEPS)
            self.counter = 0 
            self.th=th


            
            
    class EmptyControl:
        
        
        def step_callback(self, x_t):
            return x_t
        
        def between_steps(self):
            return
        
        def __call__(self, attn, is_cross: bool, place_in_unet: str):
            return attn

        
    class AttentionControl(abc.ABC):
        
        def step_callback(self, x_t):
            return x_t
        
        def between_steps(self):
            return
        
        @property
        def num_uncond_att_layers(self):
            return self.num_att_layers if LOW_RESOURCE else 0
        
        @abc.abstractmethod
        def forward (self, attn, is_cross: bool, place_in_unet: str):
            raise NotImplementedError

        def __call__(self, attn, is_cross: bool, place_in_unet: str):
            if self.cur_att_layer >= self.num_uncond_att_layers:
                if LOW_RESOURCE:
                    attn = self.forward(attn, is_cross, place_in_unet)
                else:
                    h = attn.shape[0]
                    attn[h // 2:] = self.forward(attn[h // 2:], is_cross, place_in_unet)
            self.cur_att_layer += 1
            if self.cur_att_layer == self.num_att_layers + self.num_uncond_att_layers:
                self.cur_att_layer = 0
                self.cur_step += 1
                self.between_steps()
            return attn
        
        def reset(self):
            self.cur_step = 0
            self.cur_att_layer = 0

        def __init__(self):
            self.cur_step = 0
            self.num_att_layers = -1
            self.cur_att_layer = 0

    class SpatialReplace(EmptyControl):
        
        def step_callback(self, x_t):
            if self.cur_step < self.stop_inject:
                b = x_t.shape[0]
                x_t = x_t[:1].expand(b, *x_t.shape[1:])
            return x_t

        def __init__(self, stop_inject: float):
            super(SpatialReplace, self).__init__()
            self.stop_inject = int((1 - stop_inject) * NUM_DDIM_STEPS)
            

    class AttentionStore(AttentionControl):

        @staticmethod
        def get_empty_store():
            return {"down_cross": [], "mid_cross": [], "up_cross": [],
                    "down_self": [],  "mid_self": [],  "up_self": []}

        def forward(self, attn, is_cross: bool, place_in_unet: str):
            key = f"{place_in_unet}_{'cross' if is_cross else 'self'}"
            if attn.shape[1] <= 32 ** 2:  # avoid memory overhead
                self.step_store[key].append(attn) # 'down_self' torch.Size([32768, 8, 8])
            return attn

        def between_steps(self):
            if len(self.attention_store) == 0:
                self.attention_store = self.step_store
            else:
                for key in self.attention_store:
                    for i in range(len(self.attention_store[key])):
                        self.attention_store[key][i] += self.step_store[key][i]
            self.step_store = self.get_empty_store()

        def get_average_attention(self):
            average_attention = {key: [item / self.cur_step for item in self.attention_store[key]] for key in self.attention_store}
            return average_attention


        def reset(self):
            super(AttentionStore, self).reset()
            self.step_store = self.get_empty_store()
            self.attention_store = {}

        def __init__(self):
            super(AttentionStore, self).__init__()
            self.step_store = self.get_empty_store()
            self.attention_store = {}

            
    class AttentionControlEdit(AttentionStore, abc.ABC):
        
        def step_callback(self, x_t):
            if self.local_blend is not None:
                x_t = self.local_blend(x_t, self.attention_store, self.cur_step)
            return x_t
            
        def replace_self_attention(self, attn_base, att_replace, place_in_unet):
            if att_replace.shape[2] <= 32 ** 2:
                attn_base = attn_base.unsqueeze(0).expand(att_replace.shape[0], *attn_base.shape)
                return attn_base
            else:
                return att_replace
        
        @abc.abstractmethod
        def replace_cross_attention(self, attn_base, att_replace):
            raise NotImplementedError
        
        def forward(self, attn, is_cross: bool, place_in_unet: str):
            super(AttentionControlEdit, self).forward(attn, is_cross, place_in_unet)
            if is_cross or (self.num_self_replace[0] <= self.cur_step < self.num_self_replace[1]):
                h = attn.shape[0] // (self.batch_size)
                attn = attn.reshape(self.batch_size, h, *attn.shape[1:])
                attn_base, attn_repalce = attn[0], attn[1:]
                if is_cross:
                    alpha_words = self.cross_replace_alpha[self.cur_step]
                    attn_repalce_new = self.replace_cross_attention(attn_base, attn_repalce) * alpha_words + (1 - alpha_words) * attn_repalce
                    attn[1:] = attn_repalce_new
                else:
                    attn[1:] = self.replace_self_attention(attn_base, attn_repalce, place_in_unet)
                attn = attn.reshape(self.batch_size * h, *attn.shape[2:])
            return attn
        
        def __init__(self, prompts, num_steps: int,
                    cross_replace_steps: Union[float, Tuple[float, float], Dict[str, Tuple[float, float]]],
                    self_replace_steps: Union[float, Tuple[float, float]],
                    local_blend: Optional[LocalBlend]):
            super(AttentionControlEdit, self).__init__()
            self.batch_size = len(prompts)
            self.cross_replace_alpha = ptp_utils.get_time_words_attention_alpha(prompts, num_steps, cross_replace_steps, tokenizer).to(device)
            if type(self_replace_steps) is float:
                self_replace_steps = 0, self_replace_steps
            self.num_self_replace = int(num_steps * self_replace_steps[0]), int(num_steps * self_replace_steps[1])
            self.local_blend = local_blend

    class AttentionReplace(AttentionControlEdit):

        def replace_cross_attention(self, attn_base, att_replace):
            return torch.einsum('hpw,bwn->bhpn', attn_base, self.mapper)
        
        def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float,
                    local_blend: Optional[LocalBlend] = None):
            super(AttentionReplace, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend)
            self.mapper = seq_aligner.get_replacement_mapper(prompts, tokenizer).to(device)
            

    class AttentionRefine(AttentionControlEdit):

        def replace_cross_attention(self, attn_base, att_replace):
            attn_base_replace = attn_base[:, :, self.mapper].permute(2, 0, 1, 3)
            attn_replace = attn_base_replace * self.alphas + att_replace * (1 - self.alphas)
            return attn_replace

        def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float,
                    local_blend: Optional[LocalBlend] = None):
            super(AttentionRefine, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend)
            self.mapper, alphas = seq_aligner.get_refinement_mapper(prompts, tokenizer)
            self.mapper, alphas = self.mapper.to(device), alphas.to(device)
            self.alphas = alphas.reshape(alphas.shape[0], 1, 1, alphas.shape[1])


    class AttentionReweight(AttentionControlEdit):

        def replace_cross_attention(self, attn_base, att_replace):
            if self.prev_controller is not None:
                attn_base = self.prev_controller.replace_cross_attention(attn_base, att_replace)
            attn_replace = attn_base[None, :, :, :] * self.equalizer[:, None, None, :]
            return attn_replace

        def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float, equalizer,
                    local_blend: Optional[LocalBlend] = None, controller: Optional[AttentionControlEdit] = None):
            super(AttentionReweight, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend)
            self.equalizer = equalizer.to(device)
            self.prev_controller = controller


    def get_equalizer(text: str, word_select: Union[int, Tuple[int, ...]], values: Union[List[float],
                    Tuple[float, ...]]):
        if type(word_select) is int or type(word_select) is str:
            word_select = (word_select,)
        equalizer = torch.ones(1, 77)
        
        for word, val in zip(word_select, values):
            inds = ptp_utils.get_word_inds(text, word, tokenizer)
            equalizer[:, inds] = val
        return equalizer

    def aggregate_attention(attention_store: AttentionStore, res: int, from_where: List[str], is_cross: bool, select: int):
        out = []
        attention_maps = attention_store.get_average_attention()
        num_pixels = res ** 2
        for location in from_where:
            for item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]:
                if item.shape[1] == num_pixels: # torch.Size([64, 256, 77]) all can pass
                    cross_maps = item.reshape(8, 8, res, res, item.shape[-1])
                    out.append(cross_maps)
        out = torch.cat(out, dim=1)
        out = out.sum(1) / out.shape[1]
        return out.cpu()


    def make_controller(prompts: List[str], is_replace_controller: bool, cross_replace_steps: Dict[str, float], self_replace_steps: float, blend_words=None, equilizer_params=None, mask_th=(.3,.3)) -> AttentionControlEdit:
        if blend_words is None:
            lb = None
        else:
            lb = LocalBlend(prompts, blend_word, th=mask_th)
        if is_replace_controller:
            controller = AttentionReplace(prompts, NUM_DDIM_STEPS, cross_replace_steps=cross_replace_steps, self_replace_steps=self_replace_steps, local_blend=lb)
        else:
            controller = AttentionRefine(prompts, NUM_DDIM_STEPS, cross_replace_steps=cross_replace_steps, self_replace_steps=self_replace_steps, local_blend=lb)
        if equilizer_params is not None:
            eq = get_equalizer(prompts[1], equilizer_params["words"], equilizer_params["values"])
            controller = AttentionReweight(prompts, NUM_DDIM_STEPS, cross_replace_steps=cross_replace_steps,
                                        self_replace_steps=self_replace_steps, equalizer=eq, local_blend=lb, controller=controller)
        return controller


    def load_512_seq(image_path, left=0, right=0, top=0, bottom=0, n_sample_frame=video_len, sampling_rate=1):
        images = []
        for file in sorted(os.listdir(image_path)):
            images.append(file)
        n_images = len(images)
        sequence_length = (n_sample_frame - 1) * sampling_rate + 1
        if n_images < sequence_length:
            raise ValueError
        frames = []
        for index in range(n_sample_frame):
            p = os.path.join(image_path, images[index])
            image = np.array(Image.open(p).convert("RGB"))
            h, w, c = image.shape
            left = min(left, w-1)
            right = min(right, w - left - 1)
            top = min(top, h - left - 1)
            bottom = min(bottom, h - top - 1)
            image = image[top:h-bottom, left:w-right]
            h, w, c = image.shape
            if h < w:
                offset = (w - h) // 2
                image = image[:, offset:offset + h]
            elif w < h:
                offset = (h - w) // 2
                image = image[offset:offset + w]
            image = np.array(Image.fromarray(image).resize((512, 512)))
            frames.append(image)
        return np.stack(frames)


    class NullInversion:
        
        def prev_step(self, model_output: Union[torch.FloatTensor, np.ndarray], timestep: int, sample: Union[torch.FloatTensor, np.ndarray]):
            prev_timestep = timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
            alpha_prod_t = self.scheduler.alphas_cumprod[timestep]
            alpha_prod_t_prev = self.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.scheduler.final_alpha_cumprod
            beta_prod_t = 1 - alpha_prod_t
            pred_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
            pred_sample_direction = (1 - alpha_prod_t_prev) ** 0.5 * model_output
            prev_sample = alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction
            return prev_sample
        
        def next_step(self, model_output: Union[torch.FloatTensor, np.ndarray], timestep: int, sample: Union[torch.FloatTensor, np.ndarray]): # doing inversion (math)
            timestep, next_timestep = min(timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps, 999), timestep
            alpha_prod_t = self.scheduler.alphas_cumprod[timestep] if timestep >= 0 else self.scheduler.final_alpha_cumprod
            alpha_prod_t_next = self.scheduler.alphas_cumprod[next_timestep]
            beta_prod_t = 1 - alpha_prod_t
            next_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
            next_sample_direction = (1 - alpha_prod_t_next) ** 0.5 * model_output
            next_sample = alpha_prod_t_next ** 0.5 * next_original_sample + next_sample_direction
            return next_sample
        
        def get_noise_pred_single(self, latents, t, context): # latents: torch.Size([1, 4, 64, 64]); t: tensor(1); context: torch.Size([1, 77, 768])
            # formats are correct for video unet input; Tune-A-Video also predicts the residual
            noise_pred = self.model.unet(latents, t, encoder_hidden_states=context)["sample"] # easy to out of mem
            return noise_pred

        def get_noise_pred(self, latents, t, is_forward=True, context=None):
            latents_input = torch.cat([latents] * 2)
            if context is None:
                context = self.context
            guidance_scale = 1 if is_forward else GUIDANCE_SCALE
            noise_pred = self.model.unet(latents_input, t, encoder_hidden_states=context)["sample"]
            noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2)
            noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond)
            if is_forward:
                latents = self.next_step(noise_pred, t, latents)
            else:
                latents = self.prev_step(noise_pred, t, latents)
            return latents

        @torch.no_grad()
        def latent2image(self, latents, return_type='np'):
            latents = 1 / 0.18215 * latents.detach()
            image = self.model.vae.decode(latents)['sample']
            if return_type == 'np':
                image = (image / 2 + 0.5).clamp(0, 1)
                image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
                image = (image * 255).astype(np.uint8)
            return image

        @torch.no_grad()
        def latent2image_video(self, latents, return_type='np'):
            latents = 1 / 0.18215 * latents.detach()
            latents = latents[0].permute(1, 0, 2, 3)
            image = self.model.vae.decode(latents)['sample']
            if return_type == 'np':
                image = (image / 2 + 0.5).clamp(0, 1)
                image = image.cpu().permute(0, 2, 3, 1).numpy()
                image = (image * 255).astype(np.uint8)
            return image

        @torch.no_grad()
        def image2latent(self, image):
            with torch.no_grad():
                if type(image) is Image:
                    image = np.array(image)
                if type(image) is torch.Tensor and image.dim() == 4:
                    latents = image
                else:
                    image = torch.from_numpy(image).float() / 127.5 - 1
                    image = image.permute(2, 0, 1).unsqueeze(0).to(device)
                    latents = self.model.vae.encode(image)['latent_dist'].mean
                    latents = latents * 0.18215
            return latents

        @torch.no_grad()
        def image2latent_video(self, image):
            with torch.no_grad():
                image = torch.from_numpy(image).float() / 127.5 - 1
                image = image.permute(0, 3, 1, 2).to(device)
                latents = self.model.vae.encode(image)['latent_dist'].mean
                latents = rearrange(latents, "(b f) c h w -> b c f h w", b=1)
                latents = latents * 0.18215
            return latents

        @torch.no_grad()
        def init_prompt(self, prompt: str):
            uncond_input = self.model.tokenizer(
                [""], padding="max_length", max_length=self.model.tokenizer.model_max_length,
                return_tensors="pt"
            )
            uncond_embeddings = self.model.text_encoder(uncond_input.input_ids.to(self.model.device))[0] # len=2, uncond_embeddings
            text_input = self.model.tokenizer(
                [prompt],
                padding="max_length",
                max_length=self.model.tokenizer.model_max_length,
                truncation=True,
                return_tensors="pt",
            )
            text_embeddings = self.model.text_encoder(text_input.input_ids.to(self.model.device))[0]
            self.context = torch.cat([uncond_embeddings, text_embeddings])
            self.prompt = prompt

        @torch.no_grad()
        def ddim_loop(self, latent):
            uncond_embeddings, cond_embeddings = self.context.chunk(2)
            all_latent = [latent]
            latent = latent.clone().detach()
            for i in range(NUM_DDIM_STEPS):
                t = self.model.scheduler.timesteps[len(self.model.scheduler.timesteps) - i - 1]
                # latent: torch.Size([1, 4, 8, 16, 16])
                # cond_embeddings: torch.Size([1, 77, 768])
                #  noise_pred: torch.Size([1, 4, 8, 16, 16])
                noise_pred = self.get_noise_pred_single(latent, t, cond_embeddings) # use a unet
                latent = self.next_step(noise_pred, t, latent)
                all_latent.append(latent)
            return all_latent

        @property
        def scheduler(self):
            return self.model.scheduler

        @torch.no_grad()
        def ddim_inversion(self, image):
            latent = self.image2latent_video(image)
            image_rec = self.latent2image_video(latent) # image: (512, 512, 3); latent: torch.Size([1, 4, 64, 64])
            ddim_latents = self.ddim_loop(latent)
            return image_rec, ddim_latents

        def null_optimization(self, latents, num_inner_steps, epsilon): # uncond_embeddings is what we what
            uncond_embeddings, cond_embeddings = self.context.chunk(2)
            uncond_embeddings_list = []
            latent_cur = latents[-1]
            bar = tqdm(total=num_inner_steps * NUM_DDIM_STEPS)
            for i in range(NUM_DDIM_STEPS):
                uncond_embeddings = uncond_embeddings.clone().detach()
                uncond_embeddings.requires_grad = True
                optimizer = Adam([uncond_embeddings], lr=1e-2 * (1. - i / 100.))
                latent_prev = latents[len(latents) - i - 2] # GT
                t = self.model.scheduler.timesteps[i]
                with torch.no_grad():
                    noise_pred_cond = self.get_noise_pred_single(latent_cur, t, cond_embeddings)
                for j in range(num_inner_steps):
                    noise_pred_uncond = self.get_noise_pred_single(latent_cur, t, uncond_embeddings)
                    noise_pred = noise_pred_uncond + GUIDANCE_SCALE * (noise_pred_cond - noise_pred_uncond)
                    latents_prev_rec = self.prev_step(noise_pred, t, latent_cur)
                    loss = nnf.mse_loss(latents_prev_rec, latent_prev)
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()
                    loss_item = loss.item()
                    bar.update()
                    if loss_item < epsilon + i * 2e-5:
                        break
                for j in range(j + 1, num_inner_steps):
                    bar.update()
                uncond_embeddings_list.append(uncond_embeddings[:1].detach())
                with torch.no_grad():
                    context = torch.cat([uncond_embeddings, cond_embeddings])
                    latent_cur = self.get_noise_pred(latent_cur, t, False, context)
            bar.close()
            return uncond_embeddings_list
        
        def invert(self, image_path: str, prompt: str, offsets=(0,0,0,0), num_inner_steps=10, early_stop_epsilon=1e-5, verbose=False):
            self.init_prompt(prompt)
            ptp_utils.register_attention_control(self.model, None)
            image_gt = load_512_seq(image_path, *offsets)
            if verbose:
                print("DDIM inversion...")
            image_rec, ddim_latents = self.ddim_inversion(image_gt) # ddim_latents is a list, like the link in Figure 3
            # image_rec refers to vq-autoencoder reconstruction
            if verbose:
                print("Null-text optimization...")
            uncond_embeddings = self.null_optimization(ddim_latents, num_inner_steps, early_stop_epsilon) # ddim_latents serve as GT; easy to out of mem
            return (image_gt, image_rec), ddim_latents[-1], uncond_embeddings

        def invert_(self, image_path: str, prompt: str, offsets=(0,0,0,0), num_inner_steps=10, early_stop_epsilon=1e-5, verbose=False):
            self.init_prompt(prompt)
            ptp_utils.register_attention_control(self.model, None)
            image_gt = load_512_seq(image_path, *offsets)
            if verbose:
                print("DDIM inversion...")
            image_rec, ddim_latents = self.ddim_inversion(image_gt) # ddim_latents is a list, like the link in Figure 3
            # image_rec refers to vq-autoencoder reconstruction
            if verbose:
                print("Null-text optimization...")
            return (image_gt, image_rec), ddim_latents[-1], None
        
        def __init__(self, model):
            scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False,
                                    set_alpha_to_one=False)
            self.model = model
            self.tokenizer = self.model.tokenizer
            self.model.scheduler.set_timesteps(NUM_DDIM_STEPS)
            self.prompt = None
            self.context = None

    null_inversion = NullInversion(ldm_stable)


    @torch.no_grad()
    def text2image_ldm_stable(
        model,
        prompt:  List[str],
        controller,
        num_inference_steps: int = 50,
        guidance_scale: Optional[float] = 7.5,
        generator: Optional[torch.Generator] = None,
        latent: Optional[torch.FloatTensor] = None,
        uncond_embeddings=None,
        start_time=50,
        return_type='image'
    ):
        batch_size = len(prompt)
        ptp_utils.register_attention_control(model, controller)
        height = width = 512
        
        text_input = model.tokenizer(
            prompt,
            padding="max_length",
            max_length=model.tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt",
        )
        text_embeddings = model.text_encoder(text_input.input_ids.to(model.device))[0]

        max_length = text_input.input_ids.shape[-1]
        if uncond_embeddings is None:
            uncond_input = model.tokenizer(
                [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
            )
            uncond_embeddings_ = model.text_encoder(uncond_input.input_ids.to(model.device))[0]
        else:
            uncond_embeddings_ = None

        model.scheduler.set_timesteps(num_inference_steps)
        for i, t in enumerate(tqdm(model.scheduler.timesteps[-start_time:])):
            if uncond_embeddings_ is None:
                context = torch.cat([uncond_embeddings[i].expand(*text_embeddings.shape), text_embeddings])
            else:
                context = torch.cat([uncond_embeddings_, text_embeddings])
            latents = latent
            latents = ptp_utils.diffusion_step(model, controller, latents, context, t, guidance_scale, low_resource=False)
        if return_type == 'image':
            image = ptp_utils.latent2image_video(model.vae, latents)
        else:
            image = latents
        return image, latent


    ###############
    # Custom APIs:

    ldm_stable.enable_xformers_memory_efficient_attention()

    (image_gt, image_enc), x_t, uncond_embeddings = null_inversion.invert(image_path, prompt, offsets=(0,0,0,0), verbose=True)

    ##### load uncond #####
    # uncond_embeddings_load = np.load(uncond_embeddings_path)
    # uncond_embeddings = []
    # for i in range(uncond_embeddings_load.shape[0]):
    #     uncond_embeddings.append(torch.from_numpy(uncond_embeddings_load[i]).to(device))
    #######################

    ##### save uncond #####
    # uncond_embeddings = torch.cat(uncond_embeddings)
    # uncond_embeddings = uncond_embeddings.cpu().numpy()
    #######################

    print("Start Video-P2P!")

    controller = make_controller(prompts, IRC, cross_replace_steps, self_replace_steps, blend_word, eq_params, mask_th=mask_th)
    ptp_utils.register_attention_control(ldm_stable, controller)
    generator = torch.Generator(device=device)
    with torch.no_grad():
        sequence = ldm_stable(
            prompts,
            generator=generator,
            latents=x_t,
            uncond_embeddings_pre=uncond_embeddings,
            controller = controller,
            video_length=video_len,
            simple=True,
        ).videos
    sequence1 = rearrange(sequence[0], "c t h w -> t h w c")
    sequence2 = rearrange(sequence[1], "c t h w -> t h w c")
    inversion = []
    videop2p = []
    for i in range(sequence1.shape[0]):
        inversion.append( Image.fromarray((sequence1[i] * 255).numpy().astype(np.uint8)) )
        videop2p.append( Image.fromarray((sequence2[i] * 255).numpy().astype(np.uint8)) )

    inversion[0].save(gif_name_1.replace('name', 'inversion'), save_all=True, append_images=inversion[1:], optimize=False, loop=0, duration=250)
    videop2p[0].save(gif_name_2.replace('name', 'p2p'), save_all=True, append_images=videop2p[1:], optimize=False, loop=0, duration=250)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, default="./configs/videop2p.yaml")
    args = parser.parse_args()

    main(**OmegaConf.load(args.config))