nodes.py

import os
import torch
import comfy
import yaml
import folder_paths

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, \
    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)

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}),
                "fps": ("INT", {"default": 15, "min": 1, "max": 30, "step": 1}, ),
                "frames": ("INT", {"default": 16}),
                "scale_latents": ("BOOLEAN", {"default": False})
            },
        }
        
    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

    # 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(self, 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)

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

        # forward_dict
        fd = args[1]
        
        x, t, model_in_kwargs, _ =  fd['input'], fd['timestep'], fd['c'], fd['cond_or_uncond']
        
        c_crossattn = model_in_kwargs.pop("c_crossattn")
        c_concat = conditioning['c_concat']
        num_video_frames = conditioning['num_video_frames']
        fs = conditioning['fs']

        
        original_num_frames = num_video_frames

        # Better way to determine if we're using CFG
        # The cond batch will always be num_frames >= 2 since we're doing video,
        # so we need get this condition differently here. 
        if x.shape[0] > num_video_frames:
            num_video_frames *= 2
            batch_size = 2
            use_cfg = True
        else:
            use_cfg = False
            batch_size = 1
        
        if use_cfg:
            c_concat = torch.cat([c_concat] * 2)

        self.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 = self.get_conditioning_pair(c_crossattn, use_cfg)
        c_crossattn = c_crossattn[:2] if use_cfg else c_crossattn[:1]
        context_in = c_crossattn

        img_embs = conditioning['image_emb']

        if use_cfg:
            img_emb_uncond = conditioning['image_emb_uncond']
            img_embs = torch.cat([img_embs, img_emb_uncond])

        fs = torch.cat([fs] * x_in.shape[0])
        
        outs = []
        for i in range(batch_size):
            model_in_kwargs['transformer_options']['cond_idx'] = i
            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 batch_size > 1 else num_video_frames, 
                    img_emb=img_embs[[i]],
                    fs=fs[[i]],
                    **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 assign_forward_args(
        self, 
        model, 
        c_concat, 
        image_emb, 
        image_emb_uncond, 
        fs, 
        frames,
    ):
        model.model_options['transformer_options']['conditioning'] = {
            "c_concat": c_concat,
            "image_emb": image_emb,
            'image_emb_uncond': image_emb_uncond,
            "fs": fs,
            "num_video_frames": frames,
        }

    def validate_forwardable_latent(self, 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}"
            )

    def process_image_conditioning(
        self, 
        model, 
        clip_vision, 
        vae, 
        image_proj_model, 
        images, 
        use_interpolate,
        fps: int,
        frames: int,
        scale_latents: bool
    ):
        self.model_patcher = model
        encoded_latent = vae.encode(images[:, :, :, :3])

        encoded_image = clip_vision.encode_image(images[:1])['last_hidden_state']
        image_emb = image_proj_model(encoded_image)

        encoded_image_uncond  = clip_vision.encode_image(torch.zeros_like(images)[:1])['last_hidden_state']
        image_emb_uncond = image_proj_model(encoded_image_uncond)

        c_concat = encoded_latent
        
        if scale_latents:
            vae_process_input = vae.process_input
            vae.process_input = lambda image: (image - .5) * 2
            c_concat = vae.encode(images[:, :, :, :3])
            vae.process_input = vae_process_input
            c_concat = model.model.process_latent_in(c_concat) * 1.3
        else:
            c_concat = model.model.process_latent_in(c_concat)
            
        fs = torch.tensor([fps], dtype=torch.long, device=model_management.intermediate_device())
        
        model.set_model_unet_function_wrapper(self._forward)

        used_interpolate_processing = False

        if use_interpolate and frames > 16:
            raise ValueError(
                "When using interpolation mode, the maximum amount of frames are 16."
                "If you're doing long video generation, consider using the last frame\
                     from the first generation for the next one (autoregressive)."
            )
        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]
                c_concat = mask

                used_interpolate_processing = True
        else:
            if use_interpolate and c_concat.shape[0] in [2, 3]:
                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

                # TODO
                speed = 1.0
                if speed < 1.0:
                    possible_speeds = list(torch.linspace(0, 1.0, c_concat.shape[0]))
                    speed_from_frames = enumerate(possible_speeds)
                    speed_idx = min(speed_from_frames, key=lambda n: n[1] - speed)[0]
                    end_frame_idx = speed_idx

                # End frame
                masked_frames[-1:] = c_concat[[end_frame_idx]]
                
                # Possible middle frame, but not working at the moment.
                if input_frame_count == 3:
                    middle_idx = masked_frames.shape[0] // 2
                    middle_idx_frame = c_concat.shape[0] // 2
                    masked_frames[[middle_idx]] = c_concat[[middle_idx_frame]]

                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:
                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 = torch.cat(
                [encoded_latent] + [encoded_latent[-1:]] * abs(encoded_latent.shape[0] - frames)
            )[: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(model, c_concat, image_emb, image_emb_uncond, fs, frames)
        
        return (model, {"samples": torch.zeros_like(c_concat)}, {"samples": encoded_latent},)

class DynamiCrafterLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model_path": (get_models_directory(os.listdir(MODEL_DIR_PATH)), ),
            },
        }
        
    CATEGORY = "Native_DynamiCrafter/Loaders"
    RETURN_TYPES = ("MODEL", "IMAGE_PROJ_MODEL", )
    RETURN_NAMES = ("model", "image_proj_model", )
    FUNCTION = "load_dynamicrafter"
    
    def load_model_sicts(self, model_path: str):
        model_state_dict = comfy.utils.load_torch_file(model_path)
        dynamicrafter_dict = load_dynamicrafter_dict(model_state_dict)
        image_proj_dict = load_image_proj_dict(model_state_dict)

        return dynamicrafter_dict, image_proj_dict

    def get_prediction_type(self, is_eps: bool, model_config):
        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(self, 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(self, state_dict: dict):
        left_over = state_dict.keys()
        if len(left_over) > 0:
            print("left over keys:", left_over)

    def load_dynamicrafter(self, model_path):
        model_path = os.path.join(MODEL_DIR_PATH, model_path)
        
        if os.path.exists(model_path):
            dynamicrafter_dict, image_proj_dict = self.load_model_sicts(model_path)
            model_config = DynamiCrafterBase(DYNAMICRAFTER_CONFIG)

            dynamicrafter_dict, is_eps = model_config.process_dict_version(state_dict=dynamicrafter_dict)

            MODEL_TYPE = self.get_prediction_type(is_eps, model_config)
            load_device, inital_load_device = self.handle_model_management(dynamicrafter_dict, model_config)

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

            image_proj_model = get_image_proj_model(image_proj_dict)
            model.load_model_weights(dynamicrafter_dict, "model.diffusion_model.")
            self.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
            )
    
        return (model_patcher, image_proj_model, )

NODE_CLASS_MAPPINGS = {
    "DynamiCrafterLoader": DynamiCrafterLoader,
    "DynamiCrafterProcessor": DynamiCrafterProcessor,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "DynamiCrafterLoader": "Load a DynamiCrafter Checkpoint",    
    "DynamiCrafterProcessor": "Apply DynamiCrafter",
}