Merge PR #533 from Kosinkadink/develop - Ancestral Options

Add Ancestral Options to control ancestral sampler noise

Author: Kosinkadink

animatediff/nodes.py

@@ -26,7 +26,7 @@
 from .nodes_sample import (FreeInitOptionsNode, NoiseLayerAddWeightedNode, NoiseLayerNormalizedSumNode, SampleSettingsNode, NoiseLayerAddNode, NoiseLayerReplaceNode, IterationOptionsNode,
                            CustomCFGNode, CustomCFGSimpleNode, CustomCFGKeyframeNode, CustomCFGKeyframeSimpleNode, CustomCFGKeyframeInterpolationNode, CustomCFGKeyframeFromListNode,
                            CFGExtrasPAGNode, CFGExtrasPAGSimpleNode, CFGExtrasRescaleCFGNode, CFGExtrasRescaleCFGSimpleNode,
-                           NoisedImageInjectionNode, NoisedImageInjectOptionsNode, NoiseCalibrationNode)
+                           NoisedImageInjectionNode, NoisedImageInjectOptionsNode, NoiseCalibrationNode, AncestralOptionsNode)
 from .nodes_sigma_schedule import (SigmaScheduleNode, RawSigmaScheduleNode, WeightedAverageSigmaScheduleNode, InterpolatedWeightedAverageSigmaScheduleNode, SplitAndCombineSigmaScheduleNode, SigmaScheduleToSigmasNode)
 from .nodes_context import (LegacyLoopedUniformContextOptionsNode, LoopedUniformContextOptionsNode, LoopedUniformViewOptionsNode, StandardUniformContextOptionsNode, StandardStaticContextOptionsNode, BatchedContextOptionsNode,
                             StandardStaticViewOptionsNode, StandardUniformViewOptionsNode, ViewAsContextOptionsNode,
@@ -158,6 +158,7 @@
     "ADE_SigmaScheduleToSigmas": SigmaScheduleToSigmasNode,
     "ADE_NoisedImageInjection": NoisedImageInjectionNode,
     "ADE_NoisedImageInjectOptions": NoisedImageInjectOptionsNode,
+    "ADE_AncestralOptions": AncestralOptionsNode,
     #"ADE_NoiseCalibration": NoiseCalibrationNode,
     # Scheduling
     PromptSchedulingNode.NodeID: PromptSchedulingNode,
@@ -338,6 +339,7 @@
     "ADE_NoisedImageInjection": "Image Injection 🎭🅐🅓",
     "ADE_NoisedImageInjectOptions": "Image Injection Options 🎭🅐🅓",
     "ADE_NoiseCalibration": "Noise Calibration 🎭🅐🅓",
+    "ADE_AncestralOptions": "Ancestral Options 🎭🅐🅓",
     # Scheduling
     PromptSchedulingNode.NodeID: PromptSchedulingNode.NodeName,
     PromptSchedulingLatentsNode.NodeID: PromptSchedulingLatentsNode.NodeName,

animatediff/nodes_sample.py

@@ -5,9 +5,9 @@
 from comfy.sd import VAE
 
 from .freeinit import FreeInitFilter
-from .sample_settings import (FreeInitOptions, IterationOptions,
+from .sample_settings import (FreeInitOptions, IterationOptions, AncestralOptions,
                               NoiseLayerAdd, NoiseLayerAddWeighted, NoiseLayerNormalizedSum, NoiseLayerGroup, NoiseLayerReplace, NoiseLayerType,
-                              SeedNoiseGeneration, SampleSettings, NoiseCalibration,
+                              SeedNoiseGeneration, SampleSettings, NoiseCalibration, NoiseDeterminism,
                               CustomCFGKeyframeGroup, CustomCFGKeyframe, CFGExtrasGroup, CFGExtras,
                               NoisedImageToInjectGroup, NoisedImageToInject, NoisedImageInjectOptions)
 from .utils_model import BIGMIN, BIGMAX, MAX_RESOLUTION, SigmaSchedule, InterpolationMethod
@@ -28,11 +28,12 @@ def INPUT_TYPES(s):
             "optional": {
                 "noise_layers": ("NOISE_LAYERS",),
                 "iteration_opts": ("ITERATION_OPTS",),
-                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
+                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "defaultInput": True}),
                 "adapt_denoise_steps": ("BOOLEAN", {"default": False},),
                 "custom_cfg": ("CUSTOM_CFG",),
                 "sigma_schedule": ("SIGMA_SCHEDULE",),
                 "image_inject": ("IMAGE_INJECT",),
+                "ancestral_opts": ("ANCESTRAL_OPTS",),
                 #"noise_calib": ("NOISE_CALIBRATION",), # TODO: bring back once NoiseCalibration is working
             },
             "hidden": {
@@ -48,13 +49,43 @@ def INPUT_TYPES(s):
     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, adapt_denoise_steps=False,
                         custom_cfg: CustomCFGKeyframeGroup=None, sigma_schedule: SigmaSchedule=None, image_inject: NoisedImageToInjectGroup=None,
-                        noise_calib: NoiseCalibration=None):
+                        noise_calib: NoiseCalibration=None, ancestral_opts=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, adapt_denoise_steps=adapt_denoise_steps,
-                                           custom_cfg=custom_cfg, sigma_schedule=sigma_schedule, image_injection=image_inject, noise_calibration=noise_calib)
+                                           custom_cfg=custom_cfg, sigma_schedule=sigma_schedule, image_injection=image_inject, noise_calibration=noise_calib,
+                                           ancestral_opts=ancestral_opts)
         return (sampling_settings,)
 
 
