"""
This script demonstrates how to generate a video using the CogVideoX model with the Hugging Face `diffusers` pipeline.
The script supports different types of video generation, including text-to-video (t2v), image-to-video (i2v),
and video-to-video (v2v), depending on the input data and different weight.

- text-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b
- video-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b
- image-to-video: THUDM/CogVideoX-5b-I2V or THUDM/CogVideoX1.5-5b-I2V

Running the Script:
To run the script, use the following command with appropriate arguments:

```bash
$ python cli_demo.py --prompt "A girl riding a bike." --model_path THUDM/CogVideoX1.5-5b --generate_type "t2v"
```

You can change `pipe.enable_sequential_cpu_offload()` to `pipe.enable_model_cpu_offload()` to speed up inference, but this will use more GPU memory

Additional options are available to specify the model path, guidance scale, number of inference steps, video generation type, and output paths.

"""

import os
import json
import types
import argparse
import logging
from typing import Literal, Optional, Any

import torch

from diffusers import (
    CogVideoXDPMScheduler,
    CogVideoXImageToVideoPipeline,
    CogVideoXPipeline,
    CogVideoXVideoToVideoPipeline,
    CogVideoXTransformer3DModel,
    AutoencoderKLCogVideoX,
)
from diffusers.utils import export_to_video, load_image, load_video


logging.basicConfig(level=logging.INFO)

# Recommended resolution for each model (width, height)
RESOLUTION_MAP = {
    # cogvideox1.5-*
    "cogvideox1.5-5b-i2v": (768, 1360),
    "cogvideox1.5-5b": (768, 1360),
    # cogvideox-*
    "cogvideox-5b-i2v": (480, 720),
    "cogvideox-5b": (480, 720),
    "cogvideox-2b": (480, 720),
}
from typing import Dict, List, Optional, Tuple, Union
from diffusers.schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler
from diffusers.schedulers.scheduling_dpm_cogvideox import DDIMSchedulerOutput
from diffusers import (
    CogVideoXDPMScheduler,
    CogVideoXImageToVideoPipeline,
    CogVideoXPipeline,
    CogVideoXVideoToVideoPipeline,
)
from diffusers.utils.torch_utils import randn_tensor

