Implemented ADKeyframes

animatediff/model_injection.py

@@ -9,10 +9,11 @@ import torch
 import comfy.model_management
 import comfy.utils
 from comfy.model_patcher import ModelPatcher
+from comfy.model_base import BaseModel
 
 from .motion_module_ad import AnimateDiffModel, has_mid_block, normalize_ad_state_dict
 from .logger import logger
-from .motion_utils import MotionCompatibilityError, normalize_min_max
+from .motion_utils import ADKeyframe, ADKeygrameGroup, MotionCompatibilityError, get_combined_multival, normalize_min_max
 from .motion_lora import MotionLoraInfo, MotionLoraList
 from .model_utils import get_motion_lora_path, get_motion_model_path, get_sd_model_type
 from .sample_settings import SampleSettings, SeedNoiseGeneration
@@ -91,11 +92,19 @@ class MotionModelPatcher(ModelPatcher):
         super().__init__(*args, **kwargs)
         self.model: AnimateDiffModel = self.model
         self.timestep_percent_range = (0.0, 1.0)
-        self.timestep_range = None
-        self.timestep_keyframes = None
+        self.timestep_range: tuple[float, float] = None
+        self.keyframes: ADKeygrameGroup = ADKeygrameGroup()
 
         self.scale_multival = None
         self.effect_multival = None
+        # temporary variables
+        self.current_keyframe: ADKeyframe = None
+        self.current_index = -1
+        self.current_scale: Union[float, Tensor] = None
+        self.current_effect: Union[float, Tensor] = None
+        self.combined_scale: Union[float, Tensor] = None
+        self.combined_effect: Union[float, Tensor] = None
+        self.was_within_range = False
 
     def patch_model(self, *args, **kwargs):
         # patch as normal, but prepare_weights so that lowvram meta device works properly
@@ -106,6 +115,7 @@ class MotionModelPatcher(ModelPatcher):
     def prepare_weights(self):
         # in case lowvram is active and meta device is used, need to convert weights
         # otherwise, will get exceptions thrown related to meta device
+        # TODO: with new comfy lowvram system, this is unnecessary
         state_dict = self.model.state_dict()
         for key in state_dict:
             weight = comfy.model_management.resolve_lowvram_weight(state_dict[key], self.model, key)
@@ -115,21 +125,76 @@ class MotionModelPatcher(ModelPatcher):
                 pass
 
     def pre_run(self, model: ModelPatcherAndInjector):
+        self.cleanup()
+        self.model.reset()
         # just in case, prepare_weights before every run
         self.prepare_weights()
-        # TODO: prepare timestep kf percents
-        # TODO: make this combined with timestep kfs
-        self.model.reset()
         self.model.set_scale(self.scale_multival)
         self.model.set_effect(self.effect_multival)
 
