import os
import torch
import comfy
import yaml
import folder_paths
import random

from einops import rearrange
from comfy import model_base, model_management, model_detection, latent_formats, model_sampling
from .lvdm.modules.networks.openaimodel3d import UNetModel as DynamiCrafterUNetModel

from .utils.model_utils import DynamiCrafterBase, DYNAMICRAFTER_CONFIG, VIDEOCRAFTER_CONFIG, \
    load_image_proj_dict, load_dynamicrafter_dict, get_image_proj_model, load_vae_dict

from .utils.utils import get_models_directory

MODEL_DIR= "dynamicrafter_models"
MODEL_DIR_PATH = os.path.join(folder_paths.models_dir, MODEL_DIR)


def handle_cfg(x, num_video_frames):
    use_cfg = x.shape[0] > num_video_frames
    num_video_frames *= 2 if use_cfg else 1
    batch_size = 1 if not use_cfg else 2

    return num_video_frames, use_cfg, batch_size

# There is probably a better way to do this, but with the apply_model callback, this seems necessary.
# The model gets wrapped around a CFG Denoiser class, and handles the conditioning parts there.
# We cannot access it, so we must find the conditioning according to how ComfyUI handles it.
def get_conditioning_pair(c_crossattn, use_cfg: bool):
    if not use_cfg:
        return c_crossattn

    conditioning_group = []

    for i in range(c_crossattn.shape[0]):
        # Get the positive and negative conditioning.
        positive_idx = i + 1
        negative_idx = i

        if positive_idx >= c_crossattn.shape[0]:
            break

        if not torch.equal(c_crossattn[[positive_idx]], c_crossattn[[negative_idx]]):
            conditioning_group = [
                c_crossattn[[positive_idx]], 
                c_crossattn[[negative_idx]]
            ]
            break

    if len(conditioning_group) == 0:
        raise ValueError("Could not get the appropriate conditioning group.")

    return torch.cat(conditioning_group)

def get_conditioning_args(in_dict: dict, required_args: list ):
    out_args = []
    for k in required_args:
        del_key = False
        if isinstance(k, tuple):
            if k[-1] == 'del':
                del_key = True
            k = k[0]
        out_args.append(in_dict.get(k, None))
        if del_key:
            del in_dict[k]
        return out_args

def expand_frames_tensor(tens: torch.Tensor, max_frames: int):
    """
    Expands the frame tensors by repeating the last tensor to the desired length.
    """
    return torch.cat(
            [tens] + [tens[-1:]] * abs(tens.shape[0] - max_frames)
        )[:max_frames]

def create_advanced_options( 
        fps_as_max_random: bool = False, 
        all_frames_as_context: bool = False, 
        force_start_frame: bool = False,
        reference_frames = None
    ):
    return {
            "fps_as_max_random": fps_as_max_random,
            "all_frames_as_context": all_frames_as_context,
            "force_start_frame": force_start_frame,
            "reference_frames": reference_frames,
        }  

def build_advanced_options(advanced_options: dict = None):
    if advanced_options is not None:
        return advanced_options.values()  
    else:
        return create_advanced_options().values()

def load_model_dicts(model_path: str, videocrafter_mode: False):
    model_state_dict = comfy.utils.load_torch_file(model_path)
    dynamicrafter_dict = load_dynamicrafter_dict(model_state_dict)

    if not videocrafter_mode:
        image_proj_dict = load_image_proj_dict(model_state_dict)
    else:
        image_proj_dict = None
        
    return dynamicrafter_dict, image_proj_dict

def get_prediction_type(is_eps: bool, model_config, video_crafter_mode: bool):
    if video_crafter_mode:
        model_config.unet_config["image_cross_attention_scale_learnable"] = False
        return model_base.ModelType.EPS

    if not is_eps and "image_cross_attention_scale_learnable" in model_config.unet_config.keys():
            model_config.unet_config["image_cross_attention_scale_learnable"] = False

    return model_base.ModelType.EPS if is_eps else model_base.ModelType.V_PREDICTION

