diff --git a/animatediff/context.py b/animatediff/context.py index b44fa42..50af0d8 100644 --- a/animatediff/context.py +++ b/animatediff/context.py @@ -11,8 +11,10 @@ import comfy.samplers from comfy.model_base import BaseModel from comfy.model_patcher import ModelPatcher +from .context_extras import ContextExtrasGroup from .utils_motion import get_sorted_list_via_attr + class ContextFuseMethod: FLAT = "flat" PYRAMID = "pyramid" @@ -76,6 +78,7 @@ class ContextOptions: class ContextOptionsGroup: def __init__(self): self.contexts: list[ContextOptions] = [] + self.extras = ContextExtrasGroup() self._current_context: ContextOptions = None self._current_used_steps: int = 0 self._current_index: int = 0 @@ -121,9 +124,10 @@ class ContextOptionsGroup: def is_empty(self) -> bool: return len(self.contexts) == 0 - + def clone(self): cloned = ContextOptionsGroup() + cloned.extras = self.extras.clone() for context in self.contexts: cloned.contexts.append(context) cloned._set_first_as_current() @@ -132,6 +136,11 @@ class ContextOptionsGroup: def initialize_timesteps(self, model: BaseModel): for context in self.contexts: context.start_t = model.model_sampling.percent_to_sigma(context.start_percent) + self.extras.initialize_timesteps(model) + + def prepare_current(self, t: Tensor): + self.prepare_current_context(t) + self.extras.prepare_current(t) def prepare_current_context(self, t: Tensor): curr_t: float = t[0] @@ -620,7 +629,7 @@ def generate_context_visualization(context_opts: ContextOptionsGroup, model: Mod for i, t in enumerate(sigmas): # make context_opts reflect current step/sigma - context_opts.prepare_current_context([t]) + context_opts.prepare_current([t]) context_opts.step = start_step+i # check if context should even be active in this case diff --git a/animatediff/context_extras.py b/animatediff/context_extras.py new file mode 100644 index 0000000..519e0e3 --- /dev/null +++ b/animatediff/context_extras.py @@ -0,0 +1,112 @@ +from torch import Tensor + +from comfy.model_base import BaseModel + + +class ContextExtra: + def __init__(self, start_percent: float, end_percent: float): + # scheduling + self.start_percent = float(start_percent) + self.start_t = 999999999.9 + self.end_percent = float(end_percent) + self.end_t = 0.0 + self.curr_t = 999999999.9 + + def initialize_timesteps(self, model: BaseModel): + self.start_t = model.model_sampling.percent_to_sigma(self.start_percent) + self.end_t = model.model_sampling.percent_to_sigma(self.end_percent) + + def prepare_current(self, t: Tensor): + self.curr_t = t[0] + + def should_run(self): + if self.curr_t > self.start_t or self.curr_t < self.end_t: + return False + return True + + +################################ +# Context Ref +class ContextRefParams: + def __init__(self, + attn_style_fidelity=0.0, attn_ref_weight=0.0, attn_atrength=0.0, + adain_style_fidelity=0.0, adain_ref_weight=0.0, adain_strength=0.0): + # attn1 + self.attn_style_fidelity = attn_style_fidelity + self.attn_ref_weight = attn_ref_weight + self.attn_strength = attn_atrength + # adain + self.adain_style_fidelity = adain_style_fidelity + self.adain_ref_weight = adain_ref_weight + self.adain_strength = adain_strength + + +class ContextRef(ContextExtra): + def __init__(self, start_percent: float, end_percent: float, params: ContextRefParams): + super().__init__(start_percent=start_percent, end_percent=end_percent) + self.params = params + + def should_run(self): + return super().should_run() +#-------------------------------- + + +################################ +# NaiveReuse +class NaiveReuse(ContextExtra): + def __init__(self, start_percent: float, end_percent: float, weighted_mean: float, mask_opt: Tensor=None): + super().__init__(start_percent=start_percent, end_percent=end_percent) + self.weighted_mean = weighted_mean + self.mask_opt = mask_opt + + def should_run(self): + to_return = super().should_run() + # if weighted_mean is 0.0, then reuse will take no effect anyway + return to_return and self.weighted_mean > 0.0 +#-------------------------------- + + +class ContextExtrasGroup: + def __init__(self): + self.context_ref: ContextRef = None + self.naive_reuse: NaiveReuse = None + + def get_extras_list(self) -> list[ContextExtra]: + extras_list = [] + if self.context_ref is not None: + extras_list.append(self.context_ref) + if self.naive_reuse is not None: + extras_list.append(self.naive_reuse) + return extras_list + + def initialize_timesteps(self, model: BaseModel): + for extra in self.get_extras_list(): + extra.initialize_timesteps(model) + + def prepare_current(self, t: Tensor): + for extra in self.get_extras_list(): + extra.prepare_current(t) + + def should_run_context_ref(self): + if not self.context_ref: + return False + return self.context_ref.should_run() + + def should_run_naive_reuse(self): + if not self.naive_reuse: + return False + return self.naive_reuse.should_run() + + def add(self, extra: ContextExtra): + if type(extra) == ContextRef: + self.context_ref = extra + elif type(extra) == NaiveReuse: + self.naive_reuse = extra + else: + raise Exception(f"Unrecognized ContextExtras type: {type(extra)}") + + def clone(self): + cloned = ContextExtrasGroup() + cloned.context_ref = self.context_ref + cloned.naive_reuse = self.naive_reuse + return cloned diff --git a/animatediff/nodes.py b/animatediff/nodes.py index 07ba149..e14e9ca 100644 --- a/animatediff/nodes.py +++ b/animatediff/nodes.py @@ -28,7 +28,8 @@ from .nodes_sample import (FreeInitOptionsNode, NoiseLayerAddWeightedNode, Sampl from .nodes_sigma_schedule import (SigmaScheduleNode, RawSigmaScheduleNode, WeightedAverageSigmaScheduleNode, InterpolatedWeightedAverageSigmaScheduleNode, SplitAndCombineSigmaScheduleNode, SigmaScheduleToSigmasNode) from .nodes_context import (LegacyLoopedUniformContextOptionsNode, LoopedUniformContextOptionsNode, LoopedUniformViewOptionsNode, StandardUniformContextOptionsNode, StandardStaticContextOptionsNode, BatchedContextOptionsNode, StandardStaticViewOptionsNode, StandardUniformViewOptionsNode, ViewAsContextOptionsNode, - VisualizeContextOptionsK, VisualizeContextOptionsKAdv, VisualizeContextOptionsSCustom) + VisualizeContextOptionsK, VisualizeContextOptionsKAdv, VisualizeContextOptionsSCustom, + SetContextExtrasOnContextOptions, ContextExtras_NaiveReuse, ContextExtras_ContextRef) from .nodes_ad_settings import (AnimateDiffSettingsNode, ManualAdjustPENode, SweetspotStretchPENode, FullStretchPENode, WeightAdjustAllAddNode, WeightAdjustAllMultNode, WeightAdjustIndivAddNode, WeightAdjustIndivMultNode, WeightAdjustIndivAttnAddNode, WeightAdjustIndivAttnMultNode) @@ -54,13 +55,15 @@ NODE_CLASS_MAPPINGS = { "ADE_MultivalDynamicFloatInput": MultivalDynamicFloatInputNode, "ADE_MultivalScaledMask": MultivalScaledMaskNode, "ADE_MultivalConvertToMask": MultivalConvertToMaskNode, + ############################################################################### + #------------------------------------------------------------------------------ # Context Opts "ADE_StandardStaticContextOptions": StandardStaticContextOptionsNode, "ADE_StandardUniformContextOptions": StandardUniformContextOptionsNode, "ADE_LoopedUniformContextOptions": LoopedUniformContextOptionsNode, "ADE_ViewsOnlyContextOptions": ViewAsContextOptionsNode, "ADE_BatchedContextOptions": BatchedContextOptionsNode, - "ADE_AnimateDiffUniformContextOptions": LegacyLoopedUniformContextOptionsNode, # Legacy + "ADE_AnimateDiffUniformContextOptions": LegacyLoopedUniformContextOptionsNode, # Legacy/Deprecated "ADE_VisualizeContextOptionsK": VisualizeContextOptionsK, "ADE_VisualizeContextOptionsKAdv": VisualizeContextOptionsKAdv, "ADE_VisualizeContextOptionsSCustom": VisualizeContextOptionsSCustom, @@ -68,6 +71,12 @@ NODE_CLASS_MAPPINGS = { "ADE_StandardStaticViewOptions": StandardStaticViewOptionsNode, "ADE_StandardUniformViewOptions": StandardUniformViewOptionsNode, "ADE_LoopedUniformViewOptions": LoopedUniformViewOptionsNode, + # Context Extras + "ADE_ContextExtras_Set": SetContextExtrasOnContextOptions, + "ADE_ContextExtras_ContextRef": ContextExtras_ContextRef, + "ADE_ContextExtras_NaiveReuse": ContextExtras_NaiveReuse, + #------------------------------------------------------------------------------ + ############################################################################### # Iteration Opts "ADE_IterationOptsDefault": IterationOptionsNode, "ADE_IterationOptsFreeInit": FreeInitOptionsNode, @@ -184,13 +193,15 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ADE_MultivalDynamicFloatInput": "Multival [Float List] ๐ŸŽญ๐Ÿ…๐Ÿ…“", "ADE_MultivalScaledMask": "Multival Scaled Mask ๐ŸŽญ๐Ÿ…๐Ÿ…“", "ADE_MultivalConvertToMask": "Multival to Mask ๐ŸŽญ๐Ÿ…๐Ÿ…“", + ############################################################################### + #------------------------------------------------------------------------------ # Context Opts "ADE_StandardStaticContextOptions": "Context Optionsโ—†Standard Static ๐ŸŽญ๐Ÿ…๐Ÿ…“", "ADE_StandardUniformContextOptions": "Context Optionsโ—†Standard Uniform ๐ŸŽญ๐Ÿ…๐Ÿ…“", "ADE_LoopedUniformContextOptions": "Context Optionsโ—†Looped Uniform ๐ŸŽญ๐Ÿ…๐Ÿ…“", "ADE_ViewsOnlyContextOptions": "Context Optionsโ—†Views Only [VRAMโ‡ˆ] ๐ŸŽญ๐Ÿ…๐Ÿ…“", "ADE_BatchedContextOptions": "Context Optionsโ—†Batched [Non-AD] ๐ŸŽญ๐Ÿ…๐Ÿ…“", - "ADE_AnimateDiffUniformContextOptions": "Context Optionsโ—†Looped Uniform ๐ŸŽญ๐Ÿ…๐Ÿ…“", # Legacy + "ADE_AnimateDiffUniformContextOptions": "Context Optionsโ—†Looped Uniform ๐ŸŽญ๐Ÿ…๐Ÿ…“", # Legacy/Deprecated "ADE_VisualizeContextOptionsK": "Visualize Context Options (K.) ๐ŸŽญ๐Ÿ…๐Ÿ…“", "ADE_VisualizeContextOptionsKAdv": "Visualize Context Options (K.Adv.) ๐ŸŽญ๐Ÿ…๐Ÿ…“", "ADE_VisualizeContextOptionsSCustom": "Visualize Context Options (S.Cus.) ๐ŸŽญ๐Ÿ…๐Ÿ…“", @@ -198,6 +209,12 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ADE_StandardStaticViewOptions": "View Optionsโ—†Standard Static ๐ŸŽญ๐Ÿ…๐Ÿ…“", "ADE_StandardUniformViewOptions": "View Optionsโ—†Standard Uniform ๐ŸŽญ๐Ÿ…๐Ÿ…“", "ADE_LoopedUniformViewOptions": "View Optionsโ—†Looped Uniform ๐ŸŽญ๐Ÿ…๐Ÿ…“", + # Context Extras + "ADE_ContextExtras_Set": "Set Context Extras ๐ŸŽญ๐Ÿ…๐Ÿ…“", + "ADE_ContextExtras_ContextRef": "Context Extrasโ—†ContextRef ๐ŸŽญ๐Ÿ…๐Ÿ…“", + "ADE_ContextExtras_NaiveReuse": "Context Extrasโ—†NaiveReuse ๐ŸŽญ๐Ÿ…๐Ÿ…“", + #------------------------------------------------------------------------------ + ############################################################################### # Iteration Opts "ADE_IterationOptsDefault": "Default Iteration Options ๐ŸŽญ๐Ÿ…๐Ÿ…“", "ADE_IterationOptsFreeInit": "FreeInit Iteration Options ๐ŸŽญ๐Ÿ…๐Ÿ…“", diff --git a/animatediff/nodes_context.py b/animatediff/nodes_context.py index a4ca94c..38c194d 100644 --- a/animatediff/nodes_context.py +++ b/animatediff/nodes_context.py @@ -6,6 +6,7 @@ from comfy.model_patcher import ModelPatcher from .context import (ContextFuseMethod, ContextOptions, ContextOptionsGroup, ContextSchedules, generate_context_visualization) +from .context_extras import ContextExtrasGroup, ContextRef, ContextRefParams, NaiveReuse from .utils_model import BIGMAX, MAX_RESOLUTION @@ -440,3 +441,90 @@ class VisualizeContextOptionsSCustom: images = generate_context_visualization(context_opts=context_opts, model=model, width=visual_width, video_length=latents_length, sigmas=sigmas) return (images,) + + +######################### +# Context Extras +class SetContextExtrasOnContextOptions: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "context_opts": ("CONTEXT_OPTIONS",), + "context_extras": ("CONTEXT_EXTRAS",), + }, + "optional": { + "autosize": ("ADEAUTOSIZE", {"padding": 0}), + } + } + + RETURN_TYPES = ("CONTEXT_OPTIONS",) + RETURN_NAMES = ("CONTEXT_OPTS",) + CATEGORY = "Animate Diff ๐ŸŽญ๐Ÿ…๐Ÿ…“/context opts/context extras" + FUNCTION = "set_context_extras" + + def set_context_extras(self, context_opts: ContextOptionsGroup, context_extras: ContextExtrasGroup): + context_opts = context_opts.clone() + context_opts.extras = context_extras.clone() + return (context_opts,) + + +class ContextExtras_NaiveReuse: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + }, + "optional": { + "prev_extras": ("CONTEXT_EXTRAS",), + "mask_opt": ("MASK",), + "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}), + "end_percent": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.001}), + "weighted_mean": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.001}), + "autosize": ("ADEAUTOSIZE", {"padding": 55}), + } + } + + RETURN_TYPES = ("CONTEXT_EXTRAS",) + CATEGORY = "Animate Diff ๐ŸŽญ๐Ÿ…๐Ÿ…“/context opts/context extras" + FUNCTION = "create_context_extra" + + def create_context_extra(self, start_percent=0.0, end_percent=0.1, weighted_mean=0.95, mask_opt: Tensor=None, prev_extras: ContextExtrasGroup=None): + if prev_extras is None: + prev_extras = prev_extras = ContextExtrasGroup() + prev_extras = prev_extras.clone() + # create extra + naive_reuse = NaiveReuse(start_percent=start_percent, end_percent=end_percent, weighted_mean=weighted_mean, mask_opt=mask_opt) + prev_extras.add(naive_reuse) + return (prev_extras,) + + +class ContextExtras_ContextRef: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + }, + "optional": { + "prev_extras": ("CONTEXT_EXTRAS",), + "mask_opt": ("MASK",), + "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}), + "end_percent": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.001}), + "autosize": ("ADEAUTOSIZE", {"padding": 55}), + } + } + + RETURN_TYPES = ("CONTEXT_EXTRAS",) + CATEGORY = "Animate Diff ๐ŸŽญ๐Ÿ…๐Ÿ…“/context opts/context extras" + FUNCTION = "create_context_extra" + + def create_context_extra(self, start_percent=0.0, end_percent=0.1, mask_opt: Tensor=None, prev_extras: ContextExtrasGroup=None): + if prev_extras is None: + prev_extras = prev_extras = ContextExtrasGroup() + prev_extras = prev_extras.clone() + # create extra + # TODO: make customizable, and allow mask input + params = ContextRefParams(attn_style_fidelity=1.0, attn_ref_weight=1.0, attn_atrength=1.0) + context_ref = ContextRef(start_percent=start_percent, end_percent=end_percent, params=params) + prev_extras.add(context_ref) + return (prev_extras,) diff --git a/animatediff/sampling.py b/animatediff/sampling.py index 02a2267..ddbde56 100644 --- a/animatediff/sampling.py +++ b/animatediff/sampling.py @@ -71,7 +71,7 @@ class AnimateDiffHelper_GlobalState: if self.motion_models is not None: self.motion_models.prepare_current_keyframe(x=x, t=timestep) if self.params.context_options is not None: - self.params.context_options.prepare_current_context(t=timestep) + self.params.context_options.prepare_current(t=timestep) if self.sample_settings.custom_cfg is not None: self.sample_settings.custom_cfg.prepare_current_keyframe(t=timestep) @@ -617,8 +617,10 @@ def evolved_sampling_function(model, x: Tensor, timestep: Tensor, uncond, cond, # add AD/evolved-sampling params to model_options (transformer_options) model_options = model_options.copy() - if "tranformer_options" not in model_options: - model_options["tranformer_options"] = {} + if "transformer_options" not in model_options: + model_options["transformer_options"] = {} + else: + model_options["transformer_options"] = model_options["transformer_options"].copy() model_options["transformer_options"]["ad_params"] = ADGS.create_exposed_params() if not ADGS.is_using_sliding_context(): @@ -798,7 +800,25 @@ def sliding_calc_conds_batch(model, conds, x_in: Tensor, timestep, model_options counts_final = [torch.zeros((x_in.shape[0], 1, 1, 1), device=x_in.device) for _ in conds] biases_final = [([0.0] * x_in.shape[0]) for _ in conds] - # perform calc_conds_batch per context window + CONTEXTREF_ATTN_MACHINE_STATE = "contextref_attn_machine_state" + CONTEXTREF_ADAIN_MACHINE_STATE = "contextref_adain_machine_state" + #context_ref_steps = [0,1,2,3,4,5]#[0] + context_ref = False + first_context = False + if ADGS.params.context_options.extras.should_run_context_ref(): + context_ref = True + first_context = True + + #naive_steps = [0,1]#[0,1,2]#[0,1,2,3] + naive_counts_mult = 50 + naive_init = False + cached_naive_conds = None + cached_naive_ctx_idxs = None + if ADGS.params.context_options.extras.should_run_naive_reuse(): + cached_naive_conds = [torch.zeros_like(x_in) for _ in conds] + #cached_naive_counts = [torch.zeros((x_in.shape[0], 1, 1, 1), device=x_in.device) for _ in conds] + naive_init = True + # perform calc_conds_batch per context window for ctx_idxs in context_windows: ADGS.params.sub_idxs = ctx_idxs if ADGS.motion_models is not None: @@ -817,6 +837,18 @@ def sliding_calc_conds_batch(model, conds, x_in: Tensor, timestep, model_options sub_timestep = timestep[full_idxs] sub_conds = [get_resized_cond(cond, full_idxs, len(ctx_idxs)) for cond in conds] + if context_ref: + if first_context: + first_context = False + model_options["transformer_options"][CONTEXTREF_ATTN_MACHINE_STATE] = "write" + model_options["transformer_options"][CONTEXTREF_ADAIN_MACHINE_STATE] = "write" + else: + model_options["transformer_options"][CONTEXTREF_ATTN_MACHINE_STATE] = "read" + model_options["transformer_options"][CONTEXTREF_ADAIN_MACHINE_STATE] = "read" + else: + model_options["transformer_options"][CONTEXTREF_ATTN_MACHINE_STATE] = "off" + model_options["transformer_options"][CONTEXTREF_ADAIN_MACHINE_STATE] = "off" + sub_conds_out = calc_cond_uncond_batch_wrapper(model, sub_conds, sub_x, sub_timestep, model_options) if ADGS.params.context_options.fuse_method == ContextFuseMethod.RELATIVE: @@ -841,12 +873,34 @@ def sliding_calc_conds_batch(model, conds, x_in: Tensor, timestep, model_options for i in range(len(sub_conds_out)): conds_final[i][full_idxs] += sub_conds_out[i] * weights_tensor counts_final[i][full_idxs] += weights_tensor + if naive_init: + cached_naive_ctx_idxs = ctx_idxs + for i in range(len(sub_conds)): + cached_naive_conds[i][full_idxs] = conds_final[i][full_idxs] / counts_final[i][full_idxs] + #cached_naive_conds[i][full_idxs] = conds_final[i][full_idxs] + #cached_naive_counts[i][full_idxs] = counts_final[i][full_idxs] + naive_init = False if ADGS.params.context_options.fuse_method == ContextFuseMethod.RELATIVE: # already normalized, so return as is del counts_final return conds_final else: + if cached_naive_conds is not None: + #start_idx = cached_naive_ctx_idxs[-1] + 1 + start_idx = cached_naive_ctx_idxs[0] + for z in range(start_idx, ADGS.params.full_length, len(cached_naive_ctx_idxs)): + for i in range(len(cached_naive_conds)): + new_ctx_idxs = [zz for zz in list(range(z, z+len(cached_naive_ctx_idxs))) if zz < ADGS.params.full_length] + # make sure when getting cached_naive idxs, they are adjusted for actual length leftover length + adjusted_cnaive_ctx_idxs = cached_naive_ctx_idxs[:len(new_ctx_idxs)] + weighted_mean = ADGS.params.context_options.extras.naive_reuse.weighted_mean + conds_final[i][new_ctx_idxs] = (weighted_mean * (cached_naive_conds[i][adjusted_cnaive_ctx_idxs]*counts_final[i][new_ctx_idxs])) + ((1.-weighted_mean) * conds_final[i][new_ctx_idxs]) + #conds_final[i][new_idxs] += (cached_naive_conds[i][cached_naive_full_idxs] / cached_naive_counts[i][cached_naive_full_idxs]) * counts + #counts = counts_final[i][new_idxs] * naive_counts_mult# / 2 + #conds_final[i][new_idxs] += (cached_naive_conds[i][cached_naive_full_idxs] / cached_naive_counts[i][cached_naive_full_idxs]) * counts + #counts_final[i][new_idxs] += counts# * 10#counts_final[i][full_idxs] + del cached_naive_conds # normalize conds via division by context usage counts for i in range(len(conds_final)): conds_final[i] /= counts_final[i] diff --git a/animatediff/scheduling.py b/animatediff/scheduling.py new file mode 100644 index 0000000..e69de29