-    def prepare_current_timestep(self, t: Tensor, batched_number: int):
-        # TODO: dynamically 
-        pass
+    def initialize_timesteps(self, model: BaseModel):
+        self.timestep_range = (model.model_sampling.percent_to_sigma(self.timestep_percent_range[0]),
+                               model.model_sampling.percent_to_sigma(self.timestep_percent_range[1]))
+        if self.keyframes is not None:
+            for keyframe in self.keyframes.keyframes:
+                keyframe.start_t = model.model_sampling.percent_to_sigma(keyframe.start_percent)
+
+    def prepare_current_keyframe(self, t: Tensor):
+        curr_t: float = t[0]
+        prev_index = self.current_index
+        # if has next index, loop through and see if need to switch
+        if self.keyframes.has_index(self.current_index+1):
+            for i in range(self.current_index+1, len(self.keyframes)):
+                eval_kf = self.keyframes[i]
+                # check if start_t is greater or equal to curr_t
+                # NOTE: t is in terms of sigmas, not percent, so smaller number = later step in sampling
+                if eval_kf.start_t >= curr_t:
+                    self.current_index = i
+                    self.current_keyframe = eval_kf
+                    # keep track of scale and effect multivals, accounting for inherit_missing
+                    if self.current_keyframe.has_scale():
+                        self.current_scale = self.current_keyframe.scale_multival
+                    elif not self.current_keyframe.inherit_missing:
+                        self.current_scale = None
+                    if self.current_keyframe.has_effect():
+                        self.current_effect = self.current_keyframe.effect_multival
+                    elif not self.current_keyframe.inherit_missing:
+                        self.current_effect = None
+                    # if guarantee_usage, stop searching for other keyframes
+                    if self.current_keyframe.guarantee_usage:
+                        break
+                # if eval_kf is outside the percent range, stop looking further
+                else:
+                    break
+        # if index changed, apply new combined values
+        if prev_index != self.current_index:
+            # combine model's scale and effect with keyframe's scale and effect
+            self.combined_scale = get_combined_multival(self.scale_multival, self.current_scale)
+            self.combined_effect = get_combined_multival(self.effect_multival, self.current_effect)
+            # apply scale and effect
+            self.model.set_scale(self.combined_scale)
+            self.model.set_effect(self.combined_effect)
+        # apply effect - if not within range, set effect to 0, effectively turning model off
+        if curr_t > self.timestep_range[0] or curr_t < self.timestep_range[1]:
+            self.model.set_effect(0.0)
+            self.was_within_range = False
+        else:
+            # if was not in range last step, apply effect to toggle AD status
+            if not self.was_within_range:
+                self.model.set_effect(self.combined_effect)
+                self.was_within_range = True            
+
 
     def cleanup(self):
         if self.model is not None:
             self.model.cleanup()
+        self.current_keyframe = None
+        self.current_index = -1
+        self.current_scale = None
+        self.current_effect = None
+        self.combined_scale = None
+        self.combined_effect = None
+        self.was_within_range = False
 
     def clone(self):
         # normal ModelPatcher clone actions
@@ -143,7 +208,8 @@ class MotionModelPatcher(ModelPatcher):
         n.model_keys = self.model_keys
         # extra cloned params
         n.timestep_percent_range = self.timestep_percent_range
-        n.timestep_keyframes = self.timestep_keyframes
+        n.timestep_range = self.timestep_range
+        n.keyframes = self.keyframes.clone()
         n.scale_multival = self.scale_multival
         n.effect_multival = self.effect_multival
         return n
@@ -182,19 +248,17 @@ class MotionModelGroup:
         for motion_model in self.models:
             motion_model.model.set_video_length(video_length=video_length, full_length=full_length)
     
-    def set_scale(self, multival: Union[float, Tensor]):
-        pass
-
-    def set_strength(self, multival: Union[float, Tensor]):
-        pass
+    def initialize_timesteps(self, model: BaseModel):
+        for motion_model in self.models:
+            motion_model.initialize_timesteps(model)
 
     def pre_run(self, model: ModelPatcherAndInjector):
         for motion_model in self.models:
             motion_model.pre_run(model)
     
-    def prepare_current_timestep(self, t: Tensor, batched_number: int):
+    def prepare_current_keyframe(self, t: Tensor):
         for motion_model in self.models:
-            motion_model.prepare_current_timestep(t=t, batched_number=batched_number)
+            motion_model.prepare_current_keyframe(t=t)
 
     def get_name_string(self, show_version=False):
         identifiers = []

animatediff/motion_module_ad.py

@@ -540,6 +540,7 @@ class TemporalTransformer3DModel(nn.Module):
 
     def set_scale_mask(self, mask: Tensor):
         self.raw_scale_mask = mask
+        self.temp_scale_mask = None
 
     def set_sub_idxs(self, sub_idxs: list[int]):
         self.sub_idxs = sub_idxs

animatediff/motion_utils.py

@@ -1,9 +1,11 @@
+from typing import Union
 import torch
 import torch.nn.functional as F
 from torch import Tensor, nn
 
 import comfy.model_management as model_management
 import comfy.ops