+class AncestralOptionsNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                #"batch_offset": ("INT", {"default": 0, "min": 0, "max": BIGMAX}),
+                "noise_type": (NoiseLayerType.LIST_ANCESTRAL,),
+                #"determinism": (NoiseDeterminism._LIST,),
+                "seed_offset": ("INT", {"default": 0, "min": BIGMIN, "max": BIGMAX}),
+                #"seed_gen_override": (SeedNoiseGeneration.LIST_WITH_OVERRIDE,),
+            },
+            "optional": {
+                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "defaultInput": True}),
+            },
+            "hidden": {
+                "autosize": ("ADEAUTOSIZE", {"padding": 0}),
+            }
+        }
+
+    RETURN_TYPES = ("ANCESTRAL_OPTS",)
+    CATEGORY = "Animate Diff 🎭🅐🅓/sample settings"
+    FUNCTION = "create_ancestral_opts"
+
+    def create_ancestral_opts(self, noise_type: str, seed_offset: int, determinism: str=NoiseDeterminism.DEFAULT, seed_override: int=None):
+        if isinstance(seed_override, Iterable):
+            raise Exception("Passing in a list of seeds for Ancestral Options is not supported at this time.")
+        return (AncestralOptions(noise_type=noise_type, determinism=determinism, seed_offset=seed_offset, seed_override=seed_override),)
+
+
 class NoiseLayerReplaceNode:
     @classmethod
     def INPUT_TYPES(s):
