t2v_enhanced/model/video_ldm.py

from pathlib import Path
from typing import Any, Optional, Union, Callable

import pytorch_lightning as pl
import torch
from diffusers import DDPMScheduler, DiffusionPipeline, AutoencoderKL, DDIMScheduler
from diffusers.utils.import_utils import is_xformers_available
from einops import rearrange, repeat

from transformers import CLIPTextModel, CLIPTokenizer
from t2v_enhanced.utils.video_utils import ResultProcessor, save_videos_grid, video_naming

from t2v_enhanced.model import pl_module_params_controlnet

from t2v_enhanced.model.diffusers_conditional.models.controlnet.controlnet import ControlNetModel
from t2v_enhanced.model.diffusers_conditional.models.controlnet.unet_3d_condition import UNet3DConditionModel
from t2v_enhanced.model.diffusers_conditional.models.controlnet.pipeline_text_to_video_w_controlnet_synth import TextToVideoSDPipeline

from t2v_enhanced.model.diffusers_conditional.models.controlnet.processor import set_use_memory_efficient_attention_xformers
from t2v_enhanced.model.diffusers_conditional.models.controlnet.mask_generator import MaskGenerator

import warnings
# from warnings import warn
from t2v_enhanced.utils.iimage import IImage
from t2v_enhanced.utils.object_loader import instantiate_object
from t2v_enhanced.utils.object_loader import get_class