+import comfy.utils
 from comfy.cli_args import args
 from comfy.ldm.modules.attention import attention_basic, attention_pytorch, attention_split, attention_sub_quad, default
 
@@ -93,7 +95,8 @@ def prepare_mask_batch(mask: Tensor, shape: Tensor, multiplier: int=1, match_dim
         mask = torch.cat([mask] * shape[1], dim=1)
     return mask
 
-def extend_to_batch_size(tensor, batch_size):
+
+def extend_to_batch_size(tensor: Tensor, batch_size: int):
     if tensor.shape[0] > batch_size:
         return tensor[:batch_size]
     elif tensor.shape[0] < batch_size:
@@ -104,3 +107,93 @@ def extend_to_batch_size(tensor, batch_size):
 
 class MotionCompatibilityError(ValueError):
     pass
+
+
+def get_combined_multival(multivalA: Union[float, Tensor], multivalB: Union[float, Tensor]) -> Union[float, Tensor]:
+    # if one is None, use the other
+    if multivalA == None:
+        return multivalB
+    elif multivalB == None:
+        return multivalA 
+    # both have a value - combine them based on type
+    # if both are Tensors, make dims match before multiplying
+    if type(multivalA) == Tensor and type(multivalB) == Tensor:
+        areaA = multivalA.shape[1]*multivalA.shape[2]
+        areaB = multivalB.shape[1]*multivalB.shape[2]
+        # match height/width to mask with larger area
+        leader,follower = multivalA,multivalB if areaA >= areaB else multivalB,multivalA
+        batch_size = multivalA.shape[0] if multivalA.shape[0] >= multivalB.shape[0] else multivalB.shape[0]
+        # make follower same dimensions as leader
+        follower = torch.unsqueeze(follower, 1)
+        follower = comfy.utils.common_upscale(follower, leader.shape[2], leader.shape[1], "bilinear", "center")
+        follower = torch.squeeze(follower, 1)
+        # make sure batch size will match
+        leader = extend_to_batch_size(leader, batch_size)
+        follower = extend_to_batch_size(follower, batch_size)
+        return leader * follower
+    # otherwise, just multiply them together - one of them is a float
+    return multivalA * multivalB
+
+
+class ADKeyframe:
+    def __init__(self,
+                 start_percent: float = 0.0,
+                 scale_multival: Union[float, Tensor]=None,
+                 effect_multival: Union[float, Tensor]=None,
+                 inherit_missing: bool=True,
+                 guarantee_usage: bool=True,
+                 ):
+        self.start_percent = start_percent
+        self.start_t = 999999999.9
+        self.scale_multival = scale_multival
+        self.effect_multival = effect_multival
+        self.inherit_missing = inherit_missing
+        self.guarantee_usage = guarantee_usage
+    
+    def has_scale(self):
+        return self.scale_multival is not None
+    
+    def has_effect(self):
+        return self.effect_multival is not None
+
+
+class ADKeygrameGroup:
+    def __init__(self):
+        self.keyframes: list[ADKeyframe] = []
+        self.keyframes.append(ADKeyframe())
+    
+    def add(self, keyframe: ADKeyframe):
+        added = False
+        # replace existing keyframe if same start_percent
+        for i in range(len(self.keyframes)):
+            if self.keyframes[i].start_percent == keyframe.start_percent:
+                self.keyframes[i] = keyframe
+                added = True
+                break
+        if not added:
+            self.keyframes.append(keyframe)
+        self.keyframes.sort(key=lambda k: k.start_percent)
+    
+    def get_index(self, index: int) -> Union[ADKeyframe, None]:
+        try:
+            return self.keyframes[index]
+        except IndexError:
+            return None
+    
+    def has_index(self, index: int) -> int:
+        return index >=0 and index < len(self.keyframes)
+
+    def __getitem__(self, index) -> ADKeyframe:
+        return self.keyframes[index]
+
+    def __len__(self) -> int:
+        return len(self.keyframes)
+
+    def is_empty(self) -> bool:
+        return len(self.keyframes) == 0
+
+    def clone(self) -> 'ADKeygrameGroup':
+        cloned = ADKeygrameGroup()
+        for tk in self.keyframes:
+            cloned.add(tk)
+        return cloned

animatediff/nodes.py

@@ -13,7 +13,7 @@ from .sample_settings import SampleSettings, SeedNoiseGeneration
 from .sampling import motion_sample_factory
 
 from .nodes_gen1 import AnimateDiffLoaderWithContext
-from .nodes_gen2 import UseEvolvedSamplingNode, ApplyAnimateDiffModelNode, ApplyAnimateDiffModelBasicNode, LoadAnimateDiffModelNode
+from .nodes_gen2 import UseEvolvedSamplingNode, ApplyAnimateDiffModelNode, ApplyAnimateDiffModelBasicNode, LoadAnimateDiffModelNode, ADKeyframeNode
 from .nodes_multival import MultivalDynamicNode, MultivalFloatNode, MultivalScaledMaskNode
 from .nodes_sample import FreeInitOptionsNode, NoiseLayerAddWeightedNode, SampleSettingsNode, NoiseLayerAddNode, NoiseLayerReplaceNode, IterationOptionsNode
 from .nodes_extras import AnimateDiffUnload, EmptyLatentImageLarge, CheckpointLoaderSimpleWithNoiseSelect
@@ -121,9 +121,10 @@ NODE_CLASS_MAPPINGS = {
     "ADE_ApplyAnimateDiffModel": ApplyAnimateDiffModelNode,
     "ADE_ApplyAnimateDiffModelSimple": ApplyAnimateDiffModelBasicNode,
     "ADE_LoadAnimateDiffModel": LoadAnimateDiffModelNode,
+    "ADE_AnimateDiffLoRALoader": AnimateDiffLoraLoader,
     "ADE_AnimateDiffUniformContextOptions": AnimateDiffUniformContextOptions,
     "ADE_AnimateDiffSamplingSettings": SampleSettingsNode,
-    "ADE_AnimateDiffLoRALoader": AnimateDiffLoraLoader,
+    "ADE_AnimateDiffKeyframe": ADKeyframeNode,
     # Multival Nodes
     "ADE_MultivalDynamic": MultivalDynamicNode,
     "ADE_MultivalScaledMask": MultivalScaledMaskNode,
@@ -155,9 +156,10 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ADE_ApplyAnimateDiffModel": "Apply AnimateDiff Model (Adv.) 🎭🅐🅓②",
     "ADE_ApplyAnimateDiffModelSimple": "Apply AnimateDiff Model (Basic) 🎭🅐🅓②",
     "ADE_LoadAnimateDiffModel": "Load AnimateDiff Model 🎭🅐🅓②",
+    "ADE_AnimateDiffLoRALoader": "Load AnimateDiff LoRA 🎭🅐🅓",
     "ADE_AnimateDiffUniformContextOptions": "Uniform Context Options 🎭🅐🅓",
     "ADE_AnimateDiffSamplingSettings": "Sample Settings 🎭🅐🅓",
-    "ADE_AnimateDiffLoRALoader": "Load AnimateDiff LoRA 🎭🅐🅓",
+    "ADE_AnimateDiffKeyframe": "AnimateDiff Keyframe 🎭🅐🅓",
     # Multival Nodes
     "ADE_MultivalDynamic": "Multival Dynamic 🎭🅐🅓",
     "ADE_MultivalScaledMask": "Multival Scaled Mask 🎭🅐🅓",

animatediff/nodes_gen2.py

@@ -7,8 +7,10 @@ from comfy.model_patcher import ModelPatcher
 from .context import ContextOptions, ContextSchedules, UniformContextOptions
 from .logger import logger
 from .model_utils import BetaSchedules, get_available_motion_loras, get_available_motion_models, get_motion_lora_path
+from .motion_utils import ADKeygrameGroup, ADKeyframe
 from .motion_lora import MotionLoraInfo, MotionLoraList
-from .model_injection import InjectionParams, ModelPatcherAndInjector, MotionModelGroup, MotionModelPatcher, MotionModelSettings, load_motion_module, load_motion_module_gen2, load_motion_lora_as_patches, validate_model_compatibility_gen2
+from .model_injection import (InjectionParams, ModelPatcherAndInjector, MotionModelGroup, MotionModelPatcher, MotionModelSettings,
+                              load_motion_module, load_motion_module_gen2, load_motion_lora_as_patches, validate_model_compatibility_gen2)
 from .sample_settings import SampleSettings, SeedNoiseGeneration
 from .sampling import motion_sample_factory
 
@@ -85,9 +87,9 @@ class ApplyAnimateDiffModelNode:
             },
             "optional": {
                 "motion_lora": ("MOTION_LORA",),
-                "mm_timestep_kf": ("MM_TIMESTEP_KF",),
                 "scale_multival": ("MULTIVAL",),
                 "effect_multival": ("MULTIVAL",),
+                "ad_keyframes": ("AD_KEYFRAMES",),
                 "prev_m_models": ("M_MODELS",),
             }
         }