def handle_model_management(dynamicrafter_dict: dict, model_config):
    parameters = comfy.utils.calculate_parameters(dynamicrafter_dict, "model.diffusion_model.")
    load_device = model_management.get_torch_device()
    unet_dtype = model_management.unet_dtype(
        model_params=parameters, 
        supported_dtypes=model_config.supported_inference_dtypes
    )
    manual_cast_dtype = model_management.unet_manual_cast(
        unet_dtype, 
        load_device, 
        model_config.supported_inference_dtypes
    )
    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
    inital_load_device = model_management.unet_inital_load_device(parameters, unet_dtype)
    offload_device = model_management.unet_offload_device()

    return load_device, inital_load_device

def check_leftover_keys(state_dict: dict):
    left_over = state_dict.keys()
    if len(left_over) > 0:
        print("left over keys:", left_over)

def get_fps_cond(fps: int =16, random_fps: bool = False):
    fs = torch.tensor([fps], dtype=torch.long,  device=model_management.intermediate_device())
    return fs if not random_fps else fps

class LatentConditionProcessor(object):
    def __init__(self, model, clip_vision, vae, image_proj_model=None, scale_latents=False):
        self.model = model
        self.clip_vision = clip_vision
        self.vae = vae
        self.image_proj_model = image_proj_model
        self.scale_latents = scale_latents

    def get_latent_emb(
        self, 
        images=None,
        c_concat=None, 
        norm_latent: bool = True, 
        norm_image: bool = False,
        frame_wise_emb: bool = False,
        strict: bool = False,
        all_frames: bool = False
    ):  
        # Slower, but complies with the vae processing and normalization.
        # c_concat is assumed to be an encoded latent natively through Comfy
        if strict and c_concat is not None:
            images = self.vae.decode(c_concat)
            c_concat = self.vae.encode(images[:, :, :, :3])
            images = images[:1, :, :, :3]
        else:
            if c_concat is None and images is not None:
                c_concat = self.vae.encode(images[:, :, :, :3])
            
            if images is None and c_concat is not None:
                images = self.vae.decode(c_concat[:, :, :, :3])

        if self.image_proj_model is not None:
            if norm_image:
                images = (images / 255.) * 2 - 1.0     
            
            encoded_images = []

            for idx in range(images.shape[0]):
                encoded_image = self.clip_vision.encode_image(images[[idx]])['last_hidden_state']
                encoded_images.append(encoded_image)

                if not all_frames and idx == 0:
                    break

            img_embs = self.image_proj_model(torch.cat(encoded_images))
            del encoded_images
            
        else:
            img_embs = None

        if self.scale_latents:
            if norm_latent and c_concat is None:
                vae_process_input = vae.process_input
                vae.process_input = lambda image: 2 * image - 1
                c_concat = vae.encode(images[:, :, :, :3])
                vae.process_input = vae_process_input
            c_concat = self.model.model.process_latent_in(c_concat) * 1.3
        else:
            c_concat = self.model.model.process_latent_in(c_concat)

        return c_concat, img_embs

class ReferenceProcessor:
    def __init__(
        self, 
        c_concat_ref: torch.Tensor = None, 
        img_emb_ref: torch.Tensor = None, 
        current_step: int = 0, 
        start_step: int = 1, 
        end_step: int = 20,
        strength: float = .25
    ):
        self.c_concat_ref = c_concat_ref
        self.img_emb_ref = img_emb_ref
        self.current_step = current_step
        self.start_step = start_step
        self.end_step = end_step
        self.strength = strength

    def step(self):
        self.current_step += 1

    def process(self):
        c_concat_ref = None
        img_emb_ref = None

        if self.current_step >= self.start_step and \
            self.current_step <= self.end_step:
            c_concat_ref = self.c_concat_ref
            img_emb_ref = self.img_emb_ref

        return c_concat_ref, img_emb_ref

