Merge PR #262 from Kosinkadink/develop - FreeInit, Sample Settings, Noise controls

Sample Settings + Noise controls + initial FreeInit implementation

animatediff/context.py

@@ -3,8 +3,6 @@ from typing import Callable, Optional
 
 import numpy as np
 
-from .motion_utils import NoiseType
-
 
 class ContextType:
     UNIFORM_WINDOW = "uniform window"
@@ -27,13 +25,9 @@ class UniformContextOptions(ContextOptions):
         self.context_schedule = context_schedule
         self.closed_loop = closed_loop
         self.sync_context_to_pe = False
-        self.noise_type = NoiseType.DEFAULT
     
     def set_sync_context_to_pe(self, sync_context_to_pe: bool):
         self.sync_context_to_pe = sync_context_to_pe
-    
-    def set_noise_type(self, noise_type: str):
-        self.noise_type = noise_type
 
 
 class ContextSchedules:

animatediff/freeinit.py

@@ -0,0 +1,162 @@
+# S-Lab License 1.0
+
+# Copyright 2023 S-Lab
+# Redistribution and use for non-commercial purpose in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+# 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+# 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+# 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+# 4. In the event that redistribution and/or use for commercial purpose in source or binary forms, with or without modification is required, please contact the contributor(s) of the work.
+
+# Code has been modified from https://github.com/TianxingWu/FreeInit
+
+import torch
+import torch.fft as fft
+import math
+
+
+class FreeInitFilter:
+    GAUSSIAN = "gaussian"
+    IDEAL = "ideal"
+    BOX = "box"
+    BUTTERWORTH = "butterworth"
+
+    LIST = [GAUSSIAN, BUTTERWORTH, IDEAL, BOX]
+
+
+def freq_mix_3d(x, noise, LPF):
+    """
+    Noise reinitialization.
+
+    Args:
+        x: diffused latent
+        noise: randomly sampled noise
+        LPF: low pass filter
+    """
+    # FFT
+    x_freq = fft.fftn(x, dim=(-4, -2, -1))
+    x_freq = fft.fftshift(x_freq, dim=(-4, -2, -1))
+    noise_freq = fft.fftn(noise, dim=(-4, -2, -1))
+    noise_freq = fft.fftshift(noise_freq, dim=(-4, -2, -1))
+
+    # frequency mix
+    HPF = 1 - LPF
+    x_freq_low = x_freq * LPF
+    noise_freq_high = noise_freq * HPF
+    x_freq_mixed = x_freq_low + noise_freq_high # mix in freq domain
+
+    # IFFT
+    x_freq_mixed = fft.ifftshift(x_freq_mixed, dim=(-4, -2, -1))
+    x_mixed = fft.ifftn(x_freq_mixed, dim=(-4, -2, -1)).real
+
+    return x_mixed
+
+
+def get_freq_filter(shape, device, filter_type, n, d_s, d_t):
+    """
+    Form the frequency filter for noise reinitialization.
+
+    Args:
+        shape: shape of latent (T, C, H, W)
+        filter_type: type of the freq filter
+        n: (only for butterworth) order of the filter, larger n ~ ideal, smaller n ~ gaussian
+        d_s: normalized stop frequency for spatial dimensions (0.0-1.0)
+        d_t: normalized stop frequency for temporal dimension (0.0-1.0)
+    """
+    if filter_type == FreeInitFilter.GAUSSIAN:
+        return gaussian_low_pass_filter(shape=shape, d_s=d_s, d_t=d_t).to(device)
+    elif filter_type == FreeInitFilter.IDEAL:
+        return ideal_low_pass_filter(shape=shape, d_s=d_s, d_t=d_t).to(device)
+    elif filter_type == FreeInitFilter.BOX:
+        return box_low_pass_filter(shape=shape, d_s=d_s, d_t=d_t).to(device)
+    elif filter_type == FreeInitFilter.BUTTERWORTH:
+        return butterworth_low_pass_filter(shape=shape, n=n, d_s=d_s, d_t=d_t).to(device)
+    else:
+        raise NotImplementedError
+
+def gaussian_low_pass_filter(shape, d_s=0.25, d_t=0.25):
+    """
+    Compute the gaussian low pass filter mask.
+
+    Args:
+        shape: shape of the filter (volume)
+        d_s: normalized stop frequency for spatial dimensions (0.0-1.0)
+        d_t: normalized stop frequency for temporal dimension (0.0-1.0)
+    """
+    T, H, W = shape[-4], shape[-2], shape[-1]
+    mask = torch.zeros(shape)
+    if d_s==0 or d_t==0:
+        return mask
+    for t in range(T):
+        for h in range(H):
+            for w in range(W):
+                d_square = (((d_s/d_t)*(2*t/T-1))**2 + (2*h/H-1)**2 + (2*w/W-1)**2)
+                mask[t, ..., h,w] = math.exp(-1/(2*d_s**2) * d_square)
+    return mask
+
+
+def butterworth_low_pass_filter(shape, n=4, d_s=0.25, d_t=0.25):
+    """
+    Compute the butterworth low pass filter mask.
+
+    Args:
+        shape: shape of the filter (volume)
+        n: order of the filter, larger n ~ ideal, smaller n ~ gaussian
+        d_s: normalized stop frequency for spatial dimensions (0.0-1.0)
+        d_t: normalized stop frequency for temporal dimension (0.0-1.0)
+    """
+    T, H, W = shape[-4], shape[-2], shape[-1]
+    mask = torch.zeros(shape)
+    if d_s==0 or d_t==0:
+        return mask
+    for t in range(T):
+        for h in range(H):
+            for w in range(W):
+                d_square = (((d_s/d_t)*(2*t/T-1))**2 + (2*h/H-1)**2 + (2*w/W-1)**2)
+                mask[t, ..., h,w] = 1 / (1 + (d_square / d_s**2)**n)
+    return mask
+
+
+def ideal_low_pass_filter(shape, d_s=0.25, d_t=0.25):
+    """
+    Compute the ideal low pass filter mask.
+
+    Args:
+        shape: shape of the filter (volume)
+        d_s: normalized stop frequency for spatial dimensions (0.0-1.0)
+        d_t: normalized stop frequency for temporal dimension (0.0-1.0)
+    """
+    T, H, W = shape[-4], shape[-2], shape[-1]
+    mask = torch.zeros(shape)
+    if d_s==0 or d_t==0:
+        return mask
+    for t in range(T):
+        for h in range(H):
+            for w in range(W):
+                d_square = (((d_s/d_t)*(2*t/T-1))**2 + (2*h/H-1)**2 + (2*w/W-1)**2)
+                mask[t, ...,h,w] =  1 if d_square <= d_s*2 else 0
+    return mask
+
+
+def box_low_pass_filter(shape, d_s=0.25, d_t=0.25):
+    """
+    Compute the ideal low pass filter mask (approximated version).
+
+    Args:
+        shape: shape of the filter (volume)
+        d_s: normalized stop frequency for spatial dimensions (0.0-1.0)
+        d_t: normalized stop frequency for temporal dimension (0.0-1.0)
+    """
+    T, H, W = shape[-4], shape[-2], shape[-1]
+    mask = torch.zeros(shape)
+    if d_s==0 or d_t==0:
+        return mask
+
+    threshold_s = round(int(H // 2) * d_s)
+    threshold_t = round(T // 2 * d_t)
+
+    cframe, crow, ccol = T // 2, H // 2, W //2
+    mask[cframe - threshold_t:cframe + threshold_t, ..., crow - threshold_s:crow + threshold_s, ccol - threshold_s:ccol + threshold_s] = 1.0
+
+    return mask
+

animatediff/model_injection.py

@@ -11,9 +11,10 @@ from comfy.model_patcher import ModelPatcher
 
 from .motion_module_ad import AnimateDiffModel, has_mid_block, normalize_ad_state_dict
 from .logger import logger
-from .motion_utils import MotionCompatibilityError, NoiseType, normalize_min_max
+from .motion_utils import MotionCompatibilityError, 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
 
 
 # some motion_model casts here might fail if model becomes metatensor or is not castable;
@@ -32,8 +33,8 @@ class ModelPatcherAndInjector(ModelPatcher):
 
         # injection stuff
         self.motion_injection_params: InjectionParams = None
+        self.sample_settings: SampleSettings = SampleSettings()
         self.motion_model: MotionModelPatcher = None
-        self.motion_model_sampling = None
     
     def model_patches_to(self, device):
         super().model_patches_to(device)
@@ -76,8 +77,8 @@ class ModelPatcherAndInjector(ModelPatcher):
     def clone(self):
         cloned = ModelPatcherAndInjector(self)
         cloned.motion_model = self.motion_model
+        cloned.sample_settings = self.sample_settings
         cloned.motion_injection_params = self.motion_injection_params
-        cloned.motion_model_sampling = self.motion_model_sampling
         return cloned
 
 
@@ -322,9 +323,12 @@ class InjectionParams:
         self.sync_context_to_pe = False
         self.loras: MotionLoraList = None
         self.motion_model_settings = MotionModelSettings()
-        self.noise_type: str = NoiseType.DEFAULT
+        self.sampling_settings = None
+        self.noise_type: str = SeedNoiseGeneration.COMFY
         self.sub_idxs = None  # value should NOT be included in clone, so it will auto reset
     
+    def set_noise_extra_args(self, noise_extra_args: dict):
+        noise_extra_args["context_length"] = self.context_length
 
     def set_context(self, context_length: int, context_stride: int, context_overlap: int, context_schedule: str, closed_loop: bool, sync_context_to_pe: bool=False):
         self.context_length = context_length

animatediff/motion_module_ad.py

@@ -10,8 +10,10 @@ from comfy.ldm.modules.attention import FeedForward, SpatialTransformer
 from comfy.model_patcher import ModelPatcher
 from comfy.ldm.modules.diffusionmodules import openaimodel
 from comfy.ldm.modules.diffusionmodules.openaimodel import SpatialTransformer
+from comfy.controlnet import broadcast_image_to
+from comfy.utils import repeat_to_batch_size
 
-from .motion_utils import GroupNormAD, BlockType, CrossAttentionMM, MotionCompatibilityError, TemporalTransformerGeneric
+from .motion_utils import GroupNormAD, BlockType, CrossAttentionMM, MotionCompatibilityError, prepare_mask_batch
 from .model_utils import ModelTypeSD
 
 
@@ -390,7 +392,7 @@ class VanillaTemporalModule(nn.Module):
         return output_tensor
 
 
-class TemporalTransformer3DModel(nn.Module, TemporalTransformerGeneric):
+class TemporalTransformer3DModel(nn.Module):
     def __init__(
         self,
         in_channels,
@@ -412,7 +414,15 @@ class TemporalTransformer3DModel(nn.Module, TemporalTransformerGeneric):
         temporal_position_encoding_max_len=24,
     ):
         super().__init__()
-        super().temporal_transformer_init(default_length=16)
+        self.video_length = 16
+        self.full_length = 16
+        self.scale_min = 1.0
+        self.scale_max = 1.0
+        self.raw_scale_mask: Union[Tensor, None] = None
+        self.temp_scale_mask: Union[Tensor, None] = None
+        self.sub_idxs: Union[list[int], None] = None
+        self.prev_hidden_states_batch = 0
+
 
         inner_dim = num_attention_heads * attention_head_dim
 
@@ -461,6 +471,54 @@ class TemporalTransformer3DModel(nn.Module, TemporalTransformerGeneric):
         for block in self.transformer_blocks:
             block.set_sub_idxs(sub_idxs)
 
+    def reset_temp_vars(self):
+        del self.temp_scale_mask
+        self.temp_scale_mask = None
+        self.prev_hidden_states_batch = 0
+
+    def get_scale_mask(self, hidden_states: Tensor) -> Union[Tensor, None]:
+        # if no raw mask, return None
+        if self.raw_scale_mask is None:
+            return None
+        shape = hidden_states.shape
+        batch, channel, height, width = shape
+        # if temp mask already calculated, return it
+        if self.temp_scale_mask != None:
+            # check if hidden_states batch matches
+            if batch == self.prev_hidden_states_batch:
+                if self.sub_idxs is not None:
+                    return self.temp_scale_mask[:, self.sub_idxs, :]
+                return self.temp_scale_mask
+            # if does not match, reset cached temp_scale_mask and recalculate it
+            del self.temp_scale_mask
+            self.temp_scale_mask = None
+        # otherwise, calculate temp mask
+        self.prev_hidden_states_batch = batch
+        mask = prepare_mask_batch(self.raw_scale_mask, shape=(self.full_length, 1, height, width))
+        mask = repeat_to_batch_size(mask, self.full_length)
+        # if mask not the same amount length as full length, make it match
+        if self.full_length != mask.shape[0]:
+            mask = broadcast_image_to(mask, self.full_length, 1)
+        # reshape mask to attention K shape (h*w, latent_count, 1)
+        batch, channel, height, width = mask.shape
+        # first, perform same operations as on hidden_states,
+        # turning (b, c, h, w) -> (b, h*w, c)
+        mask = mask.permute(0, 2, 3, 1).reshape(batch, height*width, channel)
+        # then, make it the same shape as attention's k, (h*w, b, c)
+        mask = mask.permute(1, 0, 2)
+        # make masks match the expected length of h*w
+        batched_number = shape[0] // self.video_length
+        if batched_number > 1:
+            mask = torch.cat([mask] * batched_number, dim=0)
+        # cache mask and set to proper device
+        self.temp_scale_mask = mask
+        # move temp_scale_mask to proper dtype + device
+        self.temp_scale_mask = self.temp_scale_mask.to(dtype=hidden_states.dtype, device=hidden_states.device)
+        # return subset of masks, if needed
+        if self.sub_idxs is not None:
+            return self.temp_scale_mask[:, self.sub_idxs, :]
+        return self.temp_scale_mask
+
     def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None):
         batch, channel, height, width = hidden_states.shape
         residual = hidden_states

animatediff/motion_utils.py

@@ -1,6 +1,3 @@
-from abc import ABC, abstractmethod
-from typing import Union
-
 import torch
 import torch.nn.functional as F
 from torch import Tensor, nn
@@ -9,10 +6,7 @@ import comfy.model_management as model_management
 import comfy.ops
 from comfy.cli_args import args
 from comfy.ldm.modules.attention import attention_basic, attention_pytorch, attention_split, attention_sub_quad, default
-from comfy.controlnet import broadcast_image_to
-from comfy.utils import repeat_to_batch_size
 
-from .motion_lora import MotionLoraInfo
 from .logger import logger
 
 
@@ -72,123 +66,12 @@ class CrossAttentionMM(nn.Module):
         return self.to_out(out)
 
 
-# super class to TemporalTransformer-like classes
-class TemporalTransformerGeneric:
-    def temporal_transformer_init(self, default_length: int):
-        self.video_length = default_length
-        self.full_length = default_length
-        self.scale_min = 1.0
-        self.scale_max = 1.0
-        self.raw_scale_mask: Union[Tensor, None] = None
-        self.temp_scale_mask: Union[Tensor, None] = None
-        self.sub_idxs: Union[list[int], None] = None
-        self.prev_hidden_states_batch = 0
-
-    def reset_temp_vars(self):
-        del self.temp_scale_mask
-        self.temp_scale_mask = None
-        self.prev_hidden_states_batch = 0
-
-    def get_scale_mask(self, hidden_states: Tensor) -> Union[Tensor, None]:
-        # if no raw mask, return None
-        if self.raw_scale_mask is None:
-            return None
-        shape = hidden_states.shape
-        batch, channel, height, width = shape
-        # if temp mask already calculated, return it
-        if self.temp_scale_mask != None:
-            # check if hidden_states batch matches
-            if batch == self.prev_hidden_states_batch:
-                if self.sub_idxs is not None:
-                    return self.temp_scale_mask[:, self.sub_idxs, :]
-                return self.temp_scale_mask
-            # if does not match, reset cached temp_scale_mask and recalculate it
-            del self.temp_scale_mask
-            self.temp_scale_mask = None
-        # otherwise, calculate temp mask
-        self.prev_hidden_states_batch = batch
-        mask = prepare_mask_batch(self.raw_scale_mask, shape=(self.full_length, 1, height, width))
-        mask = repeat_to_batch_size(mask, self.full_length)
-        # if mask not the same amount length as full length, make it match
-        if self.full_length != mask.shape[0]:
-            mask = broadcast_image_to(mask, self.full_length, 1)
-        # reshape mask to attention K shape (h*w, latent_count, 1)
-        batch, channel, height, width = mask.shape
-        # first, perform same operations as on hidden_states,
-        # turning (b, c, h, w) -> (b, h*w, c)
-        mask = mask.permute(0, 2, 3, 1).reshape(batch, height*width, channel)
-        # then, make it the same shape as attention's k, (h*w, b, c)
-        mask = mask.permute(1, 0, 2)
-        # make masks match the expected length of h*w
-        batched_number = shape[0] // self.video_length
-        if batched_number > 1:
-            mask = torch.cat([mask] * batched_number, dim=0)
-        # cache mask and set to proper device
-        self.temp_scale_mask = mask
-        # move temp_scale_mask to proper dtype + device
-        self.temp_scale_mask = self.temp_scale_mask.to(dtype=hidden_states.dtype, device=hidden_states.device)
-        # return subset of masks, if needed
-        if self.sub_idxs is not None:
-            return self.temp_scale_mask[:, self.sub_idxs, :]
-        return self.temp_scale_mask
-
-
 class BlockType:
     UP = "up"
     DOWN = "down"
     MID = "mid"
 
 
-class GenericMotionWrapper(nn.Module, ABC):
-    def __init__(self, mm_hash: str, mm_name: str, loras: list[MotionLoraInfo]):
-        super().__init__()
-        self.down_blocks: nn.ModuleList = None
-        self.up_blocks: nn.ModuleList = None
-        self.mid_block = None
-        self.mm_hash = mm_hash
-        self.mm_name = mm_name
-        self.version = "FILLTHISIN"
-        self.injector_version = "VERYIMPORTANT_FILLTHISIN"
-        self.AD_video_length: int = 0
-        self.loras = loras
-
-    def has_loras(self) -> bool:
-        # TODO: fix this to return False if has an empty list as well
-        # but only after implementing a fix for lowvram loading
-        return self.loras is not None
-
-    @abstractmethod
-    def set_video_length(self, video_length: int, full_length: int):
-        pass
-
-    @abstractmethod
-    def set_scale_multiplier(self, multiplier: Union[float, None]):
-        pass
-
-    @abstractmethod
-    def set_masks(self, masks: Tensor, min_val: float, max_val: float):
-        pass
-
-    @abstractmethod
-    def set_sub_idxs(self, sub_idxs: list[int]):
-        pass
-    
-    @abstractmethod
-    def reset_temp_vars(self):
-        pass
-
-    def reset_scale_multiplier(self):
-        self.set_scale_multiplier(None)
-
-    def reset_sub_idxs(self):
-        self.set_sub_idxs(None)
-
-    def reset(self):
-        self.reset_sub_idxs()
-        self.reset_scale_multiplier()
-        self.reset_temp_vars()
-
-
 class GroupNormAD(torch.nn.GroupNorm):
     def __init__(self, num_groups: int, num_channels: int, eps: float = 1e-5, affine: bool = True,
                  device=None, dtype=None) -> None:
@@ -215,60 +98,5 @@ def prepare_mask_batch(mask: Tensor, shape: Tensor, multiplier: int=1, match_dim
     return mask
 
 
-class NoiseType:
-    DEFAULT = "default"
-    REPEATED = "repeated"
-    CONSTANT = "constant"
-    AUTO1111 = "auto1111"
-
-    LIST = [DEFAULT, REPEATED, CONSTANT, AUTO1111]
-
-    @classmethod
-    def prepare_noise(cls, noise_type: str, latents: Tensor, noise: Tensor, context_length: int, seed: int):
-        if noise_type == cls.DEFAULT:
-            return noise
-        elif noise_type == cls.REPEATED:
-            return cls.prepare_noise_repeated(latents, noise, context_length, seed)
-        elif noise_type == cls.CONSTANT:
-            return cls.prepare_noise_constant(latents, noise, context_length, seed)
-        elif noise_type == cls.AUTO1111:
-            return cls.prepare_noise_auto1111(latents, noise, context_length, seed)
-        logger.warning(f"Noise type {noise_type} not recognized, proceeding with default noise.")
-        return noise
-
-    @classmethod
-    def prepare_noise_repeated(cls, latents: Tensor, noise: Tensor, context_length: int, seed: int):
-        if not context_length:
-            return noise
-        length = latents.shape[0]
-        generator = torch.manual_seed(seed)
-        noise = torch.randn(latents.size(), dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu")
-        noise_set = noise[:context_length]
-        cat_count = (length // context_length) + 1
-        noise_set = torch.cat([noise_set] * cat_count, dim=0)
-        noise_set = noise_set[:length]
-        return noise_set
-
-    @classmethod
-    def prepare_noise_constant(cls, latents: Tensor, noise: Tensor, context_length: int, seed: int):
-        length = latents.shape[0]
-        single_shape = (1, latents.shape[1], latents.shape[2], latents.shape[3])
-        generator = torch.manual_seed(seed)
-        noise = torch.randn(single_shape, dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu")
-        return torch.cat([noise] * length, dim=0)
-
-    @classmethod
-    def prepare_noise_auto1111(cls, latents: Tensor, noise: Tensor, context_length: int, seed: int):
-        # auto1111 applies growing seeds for a batch
-        length = latents.shape[0]
-        single_shape = (1, latents.shape[1], latents.shape[2], latents.shape[3])
-        all_noises = []
-        # i starts at 0
-        for i in range(length):
-            generator = torch.manual_seed(seed+i)
-            all_noises.append(torch.randn(single_shape, dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu"))
-        return torch.cat(all_noises, dim=0)
-
-
 class MotionCompatibilityError(ValueError):
     pass

animatediff/nodes.py

@@ -7,18 +7,19 @@ 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 NoiseType
 from .motion_lora import MotionLoraInfo, MotionLoraList
 from .model_injection import InjectionParams, ModelPatcherAndInjector, MotionModelSettings, load_motion_module
+from .sample_settings import SampleSettings, SeedNoiseGeneration
 from .sampling import motion_sample_factory
 
+from .nodes_sample import FreeInitOptionsNode, NoiseLayerAddWeightedNode, SampleSettingsNode, NoiseLayerAddNode, NoiseLayerReplaceNode, IterationOptionsNode
 from .nodes_extras import AnimateDiffUnload, EmptyLatentImageLarge, CheckpointLoaderSimpleWithNoiseSelect
 from .nodes_experimental import AnimateDiffModelSettingsSimple, AnimateDiffModelSettingsAdvanced, AnimateDiffModelSettingsAdvancedAttnStrengths
 from .nodes_deprecated import AnimateDiffLoader_Deprecated, AnimateDiffLoaderAdvanced_Deprecated, AnimateDiffCombine_Deprecated
 
 # override comfy_sample.sample with animatediff-support version
 comfy_sample.sample = motion_sample_factory(comfy_sample.sample)
-comfy_sample.sample_custom = motion_sample_factory(comfy_sample.sample_custom)
+comfy_sample.sample_custom = motion_sample_factory(comfy_sample.sample_custom, is_custom=True)
 
 
 class AnimateDiffModelSettings:
@@ -95,6 +96,7 @@ class AnimateDiffLoaderWithContext:
                 "context_options": ("CONTEXT_OPTIONS",),
                 "motion_lora": ("MOTION_LORA",),
                 "motion_model_settings": ("MOTION_MODEL_SETTINGS",),
+                "sample_settings": ("sample_settings",),
                 "motion_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
                 "apply_v2_models_properly": ("BOOLEAN", {"default": True}),
             }
@@ -109,7 +111,7 @@ class AnimateDiffLoaderWithContext:
         model: ModelPatcher,
         model_name: str, beta_schedule: str,# apply_mm_groupnorm_hack: bool,
         context_options: ContextOptions=None, motion_lora: MotionLoraList=None, motion_model_settings: MotionModelSettings=None,
-        motion_scale: float=1.0, apply_v2_models_properly: bool=False,
+        sample_settings: SampleSettings=None, motion_scale: float=1.0, apply_v2_models_properly: bool=False,
     ):
         # load motion module
         motion_model = load_motion_module(model_name, model, motion_lora=motion_lora, motion_model_settings=motion_model_settings)
@@ -133,7 +135,6 @@ class AnimateDiffLoaderWithContext:
                         closed_loop=context_options.closed_loop,
                         sync_context_to_pe=context_options.sync_context_to_pe,
                 )
-                params.noise_type = context_options.noise_type
         if motion_lora:
             params.set_loras(motion_lora)
         # set motion_scale and motion_model_settings
@@ -154,6 +155,7 @@ class AnimateDiffLoaderWithContext:
 
         model = ModelPatcherAndInjector(model)
         model.motion_model = motion_model
+        model.sample_settings = sample_settings if sample_settings is not None else SampleSettings()
         model.motion_injection_params = params
 
         # save model sampling from BetaSchedule as object patch
@@ -195,45 +197,19 @@ class AnimateDiffUniformContextOptions:
         return (context_options,)
 
 
-class AnimateDiffUniformContextOptionsExperimental:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "context_length": ("INT", {"default": 16, "min": 1, "max": 32}), # keep an eye on these max values
-                "context_stride": ("INT", {"default": 1, "min": 1, "max": 32}),  # would need to be updated
-                "context_overlap": ("INT", {"default": 0, "min": 0, "max": 32}), # if new motion modules come out
-                "context_schedule": (ContextSchedules.CONTEXT_SCHEDULE_LIST,),
-                "closed_loop": ("BOOLEAN", {"default": False},),
-                "noise_override": (NoiseType.LIST,)
-                #"sync_context_to_pe": ("BOOLEAN", {"default": False},),
-            },
-        }
-    
-    RETURN_TYPES = ("CONTEXT_OPTIONS",)
-    CATEGORY = "Animate Diff 🎭🅐🅓"
-    FUNCTION = "create_options"
-
-    def create_options(self, context_length: int, context_stride: int, context_overlap: int, context_schedule: int, closed_loop: bool,
-                       noise_override: str):
-        context_options = UniformContextOptions(
-            context_length=context_length,
-            context_stride=context_stride,
-            context_overlap=context_overlap,
-            context_schedule=context_schedule,
-            closed_loop=closed_loop,
-            )
-        context_options.set_noise_type(noise_override)
-        #context_options.set_sync_context_to_pe(sync_context_to_pe)
-        return (context_options,)
-
-
 NODE_CLASS_MAPPINGS = {
     "ADE_AnimateDiffUniformContextOptions": AnimateDiffUniformContextOptions,
-    #"ADE_AnimateDiffUniformContextOptionsExperimental": AnimateDiffUniformContextOptionsExperimental,
+    "ADE_AnimateDiffSamplingSettings": SampleSettingsNode,
     "ADE_AnimateDiffLoaderWithContext": AnimateDiffLoaderWithContext,
     "ADE_AnimateDiffLoRALoader": AnimateDiffLoraLoader,
     "ADE_AnimateDiffModelSettings_Release": AnimateDiffModelSettings,
+    # Noise Layer Nodes
+    "ADE_NoiseLayerAdd": NoiseLayerAddNode,
+    "ADE_NoiseLayerAddWeighted": NoiseLayerAddWeightedNode,
+    "ADE_NoiseLayerReplace": NoiseLayerReplaceNode,
+    # Iteration Opts
+    "ADE_IterationOptsDefault": IterationOptionsNode,
+    "ADE_IterationOptsFreeInit": FreeInitOptionsNode,
     # Experimental Nodes
     "ADE_AnimateDiffModelSettingsSimple": AnimateDiffModelSettingsSimple,
     "ADE_AnimateDiffModelSettings": AnimateDiffModelSettingsAdvanced,
@@ -249,10 +225,17 @@ NODE_CLASS_MAPPINGS = {
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "ADE_AnimateDiffUniformContextOptions": "Uniform Context Options 🎭🅐🅓",
-    #"ADE_AnimateDiffUniformContextOptionsExperimental": "EXPERIMENTAL Uniform Context Options 🎭🅐🅓",
+    "ADE_AnimateDiffSamplingSettings": "Sample Settings 🎭🅐🅓",
     "ADE_AnimateDiffLoaderWithContext": "AnimateDiff Loader 🎭🅐🅓",
     "ADE_AnimateDiffLoRALoader": "AnimateDiff LoRA Loader 🎭🅐🅓",
     "ADE_AnimateDiffModelSettings_Release": "Motion Model Settings 🎭🅐🅓",
+    # Noise Layer Nodes
+    "ADE_NoiseLayerAdd": "Noise Layer [Add] 🎭🅐🅓",
+    "ADE_NoiseLayerAddWeighted": "Noise Layer [Add Weighted] 🎭🅐🅓",
+    "ADE_NoiseLayerReplace": "Noise Layer [Replace] 🎭🅐🅓",
+    # Iteration Opts
+    "ADE_IterationOptsDefault": "Default Iteration Options 🎭🅐🅓",
+    "ADE_IterationOptsFreeInit": "FreeInit Iteration Options 🎭🅐🅓",
     # Experimental Nodes
     "ADE_AnimateDiffModelSettingsSimple": "EXP Motion Model Settings (Simple) 🎭🅐🅓",
     "ADE_AnimateDiffModelSettings": "EXP Motion Model Settings (Advanced) 🎭🅐🅓",

animatediff/nodes_extras.py

@@ -37,19 +37,25 @@ class CheckpointLoaderSimpleWithNoiseSelect:
                 "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
                 "beta_schedule": (BetaSchedules.ALIAS_LIST, {"default": BetaSchedules.LINEAR}, )
             },
+            "optional": {
+                "use_custom_scale_factor": ("BOOLEAN", {"default": False}),
+                "scale_factor": ("FLOAT", {"default": 0.18215, "min": 0.0, "max": 1.0, "step": 0.00001})
+            }
         }
     RETURN_TYPES = ("MODEL", "CLIP", "VAE")
     FUNCTION = "load_checkpoint"
 
     CATEGORY = "Animate Diff 🎭🅐🅓/extras"
 
-    def load_checkpoint(self, ckpt_name, beta_schedule, output_vae=True, output_clip=True):
+    def load_checkpoint(self, ckpt_name, beta_schedule, output_vae=True, output_clip=True, use_custom_scale_factor=False, scale_factor=0.18215):
         ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
         out = load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
         # register chosen beta schedule on model - convert to beta_schedule name recognized by ComfyUI
         new_model_sampling = BetaSchedules.to_model_sampling(beta_schedule, out[0])
         if new_model_sampling is not None:
             out[0].model.model_sampling = new_model_sampling
+        if use_custom_scale_factor:
+            out[0].model.latent_format.scale_factor = scale_factor
         return out
 
 

animatediff/nodes_sample.py

@@ -0,0 +1,186 @@
+from torch import Tensor
+
+from .freeinit import FreeInitFilter
+from .sample_settings import FreeInitOptions, IterationOptions, NoiseLayerAdd, NoiseLayerAddWeighted, NoiseLayerGroup, NoiseLayerReplace, NoiseLayerType, SeedNoiseGeneration, SampleSettings
+
+
+class SampleSettingsNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "batch_offset": ("INT", {"default": 0, "min": 0}),
+                "noise_type": (NoiseLayerType.LIST,),
+                "seed_gen": (SeedNoiseGeneration.LIST,),
+                "seed_offset": ("INT", {"default": 0, "min": -999999999999999}),
+            },
+            "optional": {
+                "noise_layers": ("NOISE_LAYERS",),
+                "iteration_opts": ("ITERATION_OPTS",),
+                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
+            }
+        }
+    
+    RETURN_TYPES = ("SAMPLE_SETTINGS",)
+    RETURN_NAMES = ("settings",)
+    CATEGORY = "Animate Diff 🎭🅐🅓"
+    FUNCTION = "create_settings"
+
+    def create_settings(self, batch_offset: int, noise_type: str, seed_gen: str, seed_offset: int, noise_layers: NoiseLayerGroup=None, iteration_opts: IterationOptions=None, seed_override: int=None,):
+        sampling_settings = SampleSettings(batch_offset=batch_offset, noise_type=noise_type, seed_gen=seed_gen, seed_offset=seed_offset, noise_layers=noise_layers, iteration_opts=iteration_opts, seed_override=seed_override)
+        return (sampling_settings,)
+
+
+class NoiseLayerReplaceNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "batch_offset": ("INT", {"default": 0, "min": 0}),
+                "noise_type": (NoiseLayerType.LIST,),
+                "seed_gen_override": (SeedNoiseGeneration.LIST_WITH_OVERRIDE,),
+                "seed_offset": ("INT", {"default": 0, "min": -999999999999999}),
+            },
+            "optional": {
+                "prev_noise_layers": ("NOISE_LAYERS",),
+                "mask_optional": ("MASK",),
+                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
+            }
+        }
+    
+    RETURN_TYPES = ("NOISE_LAYERS",)
+    CATEGORY = "Animate Diff 🎭🅐🅓/noise layers"
+    FUNCTION = "create_layers"
+
+    def create_layers(self, batch_offset: int, noise_type: str, seed_gen_override: str, seed_offset: int,
+                      prev_noise_layers: NoiseLayerGroup=None, mask_optional: Tensor=None, seed_override: int=None,):
+        # prepare prev_noise_layers
+        if prev_noise_layers is None:
+            prev_noise_layers = NoiseLayerGroup()
+        prev_noise_layers = prev_noise_layers.clone()
+        # create layer
+        layer = NoiseLayerReplace(noise_type=noise_type, batch_offset=batch_offset, seed_gen_override=seed_gen_override, seed_offset=seed_offset,
+                                  seed_override=seed_override, mask=mask_optional)
+        prev_noise_layers.add_to_start(layer)
+        return (prev_noise_layers,)
+
+
+class NoiseLayerAddNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "batch_offset": ("INT", {"default": 0, "min": 0}),
+                "noise_type": (NoiseLayerType.LIST,),
+                "seed_gen_override": (SeedNoiseGeneration.LIST_WITH_OVERRIDE,),
+                "seed_offset": ("INT", {"default": 0, "min": -999999999999999}),
+                "noise_weight": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 10.0, "step": 0.001}),
+            },
+            "optional": {
+                "prev_noise_layers": ("NOISE_LAYERS",),
+                "mask_optional": ("MASK",),
+                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
+            }
+        }
+    
+    RETURN_TYPES = ("NOISE_LAYERS",)
+    CATEGORY = "Animate Diff 🎭🅐🅓/noise layers"
+    FUNCTION = "create_layers"
+
+    def create_layers(self, batch_offset: int, noise_type: str, seed_gen_override: str, seed_offset: int,
+                      noise_weight: float,
+                      prev_noise_layers: NoiseLayerGroup=None, mask_optional: Tensor=None, seed_override: int=None,):
+        # prepare prev_noise_layers
+        if prev_noise_layers is None:
+            prev_noise_layers = NoiseLayerGroup()
+        prev_noise_layers = prev_noise_layers.clone()
+        # create layer
+        layer = NoiseLayerAdd(noise_type=noise_type, batch_offset=batch_offset, seed_gen_override=seed_gen_override, seed_offset=seed_offset,
+                              seed_override=seed_override, mask=mask_optional,
+                              noise_weight=noise_weight)
+        prev_noise_layers.add_to_start(layer)
+        return (prev_noise_layers,)
+
+
+class NoiseLayerAddWeightedNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "batch_offset": ("INT", {"default": 0, "min": 0}),
+                "noise_type": (NoiseLayerType.LIST,),
+                "seed_gen_override": (SeedNoiseGeneration.LIST_WITH_OVERRIDE,),
+                "seed_offset": ("INT", {"default": 0, "min": -999999999999999}),
+                "noise_weight": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 10.0, "step": 0.001}),
+                "balance_multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "step": 0.001}),
+            },
+            "optional": {
+                "prev_noise_layers": ("NOISE_LAYERS",),
+                "mask_optional": ("MASK",),
+                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
+            }
+        }
+    
+    RETURN_TYPES = ("NOISE_LAYERS",)
+    CATEGORY = "Animate Diff 🎭🅐🅓/noise layers"
+    FUNCTION = "create_layers"
+
+    def create_layers(self, batch_offset: int, noise_type: str, seed_gen_override: str, seed_offset: int,
+                      noise_weight: float, balance_multiplier: float,
+                      prev_noise_layers: NoiseLayerGroup=None, mask_optional: Tensor=None, seed_override: int=None,):
+        # prepare prev_noise_layers
+        if prev_noise_layers is None:
+            prev_noise_layers = NoiseLayerGroup()
+        prev_noise_layers = prev_noise_layers.clone()
+        # create layer
+        layer = NoiseLayerAddWeighted(noise_type=noise_type, batch_offset=batch_offset, seed_gen_override=seed_gen_override, seed_offset=seed_offset,
+                              seed_override=seed_override, mask=mask_optional,
+                              noise_weight=noise_weight, balance_multiplier=balance_multiplier)
+        prev_noise_layers.add_to_start(layer)
+        return (prev_noise_layers,)
+
+
+class IterationOptionsNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "iterations": ("INT", {"default": 1, "min": 1}),
+            }
+        }
+    
+    RETURN_TYPES = ("ITERATION_OPTS",)
+    CATEGORY = "Animate Diff 🎭🅐🅓/iteration opts"
+    FUNCTION = "create_iter_opts"
+
+    def create_iter_opts(self, iterations: int):
+        iter_opts = IterationOptions(iterations=iterations)
+        return (iter_opts,)
+
+
+class FreeInitOptionsNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "iterations": ("INT", {"default": 2, "min": 1}),
+                "filter": (FreeInitFilter.LIST,),
+                "d_s": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.001}),
+                "d_t": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.001}),
+                "n_butterworth": ("INT", {"default": 4, "min": 1, "max": 100},),
+                "sigma_step": ("INT", {"default": 999, "min": 1, "max": 999}),
+                "apply_to_1st_iter": ("BOOLEAN", {"default": False}),
+                "init_type": (FreeInitOptions.LIST,)
+            }
+        }
+    
+    RETURN_TYPES = ("ITERATION_OPTS",)
+    CATEGORY = "Animate Diff 🎭🅐🅓/iteration opts"
+    FUNCTION = "create_iter_opts"
+
+    def create_iter_opts(self, iterations: int, filter: str, d_s: float, d_t: float, n_butterworth: int,
+                         sigma_step: int, apply_to_1st_iter: bool, init_type: str):
+        # init_type does nothing for now, not until I add more methods of applying low+high freq noise
+        iter_opts = FreeInitOptions(iterations=iterations, step=sigma_step, apply_to_1st_iter=apply_to_1st_iter,
+                                    filter=filter, d_s=d_s, d_t=d_t, n=n_butterworth)
+        return (iter_opts,)