@@ -97,7 +99,7 @@ class ApplyAnimateDiffModelNode:
     FUNCTION = "apply_motion_model"
 
     def apply_motion_model(self, motion_model: MotionModelPatcher, start_percent: float=0.0, end_percent: float=1.0,
-                           motion_lora: MotionLoraList=None, mm_timestep_kf=None,
+                           motion_lora: MotionLoraList=None, ad_keyframes: ADKeygrameGroup=None,
                            scale_multival=None, effect_multival=None,
                            prev_m_models: MotionModelGroup=None,):
         # set up motion models list
@@ -109,9 +111,10 @@ class ApplyAnimateDiffModelNode:
         if motion_lora is not None:
             for lora in motion_lora.loras:
                 load_motion_lora_as_patches(motion_model, lora)
-        # TODO: handle mm_timestep_kf
         motion_model.scale_multival = scale_multival
         motion_model.effect_multival = effect_multival
+        motion_model.keyframes = ad_keyframes.clone() if ad_keyframes else ADKeygrameGroup()
+        motion_model.timestep_percent_range = (start_percent, end_percent)
         # add to beginning, so that after injection, it will be the earliest of prev_m_models to be run
         prev_m_models.add_to_start(mm=motion_model)
         return (prev_m_models,)
@@ -122,11 +125,13 @@ class ApplyAnimateDiffModelBasicNode:
     def INPUT_TYPES(s):
         return {
             "required": {
-                "motion_model": ("MOTION_M",),
+                "motion_model": ("MOTION_MODEL_ADE",),
             },
             "optional": {
                 "motion_lora": ("MOTION_LORA",),
                 "scale_multival": ("MULTIVAL",),
+                "effect_multival": ("MULTIVAL",),
+                "ad_keyframes": ("AD_KEYFRAMES",),
             }
         }
     