class DynamiCrafterReferenceProcessor:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE", ),
                "start_step": ("INT", {"default": 1}),
                "end_step": ("INT", {"default": 20}),
                "strength": ("FLOAT", {
                    "default": 0.25, 
                    "min": 0., 
                    "max": 1.,
                    "step": 1e-2
                })
            },
        }
        
    CATEGORY = "Native_DynamiCrafter/Conditioning"
    RETURN_TYPES = ("DYNAMI_REFERENCE_PROCESSOR", )
    RETURN_NAMES = ("reference_frames", )

    FUNCTION = "initialize_reference_processor"
    
    def __init__(self):
        self.images = None
        self.start_step = None
        self.end_step = None
        self.strength = None

    def preprocess(self, model, vae, clip_vision, image_proj_model):
        c_concat_ref = vae.encode(self.images[:, :, :, :3])
        frame_wise_img_ref = []
        
        for i in range(self.images.shape[0]):
            encoded_image_ref = clip_vision.encode_image(self.images[[i]])['last_hidden_state']
            frame_wise_img_ref.append(encoded_image_ref)

        image_emb_ref = image_proj_model(torch.cat(frame_wise_img_ref))
        c_concat_ref = model.model.process_latent_in(c_concat_ref)
        
        reference_processor = ReferenceProcessor(
            c_concat_ref, 
            image_emb_ref,
            start_step=self.start_step,
            end_step=self.end_step,
            strength=self.strength
        )
        return reference_processor

    def initialize_reference_processor(self, images, start_step, end_step, strength):
        keys = ("images", "start_step", "end_step", "strength")
        values = (images, start_step, end_step, strength)

        for k, v in zip(keys, values):
            setattr(self, k, v)

        return ({"preprocess": self.preprocess}, )

def validate_forwardable_latent(latent, c_concat, num_video_frames, use_cfg):
    check_no_cfg = latent.shape[0] != num_video_frames
    check_with_cfg = latent.shape[0] != (num_video_frames * 2)

    latent_batch_size = latent.shape[0] if not use_cfg else latent.shape[0] // 2
    num_frames = num_video_frames if not use_cfg else num_video_frames // 2

    if all([check_no_cfg, check_with_cfg]):
        raise ValueError(
            "Please make sure your latent inputs match the number of frames in the DynamiCrafter Processor."
            f"Got a latent batch size of ({latent_batch_size}) with number of frames being ({num_frames})."
        )
    
    latent_h, latent_w = latent.shape[-2:]
    c_concat_h, c_concat_w = c_concat.shape[-2:]

    if not all([latent_h == c_concat_h, latent_w == c_concat_w]):
        raise ValueError(
            "Please make sure that your input latent and image frames are the same height and width.",
            f"Image Size: {c_concat_w * 8}, {c_concat_h * 8}, Latent Size: {latent_h * 8}, {latent_w * 8}"
        )

class DynamiAdvancedOpts:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "fps_as_max_random": ("BOOLEAN", {"default": False}),
                "all_frames_as_context": ("BOOLEAN", {"default": False}),
                "force_start_frame": ("BOOLEAN", {"default": False}),
                           },
            "optional": {
                "reference_frames": ("DYNAMI_REFERENCE_PROCESSOR", ),
            }
        }
    CATEGORY = "Native_DynamiCrafter/Options"
    RETURN_TYPES = ("DYNAMI_ADVANCED_OPTS", )
    RETURN_NAMES = ("advanced_options", )

    FUNCTION = "get_advanced_options"


    def get_advanced_options(
        self, 
        fps_as_max_random, 
        all_frames_as_context, 
        force_start_frame,
        reference_frames: DynamiCrafterReferenceProcessor.preprocess = None
    ):

        return ({
        "options": create_advanced_options(
            fps_as_max_random, 
            all_frames_as_context, 
            force_start_frame, 
            reference_frames
        )
        },)