@@ -68,7 +99,7 @@ def INPUT_TYPES(s):
             "optional": {
                 "prev_noise_layers": ("NOISE_LAYERS",),
                 "mask_optional": ("MASK",),
-                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
+                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "defaultInput": True}),
             },
             "hidden": {
                 "autosize": ("ADEAUTOSIZE", {"padding": 0}),
@@ -106,7 +137,7 @@ def INPUT_TYPES(s):
             "optional": {
                 "prev_noise_layers": ("NOISE_LAYERS",),
                 "mask_optional": ("MASK",),
-                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
+                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "defaultInput": True}),
             },
             "hidden": {
                 "autosize": ("ADEAUTOSIZE", {"padding": 0}),
@@ -147,7 +178,7 @@ def INPUT_TYPES(s):
             "optional": {
                 "prev_noise_layers": ("NOISE_LAYERS",),
                 "mask_optional": ("MASK",),
-                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
+                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "defaultInput": True}),
             },
             "hidden": {
                 "autosize": ("ADEAUTOSIZE", {"padding": 0}),
@@ -187,7 +218,7 @@ def INPUT_TYPES(s):
             "optional": {
                 "prev_noise_layers": ("NOISE_LAYERS",),
                 "mask_optional": ("MASK",),
-                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}),
+                "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "defaultInput": True}),
             },
             "hidden": {
                 "autosize": ("ADEAUTOSIZE", {"padding": 0}),

animatediff/sample_settings.py

@@ -1,13 +1,16 @@
+from __future__ import annotations
 from collections.abc import Iterable
 from typing import Union, Callable
 import torch
 from torch import Tensor
 import torch.fft as fft
 from einops import rearrange
 
+import comfy.k_diffusion.sampling
 import comfy.sample
 import comfy.samplers
 import comfy.model_management
+from comfy.patcher_extension import WrappersMP, add_wrapper_with_key
 from comfy.model_patcher import ModelPatcher
 from comfy.model_base import BaseModel
 from comfy.sd import VAE
@@ -29,6 +32,13 @@ def prepare_mask_ad(noise_mask, shape, device):
     return noise_mask
 
 
+class NoiseDeterminism:
+    DEFAULT = "default"
+    DETERMINISTIC = "deterministic"
+
+    _LIST = [DEFAULT, DETERMINISTIC]
+
+
 class NoiseLayerType:
     DEFAULT = "default"
     CONSTANT = "constant"
@@ -37,6 +47,7 @@ class NoiseLayerType:
     FREENOISE = "FreeNoise"
 
     LIST = [DEFAULT, CONSTANT, EMPTY, REPEATED_CONTEXT, FREENOISE]
+    LIST_ANCESTRAL = [DEFAULT, CONSTANT]
 
 
 class NoiseApplication:
@@ -56,10 +67,10 @@ class NoiseNormalize:
 
 
 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,
+    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, adapt_denoise_steps: bool=False,
-                 custom_cfg: 'CustomCFGKeyframeGroup'=None, sigma_schedule: SigmaSchedule=None, image_injection: 'NoisedImageToInjectGroup'=None,
-                 noise_calibration: 'NoiseCalibration'=None):
+                 custom_cfg: CustomCFGKeyframeGroup=None, sigma_schedule: SigmaSchedule=None, image_injection: NoisedImageToInjectGroup=None,
+                 noise_calibration: NoiseCalibration=None, ancestral_opts: AncestralOptions=None):
         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