animatediff/sample_settings.py

@@ -0,0 +1,366 @@
+from collections.abc import Iterable
+import torch
+from torch import Tensor
+
+import comfy.sample
+from comfy.model_patcher import ModelPatcher
+
+from .freeinit import FreeInitFilter, freq_mix_3d, get_freq_filter
+from .logger import logger
+
+
+def prepare_mask_ad(noise_mask, shape, device):
+    """ensures noise mask is of proper dimensions"""
+    noise_mask = torch.nn.functional.interpolate(noise_mask.reshape((-1, 1, noise_mask.shape[-2], noise_mask.shape[-1])), size=(shape[2], shape[3]), mode="bilinear")
+    #noise_mask = noise_mask.round()
+    noise_mask = torch.cat([noise_mask] * shape[1], dim=1)
+    noise_mask = comfy.utils.repeat_to_batch_size(noise_mask, shape[0])
+    noise_mask = noise_mask.to(device)
+    return noise_mask
+
+
+class NoiseLayerType:
+    DEFAULT = "default"
+    CONSTANT = "constant"
+    EMPTY = "empty"
+    REPEATED_CONTEXT = "repeated_context"
+
+    LIST = [DEFAULT, CONSTANT, EMPTY, REPEATED_CONTEXT]
+
+
+class NoiseApplication:
+    ADD = "add"
+    ADD_WEIGHTED = "add_weighted"
+    REPLACE = "replace"
+    
+    LIST = [ADD, ADD_WEIGHTED, REPLACE]
+
+
+class NoiseNormalize:
+    DISABLE = "disable"
+    NORMAL = "normal"
+
+    LIST = [DISABLE, NORMAL]
+
+
+class SampleSettings:
+    def __init__(self, batch_offset: int=0, noise_type: str=None, seed_gen: str=None, seed_offset: int=0, noise_layers: 'NoiseLayerGroup'=None, iteration_opts=None, seed_override:int=None, negative_cond_flipflop=False):
+        self.batch_offset = batch_offset
+        self.noise_type = noise_type if noise_type is not None else NoiseLayerType.DEFAULT
+        self.seed_gen = seed_gen if seed_gen is not None else SeedNoiseGeneration.COMFY
+        self.noise_layers = noise_layers if noise_layers else NoiseLayerGroup()
+        self.iteration_opts = iteration_opts if iteration_opts else IterationOptions()
+        self.seed_offset = seed_offset
+        self.seed_override = seed_override
+        self.negative_cond_flipflop = negative_cond_flipflop
+    
+    def prepare_noise(self, seed: int, latents: Tensor, noise: Tensor, extra_seed_offset=0, extra_args:dict={}, force_create_noise=True):
+        if self.seed_override is not None:
+            seed = self.seed_override
+        # if seed is iterable, attempt to do per-latent noises
+        if isinstance(seed, Iterable):
+            noise = SeedNoiseGeneration.create_noise_individual_seeds(seeds=seed, latents=latents, seed_offset=self.seed_offset+extra_seed_offset, extra_args=extra_args)
+            seed = seed[0]+self.seed_offset
+        else:
+            seed += self.seed_offset
+            # replace initial noise if not batch_offset 0 or Comfy seed_gen or not NoiseType default
+            if self.batch_offset != 0 or self.seed_offset != 0 or self.noise_type != NoiseLayerType.DEFAULT or self.seed_gen != SeedNoiseGeneration.COMFY or force_create_noise:
+                noise = SeedNoiseGeneration.create_noise(seed=seed+extra_seed_offset, latents=latents, existing_seed_gen=self.seed_gen, seed_gen=self.seed_gen,
+                                                        noise_type=self.noise_type, batch_offset=self.batch_offset, extra_args=extra_args)
+        # apply noise layers
+        for noise_layer in self.noise_layers.layers:
+            # first, generate new noise matching seed gen override
+            layer_noise = noise_layer.create_layer_noise(existing_seed_gen=self.seed_gen, seed=seed, latents=latents,
+                                                         extra_seed_offset=extra_seed_offset, extra_args=extra_args)
+            # next, get noise after applying layer
+            noise = noise_layer.apply_layer_noise(new_noise=layer_noise, old_noise=noise)
+        # noise prepared now
+        return noise
+
+
+class NoiseLayer:
+    def __init__(self, noise_type: str, batch_offset: int, seed_gen_override: str, seed_offset: int, seed_override: int=None, mask: Tensor=None):
+        self.application: str = NoiseApplication.REPLACE
+        self.noise_type = noise_type
+        self.batch_offset = batch_offset
+        self.seed_gen_override = seed_gen_override
+        self.seed_offset = seed_offset
+        self.seed_override = seed_override
+        self.mask = mask
+    
+    def create_layer_noise(self, existing_seed_gen: str, seed: int, latents: Tensor, extra_seed_offset=0, extra_args:dict={}) -> Tensor:
+        if self.seed_override is not None:
+            seed = self.seed_override
+         # if seed is iterable, attempt to do per-latent noises
+        if isinstance(seed, Iterable):
+            return SeedNoiseGeneration.create_noise_individual_seeds(seeds=seed, latents=latents, seed_offset=self.seed_offset+extra_seed_offset, extra_args=extra_args)
+        seed += self.seed_offset + extra_seed_offset
+        return SeedNoiseGeneration.create_noise(seed=seed, latents=latents, existing_seed_gen=existing_seed_gen, seed_gen=self.seed_gen_override,
+                                                noise_type=self.noise_type, batch_offset=self.batch_offset, extra_args=extra_args)
+
+    def apply_layer_noise(self, new_noise: Tensor, old_noise: Tensor) -> Tensor:
+        return old_noise
+    
+    def get_noise_mask(self, noise: Tensor) -> Tensor:
+        if self.mask is None:
+            return 1
+        noise_mask = self.mask.reshape((-1, 1, self.mask.shape[-2], self.mask.shape[-1]))
+        return prepare_mask_ad(noise_mask, noise.shape, noise.device)
+
+
+class NoiseLayerReplace(NoiseLayer):
+    def __init__(self, noise_type: str, batch_offset: int, seed_gen_override: str, seed_offset: int, seed_override: int=None, mask: Tensor=None):
+        super().__init__(noise_type, batch_offset, seed_gen_override, seed_offset, seed_override, mask)
+        self.application = NoiseApplication.REPLACE
+
+    def apply_layer_noise(self, new_noise: Tensor, old_noise: Tensor) -> Tensor:
+        noise_mask = self.get_noise_mask(old_noise)
+        return (1-noise_mask)*old_noise + noise_mask*new_noise
+
+
+class NoiseLayerAdd(NoiseLayer):
+    def __init__(self, noise_type: str, batch_offset: int, seed_gen_override: str, seed_offset: int, seed_override: int=None, mask: Tensor=None,
+                 noise_weight=1.0):
+        super().__init__(noise_type, batch_offset, seed_gen_override, seed_offset, seed_override, mask)
+        self.noise_weight = noise_weight
+        self.application = NoiseApplication.ADD
+
+    def apply_layer_noise(self, new_noise: Tensor, old_noise: Tensor) -> Tensor:
+        noise_mask = self.get_noise_mask(old_noise)
+        return (1-noise_mask)*old_noise + noise_mask*(old_noise + new_noise * self.noise_weight)
+
+
+class NoiseLayerAddWeighted(NoiseLayerAdd):
+    def __init__(self, noise_type: str, batch_offset: int, seed_gen_override: str, seed_offset: int, seed_override: int=None, mask: Tensor=None,
+                 noise_weight=1.0, balance_multiplier=1.0):
+        super().__init__(noise_type, batch_offset, seed_gen_override, seed_offset, seed_override, mask, noise_weight)
+        self.balance_multiplier = balance_multiplier
+        self.application = NoiseApplication.ADD_WEIGHTED
+
+    def apply_layer_noise(self, new_noise: Tensor, old_noise: Tensor) -> Tensor:
+        noise_mask = self.get_noise_mask(old_noise)
+        return (1-noise_mask)*old_noise + noise_mask*(old_noise * (1.0-(self.noise_weight*self.balance_multiplier)) + new_noise * self.noise_weight)
+
+
+class NoiseLayerGroup:
+    def __init__(self):
+        self.layers: list[NoiseLayer] = []
+    
+    def add(self, layer: NoiseLayer) -> None:
+        # add to the end of list
+        self.layers.append(layer)
+
+    def add_to_start(self, layer: NoiseLayer) -> None:
+        # add to the beginning of list
+        self.layers.insert(0, layer)
+
+    def __getitem__(self, index) -> NoiseLayer:
+        return self.layers[index]
+    
+    def is_empty(self) -> bool:
+        return len(self.layers) == 0
+    
+    def clone(self) -> 'NoiseLayerGroup':
+        cloned = NoiseLayerGroup()
+        for layer in self.layers:
+            cloned.add(layer)
+        return cloned
+
+class SeedNoiseGeneration:
+    COMFY = "comfy"
+    AUTO1111 = "auto1111"
+    USE_EXISTING = "use existing"
+
+    LIST = [COMFY, AUTO1111]
+    LIST_WITH_OVERRIDE = [USE_EXISTING, COMFY, AUTO1111]
+
+    @classmethod
+    def create_noise(cls, seed: int, latents: Tensor, existing_seed_gen: str=COMFY, seed_gen: str=USE_EXISTING, noise_type: str=NoiseLayerType.DEFAULT, batch_offset: int=0, extra_args: dict={}):
+        # determine if should use existing type
+        if seed_gen == cls.USE_EXISTING:
+            seed_gen = existing_seed_gen
+        if seed_gen == cls.COMFY:
+            return cls.create_noise_comfy(seed, latents, noise_type, batch_offset, extra_args)
+        elif seed_gen == cls.AUTO1111:
+            return cls.create_noise_auto1111(seed, latents, noise_type, batch_offset, extra_args)
+        raise ValueError(f"Noise seed_gen {seed_gen} is not recognized.")
+
+    @staticmethod
+    def create_noise_comfy(seed: int, latents: Tensor, noise_type: str=NoiseLayerType.DEFAULT, batch_offset: int=0, extra_args: dict={}):
+        common_noise = SeedNoiseGeneration._create_common_noise(seed, latents, noise_type, batch_offset, extra_args)
+        if common_noise is not None:
+            return common_noise
+        if noise_type == NoiseLayerType.CONSTANT:
+            generator = torch.manual_seed(seed)
+            length = latents.shape[0]
+            single_shape = (1 + batch_offset, latents.shape[1], latents.shape[2], latents.shape[3])
+            single_noise = torch.randn(single_shape, dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu")
+            return torch.cat([single_noise[batch_offset:]] * length, dim=0)
+        # comfy creates noise with a single seed for the entire shape of the latents batched tensor
+        generator = torch.manual_seed(seed)
+        offset_shape = (latents.shape[0] + batch_offset, latents.shape[1], latents.shape[2], latents.shape[3])
+        final_noise = torch.randn(offset_shape, dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu")
+        final_noise = final_noise[batch_offset:]
+        # convert to derivative noise type, if needed
+        derivative_noise = SeedNoiseGeneration._create_derivative_noise(final_noise, noise_type=noise_type, extra_args=extra_args)
+        if derivative_noise is not None:
+            return derivative_noise
+        return final_noise
+    
+    @staticmethod
+    def create_noise_auto1111(seed: int, latents: Tensor, noise_type: str=NoiseLayerType.DEFAULT, batch_offset: int=0, extra_args: dict={}):
+        common_noise = SeedNoiseGeneration._create_common_noise(seed, latents, noise_type, batch_offset, extra_args)
+        if common_noise is not None:
+            return common_noise
+        if noise_type == NoiseLayerType.CONSTANT:
+            generator = torch.manual_seed(seed+batch_offset)
+            length = latents.shape[0]
+            single_shape = (1, latents.shape[1], latents.shape[2], latents.shape[3])
+            single_noise = torch.randn(single_shape, dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu")
+            return torch.cat([single_noise] * length, dim=0)
+        # auto1111 applies growing seeds for a batch
+        length = latents.shape[0]
+        single_shape = (1, latents.shape[1], latents.shape[2], latents.shape[3])
+        all_noises = []
+        # i starts at 0
+        for i in range(length):
+            generator = torch.manual_seed(seed+i+batch_offset)
+            all_noises.append(torch.randn(single_shape, dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu"))
+        final_noise = torch.cat(all_noises, dim=0)
+        # convert to derivative noise type, if needed
+        derivative_noise = SeedNoiseGeneration._create_derivative_noise(final_noise, noise_type=noise_type, extra_args=extra_args)
+        if derivative_noise is not None:
+            return derivative_noise
+        return final_noise
+    
+    @staticmethod
+    def create_noise_individual_seeds(seeds: list[int], latents: Tensor, seed_offset: int=0, extra_args: dict={}):
+        length = latents.shape[0]
+        if len(seeds) < length:
+            raise ValueError(f"{len(seeds)} seeds in seed_override were provided, but at least {length} are required to work with the current latents.")
+        seeds = seeds[:length]
+        single_shape = (1, latents.shape[1], latents.shape[2], latents.shape[3])
+        all_noises = []
+        for seed in seeds:
+            generator = torch.manual_seed(seed+seed_offset)
+            all_noises.append(torch.randn(single_shape, dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu"))
+        return torch.cat(all_noises, dim=0)
+
+    @staticmethod
+    def _create_common_noise(seed: int, latents: Tensor, noise_type: str=NoiseLayerType.DEFAULT, batch_offset: int=0, extra_args: dict={}):
+        if noise_type == NoiseLayerType.EMPTY:
+            return torch.zeros_like(latents)
+        return None
+    
+    @staticmethod
+    def _create_derivative_noise(noise: Tensor, noise_type: str, extra_args: dict):
+        derivative_func = DERIVATIVE_NOISE_FUNC_MAP.get(noise_type, None)
+        if derivative_func is None:
+            return None
+        return derivative_func(noise=noise, extra_args=extra_args)
+
+    @staticmethod
+    def _convert_to_repeated_context(noise: Tensor, extra_args: dict):
+        # if no context_length, return unmodified noise
+        context_length: int = extra_args["context_length"]
+        if context_length is None:
+            return noise
+        length = noise.shape[0]
+        noise = noise[:context_length]
+        cat_count = (length // context_length) + 1
+        return torch.cat([noise] * cat_count, dim=0)[:length]
+
+
+DERIVATIVE_NOISE_FUNC_MAP = {NoiseLayerType.REPEATED_CONTEXT: SeedNoiseGeneration._convert_to_repeated_context,}
+
+
+class IterationOptions:
+    SAMPLER = "sampler"
+
+    def __init__(self, iterations: int=1, cache_init_noise=False, cache_init_latents=False):
+        self.iterations = iterations
+        self.cache_init_noise = cache_init_noise
+        self.cache_init_latents = cache_init_latents
+        self.need_sampler = False
+
+    def get_sigma(self, model: ModelPatcher, step: int):
+        model_sampling = model.model.model_sampling
+        if "model_sampling" in model.object_patches:
+            model_sampling = model.object_patches["model_sampling"]
+        return model_sampling.sigmas[step]
+
+    def initialize(self, latents: Tensor):
+        pass
+
+    def preprocess_latents(self, curr_i: int, model: ModelPatcher, latents: Tensor, noise: Tensor, **kwargs):
+        return latents, noise
+
+
+class FreeInitOptions(IterationOptions):
+    DINKINIT_V1 = "DinkInit_v1"
+
+    LIST = [DINKINIT_V1]
+
+    def __init__(self, iterations: int, step: int=999, apply_to_1st_iter: bool=False,
+                 filter=FreeInitFilter.GAUSSIAN, d_s=0.25, d_t=0.25, n=4):
+        super().__init__(iterations=iterations, cache_init_noise=True, cache_init_latents=True)
+        self.apply_to_1st_iter = apply_to_1st_iter
+        self.step = step
+        self.filter = filter
+        self.d_s = d_s
+        self.d_t = d_t
+        self.n = n
+        self.freq_filter = None
+        self.need_sampler = True
+
+    def initialize(self, latents: Tensor):
+        self.freq_filter = get_freq_filter(latents.shape, device=latents.device, filter_type=self.filter,
+                                           n=self.n, d_s=self.d_s, d_t=self.d_t)
+    
+    def preprocess_latents(self, curr_i: int, model: ModelPatcher, latents: Tensor, noise: Tensor, cached_latents: Tensor, cached_noise: Tensor,
+                           seed:int, sample_settings: SampleSettings, noise_extra_args: dict, **kwargs):
+        # if first iter and should not apply, do nothing
+        if curr_i == 0 and not self.apply_to_1st_iter:
+            return latents, noise
+        # otherwise, do FreeInit stuff - DinkInit_v1
+        # 1. apply initial noise with appropriate step sigma
+        sigma = self.get_sigma(model, self.step).to(latents.device)
+        alpha_cumprod = 1 / ((sigma * sigma) + 1) #1 / ((sigma * sigma)) # 1 / ((sigma * sigma) + 1)
+        noised_latents = (latents + (cached_noise * sigma)) * alpha_cumprod
+        # 2. create random noise z_rand for high frequency
+        z_rand = sample_settings.prepare_noise(seed=seed, latents=latents, noise=None, extra_seed_offset=curr_i,
+                                                extra_args=noise_extra_args, force_create_noise=True)
+        ####z_rand = torch.randn_like(latents, dtype=latents.dtype, device=latents.device)
+        # 3. noise reinitialization - combines low freq. noise from noised_latents and high freq. noise from z_rand
+        noise = freq_mix_3d(x=noised_latents, noise=z_rand.to(dtype=latents.dtype, device=latents.device), LPF=self.freq_filter)
+        return cached_latents, noise
+    
+        #noised_latents, noise
+
+
+        # otherwise, do FreeInit stuff
+        # 1. apply initial noise with appropriate step sigma
+        sigma = self.get_sigma(model, self.step).to(latents.device)
+        #####latents = model.model.process_latent_in(latents)
+        latents += noise * sigma
+        # 2. create random noise z_rand for high frequency
+        # Note: original implementation does not use a generator for this... could this cause repeatability issues?
+        z_rand = torch.randn_like(latents, dtype=latents.dtype, device=latents.device)
+        # 3. noise reinitialization
+        #latents = freq_mix_3d(x=latents, noise=z_rand, LPF=self.freq_filter)
+        noise = freq_mix_3d(x=noise, noise=z_rand, LPF=self.freq_filter)
+        # treat latents as empty, and freq-mixed latents as noise
+        #####latents = model.model.process_latent_out(latents)
+        return latents, noise
+        ######return torch.zeros_like(latents), latents
+
+        # otherwise, do FreeInit stuff
+        # 1. apply initial noise with appropriate step sigma (default=999)
+        sigma = self.get_sigma(model, self.step).to(latents.device)
+        latents += noise * sigma
+        # 2. create random noise z_rand for high frequency
+        # Note: original implementation does not use a generator for this... could this cause repeatability issues?
+        z_rand = torch.randn_like(latents, dtype=latents.dtype, device=latents.device)
+        # 3. noise reinitialization
+        latents = freq_mix_3d(x=latents, noise=z_rand, LPF=self.freq_filter)
+        return latents, noise

animatediff/sampling.py

@@ -9,13 +9,14 @@ from einops import rearrange
 import comfy.ldm.modules.attention as attention
 from comfy.ldm.modules.diffusionmodules import openaimodel
 import comfy.model_management as model_management
-import comfy.samplers as comfy_samplers
-import comfy.sample as comfy_sample
+import comfy.samplers
+import comfy.sample
 import comfy.utils
 from comfy.controlnet import ControlBase
 
 from .context import get_context_scheduler
-from .motion_utils import GroupNormAD, NoiseType
+from .sample_settings import IterationOptions, SeedNoiseGeneration, prepare_mask_ad
+from .motion_utils import GroupNormAD
 from .model_utils import ModelTypeSD, wrap_function_to_inject_xformers_bug_info
 from .model_injection import InjectionParams, ModelPatcherAndInjector, MotionModelPatcher
 from .motion_module_ad import AnimateDiffFormat, AnimateDiffInfo, AnimateDiffVersion, VanillaTemporalModule
@@ -36,14 +37,6 @@ class AnimateDiffHelper_GlobalState:
         self.last_step: int = 0
         self.current_step: int = 0
         self.total_steps: int = 0
-        self.video_length: int  = 0
-        self.context_frames: int = None
-        self.context_stride: int = None
-        self.context_overlap: int = None
-        self.context_schedule: str = None
-        self.closed_loop: bool = False
-        self.sync_context_to_pe: bool = False
-        self.sub_idxs: list = None
         if self.motion_model is not None:
             del self.motion_model
             self.motion_model = None
@@ -52,26 +45,19 @@ class AnimateDiffHelper_GlobalState:
             self.params = None
     
     def update_with_inject_params(self, params: InjectionParams):
-        self.video_length = params.video_length
-        self.context_frames = params.context_length
-        self.context_stride = params.context_stride
-        self.context_overlap = params.context_overlap
-        self.context_schedule = params.context_schedule
-        self.closed_loop = params.closed_loop
-        self.sync_context_to_pe = params.sync_context_to_pe
         self.params = params
 
     def is_using_sliding_context(self):
-        return self.context_frames is not None
+        return self.params is not None and self.params.context_length is not None
     
     def create_exposed_params(self):
         # This dict will be exposed to be used by other extensions
         # DO NOT change any of the key names
         # or I will find you 👁.👁
         return {
-            "full_length": self.video_length,
-            "context_length": self.context_frames,
-            "sub_idxs": self.sub_idxs,
+            "full_length": self.params.video_length,
+            "context_length": self.params.context_length,
+            "sub_idxs": self.params.sub_idxs,
         }
 
 ADGS = AnimateDiffHelper_GlobalState()
@@ -161,16 +147,6 @@ def get_additional_models_factory(orig_get_additional_models: Callable, motion_m
 ##################################################################################
 
 
-def prepare_mask_ad(noise_mask, shape, device):
-    """ensures noise mask is of proper dimensions"""
-    noise_mask = torch.nn.functional.interpolate(noise_mask.reshape((-1, 1, noise_mask.shape[-2], noise_mask.shape[-1])), size=(shape[2], shape[3]), mode="bilinear")
-    #noise_mask = noise_mask.round()
-    noise_mask = torch.cat([noise_mask] * shape[1], dim=1)
-    noise_mask = comfy.utils.repeat_to_batch_size(noise_mask, shape[0])
-    noise_mask = noise_mask.to(device)
-    return noise_mask
-
-
 def apply_params_to_motion_model(motion_model: MotionModelPatcher, params: InjectionParams):
     if params.context_length and params.video_length > params.context_length:
         logger.info(f"Sliding context window activated - latents passed in ({params.video_length}) greater than context_length {params.context_length}.")
@@ -201,9 +177,9 @@ class FunctionInjectionHolder:
         self.orig_memory_required = model.model.memory_required # allows for "unlimited area hack" to prevent halving of conds/unconds
         self.orig_groupnorm_forward = torch.nn.GroupNorm.forward # used to normalize latents to remove "flickering" of colors/brightness between frames
         self.orig_groupnormad_forward = GroupNormAD.forward
-        self.orig_sampling_function = comfy_samplers.sampling_function # used to support sliding context windows in samplers
-        self.orig_prepare_mask = comfy_sample.prepare_mask
-        self.orig_get_additional_models = comfy_sample.get_additional_models
+        self.orig_sampling_function = comfy.samplers.sampling_function # used to support sliding context windows in samplers
+        self.orig_prepare_mask = comfy.sample.prepare_mask
+        self.orig_get_additional_models = comfy.sample.get_additional_models
         # Inject Functions
         openaimodel.forward_timestep_embed = forward_timestep_embed_factory()
         if params.unlimited_area_hack:
@@ -215,9 +191,9 @@ class FunctionInjectionHolder:
             torch.nn.GroupNorm.forward = groupnorm_mm_factory(params)
             if params.apply_mm_groupnorm_hack:
                 GroupNormAD.forward = groupnorm_mm_factory(params)
-        comfy_samplers.sampling_function = sliding_sampling_function
-        comfy_sample.prepare_mask = prepare_mask_ad
-        comfy_sample.get_additional_models = get_additional_models_factory(self.orig_get_additional_models, model.motion_model)
+        comfy.samplers.sampling_function = sliding_sampling_function
+        comfy.sample.prepare_mask = prepare_mask_ad
+        comfy.sample.get_additional_models = get_additional_models_factory(self.orig_get_additional_models, model.motion_model)
         del info
 
     def restore_functions(self, model: ModelPatcherAndInjector):
@@ -227,27 +203,29 @@ class FunctionInjectionHolder:
             openaimodel.forward_timestep_embed = self.orig_forward_timestep_embed
             torch.nn.GroupNorm.forward = self.orig_groupnorm_forward
             GroupNormAD.forward = self.orig_groupnormad_forward
-            comfy_samplers.sampling_function = self.orig_sampling_function
-            comfy_sample.prepare_mask = self.orig_prepare_mask
-            comfy_sample.get_additional_models = self.orig_get_additional_models
+            comfy.samplers.sampling_function = self.orig_sampling_function
+            comfy.sample.prepare_mask = self.orig_prepare_mask
+            comfy.sample.get_additional_models = self.orig_get_additional_models
         except AttributeError:
             logger.error("Encountered AttributeError while attempting to restore functions - likely, an error occured while trying " + \
                          "to save original functions before injection, and a more specific error was thrown by ComfyUI.")
 
 
-def motion_sample_factory(orig_comfy_sample: Callable) -> Callable:
+def motion_sample_factory(orig_comfy_sample: Callable, is_custom: bool=False) -> Callable:
     def motion_sample(model: ModelPatcherAndInjector, noise: Tensor, *args, **kwargs):
         # check if model is intended for injecting
         if type(model) != ModelPatcherAndInjector:
             return orig_comfy_sample(model, noise, *args, **kwargs)
         # otherwise, injection time
         latents = None
+        cached_latents = None
+        cached_noise = None
         function_injections = FunctionInjectionHolder()
         try:
             # clone params from model
             params = model.motion_injection_params.clone()
             # get amount of latents passed in, and store in params
-            latents = args[-1]
+            latents: Tensor = args[-1]
             params.video_length = latents.size(0)
             params.full_length = latents.size(0)
             # reset global state
@@ -258,19 +236,18 @@ def motion_sample_factory(orig_comfy_sample: Callable) -> Callable:
             # apply custom noise, if needed
             disable_noise = kwargs.get("disable_noise") or False
             seed = kwargs["seed"]
-            if not disable_noise:
-                # if context asks for specific noise, do it
-                noise = NoiseType.prepare_noise(params.noise_type, latents=latents, noise=noise, context_length=params.context_length, seed=seed)
 
             # apply params to motion model
             apply_params_to_motion_model(model.motion_model, params)
 
-            # handle GLOBALSTATE vars and step tally
-            ADGS.update_with_inject_params(params)
-            ADGS.start_step = kwargs.get("start_step") or 0
-            ADGS.current_step = ADGS.start_step
-            ADGS.last_step = kwargs.get("last_step") or 0
+            # prepare noise_extra_args for noise generation purposes
+            noise_extra_args = {"disable_noise": disable_noise}
+            params.set_noise_extra_args(noise_extra_args)
+            # if noise is not disabled, do noise stuff
+            if not disable_noise:
+                noise = model.sample_settings.prepare_noise(seed, latents, noise, extra_args=noise_extra_args, force_create_noise=False)
 
+            # callback setup
             original_callback = kwargs.get("callback", None)
             def ad_callback(step, x0, x, total_steps):
                 if original_callback is not None:
@@ -280,11 +257,52 @@ def motion_sample_factory(orig_comfy_sample: Callable) -> Callable:
             kwargs["callback"] = ad_callback
             ADGS.motion_model = model.motion_model
 
-            model.motion_model.pre_run()
-            return wrap_function_to_inject_xformers_bug_info(orig_comfy_sample)(model, noise, *args, **kwargs)
+            iter_opts = IterationOptions()
+            if model.sample_settings is not None:
+                iter_opts = model.sample_settings.iteration_opts
+            iter_opts.initialize(latents)
+            # cache initial noise and latents, if needed
+            if iter_opts.cache_init_latents:
+                cached_latents = latents.clone()
+            if iter_opts.cache_init_noise:
+                cached_noise = noise.clone()
+            # prepare iter opts preprocess kwargs, if needed
+            iter_kwargs = {}
+            if iter_opts.need_sampler:
+                # -5 for sampler_name (not custom) and sampler (custom)
+                if is_custom:
+                    iter_kwargs[IterationOptions.SAMPLER] = args[-5]
+                else:
+                    iter_kwargs[IterationOptions.SAMPLER] = comfy.samplers.KSampler(
+                        model.model, steps=args[-7],
+                        device=model.current_device, sampler=args[-5],
+                        scheduler=args[-4], denoise=kwargs["denoise"],
+                        model_options=model.model_options)
+
+            args = list(args)
+            for curr_i in range(iter_opts.iterations):
+                # handle GLOBALSTATE vars and step tally
+                ADGS.update_with_inject_params(params)
+                ADGS.start_step = kwargs.get("start_step") or 0
+                ADGS.current_step = ADGS.start_step
+                ADGS.last_step = kwargs.get("last_step") or 0
+                if iter_opts.iterations > 1:
+                    logger.info(f"Iteration {curr_i+1}/{iter_opts.iterations}")
+                # perform any iter_opts preprocessing on latents
+                latents, noise = iter_opts.preprocess_latents(curr_i=curr_i, model=model, latents=latents, noise=noise,
+                                                              cached_latents=cached_latents, cached_noise=cached_noise,
+                                                              seed=seed,
+                                                              sample_settings=model.sample_settings, noise_extra_args=noise_extra_args)
+                args[-1] = latents
+
+                model.motion_model.pre_run()
+                latents = wrap_function_to_inject_xformers_bug_info(orig_comfy_sample)(model, noise, *args, **kwargs)
+            return latents
         finally:
             del latents
             del noise
+            del cached_latents
+            del cached_noise
             # reset global state
             ADGS.reset()
             # restore injected functions
@@ -533,9 +551,9 @@ def sliding_sampling_function(model, x, timestep, uncond, cond, cond_scale, mode
         # https://github.com/comfyanonymous/ComfyUI/compare/master...ashen-sensored:ComfyUI:master
         def sliding_calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options):
             # get context scheduler
-            context_scheduler = get_context_scheduler(ADGS.context_schedule)
+            context_scheduler = get_context_scheduler(ADGS.params.context_schedule)
             # figure out how input is split
-            axes_factor = x_in.size(0)//ADGS.video_length
+            axes_factor = x_in.size(0)//ADGS.params.video_length
 
             # prepare final cond, uncond, and out_count
             cond_final = torch.zeros_like(x_in)
@@ -546,8 +564,8 @@ def sliding_sampling_function(model, x, timestep, uncond, cond, cond_scale, mode
                 if control.previous_controlnet is not None:
                     prepare_control_objects(control.previous_controlnet, full_idxs)
                 control.sub_idxs = full_idxs
-                control.full_latent_length = ADGS.video_length
-                control.context_length = ADGS.context_frames
+                control.full_latent_length = ADGS.params.video_length
+                control.context_length = ADGS.params.context_length
             
             def get_resized_cond(cond_in, full_idxs) -> list:
                 # reuse or resize cond items to match context requirements
@@ -574,7 +592,7 @@ def sliding_sampling_function(model, x, timestep, uncond, cond, cond_scale, mode
                                     prepare_control_objects(control_item, full_idxs)
                                 else:
                                     raise ValueError(f"Control type {type(control_item).__name__} may not support required features for sliding context window; \
-                                                        use Control objects from Kosinkadink/Advanced-ControlNet nodes, or make sure Advanced-ControlNet is updated.")
+                                                        use Control objects from Kosinkadink/ComfyUI-Advanced-ControlNet nodes, or make sure Advanced-ControlNet is updated.")
                                 resized_actual_cond[key] = control_item
                                 del control_item
                             elif isinstance(cond_item, dict):
@@ -597,15 +615,14 @@ def sliding_sampling_function(model, x, timestep, uncond, cond, cond_scale, mode
                 return resized_cond
 
             # perform calc_cond_uncond_batch per context window
-            for ctx_idxs in context_scheduler(ADGS.current_step, ADGS.total_steps, ADGS.video_length, ADGS.context_frames, ADGS.context_stride, ADGS.context_overlap, ADGS.closed_loop):
-                ADGS.sub_idxs = ctx_idxs
-                ADGS.params.sub_idxs = ADGS.sub_idxs
-                ADGS.motion_model.model.set_sub_idxs(ADGS.sub_idxs)
+            for ctx_idxs in context_scheduler(ADGS.current_step, ADGS.total_steps, ADGS.params.video_length, ADGS.params.context_length, ADGS.params.context_stride, ADGS.params.context_overlap, ADGS.params.closed_loop):
+                ADGS.params.sub_idxs = ctx_idxs
+                ADGS.motion_model.model.set_sub_idxs(ctx_idxs)
                 # account for all portions of input frames
                 full_idxs = []
                 for n in range(axes_factor):
                     for ind in ctx_idxs:
-                        full_idxs.append((ADGS.video_length*n)+ind)
+                        full_idxs.append((ADGS.params.video_length*n)+ind)
                 # get subsections of x, timestep, cond, uncond, cond_concat
                 sub_x = x_in[full_idxs]
                 sub_timestep = timestep[full_idxs]