class DynamiCrafterProcessor:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL", ),
                "clip_vision": ("CLIP_VISION", ),
                "vae": ("VAE", ),
                "image_proj_model": ("IMAGE_PROJ_MODEL", ),
                "images": ("IMAGE", ),
                "use_interpolate": ("BOOLEAN", {"default": False}),
                "img_cfg": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 1000}),
                "fps": ("INT", {"default": 15, "min": 1, "max": 1000, "step": 1}, ),
                "frames": ("INT", {"default": 16}),
                "scale_latents": ("BOOLEAN", {"default": False}),
            },
            "optional": {"advanced_options": ("DYNAMI_ADVANCED_OPTS", ) }
        }
        
    CATEGORY = "Native_DynamiCrafter/Processing"
    RETURN_TYPES = ("MODEL", "LATENT", "LATENT", )
    RETURN_NAMES = ("model", "empty_latent", "latent_img", )

    FUNCTION = "process_image_conditioning"

    def __init__(self):
        self.model_patcher = None

 
    def get_autoregressive_concat(self, transformer_options: dict, c_concat, img_embs):
        if 'c_concat_ar' in transformer_options and transformer_options['c_concat_ar'] is not None:
            c_concat = transformer_options['c_concat_ar']
            c_concat, img_embs =  self.latent_condition_processor.get_latent_emb(c_concat=c_concat, norm_image=True)
        return c_concat, img_embs

    # apply_model, {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": cond_or_uncond}
    def _forward(self, *args):
        transformer_options = self.model_patcher.model_options['transformer_options']
        conditioning = transformer_options['conditioning']
        apply_model = args[0]
        emb_sz = 1

        # Forward_dict = args[1]
        x, t, model_in_kwargs, _ =  args[1].values()

        if 'uncond_img' in conditioning.keys():
            del conditioning['uncond_img']
            
        c_concat, c_concat_ref ,img_embs, img_emb_uncond, img_cfg, fs, num_video_frames  = conditioning.values()

        c_concat, img_embs = self.get_autoregressive_concat(transformer_options, c_concat, img_embs)
        
        num_video_frames, use_cfg, batch_size = handle_cfg(x, num_video_frames)
        
        if use_cfg:
            c_concat = torch.cat([c_concat] * 2)

        validate_forwardable_latent(x, c_concat, num_video_frames, use_cfg)

        x_in, c_concat = map(lambda xc: rearrange(xc, '(b t) c h w -> b c t h w', b=batch_size), (x, c_concat))

        # We always assume video, so there will always be batched conditionings.
        c_crossattn = get_conditioning_pair(model_in_kwargs.pop('c_crossattn'), use_cfg)
        c_crossattn = c_crossattn[:2] if use_cfg else c_crossattn[:1]
        context_in = c_crossattn
        
        if use_cfg:
            img_emb_uncond = conditioning['image_emb_uncond']
            img_embs = torch.cat([img_embs, img_emb_uncond])

        
        if img_embs.shape[0] == num_video_frames:
            emb_sz = num_video_frames // 2 if batch_size > 1 else num_video_frames

        if c_concat_ref is not None:
            process_ref = c_concat_ref.process()
            strength = c_concat_ref.strength

            if all([ref is not None for ref in process_ref]):
                ir = process_ref[1].to(c_concat.device, c_concat.dtype)

                if img_embs.shape[0] == 2:
                    if ir.shape[0] > img_embs.shape[0]:
                        emb_sz = ir.shape[0] 
                        img_embs_cond = torch.cat([img_embs[:1]] * ir.shape[0])
                        img_embs_uncond = torch.zeros_like(img_embs_cond)
                        img_embs = torch.cat([img_embs_cond, img_embs_uncond], dim=0)
                        
                    else:
                        ir = torch.cat([ir] * 2)
                else:
                    if ir.shape[0] > img_embs.shape[0]:
                        # The img_emb reference should always be conditional, so no multiples for CFG.
                        emb_sz = ir.shape[0] 
                        
                        max_emb_sz = img_embs.shape[0]
                        img_emb_pad = [img_embs[-1:]] * (ir.shape[0] - max_emb_sz)
                        img_embs_cond= torch.cat([img_embs[:max_emb_sz]] + img_emb_pad)[:emb_sz]
                        img_embs_uncond = torch.zeros_like(img_embs_cond)

                        img_embs = torch.cat([img_embs, img_embs_uncond])

                _ie, _ir = map(lambda t: rearrange(t, 'b s l -> (b s) l'), (img_embs[:emb_sz], ir[:emb_sz]))
                sim = torch.nn.CosineSimilarity(dim=1)(_ie, _ir).unsqueeze(1)
                _ie = _ie * sim
                _ie = rearrange(_ie, '(b s) l -> b s l', b=emb_sz, s=img_embs.shape[1])
                ir[:emb_sz] = _ie

                del _ie
                del _ir

                img_embs[:emb_sz] = ir[:emb_sz] * strength + (1. - strength) * img_embs[:emb_sz]
 
            c_concat_ref.step()

        if isinstance(fs, int):
            fs = torch.tensor(
                [random.randint(3, fs)], 
                dtype=torch.long, 
                device=model_management.intermediate_device()
            )     
        fs = torch.cat([fs] * x_in.shape[0])
        
        outs = []
        for i in range(batch_size):
            model_in_kwargs['transformer_options']['cond_idx'] = i
            img_emb = img_embs[:emb_sz] if i == 0 else img_embs[-emb_sz:]

            x_out = apply_model(
                    x_in[[i]], 
                    t=torch.cat([t[:1]]),
                    context_in=context_in[[i]],
                    c_crossattn=c_crossattn, 
                    cc_concat=c_concat[[i]], # "cc" is to handle naming conflict with apply_model wrapper. 
                    # We want to handle this in the UNet forward.
                    num_video_frames=num_video_frames // 2 if use_cfg else num_video_frames, 
                    img_emb=img_emb,
                    fs=fs[[i]],
                    **model_in_kwargs
                )
            outs.append(x_out)
        
        if img_cfg > 1.0 and use_cfg:
            model_in_kwargs['transformer_options']['cond_idx'] = 0
            uncond_img = apply_model(
                x_in[[0]],
                t=torch.cat([t[:1]]),
                context_in=torch.zeros_like(context_in[[0]]),
                c_crossattn=torch.zeros_like(c_crossattn),
                cc_concat=c_concat[[0]],
                num_video_frames=num_video_frames // 2 if use_cfg else num_video_frames,
                img_emb=img_embs[:emb_sz],
                fs=fs[[0]],
                **model_in_kwargs 
            )
            uncond_img = rearrange(uncond_img, 'b c t h w -> (b t) c h w')

            self.model_patcher.model_options\
                ['transformer_options']['conditioning']['uncond_img'] = uncond_img

        x_out = torch.cat(list(reversed(outs)))
        x_out = rearrange(x_out, 'b c t h w -> (b t) c h w')

        return x_out

    def cfg_with_img_cond(self, args):
        conditioning = self.model_patcher.model_options["transformer_options"]["conditioning"]
        img_cfg = conditioning["img_cfg"]
        uncond_img = conditioning.get("uncond_img", None)
        uncond_pred, cond_pred, cond_scale = args["uncond"], args["cond"], args["cond_scale"]

        if img_cfg > 1.0 and uncond_img is not None:
            uncond_img = args["input"] - uncond_img
            cfg_result = uncond_pred + img_cfg * (uncond_img - uncond_pred) + cond_scale * (cond_pred - uncond_img)
        else:
            cfg_result = uncond_pred + (cond_pred - uncond_pred) * cond_scale

        return cfg_result

    def assign_forward_args(
        self, 
        c_concat, 
        image_emb, 
        image_emb_uncond, 
        img_cfg,
        fs, 
        frames,
        c_concat_ref: ReferenceProcessor = None
    ):
        self.model_patcher.model_options['transformer_options']['conditioning'] = {
            "c_concat": c_concat,
            "c_concat_ref": c_concat_ref,
            "image_emb": image_emb ,
            'image_emb_uncond': image_emb_uncond,
            "img_cfg": img_cfg,
            "fs": fs,
            "num_video_frames": frames,
        }

    def process_image_conditioning(
        self, 
        model, 
        clip_vision, 
        vae, 
        image_proj_model, 
        images, 
        use_interpolate,
        img_cfg,
        fps: int,
        frames: int,
        scale_latents: bool,
        advanced_options: dict = {}
    ):

        fps_as_max_random, all_frames_as_context, \
            force_start_frame, reference_frames = build_advanced_options(advanced_options.get("options", None))

        self.model_patcher = model
        self.latent_condition_processor = LatentConditionProcessor(
            model, 
            clip_vision, 
            vae, 
            image_proj_model, 
            scale_latents
        )

        encoded_latent = vae.encode(images[:, :, :, :3])
        c_concat, image_emb = self.latent_condition_processor.get_latent_emb(images, all_frames=all_frames_as_context)
        _, image_emb_uncond = self.latent_condition_processor.get_latent_emb(torch.zeros_like(images), all_frames=all_frames_as_context)
        fs = get_fps_cond(fps, fps_as_max_random)

        model.set_model_unet_function_wrapper(self._forward)
        model.set_model_sampler_cfg_function(self.cfg_with_img_cond)

        if reference_frames is not None:
            preprocess = reference_frames['preprocess']
            reference_frames = preprocess(model, vae, clip_vision, image_proj_model)
 
        used_interpolate_processing = False

        if use_interpolate and frames > 16:
            raise ValueError("Interpolation mode limits frames to 16.Consider using autoregressive approach for long videos.")
        if encoded_latent.shape[0] == 1:
            c_concat = torch.cat([c_concat] * frames, dim=0)[:frames]
        
            if use_interpolate:
                mask = torch.zeros_like(c_concat)
                mask[:1] = c_concat[:1]
                mask[-1:] = c_concat[:1]
                c_concat = mask
                used_interpolate_processing = True

                print("One frame was provided with use_interpolate set to True. Using start frame for both start and end.")
        else:
            if use_interpolate:
                if c_concat.shape[0] > 2:
                    print(f"To use interpolation mode,please make sure to pass a batch of exactly two images.")
                    frame_one = c_concat[[0]]

                    if c_concat.shape[0] == 1:
                        frame_two = frame_one.clone()
                    else:
                        frame_two = c_concat[[-1]]

                    c_concat = torch.cat([frame_one, frame_two])

                input_frame_count = c_concat.shape[0]

                # We're just padding to the same type an size of the concat
                masked_frames = torch.zeros_like(torch.cat([c_concat[:1]] * frames))[:frames]

                # Start frame
                masked_frames[:1] = c_concat[:1]

                end_frame_idx = -1

                # End frame
                masked_frames[-1:] = c_concat[[end_frame_idx]]
    
                c_concat = masked_frames
                used_interpolate_processing = True

                print(f"Using interpolation mode with {input_frame_count} frames.")

            if c_concat.shape[0] < frames and not used_interpolate_processing:

                if all_frames_as_context:
                    c_concat = expand_frames_tensor(c_concat, frames)
                else:
                    print(
                        "Multiple images found, but interpolation mode is unset. Using the first frame as condition.",
                    )
                    c_concat = torch.cat([c_concat[:1]] * frames)

        c_concat = c_concat[:frames]
        
        if encoded_latent.shape[0] == 1:
            encoded_latent = torch.cat([encoded_latent] * frames)[:frames]

        if encoded_latent.shape[0] < frames and encoded_latent.shape[0] != 1:
            encoded_latent = expand_frames_tensor(encoded_latent, frames)
        
        # We could store this as a state in this Node Class Instance, but to prevent any weird edge cases,
        # this should always be passed through the 'stateless' way, and let ComfyUI handle the transformer_options state.
        self.assign_forward_args(c_concat, image_emb, image_emb_uncond, img_cfg, fs, frames, reference_frames)
        
        out_zeros, out = {"samples": torch.zeros_like(c_concat)}, {"samples": encoded_latent}

        if force_start_frame:
            mask = out_zeros['samples'][:, :1, ...].clone()
            mask[1:] = 1.
            mask[:1] = 0. # Add a little noise to initialize generation.
            out.update({"noise_mask": mask}) 

        return (model, out_zeros, out)

class DynamiCrafterLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model_path": (get_models_directory(os.listdir(MODEL_DIR_PATH)), ),
                "videocrafter_mode": ("BOOLEAN", {"default": False})
            },
        }
        
    CATEGORY = "Native_DynamiCrafter/Loaders"
    RETURN_TYPES = ("MODEL", "IMAGE_PROJ_MODEL", )
    RETURN_NAMES = ("model", "image_proj_model", )
    FUNCTION = "load_dynamicrafter"
     
    def load_dynamicrafter(self, model_path, videocrafter_mode):
        model_path = os.path.join(MODEL_DIR_PATH, model_path)
        
        if os.path.exists(model_path):
            dynamicrafter_dict, image_proj_dict = load_model_dicts(model_path, videocrafter_mode)
            CONFIG = DYNAMICRAFTER_CONFIG if not videocrafter_mode else VIDEOCRAFTER_CONFIG
            model_config = DynamiCrafterBase(CONFIG)

            dynamicrafter_dict, is_eps = model_config.process_dict_version(state_dict=dynamicrafter_dict)
            MODEL_TYPE = get_prediction_type(is_eps, model_config, videocrafter_mode)

            load_device, inital_load_device = handle_model_management(dynamicrafter_dict, model_config)

            model = model_base.BaseModel(
                model_config, 
                model_type=MODEL_TYPE, 
                device=inital_load_device, 
                unet_model=DynamiCrafterUNetModel
            )

            if not videocrafter_mode:
                image_proj_model = get_image_proj_model(image_proj_dict)
                model.load_model_weights(dynamicrafter_dict, "model.diffusion_model.")
            else:
                sd = {k.replace("model.diffusion_model.", ""): v for k, v in dynamicrafter_dict.items()}
                model.diffusion_model.load_state_dict(sd, strict=False)
                dynamicrafter_dict = {}
                image_proj_model = None
        
            check_leftover_keys(dynamicrafter_dict)

            model_patcher = comfy.model_patcher.ModelPatcher(
                model, 
                load_device=load_device, 
                offload_device=model_management.unet_offload_device(), 
                current_device=inital_load_device
            )

            model_patcher.model.diffusion_model.videocrafter_mode = videocrafter_mode
    
        return (model_patcher, image_proj_model, )