def custom_step(
    self,
    model_output: torch.Tensor,
    old_pred_original_sample: torch.Tensor,
    timestep: int,
    timestep_back: int,
    sample: torch.Tensor,
    eta: float = 0.0,
    use_clipped_model_output: bool = False,
    generator=None,
    variance_noise: Optional[torch.Tensor] = None,
    return_dict: bool = False,
) -> Union[DDIMSchedulerOutput, Tuple]:
    """
    Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
    process from the learned model outputs (most often the predicted noise).

    Args:
        model_output (`torch.Tensor`):
            The direct output from learned diffusion model.
        timestep (`float`):
            The current discrete timestep in the diffusion chain.
        sample (`torch.Tensor`):
            A current instance of a sample created by the diffusion process.
        eta (`float`):
            The weight of noise for added noise in diffusion step.
        use_clipped_model_output (`bool`, defaults to `False`):
            If `True`, computes "corrected" `model_output` from the clipped predicted original sample. Necessary
            because predicted original sample is clipped to [-1, 1] when `self.config.clip_sample` is `True`. If no
            clipping has happened, "corrected" `model_output` would coincide with the one provided as input and
            `use_clipped_model_output` has no effect.
        generator (`torch.Generator`, *optional*):
            A random number generator.
        variance_noise (`torch.Tensor`):
            Alternative to generating noise with `generator` by directly providing the noise for the variance
            itself. Useful for methods such as [`CycleDiffusion`].
        return_dict (`bool`, *optional*, defaults to `True`):
            Whether or not to return a [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] or `tuple`.

    Returns:
        [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] or `tuple`:
            If return_dict is `True`, [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] is returned, otherwise a
            tuple is returned where the first element is the sample tensor.

    """

    if self.num_inference_steps is None:
        raise ValueError(
            "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
        )

    # See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf
    # Ideally, read DDIM paper in-detail understanding

    # Notation (<variable name> -> <name in paper>
    # - pred_noise_t -> e_theta(x_t, t)
    # - pred_original_sample -> f_theta(x_t, t) or x_0
    # - std_dev_t -> sigma_t
    # - eta -> η
    # - pred_sample_direction -> "direction pointing to x_t"
    # - pred_prev_sample -> "x_t-1"

    # 1. get previous step value (=t-1)
    prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps

    if self.step_idx == self.num_inference_steps:
        self.step_idx = 0

    # 2. compute alphas, betas
    alpha_prod_t = self.alphas_cumprod[timestep]
    alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
    alpha_prod_t_back = self.alphas_cumprod[timestep_back] if timestep_back is not None else None

    beta_prod_t = 1 - alpha_prod_t

    # 3. compute predicted original sample from predicted noise also called
    # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
    # To make style tests pass, commented out `pred_epsilon` as it is an unused variable
    if self.config.prediction_type == "epsilon":
        pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
        # pred_epsilon = model_output
    elif self.config.prediction_type == "sample":
        pred_original_sample = model_output
        # pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)
    elif self.config.prediction_type == "v_prediction":
        pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
        # pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
    else:
        raise ValueError(
            f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
            " `v_prediction`"
        )
    
    h, r, lamb, lamb_next = self.get_variables(alpha_prod_t, alpha_prod_t_prev, alpha_prod_t_back)
    mult = list(self.get_mult(h, r, alpha_prod_t, alpha_prod_t_prev, alpha_prod_t_back))
    mult_noise = (1 - alpha_prod_t_prev) ** 0.5 * (1 - (-2 * h).exp()) ** 0.5

    noise = randn_tensor(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype)
    prev_sample = mult[0] * sample - mult[1] * pred_original_sample + mult_noise * noise

    if old_pred_original_sample is None or prev_timestep < 0:
        # Save a network evaluation if all noise levels are 0 or on the first step
        self.step_idx += 1
        return prev_sample, pred_original_sample
    else:
        denoised_d = mult[2] * pred_original_sample - mult[3] * old_pred_original_sample
        noise = randn_tensor(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype)
        x_advanced = mult[0] * sample - mult[1] * denoised_d + mult_noise * noise
        prev_sample = x_advanced

    if not self.__class__.switched and self.step_idx >= 5:
        current_model_outputs = prev_sample
        prev_model_output = self.__class__.prev_model_output
        diff = torch.tanh(
                    (current_model_outputs - prev_model_output).abs().mean() /
                    prev_model_output.abs().mean()
                )
        self_diff = self.__class__.prev_diff - diff
        if 0 < self_diff < self.self_diff_threshold or self.step_idx == 30:
            logging.info(f"Switched to rendering model at Step {self.step_idx}!")
            self.__class__.parent_pipe.rendering_model.use_image_rotary_emb = False
            self.__class__.parent_pipe.transformer = self.__class__.parent_pipe.rendering_model

            def patch_forward(cls, method_name, arg_name="image_rotary_emb"):
                orig_method = getattr(cls, method_name)
                def new_method(*args, **kwargs):
                    kwargs[arg_name] = None
                    return orig_method(*args, **kwargs)
                setattr(cls, method_name, new_method)

            patch_forward(self.__class__.parent_pipe, "transformer")

            self.__class__.parent_pipe.scheduler.alphas_cumprod = self.__class__.parent_pipe.rendering_scheduler.alphas_cumprod
            self.__class__.parent_pipe.scheduler.final_alpha_cumprod = self.__class__.parent_pipe.rendering_scheduler.final_alpha_cumprod
            self.__class__.switched = True
        self.__class__.prev_diff = diff
    
    self.__class__.prev_model_output = prev_sample

    if not return_dict:
        self.step_idx += 1
        return (prev_sample, pred_original_sample)

    self.step_idx += 1
    return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)