class VideoLDM(pl.LightningModule):

    def __init__(self,
                 inference_params: pl_module_params_controlnet.InferenceParams,
                 opt_params: pl_module_params_controlnet.OptimizerParams = None,
                 unet_params: pl_module_params_controlnet.UNetParams = None,
                 ):
        super().__init__()

        self.inference_generator = torch.Generator(device=self.device)

        self.opt_params = opt_params
        self.unet_params = unet_params

        print(f"Base pipeline from: {unet_params.pipeline_repo}")
        print(f"Pipeline class {unet_params.pipeline_class}")
        # load entire pipeline (unet, vq, text encoder,..)
        state_dict_control_model = None
        state_dict_fusion = None
        state_dict_base_model = None

        if len(opt_params.load_trained_controlnet_from_ckpt) > 0:
            state_dict_ckpt = torch.load(opt_params.load_trained_controlnet_from_ckpt, map_location=torch.device("cpu"))
            state_dict_ckpt = state_dict_ckpt["state_dict"]
            state_dict_control_model = dict(filter(lambda x: x[0].startswith("unet"), state_dict_ckpt.items()))
            state_dict_control_model = {k.split("unet.")[1]: v for (k, v) in state_dict_control_model.items()}

            state_dict_fusion = dict(filter(lambda x: "cross_attention_merger" in x[0], state_dict_ckpt.items()))
            state_dict_fusion = {k.split("base_model.")[1]: v for (k, v) in state_dict_fusion.items()}
            del state_dict_ckpt

        state_dict_proj = None
        state_dict_ckpt = None

        if hasattr(unet_params, "use_resampler") and unet_params.use_resampler:
            num_queries = unet_params.num_frames if unet_params.num_frames > 1 else None
            if unet_params.use_image_tokens_ctrl:
                num_queries = unet_params.num_control_input_frames
                assert unet_params.frame_expansion == "none"
            image_encoder = self.unet_params.image_encoder
            embedding_dim = image_encoder.embedding_dim

            resampler = instantiate_object(self.unet_params.resampler_cls, video_length=num_queries, embedding_dim=embedding_dim, input_tokens=image_encoder.num_tokens, num_layers=self.unet_params.resampler_merging_layers, aggregation=self.unet_params.aggregation)

            state_dict_proj = None

            self.resampler = resampler
            self.image_encoder = image_encoder


        noise_scheduler = DDPMScheduler.from_pretrained(self.unet_params.pipeline_repo, subfolder="scheduler")
        tokenizer = CLIPTokenizer.from_pretrained(self.unet_params.pipeline_repo, subfolder="tokenizer")
        text_encoder = CLIPTextModel.from_pretrained(self.unet_params.pipeline_repo, subfolder="text_encoder")
        vae = AutoencoderKL.from_pretrained(self.unet_params.pipeline_repo, subfolder="vae")
        base_model = UNet3DConditionModel.from_pretrained(self.unet_params.pipeline_repo, subfolder="unet", low_cpu_mem_usage=False, device_map=None, merging_mode=self.unet_params.merging_mode_base, use_image_embedding=unet_params.use_resampler and unet_params.use_image_tokens_main, use_fps_conditioning=self.opt_params.use_fps_conditioning, unet_params=unet_params)

        if state_dict_base_model is not None:
            miss, unex = base_model.load_state_dict(state_dict_base_model, strict=False)
            assert len(unex) == 0
            if len(miss) > 0:
                warnings.warn(f"Missing keys when loading base_mode:{miss}")
            del state_dict_base_model
        if state_dict_fusion is not None:
            miss, unex = base_model.load_state_dict(state_dict_fusion, strict=False)
            assert len(unex) == 0
            del state_dict_fusion

        print("PIPE LOADING DONE")
        self.noise_scheduler = noise_scheduler
        self.tokenizer = tokenizer
        self.text_encoder = text_encoder
        self.vae = vae

        self.unet = ControlNetModel.from_unet(
            unet=base_model,
            conditioning_embedding_out_channels=unet_params.conditioning_embedding_out_channels,
            downsample_controlnet_cond=unet_params.downsample_controlnet_cond,
            num_frames=unet_params.num_frames if (unet_params.frame_expansion != "none" or self.unet_params.use_controlnet_mask) else unet_params.num_control_input_frames,
            num_frame_conditioning=unet_params.num_control_input_frames,
            frame_expansion=unet_params.frame_expansion,
            pre_transformer_in_cond=unet_params.pre_transformer_in_cond,
            num_tranformers=unet_params.num_tranformers,
            vae=AutoencoderKL.from_pretrained(self.unet_params.pipeline_repo, subfolder="vae"),
            zero_conv_mode=unet_params.zero_conv_mode,
            merging_mode=unet_params.merging_mode,
            condition_encoder=unet_params.condition_encoder,
            use_controlnet_mask=unet_params.use_controlnet_mask,
            use_image_embedding=unet_params.use_resampler and unet_params.use_image_tokens_ctrl,
            unet_params=unet_params,
            use_image_encoder_normalization=unet_params.use_image_encoder_normalization,
        )
        if state_dict_control_model is not None:
            miss, unex = self.unet.load_state_dict(
                state_dict_control_model, strict=False)
            if len(miss) > 0:
                print("WARNING: Loading checkpoint for controlnet misses states")
                print(miss)

        if unet_params.frame_expansion == "none":
            attention_params = self.unet_params.attention_mask_params
            assert not attention_params.temporal_self_attention_only_on_conditioning and not attention_params.spatial_attend_on_condition_frames and not attention_params.temp_attend_on_neighborhood_of_condition_frames

        self.mask_generator = MaskGenerator(
            self.unet_params.attention_mask_params, num_frame_conditioning=self.unet_params.num_control_input_frames, num_frames=self.unet_params.num_frames)
        self.mask_generator_base = MaskGenerator(
            self.unet_params.attention_mask_params_base, num_frame_conditioning=self.unet_params.num_control_input_frames, num_frames=self.unet_params.num_frames)

        if state_dict_proj is not None and unet_params.use_image_tokens_main:
            if unet_params.use_image_tokens_main:
                missing, unexpected = base_model.load_state_dict(
                    state_dict_proj, strict=False)
            elif unet_params.use_image_tokens_ctrl:
                missing, unexpected = unet.load_state_dict(
                    state_dict_proj, strict=False)
            assert len(unexpected) == 0, f"Unexpected entries {unexpected}"
            print(f"Missing keys state proj = {missing}")
            del state_dict_proj

        base_model.requires_grad_(False)
        self.base_model = base_model
        self.unet.requires_grad_(False)
        self.text_encoder.requires_grad_(False)
        self.vae.requires_grad_(False)

        layers_config = opt_params.layers_config
        layers_config.set_requires_grad(self)

        print("CUSTOM XFORMERS ATTENTION USED.")
        if is_xformers_available():
            set_use_memory_efficient_attention_xformers(self.unet, num_frame_conditioning=self.unet_params.num_control_input_frames,
                                                        num_frames=self.unet_params.num_frames,
                                                        attention_mask_params=self.unet_params.attention_mask_params
                                                        )
            set_use_memory_efficient_attention_xformers(self.base_model, num_frame_conditioning=self.unet_params.num_control_input_frames,
                                                        num_frames=self.unet_params.num_frames,
                                                        attention_mask_params=self.unet_params.attention_mask_params_base)

        if len(inference_params.scheduler_cls) > 0:
            inf_scheduler_class = get_class(inference_params.scheduler_cls)
        else:
            inf_scheduler_class = DDIMScheduler

        inf_scheduler = inf_scheduler_class.from_pretrained(
            self.unet_params.pipeline_repo, subfolder="scheduler")
        inference_pipeline = TextToVideoSDPipeline(vae=self.vae,
                                                   text_encoder=self.text_encoder,
                                                   tokenizer=self.tokenizer,
                                                   unet=self.base_model,
                                                   controlnet=self.unet,
                                                   scheduler=inf_scheduler
                                                   )

        inference_pipeline.set_noise_generator(self.opt_params.noise_generator)
        inference_pipeline.enable_vae_slicing()

        inference_pipeline.set_progress_bar_config(disable=True)

        self.inference_params = inference_params
        self.inference_pipeline = inference_pipeline

        self.result_processor = ResultProcessor(fps=self.inference_params.frame_rate, n_frames=self.inference_params.video_length)

    def on_start(self):
        datamodule = self.trainer._data_connector._datahook_selector.datamodule
        pipe_id_model = self.unet_params.pipeline_repo
        for dataset_key in ["video_dataset", "image_dataset", "predict_dataset"]:
            dataset = getattr(datamodule, dataset_key, None)
            if dataset is not None and hasattr(dataset, "model_id"):
                pipe_id_data = dataset.model_id
                assert pipe_id_model == pipe_id_data, f"Model and Dataloader need the same pipeline path. Found '{pipe_id_model}' and '{dataset_key}.model_id={pipe_id_data}'. Consider setting '--data.{dataset_key}.model_id={pipe_id_data}'"
        self.result_processor.set_logger(self.logger)

    def on_predict_start(self) -> None:
        self.on_start()
        # pipe = DiffusionPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b", torch_dtype=torch.float16, variant="fp16")
        # pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
        # pipe.set_progress_bar_config(disable=True)
        # self.first_stage = pipe.to(self.device)

    def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
        cfg = self.trainer.predict_cfg

        result_file_stem = cfg["result_file_stem"]
        storage_fol = Path(cfg['predict_dir'])
        prompts = [cfg["prompt"]]

        inference_params: pl_module_params_controlnet.InferenceParams = self.inference_params
        conditioning_type = inference_params.conditioning_type
        # n_autoregressive_generations = inference_params.n_autoregressive_generations
        n_autoregressive_generations = cfg["n_autoregressive_generations"]
        mode = inference_params.mode
        start_from_real_input = inference_params.start_from_real_input
        assert isinstance(prompts, list)

        prompts = n_autoregressive_generations * prompts

        self.inference_generator.manual_seed(self.inference_params.seed)

        assert self.unet_params.num_control_input_frames == self.inference_params.video_length//2, f"currently we assume to have an equal size for and second half of the frame interval, e.g. 16 frames, and we condition on 8. Current setup: {self.unet_params.num_frame_conditioning} and {self.inference_params.video_length}"

        chunks_conditional = []
        batch_size = 1
        shape = (batch_size, self.inference_pipeline.unet.config.in_channels, self.inference_params.video_length,
                 self.inference_pipeline.unet.config.sample_size, self.inference_pipeline.unet.config.sample_size)
        for idx, prompt in enumerate(prompts):
            if idx > 0:
                content = sample*2-1
                content_latent = self.vae.encode(content).latent_dist.sample() * self.vae.config.scaling_factor
                content_latent = rearrange(content_latent, "F C W H -> 1 C F W H")
                content_latent = content_latent[:, :, self.unet_params.num_control_input_frames:].detach().clone()

            if hasattr(self.inference_pipeline, "noise_generator"):
                latents = self.inference_pipeline.noise_generator.sample_noise(shape=shape, device=self.device, dtype=self.dtype, generator=self.inference_generator, content=content_latent if idx > 0 else None)
            else:
                latents = None
            if idx == 0:
                sample = cfg["video"]
            else:
                if inference_params.conditioning_type == "fixed":
                    context = chunks_conditional[0][:self.unet_params.num_frame_conditioning]
                    context = [context]
                    context = [2*sample-1 for sample in context]

                    input_frames_conditioning = torch.cat(context).detach().clone()
                    input_frames_conditioning = rearrange(input_frames_conditioning, "F C W H -> 1 F C W H")
                elif inference_params.conditioning_type == "last_chunk":
                    input_frames_conditioning = condition_input[:, -self.unet_params.num_frame_conditioning:].detach().clone()
                elif inference_params.conditioning_type == "past":
                    context = [sample[:self.unet_params.num_control_input_frames] for sample in chunks_conditional]
                    context = [2*sample-1 for sample in context]

                    input_frames_conditioning = torch.cat(context).detach().clone()
                    input_frames_conditioning = rearrange(input_frames_conditioning, "F C W H -> 1 F C W H")
                else:
                    raise NotImplementedError()

                input_frames = condition_input[:, self.unet_params.num_control_input_frames:].detach().clone()

                sample = self(prompt, input_frames=input_frames, input_frames_conditioning=input_frames_conditioning, latents=latents)

            if hasattr(self.inference_pipeline, "reset_noise_generator_state"):
                self.inference_pipeline.reset_noise_generator_state()

            condition_input = rearrange(sample, "F C W H -> 1 F C W H")
            condition_input = (2*condition_input)-1  # range: [-1,1]

            # store first 16 frames, then always last 8 of a chunk
            chunks_conditional.append(sample)

        result_formats = self.inference_params.result_formats
        # result_formats = [gif", "mp4"]
        concat_video = self.inference_params.concat_video

        def IImage_normalized(x): return IImage(x, vmin=0, vmax=1)
        for result_format in result_formats:
            save_format = result_format.replace("eval_", "")

            merged_video = None
            for chunk_idx, (prompt, video) in enumerate(zip(prompts, chunks_conditional)):
                if chunk_idx == 0:
                    current_video = IImage_normalized(video)
                else:
                    current_video = IImage_normalized(video[self.unet_params.num_control_input_frames:])

                if merged_video is None:
                    merged_video = current_video
                else:
                    merged_video &= current_video

            if concat_video:
                filename = video_naming(prompts[0], save_format, batch_idx, 0)
                result_file_video = (storage_fol / filename).absolute().as_posix()
                result_file_video = (Path(result_file_video).parent / (result_file_stem+Path(result_file_video).suffix)).as_posix()
                self.result_processor.save_to_file(video=merged_video.torch(vmin=0, vmax=1), prompt=prompts[0], video_filename=result_file_video, prompt_on_vid=False)

    def forward(self, prompt, input_frames=None, input_frames_conditioning=None, latents=None):
        call_params = self.inference_params.to_dict()
        # print(f"INFERENCE PARAMS = {call_params}")
        call_params["prompt"] = prompt

        call_params["image"] = input_frames
        call_params["num_frames"] = self.inference_params.video_length
        call_params["return_dict"] = False
        call_params["output_type"] = "pt_t2v"
        call_params["mask_generator"] = self.mask_generator
        call_params["precision"] = "16" if self.trainer.precision.startswith("16") else "32"
        call_params["no_text_condition_control"] = self.opt_params.no_text_condition_control
        call_params["weight_control_sample"] = self.unet_params.weight_control_sample
        call_params["use_controlnet_mask"] = self.unet_params.use_controlnet_mask
        call_params["skip_controlnet_branch"] = self.opt_params.skip_controlnet_branch
        call_params["img_cond_resampler"] = self.resampler if self.unet_params.use_resampler else None
        call_params["img_cond_encoder"] = self.image_encoder if self.unet_params.use_resampler else None
        call_params["input_frames_conditioning"] = input_frames_conditioning
        call_params["cfg_text_image"] = self.unet_params.cfg_text_image
        call_params["use_of"] = self.unet_params.use_of
        if latents is not None:
            call_params["latents"] = latents

        sample = self.inference_pipeline(generator=self.inference_generator, **call_params)
        return sample