diff --git a/animatediff/model_injection.py b/animatediff/model_injection.py index a939714..686446d 100644 --- a/animatediff/model_injection.py +++ b/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 = [] diff --git a/animatediff/motion_module_ad.py b/animatediff/motion_module_ad.py index bb6199f..2f0969f 100644 --- a/animatediff/motion_module_ad.py +++ b/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 diff --git a/animatediff/motion_utils.py b/animatediff/motion_utils.py index 3073b50..d6937a2 100644 --- a/animatediff/motion_utils.py +++ b/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 diff --git a/animatediff/nodes.py b/animatediff/nodes.py index 98e0aa1..f79991f 100644 --- a/animatediff/nodes.py +++ b/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 🎭🅐🅓", diff --git a/animatediff/nodes_gen2.py b/animatediff/nodes_gen2.py index 5fb2153..f900ed2 100644 --- a/animatediff/nodes_gen2.py +++ b/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,) diff --git a/animatediff/nodes_multival.py b/animatediff/nodes_multival.py index 9e3348e..f17ae30 100644 --- a/animatediff/nodes_multival.py +++ b/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,) diff --git a/animatediff/sampling.py b/animatediff/sampling.py index b6fdf45..7fa34e1 100644 --- a/animatediff/sampling.py +++ b/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