@@ -73,6 +84,7 @@ def __init__(self, batch_offset: int=0, noise_type: str=None, seed_gen: str=None
         self.sigma_schedule = sigma_schedule
         self.image_injection = image_injection.clone() if image_injection else NoisedImageToInjectGroup()
         self.noise_calibration = noise_calibration
+        self.ancestral_opts = ancestral_opts
 
     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:
@@ -113,7 +125,84 @@ def clone(self):
         return SampleSettings(batch_offset=self.batch_offset, noise_type=self.noise_type, seed_gen=self.seed_gen, seed_offset=self.seed_offset,
                            noise_layers=self.noise_layers.clone(), iteration_opts=self.iteration_opts, seed_override=self.seed_override,
                            negative_cond_flipflop=self.negative_cond_flipflop, adapt_denoise_steps=self.adapt_denoise_steps, custom_cfg=self.custom_cfg,
-                           sigma_schedule=self.sigma_schedule, image_injection=self.image_injection, noise_calibration=self.noise_calibration)
+                           sigma_schedule=self.sigma_schedule, image_injection=self.image_injection, noise_calibration=self.noise_calibration,
+                           ancestral_opts=self.ancestral_opts)
+
+
+class AncestralOptions:
+    def __init__(self, noise_type: str, determinism: str, seed_offset: int, seed_override: int=None):
+        self.noise_type = noise_type
+        self.determinism = determinism
+        self.seed_offset = seed_offset
+        self.seed_override = seed_override
+    
+    def init_custom_noise_sampler(self, seed: int):
+        if self.seed_override is not None:
+            seed = self.seed_override
+        if isinstance(seed, Iterable):
+            raise Exception("Passing in a list of seeds for Ancestral Options is not supported at this time.")
+        seed += self.seed_offset
+        return _custom_noise_sampler_factory(real_seed=seed, noise_type=self.noise_type, determinism=self.determinism)
+
+    def add_wrapper_sampler_sample(self, model_options, seed):
+        add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, "ADE",
+                             _sampler_sample_ancestral_options_factory(self.init_custom_noise_sampler(seed)),
+                             model_options, is_model_options=True)
+
+
+def _sampler_sample_ancestral_options_factory(custom_noise_sampler: Callable):
+    def sampler_sample_ancestral_options_wrapper(executor, *args, **kwargs):
+        try:
+            # TODO: implement this as a model_options thing instead in core ComfyUI
+            orig_default_noise_sampler = comfy.k_diffusion.sampling.default_noise_sampler
+            comfy.k_diffusion.sampling.default_noise_sampler = custom_noise_sampler
+            return executor(*args, **kwargs)
+        finally:
+            comfy.k_diffusion.sampling.default_noise_sampler = orig_default_noise_sampler
+    return sampler_sample_ancestral_options_wrapper
+
+
+def _custom_noise_sampler_factory(real_seed: int, noise_type: str, determinism: str):
+    def custom_noise_sampler(x: Tensor, seed: int=None):
+        single_generator = None
+        multiple_generators = []
+        if determinism == NoiseDeterminism.DEFAULT:
+            # prepare generators
+            single_generator = torch.Generator(device=x.device)
+            single_generator.manual_seed(real_seed)
+            # create function to handle determinism type
+            def sample_default(sigma, sigma_next):
+                if noise_type == NoiseLayerType.CONSTANT:
+                    goal_shape = list(x.shape)
+                    goal_shape[0] = 1
+                    one_noise = torch.randn(goal_shape, dtype=x.dtype, layout=x.layout, device=x.device, generator=single_generator)
+                    return torch.cat([one_noise]*x.shape[0], dim=0)
+                return torch.randn(x.size(), dtype=x.dtype, layout=x.layout, device=x.device, generator=single_generator)
+            # return function
+            return sample_default
+        elif determinism == NoiseDeterminism.DETERMINISTIC:
+            # prepare generators
+            for i in range(x.size(0)):
+                generator = torch.Generator(device=x.device)
+                multiple_generators.append(generator.manual_seed(real_seed+i))
+            # create function to handle determinism type
+            def sample_deterministic(sigma, sigma_next):
+                goal_shape = list(x.shape)
+                goal_shape[0] = 1
+                if noise_type == NoiseLayerType.CONSTANT:
+                    one_noise = torch.randn(goal_shape, dtype=x.dtype, layout=x.layout, device=x.device, generator=multiple_generators[0])
+                    return torch.cat([one_noise]*x.shape[0], dim=0)
+                noises = []
+                for generator in multiple_generators:
+                    one_noise = torch.randn(goal_shape, dtype=x.dtype, layout=x.layout, device=x.device, generator=generator)
+                    noises.append(one_noise)
+                return torch.cat(noises, dim=0)
+            # return function
+            return sample_deterministic
+        else:
+            raise Exception(f"Determinism type '{determinism}' is not recognized.")
+    # return function
+    return custom_noise_sampler
 
 
 class NoiseLayer:

animatediff/sampling.py

@@ -408,6 +408,9 @@ def outer_sample_wrapper(executor: WrapperExecutor, *args, **kwargs):
         # if noise is not disabled, do noise stuff
         if not disable_noise:
             noise = helper.get_sample_settings().prepare_noise(seed, latents, noise, extra_args=noise_extra_args, force_create_noise=False)
+        # handle AncestralOptions, if present
+        if helper.get_sample_settings().ancestral_opts is not None:
+            helper.get_sample_settings().ancestral_opts.add_wrapper_sampler_sample(guider.model_options, seed)
 
         # callback setup
         original_callback = args[-3]

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-animatediff-evolved"
 description = "Improved AnimateDiff integration for ComfyUI."
-version = "1.4.2"
+version = "1.4.3"
 license = { file = "LICENSE" }
 dependencies = []