def generate_video(
    prompt: str,
    model_path: str,
    rendering_model_path: str,
    lora_path: str = None,
    lora_rank: int = 128,
    num_frames: int = 81,
    width: Optional[int] = None,
    height: Optional[int] = None,
    output_path: str = "./",
    image_or_video_path: str = "",
    num_inference_steps: int = 50,
    guidance_scale: float = 6.0,
    num_videos_per_prompt: int = 1,
    dtype: torch.dtype = torch.bfloat16,
    generate_type: str = Literal["t2v", "i2v", "v2v"],  # i2v: image to video, v2v: video to video
    seed: int = 42,
    fps: int = 16,
    self_diff_threshold: float = 0.01,
    force_regenerate: bool = False,
):
    """
    Generates a video based on the given prompt and saves it to the specified path.

    Parameters:
    - prompt (str): The description of the video to be generated.
    - model_path (str): The path of the pre-trained model to be used.
    - lora_path (str): The path of the LoRA weights to be used.
    - lora_rank (int): The rank of the LoRA weights.
    - output_path (str): The path where the generated video will be saved.
    - num_inference_steps (int): Number of steps for the inference process. More steps can result in better quality.
    - num_frames (int): Number of frames to generate. CogVideoX1.0 generates 49 frames for 6 seconds at 8 fps, while CogVideoX1.5 produces either 81 or 161 frames, corresponding to 5 seconds or 10 seconds at 16 fps.
    - width (int): The width of the generated video, applicable only for CogVideoX1.5-5B-I2V
    - height (int): The height of the generated video, applicable only for CogVideoX1.5-5B-I2V
    - guidance_scale (float): The scale for classifier-free guidance. Higher values can lead to better alignment with the prompt.
    - num_videos_per_prompt (int): Number of videos to generate per prompt.
    - dtype (torch.dtype): The data type for computation (default is torch.bfloat16).
    - generate_type (str): The type of video generation (e.g., 't2v', 'i2v', 'v2v').·
    - seed (int): The seed for reproducibility.
    - fps (int): The frames per second for the generated video.
    """

    # 1.  Load the pre-trained CogVideoX pipeline with the specified precision (bfloat16).
    # add device_map="balanced" in the from_pretrained function and remove the enable_model_cpu_offload()
    # function to use Multi GPUs.

    image = None
    video = None

    model_name = model_path.split("/")[-1].lower()
    desired_resolution = RESOLUTION_MAP[model_name]
    if width is None or height is None:
        height, width = desired_resolution
        logging.info(
            f"\033[1mUsing default resolution {desired_resolution} for {model_name}\033[0m"
        )
    elif (height, width) != desired_resolution:
        if generate_type == "i2v":
            # For i2v models, use user-defined width and height
            logging.warning(
                f"\033[1;31mThe width({width}) and height({height}) are not recommended for {model_name}. The best resolution is {desired_resolution}.\033[0m"
            )
        else:
            # Otherwise, use the recommended width and height
            logging.warning(
                f"\033[1;31m{model_name} is not supported for custom resolution. Setting back to default resolution {desired_resolution}.\033[0m"
            )
            height, width = desired_resolution

    if generate_type == "i2v":
        pipe = CogVideoXImageToVideoPipeline.from_pretrained(model_path, torch_dtype=dtype)
        image = load_image(image=image_or_video_path)
    elif generate_type == "t2v":
        pipe = CogVideoXPipeline.from_pretrained(model_path, torch_dtype=dtype)
    else:
        pipe = CogVideoXVideoToVideoPipeline.from_pretrained(model_path, torch_dtype=dtype)
        video = load_video(image_or_video_path)

    # If you're using with lora, add this code
    if lora_path:
        pipe.load_lora_weights(
            lora_path, weight_name="pytorch_lora_weights.safetensors", adapter_name="test_1"
        )
        pipe.fuse_lora(components=["transformer"], lora_scale=1.0)

    # 2. Set Scheduler.
    # Can be changed to `CogVideoXDPMScheduler` or `CogVideoXDDIMScheduler`.
    # We recommend using `CogVideoXDDIMScheduler` for CogVideoX-2B.
    # using `CogVideoXDPMScheduler` for CogVideoX-5B / CogVideoX-5B-I2V.

    # pipe.scheduler = CogVideoXDDIMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
    
    pipe.scheduler = CogVideoXDPMScheduler.from_config(
        pipe.scheduler.config, timestep_spacing="trailing"
    )
    config = dict(pipe.scheduler.config)
    config["snr_shift_scale"] = 3.0
    pipe.rendering_scheduler = CogVideoXDPMScheduler.from_config(
        config, timestep_spacing="trailing"
    )

    # 3. Enable CPU offload for the model.
    # turn off if you have multiple GPUs or enough GPU memory(such as H100) and it will cost less time in inference
    # and enable to("cuda")
    pipe.to("cuda")

    pipe.sketching_model = pipe.transformer
    pipe.rendering_model = CogVideoXTransformer3DModel.from_pretrained(rendering_model_path, 
                                    subfolder="transformer", torch_dtype=dtype)
    pipe.rendering_model.to("cuda")

    pipe.vae = AutoencoderKLCogVideoX.from_pretrained(rendering_model_path, 
                                    subfolder="vae", torch_dtype=dtype)
    pipe.vae.to("cuda")
    pipe.vae_scaling_factor_image = pipe.vae.config.scaling_factor

    # pipe.enable_model_cpu_offload()
    # pipe.enable_sequential_cpu_offload()
    # pipe.vae.enable_slicing()
    # pipe.vae.enable_tiling()

    pipe.scheduler.__class__.step = custom_step
    pipe.scheduler.__class__.self_diff_threshold = self_diff_threshold
    pipe.scheduler.__class__.parent_pipe = pipe

    out_file_names = None
    if os.path.splitext(prompt)[-1] == ".txt":
        with open(prompt, "r") as f:
            all_prompts = [l.strip() for l in f.readlines()]
    elif os.path.splitext(prompt)[-1] == ".json":
        with open(prompt, "r") as f:
            json_prompts = json.load(f)
        all_prompts = [prompt_info['new_prompt'] for prompt_info in json_prompts['prompts']]
        out_file_names = [prompt_info['orig_prompt'] for prompt_info in json_prompts['prompts']]
    else:
        all_prompts = [prompt]

    rank = int(os.getenv("RANK", 0))
    world_size = int(os.getenv("WORLD_SIZE", 1))
    local_rank = int(os.getenv("LOCAL_RANK", 0))
    nproc_per_node = int(os.getenv("NPROC_PER_NODE", 1))
    if world_size > 1 or nproc_per_node > 1:
        logging.info(f"Generating {len(all_prompts)} prompts in parallel.")
        num_per_node = (len(all_prompts)+world_size-1) // world_size
        start = num_per_node*rank
        if start >=  len(all_prompts):
            all_prompts = []
        else:
            end = start+num_per_node
            if end > len(all_prompts):
                end = len(all_prompts)
            all_prompts = all_prompts[start:end]
            if out_file_names is not None:
                out_file_names = out_file_names[start:end]

        num_per_proc = (num_per_node+nproc_per_node-1) // nproc_per_node
        start = num_per_proc*local_rank
        if start >=  len(all_prompts):
            all_prompts = []
        else:
            end = start+num_per_proc
            if end > len(all_prompts):
                end = len(all_prompts)
            all_prompts = all_prompts[start:end]
            if out_file_names is not None:
                out_file_names = out_file_names[start:end]

    for idx, prompt in enumerate(all_prompts):
        logging.info(f"Input prompt: {prompt}")
        if out_file_names is not None:
            out_file = out_file_names[idx]
        else:
            out_file = prompt
        for index in range(num_videos_per_prompt):

            logging.info(f"Generating video ... for seed {seed + index}")
            suffix = '.mp4'
            os.makedirs(output_path, exist_ok=True)
            save_file = os.path.join(output_path, f'{out_file}-{index}' + suffix)
            # prevent file name too long 
            #if len(prompt.encode()) > 200:
            if len(out_file.encode()) > 230:
                save_file = os.path.join(output_path, f'{out_file[:20]}-{index}' + suffix)

            if args.force_regenerate or (not os.path.exists(save_file)):

                pipe.scheduler.__class__.prev_model_output = None
                pipe.scheduler.__class__.switched = False
                pipe.scheduler.__class__.step_idx = 0
                pipe.scheduler.__class__.prev_diff = 0
                pipe.transformer = pipe.sketching_model 

                # 4. Generate the video frames based on the prompt.
                # `num_frames` is the Number of frames to generate.
                if generate_type == "i2v":
                    video_generate = pipe(
                        height=height,
                        width=width,
                        prompt=prompt,
                        image=image,
                        # The path of the image, the resolution of video will be the same as the image for CogVideoX1.5-5B-I2V, otherwise it will be 720 * 480
                        num_videos_per_prompt=num_videos_per_prompt,  # Number of videos to generate per prompt
                        num_inference_steps=num_inference_steps,  # Number of inference steps
                        num_frames=num_frames,  # Number of frames to generate
                        use_dynamic_cfg=True,  # This id used for DPM scheduler, for DDIM scheduler, it should be False
                        guidance_scale=guidance_scale,
                        generator=torch.Generator().manual_seed(seed),  # Set the seed for reproducibility
                    ).frames[0]
                elif generate_type == "t2v":
                    video_generate = pipe( 
                        height=height,
                        width=width,
                        prompt=prompt,
                        num_videos_per_prompt=1,
                        num_inference_steps=num_inference_steps,
                        num_frames=num_frames,
                        use_dynamic_cfg=True,
                        guidance_scale=guidance_scale,
                        generator=torch.Generator().manual_seed(seed+index),
                    ).frames[0]
                else:
                    video_generate = pipe(
                        height=height,
                        width=width,
                        prompt=prompt,
                        video=video,  # The path of the video to be used as the background of the video
                        num_videos_per_prompt=num_videos_per_prompt,
                        num_inference_steps=num_inference_steps,
                        num_frames=num_frames,
                        use_dynamic_cfg=True,
                        guidance_scale=guidance_scale,
                        generator=torch.Generator().manual_seed(seed),  # Set the seed for reproducibility
                    ).frames[0]

                export_to_video(video_generate, save_file, fps=fps)
            else:
                logging.info(f"Video {save_file} already exists, skipping generation.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Generate a video from a text prompt using CogVideoX"
    )
    parser.add_argument(
        "--prompt", type=str, required=True, help="The description of the video to be generated"
    )
    parser.add_argument(
        "--image_or_video_path",
        type=str,
        default=None,
        help="The path of the image to be used as the background of the video",
    )
    parser.add_argument(
        "--model_path",
        type=str,
        default="THUDM/CogVideoX1.5-5B",
        help="Path of the pre-trained model use",
    )
    parser.add_argument(
        "--rendering_model_path",
        type=str,
        default="THUDM/CogVideoX1.5-2B",
        help="Path of the pre-trained model use",
    )
    parser.add_argument(
        "--lora_path", type=str, default=None, help="The path of the LoRA weights to be used"
    )
    parser.add_argument("--lora_rank", type=int, default=128, help="The rank of the LoRA weights")
    parser.add_argument(
        "--output_path", type=str, default="./", help="The path save generated video"
    )
    parser.add_argument(
        "--guidance_scale", type=float, default=6.0, help="The scale for classifier-free guidance"
    )
    parser.add_argument("--num_inference_steps", type=int, default=50, help="Inference steps")
    parser.add_argument(
        "--num_frames", type=int, default=81, help="Number of steps for the inference process"
    )
    parser.add_argument("--width", type=int, default=None, help="The width of the generated video")
    parser.add_argument(
        "--height", type=int, default=None, help="The height of the generated video"
    )
    parser.add_argument(
        "--fps", type=int, default=16, help="The frames per second for the generated video"
    )
    parser.add_argument(
        "--num_videos_per_prompt",
        type=int,
        default=1,
        help="Number of videos to generate per prompt",
    )
    parser.add_argument(
        "--generate_type", type=str, default="t2v", help="The type of video generation"
    )
    parser.add_argument(
        "--self_diff_threshold", type=float, default=0.01
    )
    parser.add_argument(
        "--dtype", type=str, default="bfloat16", help="The data type for computation"
    )
    parser.add_argument("--seed", type=int, default=42, help="The seed for reproducibility")
    parser.add_argument(
        "--force_regenerate",
        action="store_true",
        default=False,
        help="turn on for force regenerate, even already find video saved in save_path")

    args = parser.parse_args()
    dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16
    generate_video(
        prompt=args.prompt,
        model_path=args.model_path,
        rendering_model_path=args.rendering_model_path,
        lora_path=args.lora_path,
        lora_rank=args.lora_rank,
        output_path=args.output_path,
        num_frames=args.num_frames,
        width=args.width,
        height=args.height,
        image_or_video_path=args.image_or_video_path,
        num_inference_steps=args.num_inference_steps,
        guidance_scale=args.guidance_scale,
        num_videos_per_prompt=args.num_videos_per_prompt,
        dtype=dtype,
        generate_type=args.generate_type,
        seed=args.seed,
        fps=args.fps,
        self_diff_threshold=args.self_diff_threshold,
        force_regenerate=args.force_regenerate,
    )