@@ -136,10 +141,11 @@ class ApplyAnimateDiffModelBasicNode:
 
     def apply_motion_model(self,
                            motion_model: MotionModelPatcher, motion_lora: MotionLoraList=None,
-                           scale_multival=None,):
+                           scale_multival=None, effect_multival=None, ad_keyframes=None):
         # just a subset of normal ApplyAnimateDiffModelNode inputs
         return ApplyAnimateDiffModelNode.apply_motion_model(self, motion_model, motion_lora=motion_lora,
-                                                            scale_multival=scale_multival)
+                                                            scale_multival=scale_multival, effect_multival=effect_multival,
+                                                            ad_keyframes=ad_keyframes)
 
 
 class LoadAnimateDiffModelNode:
@@ -150,7 +156,7 @@ class LoadAnimateDiffModelNode:
                 "model_name": (get_available_motion_models(),),
             },
             "optional": {
-                "mm_settings": ("MOTION_MODEL_SETTINGS",),
+                "ad_settings": ("MOTION_MODEL_SETTINGS",),
             }
         }
 
@@ -159,7 +165,41 @@ class LoadAnimateDiffModelNode:
     CATEGORY = "Animate Diff 🎭🅐🅓"
     FUNCTION = "load_motion_model"
 
-    def load_motion_model(self, model_name: str, mm_settings: MotionModelSettings=None):
+    def load_motion_model(self, model_name: str, ad_settings: MotionModelSettings=None):
         # load motion module and motion settings, if included