class VideoCrafterProcessor:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL", ),
                "frames": ("INT", {"default": 16, "min": 1, "max": 133333337}),
                "fps": ("INT", {"default": 15, "min": 1, "max": 1000, "step": 1}, ),
            }
        }
        
    CATEGORY = "Native_DynamiCrafter/Processing"
    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("model",)

    FUNCTION = "process_fps"
    
    def __init__(self):
        self.model = None

    def _forward(self, *args):
        apply_model = args[0]

        # forward_dict
        fd = args[1]
        
        x, t, model_in_kwargs, cond_or_uncond =  fd['input'], fd['timestep'], fd['c'], fd['cond_or_uncond']
        num_video_frames = self.model.model_options['transformer_options']['num_video_frames']
        
        num_video_frames, use_cfg, batch_size = handle_cfg(x, num_video_frames)

        x = rearrange(x, '(b t) c h w -> b c t h w', b=batch_size)
    
        c_crossattn = get_conditioning_pair(model_in_kwargs.pop('c_crossattn'), use_cfg)
        
        c_crossattn = c_crossattn[:2] if use_cfg else c_crossattn[:1]
        outs = []

        for i in range(batch_size):
            model_in_kwargs['transformer_options']['cond_idx'] = i
            nvs = num_video_frames // 2 if batch_size > 1 else num_video_frames
            x_out = apply_model(
                x[[i]], 
                torch.cat([t[:1]]), 
                c_crossattn=torch.cat([c_crossattn[[i]]] * nvs), 
                num_video_frames=nvs, 
                **model_in_kwargs
            )
            outs.append(x_out)

        x_out = torch.cat(list(reversed(outs)))
        x_out = rearrange(x_out, 'b c t h w -> (b t) c h w')

        return x_out

    def process_fps(self, model, fps, frames):
        fs_t = torch.tensor([fps], dtype=torch.long, device=model_management.intermediate_device())
        model.model_options['transformer_options']['fs_t'] = fs_t
        model.model_options['transformer_options']['num_video_frames'] = frames
        model.set_model_unet_function_wrapper(self._forward)
        self.model = model
        return (model, )

NODE_CLASS_MAPPINGS = {
    "DynamiCrafterLoader": DynamiCrafterLoader,
    "DynamiCrafterProcessor": DynamiCrafterProcessor,
    "DynamiCrafterReferenceProcessor": DynamiCrafterReferenceProcessor,
    "VideoCrafterProcessor": VideoCrafterProcessor,
    "DynamiAdvancedOpts": DynamiAdvancedOpts
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "DynamiCrafterLoader": "Load a DynamiCrafter Checkpoint",    
    "DynamiCrafterProcessor": "Apply DynamiCrafter",
    "DynamiCrafterReferenceProcessor": "Apply Reference Frames [EXPERIMENTAL]",
    "VideoCrafterProcessor": "Apply VideoCrafter",
    "DynamiAdvancedOpts": "Dynamicrafter Advanced Options"
}