diff --git a/lvdm/modules/networks/openaimodel3d.py b/lvdm/modules/networks/openaimodel3d.py index cd4e1c6..d242e7e 100644 --- a/lvdm/modules/networks/openaimodel3d.py +++ b/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) diff --git a/nodes.py b/nodes.py index 6055180..b019aec 100644 --- a/nodes.py +++ b/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" } diff --git a/utils/model_utils.py b/utils/model_utils.py index adc1be4..080b47f 100644 --- a/utils/model_utils.py +++ b/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,