-        motion_model = load_motion_module_gen2(model_name=model_name, motion_model_settings=mm_settings)
+        motion_model = load_motion_module_gen2(model_name=model_name, motion_model_settings=ad_settings)
         return (motion_model,)
+
+
+class ADKeyframeNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}, ),
+            },
+            "optional": {
+                "prev_ad_keyframes": ("AD_KEYFRAMES", ),
+                "scale_multival": ("MULTIVAL",),
+                "effect_multival": ("MULTIVAL",),
+                "inherit_missing": ("BOOLEAN", {"default": True}, ),
+                "guarantee_usage": ("BOOLEAN", {"default": True}, ),
+            }
+        }
+    
+    RETURN_TYPES = ("AD_KEYFRAMES", )
+    FUNCTION = "load_keyframe"
+
+    CATEGORY = "Animate Diff 🎭🅐🅓"
+
+    def load_keyframe(self,
+                      start_percent: float, prev_ad_keyframes=None,
+                      scale_multival: [float, torch.Tensor]=None, effect_multival: [float, torch.Tensor]=None,
+                      inherit_missing: bool=True, guarantee_usage: bool=True):
+        if not prev_ad_keyframes:
+            prev_ad_keyframes = ADKeygrameGroup()
+        prev_ad_keyframes.clone()
+        keyframe = ADKeyframe(start_percent=start_percent, scale_multival=scale_multival, effect_multival=effect_multival,
+                              inherit_missing=inherit_missing, guarantee_usage=guarantee_usage)
+        prev_ad_keyframes.add(keyframe)
+        return (prev_ad_keyframes,)

animatediff/nodes_multival.py

@@ -44,10 +44,10 @@ class MultivalDynamicNode:
         else:
             if not float_is_iterable:
                 return (float_val,)
-            # create a dummy mask of b,c,h,w=float_len,1,1,1 (1 channel, sigle pixel)
+            # create a dummy mask of b,h,w=float_len,1,1 (sigle pixel)
             # purpose is for float input to work with mask code, without special cases
             float_len = len(float_val) if float_is_iterable else 1
-            shape = (float_len,1,1,1)
+            shape = (float_len,1,1)
             mask_optional = torch.ones(shape)
             mask_optional = mask_optional * float_val
             return (mask_optional,)

animatediff/sampling.py

@@ -32,7 +32,15 @@ class AnimateDiffHelper_GlobalState:
         self.params: InjectionParams = None
         self.reset()
     
+    def initialize(self, model):
+        # this function is to be run in sampling func
+        if not self.initialized:
+            self.initialized = True
+            if self.motion_models is not None:
+                self.motion_models.initialize_timesteps(model)
+
     def reset(self):
+        self.initialized = False
         self.start_step: int = 0
         self.last_step: int = 0
         self.current_step: int = 0
@@ -320,7 +328,6 @@ def motion_sample_factory(orig_comfy_sample: Callable, is_custom: bool=False) ->
     return motion_sample
 
 
-
 def sliding_sampling_function(model, x, timestep, uncond, cond, cond_scale, model_options={}, seed=None):
         def get_area_and_mult(conds, x_in, timestep_in):
             area = (x_in.shape[2], x_in.shape[3], 0, 0)
@@ -652,6 +659,9 @@ def sliding_sampling_function(model, x, timestep, uncond, cond, cond_scale, mode
             del out_count_final
             return cond_final, uncond_final
 
+        ADGS.initialize(model)
+        if ADGS.motion_models is not None:
+            ADGS.motion_models.prepare_current_keyframe(t=timestep)
 
         if math.isclose(cond_scale, 1.0):
             uncond = None