Initial commit for v2.

Changes and workflows will be updated in a PR.

lvdm/modules/networks/openaimodel3d.py

@@ -313,18 +313,24 @@ def context_processor(context, t, img_emb=None, temporal_size=16, concat_only=Fa
         
     ## repeat t times for context [(b t) 77 768] & time embedding
     ## check if we use per-frame image conditioning
-
+    used_concat = False
     if img_emb is not None:
+        if context.shape[0] < img_emb.shape[0]:
+            context = torch.cat([context] * img_emb.shape[0])[:img_emb.shape[0]]
+            used_concat = True
         context = torch.cat([context, img_emb.to(context.device, context.dtype)], dim=1)
-        
+    
     if concat_only:
         return context
 
     b, l_context, _ = context.shape
     if l_context == 77 + t * temporal_size:
         context_text, context_img = context[:,:77,:], context[:,77:,:]
-        context_text = context_text.repeat_interleave(repeats=t, dim=0)
-        context_img = rearrange(context_img, 'b (t l) c -> (b t) l c', t=t)
+
+        if not used_concat:
+            context_text = context_text.repeat_interleave(repeats=t, dim=0)
+            context_img = rearrange(context_img, 'b (t l) c -> (b t) l c', t=t)
+            
         context = torch.cat([context_text, context_img], dim=1)
     else:
         context = context.repeat_interleave(repeats=t, dim=0)
@@ -419,6 +425,7 @@ class UNetModel(nn.Module):
                  default_fs=4,
                  fs_condition=False,
                  device=None,
+                 videocrafter_mode=False,
                  dtype=torch.float16,
                  operations=ops
                 ):
@@ -449,6 +456,7 @@ class UNetModel(nn.Module):
         self.device = device
         #self.dtype = dtype
         self.dtype = torch.float32
+        self.videocrafter_mode = videocrafter_mode
 
         ## Time embedding blocks
         self.time_embed = nn.Sequential(
@@ -699,11 +707,20 @@ class UNetModel(nn.Module):
             cond_idx=None,
             **kwargs
         ):
+        fs_t = transformer_options.get("fs_t", None)
+        if fs_t is not None:
+            fs = fs_t
+
+        null_conditions = [fs is None]
         
-        if any([fs is None, img_emb is None, cc_concat is None]):
+        if not self.videocrafter_mode:
+            null_conditions += [img_emb is None, cc_concat is None]
+
+        if any(null_conditions):
             raise ValueError("One or more of the required inputs for UNet Forward is None.")
         
         cond_idx = transformer_options.get("cond_idx", None)
+
         transformer_options['original_shape'] = list(x.shape)
         transformer_options['transformer_index'] = 0
         transformer_patches = transformer_options.get("patches", {})
@@ -713,14 +730,15 @@ class UNetModel(nn.Module):
         # We usually denote "f" as frames, but will use "t" (time) to be consistent with DynamiCrafter.
         b,_,t,_,_ = x.shape
         
-        context = context_in
-        cc_concat = cc_concat.to(x.device, x.dtype)
-        x = torch.cat([x, cc_concat], dim=1)
-
+        context = context_in if not self.videocrafter_mode else context
         fs = fs.to(x.device, x.dtype)
 
         timestep = timesteps
-        context = context_processor(context, num_video_frames, img_emb=img_emb)
+
+        if not self.videocrafter_mode:
+            cc_concat = cc_concat.to(x.device, x.dtype)
+            x = torch.cat([x, cc_concat], dim=1)
+            context = context_processor(context, num_video_frames, img_emb=img_emb)
 
         t_emb = timestep_embedding(timestep, self.model_channels, repeat_only=False, dtype=self.dtype)
         emb = self.time_embed(t_emb)
@@ -738,7 +756,6 @@ class UNetModel(nn.Module):
 
             fs_embed = self.fps_embedding(fs_emb)
             fs_embed = fs_embed.repeat_interleave(repeats=t, dim=0)
-
             emb = emb + fs_embed
 
         h = x.type(self.dtype)

nodes.py

@@ -3,12 +3,13 @@ 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, \
+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
@@ -16,6 +17,341 @@ 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):
@@ -27,10 +363,12 @@ class DynamiCrafterProcessor:
                 "image_proj_model": ("IMAGE_PROJ_MODEL", ),
                 "images": ("IMAGE", ),
                 "use_interpolate": ("BOOLEAN", {"default": False}),
