diff --git a/animatediff/nodes.py b/animatediff/nodes.py index cca9537..ffca1fc 100644 --- a/animatediff/nodes.py +++ b/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, diff --git a/animatediff/nodes_sample.py b/animatediff/nodes_sample.py index e9a6102..1d8b9d7 100644 --- a/animatediff/nodes_sample.py +++ b/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}), diff --git a/animatediff/sample_settings.py b/animatediff/sample_settings.py index 35e0a94..68acef9 100644 --- a/animatediff/sample_settings.py +++ b/animatediff/sample_settings.py @@ -1,3 +1,4 @@ +from __future__ import annotations from collections.abc import Iterable from typing import Union, Callable import torch @@ -5,9 +6,11 @@ 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: diff --git a/animatediff/sampling.py b/animatediff/sampling.py index ec2174b..92167f6 100644 --- a/animatediff/sampling.py +++ b/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] diff --git a/pyproject.toml b/pyproject.toml index 537e762..be2abee 100644 --- a/pyproject.toml +++ b/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 = []