-                "fps": ("INT", {"default": 15, "min": 1, "max": 30, "step": 1}, ),
+                "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})
+                "scale_latents": ("BOOLEAN", {"default": False}),
             },
+            "optional": {"advanced_options": ("DYNAMI_ADVANCED_OPTS", ) }
         }
         
     CATEGORY = "Native_DynamiCrafter/Processing"
@@ -42,89 +380,106 @@ class DynamiCrafterProcessor:
     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)
+ 
+    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
-        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']
+        # Forward_dict = args[1]
+        x, t, model_in_kwargs, _ =  args[1].values()
 
-        
-        original_num_frames = num_video_frames
+        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()
 
-        # 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
+        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)
 
-        self.validate_forwardable_latent(x, c_concat, num_video_frames, use_cfg)
+        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 = 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
-
-        img_embs = conditioning['image_emb']
-
+        
         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]]),
@@ -132,57 +487,70 @@ class DynamiCrafterProcessor:
                     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]],
+                    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, 
-        model, 
         c_concat, 
         image_emb, 
         image_emb_uncond, 
+        img_cfg,
         fs, 
         frames,
+        c_concat_ref: ReferenceProcessor = None
     ):
-        model.model_options['transformer_options']['conditioning'] = {
+        self.model_patcher.model_options['transformer_options']['conditioning'] = {
             "c_concat": c_concat,
-            "image_emb": image_emb,
+            "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 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, 
@@ -191,53 +559,65 @@ class DynamiCrafterProcessor:
         image_proj_model, 
         images, 
         use_interpolate,
+        img_cfg,
         fps: int,
         frames: int,
-        scale_latents: bool
+        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)
 
-        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)
+        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(
-                "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)."
-            )
+            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 and c_concat.shape[0] in [2, 3]:
+            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
@@ -248,33 +628,23 @@ class DynamiCrafterProcessor:
 
                 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)
+
+                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]
         
@@ -282,15 +652,21 @@ class DynamiCrafterProcessor:
             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]
+            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(model, c_concat, image_emb, image_emb_uncond, fs, frames)
+        self.assign_forward_args(c_concat, image_emb, image_emb_uncond, img_cfg, fs, frames, reference_frames)
         
-        return (model, {"samples": torch.zeros_like(c_concat)}, {"samples": encoded_latent},)
+        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
@@ -298,6 +674,7 @@ class DynamiCrafterLoader:
         return {
             "required": {
                 "model_path": (get_models_directory(os.listdir(MODEL_DIR_PATH)), ),
+                "videocrafter_mode": ("BOOLEAN", {"default": False})
             },
         }
         
@@ -305,54 +682,19 @@ class DynamiCrafterLoader:
     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):
+     
+    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 = self.load_model_sicts(model_path)
-            model_config = DynamiCrafterBase(DYNAMICRAFTER_CONFIG)
+            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)
 
-            MODEL_TYPE = self.get_prediction_type(is_eps, model_config)
-            load_device, inital_load_device = self.handle_model_management(dynamicrafter_dict, model_config)
+            load_device, inital_load_device = handle_model_management(dynamicrafter_dict, model_config)
 
             model = model_base.BaseModel(
                 model_config, 
@@ -361,9 +703,16 @@ class DynamiCrafterLoader:
                 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)
+            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, 
@@ -371,15 +720,86 @@ class DynamiCrafterLoader:
                 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"
 }

utils/model_utils.py

@@ -38,6 +38,31 @@ DYNAMICRAFTER_CONFIG = {
     'fs_condition': True
 }
 
+VIDEOCRAFTER_CONFIG = {
+    'in_channels': 4,
+    'out_channels': 4,
+    'model_channels': 320,
+    'attention_resolutions': [4, 2, 1],
+    'num_res_blocks': 2,
+    'channel_mult': [1, 2, 4, 4],
+    'num_head_channels': 64,
+    'transformer_depth': 1,
+    'context_dim': 1024,
+    'use_linear': True,
+    'use_checkpoint': False,
+    'temporal_conv': True,
+    'temporal_attention': True,
+    'temporal_selfatt_only': True,
+    'use_relative_position': False,
+    'use_causal_attention': False,
+    'temporal_length': 16,
+    'addition_attention': True,
+    'image_cross_attention': True,
+    'image_cross_attention_scale_learnable': True,
+    'default_fs': 3,
+    'fs_condition': True
+}
+
 IMAGE_PROJ_CONFIG = {
     "dim": 1024,
     "depth": 4,