From da46246422b2a77b01d062dbe66af2db95b38b19 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sun, 12 Nov 2023 16:27:43 -0600 Subject: [PATCH 01/12] Reenabled experimental uniform context options node, again --- animatediff/nodes.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/animatediff/nodes.py b/animatediff/nodes.py index b1f5904..21275f5 100644 --- a/animatediff/nodes.py +++ b/animatediff/nodes.py @@ -210,7 +210,7 @@ class AnimateDiffUniformContextOptionsExperimental: NODE_CLASS_MAPPINGS = { "ADE_AnimateDiffUniformContextOptions": AnimateDiffUniformContextOptions, - #"ADE_AnimateDiffUniformContextOptionsExperimental": AnimateDiffUniformContextOptionsExperimental, + "ADE_AnimateDiffUniformContextOptionsExperimental": AnimateDiffUniformContextOptionsExperimental, "ADE_AnimateDiffLoaderWithContext": AnimateDiffLoaderWithContext, "ADE_AnimateDiffLoRALoader": AnimateDiffLoRALoader, "ADE_AnimateDiffModelSettings_Release": AnimateDiffModelSettings, @@ -229,7 +229,7 @@ NODE_CLASS_MAPPINGS = { } NODE_DISPLAY_NAME_MAPPINGS = { "ADE_AnimateDiffUniformContextOptions": "Uniform Context Options 🎭🅐🅓", - #"ADE_AnimateDiffUniformContextOptionsExperimental": "EXPERIMENTAL Uniform Context Options 🎭🅐🅓", + "ADE_AnimateDiffUniformContextOptionsExperimental": "EXPERIMENTAL Uniform Context Options 🎭🅐🅓", "ADE_AnimateDiffLoaderWithContext": "AnimateDiff Loader 🎭🅐🅓", "ADE_AnimateDiffLoRALoader": "AnimateDiff LoRA Loader 🎭🅐🅓", "ADE_AnimateDiffModelSettings_Release": "Motion Model Settings 🎭🅐🅓", From 08e0f11018c1a6963ff1f3fc50de42255767a64c Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sun, 3 Dec 2023 05:02:15 -0600 Subject: [PATCH 02/12] Some in progress noise and beta_schedule passthrough testing --- animatediff/model_utils.py | 13 ++++++++++++- animatediff/motion_utils.py | 7 ++++++- animatediff/nodes_extras.py | 8 +++++++- 3 files changed, 25 insertions(+), 3 deletions(-) diff --git a/animatediff/model_utils.py b/animatediff/model_utils.py index 1299748..3343ad5 100644 --- a/animatediff/model_utils.py +++ b/animatediff/model_utils.py @@ -32,16 +32,20 @@ class BetaSchedules: SQRT_LINEAR = "sqrt_linear (AnimateDiff)" LINEAR_ADXL = "linear (AnimateDiff-SDXL)" LINEAR = "linear (HotshotXL/default)" + USE_PATCHED = "use_patched" + USE_ORIGINAL = "use_original" SQRT = "sqrt" COSINE = "cosine" SQUAREDCOS_CAP_V2 = "squaredcos_cap_v2" - ALIAS_LIST = [SQRT_LINEAR, LINEAR_ADXL, LINEAR, SQRT, COSINE, SQUAREDCOS_CAP_V2] + ALIAS_LIST = [SQRT_LINEAR, LINEAR_ADXL, LINEAR, USE_PATCHED, USE_ORIGINAL, SQRT, COSINE, SQUAREDCOS_CAP_V2] ALIAS_MAP = { SQRT_LINEAR: "sqrt_linear", LINEAR_ADXL: "linear", # also linear, but has different linear_end (0.020) LINEAR: "linear", + USE_PATCHED: None, + USE_ORIGINAL: None, SQRT: "sqrt", COSINE: "cosine", SQUAREDCOS_CAP_V2: "squaredcos_cap_v2", @@ -57,6 +61,13 @@ class BetaSchedules: @classmethod def to_model_sampling(cls, alias: str, model: ModelPatcher): + # if should use patched, try to get model_sampling patch; otherwise, use original + if alias == cls.USE_PATCHED: + ms_obj = model.object_patches.get("model_sampling", None) + return ms_obj if ms_obj is not None else model.model.model_sampling + # if should use original, use model_sampling that is already on the model + if alias == cls.USE_ORIGINAL: + return model.model.model_sampling ms_obj = model_sampling(cls.to_config(alias), model_type=model.model.model_type) if alias == cls.LINEAR_ADXL: # uses linear_end=0.020 diff --git a/animatediff/motion_utils.py b/animatediff/motion_utils.py index ec4d566..576aab6 100644 --- a/animatediff/motion_utils.py +++ b/animatediff/motion_utils.py @@ -223,9 +223,10 @@ class NoiseType: DEFAULT = "default" REPEATED = "repeated" CONSTANT = "constant" + CONSTANT_ADDED = "constant_added" AUTO1111 = "auto1111" - LIST = [DEFAULT, REPEATED, CONSTANT, AUTO1111] + LIST = [DEFAULT, REPEATED, CONSTANT, CONSTANT_ADDED, AUTO1111] @classmethod def prepare_noise(cls, noise_type: str, latents: Tensor, noise: Tensor, context_length: int, seed: int): @@ -235,6 +236,10 @@ class NoiseType: 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.CONSTANT_ADDED: + new_noise = cls.prepare_noise_constant(latents, noise, context_length, seed) + noise_weight = 0.2 + return noise * (1.0-(noise_weight/3)) + new_noise * noise_weight 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.") diff --git a/animatediff/nodes_extras.py b/animatediff/nodes_extras.py index 60fba49..d9c64b5 100644 --- a/animatediff/nodes_extras.py +++ b/animatediff/nodes_extras.py @@ -36,17 +36,23 @@ class CheckpointLoaderSimpleWithNoiseSelect: "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), "beta_schedule": (BetaSchedules.get_alias_list_with_first_element(BetaSchedules.LINEAR), ) }, + "optional": { + "use_custom_scale_factor": ("BOOLEAN", {"default": False}), + "scale_factor": ("FLOAT", {"default": 0.18215, "min": 0.0, "max": 1.0}) + } } 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 out[0].model.model_sampling = BetaSchedules.to_model_sampling(beta_schedule, out[0]) + if use_custom_scale_factor: + out[0].model.latent_format.scale_factor = scale_factor return out From 011e3034b31adf5a3dc273d71fb283ed4e49e088 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Wed, 13 Dec 2023 10:02:05 -0600 Subject: [PATCH 03/12] fix issue with merge with develop_alpha --- animatediff/model_utils.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/animatediff/model_utils.py b/animatediff/model_utils.py index 883cae9..59b925b 100644 --- a/animatediff/model_utils.py +++ b/animatediff/model_utils.py @@ -44,8 +44,6 @@ class BetaSchedules: SQRT_LINEAR: "sqrt_linear", LINEAR_ADXL: "linear", # also linear, but has different linear_end (0.020) LINEAR: "linear", - USE_PATCHED: None, - USE_ORIGINAL: None, SQRT: "sqrt", COSINE: "cosine", SQUAREDCOS_CAP_V2: "squaredcos_cap_v2", From 6e26c305feebce3587c3b290cfb5346e4633e442 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Fri, 15 Dec 2023 08:41:12 -0600 Subject: [PATCH 04/12] Started work on SamplingSettings/Noise Layers --- animatediff/context.py | 6 -- animatediff/model_injection.py | 8 +- animatediff/motion_module_ad.py | 64 +++++++++++- animatediff/motion_utils.py | 177 -------------------------------- animatediff/nodes.py | 4 +- animatediff/nodes_sample.py | 16 +++ animatediff/sample_settings.py | 155 ++++++++++++++++++++++++++++ animatediff/sampling.py | 47 +++------ 8 files changed, 254 insertions(+), 223 deletions(-) create mode 100644 animatediff/nodes_sample.py create mode 100644 animatediff/sample_settings.py diff --git a/animatediff/context.py b/animatediff/context.py index 316665b..10be918 100644 --- a/animatediff/context.py +++ b/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: diff --git a/animatediff/model_injection.py b/animatediff/model_injection.py index 02ce486..114a730 100644 --- a/animatediff/model_injection.py +++ b/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 SeedNoiseType # some motion_model casts here might fail if model becomes metatensor or is not castable; @@ -33,7 +34,6 @@ class ModelPatcherAndInjector(ModelPatcher): # injection stuff self.motion_injection_params: InjectionParams = None self.motion_model: MotionModelPatcher = None - self.motion_model_sampling = None def model_patches_to(self, device): super().model_patches_to(device) @@ -77,7 +77,6 @@ class ModelPatcherAndInjector(ModelPatcher): cloned = ModelPatcherAndInjector(self) cloned.motion_model = self.motion_model cloned.motion_injection_params = self.motion_injection_params - cloned.motion_model_sampling = self.motion_model_sampling return cloned @@ -315,7 +314,8 @@ 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 = SeedNoiseType.DEFAULT self.sub_idxs = None # value should NOT be included in clone, so it will auto reset diff --git a/animatediff/motion_module_ad.py b/animatediff/motion_module_ad.py index 69012a6..e4a3068 100644 --- a/animatediff/motion_module_ad.py +++ b/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 @@ -387,7 +389,7 @@ class VanillaTemporalModule(nn.Module): return output_tensor -class TemporalTransformer3DModel(nn.Module, TemporalTransformerGeneric): +class TemporalTransformer3DModel(nn.Module): def __init__( self, in_channels, @@ -409,7 +411,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 @@ -458,6 +468,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 diff --git a/animatediff/motion_utils.py b/animatediff/motion_utils.py index 237a079..1494baf 100644 --- a/animatediff/motion_utils.py +++ b/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,65 +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" - CONSTANT_ADDED = "constant_added" - AUTO1111 = "auto1111" - - LIST = [DEFAULT, REPEATED, CONSTANT, CONSTANT_ADDED, 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.CONSTANT_ADDED: - new_noise = cls.prepare_noise_constant(latents, noise, context_length, seed) - noise_weight = 0.2 - return noise * (1.0-(noise_weight/3)) + new_noise * noise_weight - 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 diff --git a/animatediff/nodes.py b/animatediff/nodes.py index db59a3d..9413f70 100644 --- a/animatediff/nodes.py +++ b/animatediff/nodes.py @@ -7,9 +7,9 @@ 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 SeedNoiseType from .sampling import motion_sample_factory from .nodes_extras import AnimateDiffUnload, EmptyLatentImageLarge, CheckpointLoaderSimpleWithNoiseSelect @@ -205,7 +205,7 @@ class AnimateDiffUniformContextOptionsExperimental: "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,) + "noise_override": (SeedNoiseType.LIST,) #"sync_context_to_pe": ("BOOLEAN", {"default": False},), }, } diff --git a/animatediff/nodes_sample.py b/animatediff/nodes_sample.py new file mode 100644 index 0000000..e09c77b --- /dev/null +++ b/animatediff/nodes_sample.py @@ -0,0 +1,16 @@ +from .sample_settings import SeedNoiseType + +class SampleSettingsNode: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "seed_noise": (SeedNoiseType.LIST), + }, + "optional": { + "noise_layers": ("NOISE_LAYERS", ), + } + } + + RETURN_TYPES = ("SAMPLE_SETTINGS") + diff --git a/animatediff/sample_settings.py b/animatediff/sample_settings.py new file mode 100644 index 0000000..32f66e0 --- /dev/null +++ b/animatediff/sample_settings.py @@ -0,0 +1,155 @@ +import torch +from torch import Tensor + +from .logger import logger + + +class SamplingSettings: + def __init__(self): + self.seed_noise_type: str = SeedNoiseType.DEFAULT + self.noise_layers = NoiseLayerGroup() + + +class NoiseLayerType: + DEFAULT = "default" + CONSTANT = "constant" + + +class NoiseLayer: + def __init__(self): + self.mask: Tensor = None + self.layer_type: str = NoiseLayerType.DEFAULT + self.options: dict = {} + + +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 __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 SeedNoiseType: + DEFAULT = "default" + AUTO1111 = "auto1111" + + LIST = [DEFAULT, 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.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 __OLDNoiseType: + DEFAULT = "default" + REPEATED = "repeated" + CONSTANT = "constant" + CONSTANT_ADDED = "constant_added" + AUTO1111 = "auto1111" + + LIST = [DEFAULT, REPEATED, CONSTANT, CONSTANT_ADDED, 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.CONSTANT_ADDED: + new_noise = cls.prepare_noise_constant(latents, noise, context_length, seed) + noise_weight = 0.2 + return noise * (1.0-(noise_weight/3)) + new_noise * noise_weight + 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) diff --git a/animatediff/sampling.py b/animatediff/sampling.py index 77d9972..e7f867e 100644 --- a/animatediff/sampling.py +++ b/animatediff/sampling.py @@ -15,7 +15,8 @@ import comfy.utils from comfy.controlnet import ControlBase from .context import get_context_scheduler -from .motion_utils import GroupNormAD, NoiseType +from .sample_settings import SeedNoiseType +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() @@ -260,7 +246,7 @@ def motion_sample_factory(orig_comfy_sample: Callable) -> Callable: 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) + noise = SeedNoiseType.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) @@ -533,9 +519,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 +532,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 +560,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 +583,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] From 161c87d1178357de33c9cd834ed67f77517befe8 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 26 Dec 2023 05:02:11 -0600 Subject: [PATCH 05/12] Work on Sample Settings + Noise Layer nodes --- animatediff/model_injection.py | 6 +- animatediff/nodes.py | 51 +++------- animatediff/nodes_extras.py | 2 +- animatediff/nodes_sample.py | 87 ++++++++++++++++- animatediff/sample_settings.py | 172 ++++++++++++++++++++++++--------- animatediff/sampling.py | 11 ++- 6 files changed, 235 insertions(+), 94 deletions(-) diff --git a/animatediff/model_injection.py b/animatediff/model_injection.py index 88a1010..9e6b29e 100644 --- a/animatediff/model_injection.py +++ b/animatediff/model_injection.py @@ -14,7 +14,7 @@ from .logger import logger 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 SeedNoiseType +from .sample_settings import SampleSettings, SeedNoiseGeneration # some motion_model casts here might fail if model becomes metatensor or is not castable; @@ -33,6 +33,7 @@ class ModelPatcherAndInjector(ModelPatcher): # injection stuff self.motion_injection_params: InjectionParams = None + self.sample_settings: SampleSettings = SampleSettings() self.motion_model: MotionModelPatcher = None def model_patches_to(self, device): @@ -76,6 +77,7 @@ 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 return cloned @@ -322,7 +324,7 @@ class InjectionParams: self.loras: MotionLoraList = None self.motion_model_settings = MotionModelSettings() self.sampling_settings = None - self.noise_type: str = SeedNoiseType.DEFAULT + self.noise_type: str = SeedNoiseGeneration.COMFY self.sub_idxs = None # value should NOT be included in clone, so it will auto reset diff --git a/animatediff/nodes.py b/animatediff/nodes.py index 9413f70..f1063a6 100644 --- a/animatediff/nodes.py +++ b/animatediff/nodes.py @@ -9,9 +9,10 @@ from .logger import logger from .model_utils import BetaSchedules, get_available_motion_loras, get_available_motion_models, get_motion_lora_path from .motion_lora import MotionLoraInfo, MotionLoraList from .model_injection import InjectionParams, ModelPatcherAndInjector, MotionModelSettings, load_motion_module -from .sample_settings import SeedNoiseType +from .sample_settings import SampleSettings, SeedNoiseGeneration from .sampling import motion_sample_factory +from .nodes_sample import SampleSettingsNode, NoiseLayerAddNode, NoiseLayerReplaceNode from .nodes_extras import AnimateDiffUnload, EmptyLatentImageLarge, CheckpointLoaderSimpleWithNoiseSelect from .nodes_experimental import AnimateDiffModelSettingsSimple, AnimateDiffModelSettingsAdvanced, AnimateDiffModelSettingsAdvancedAttnStrengths from .nodes_deprecated import AnimateDiffLoader_Deprecated, AnimateDiffLoaderAdvanced_Deprecated, AnimateDiffCombine_Deprecated @@ -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,15 @@ 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": (SeedNoiseType.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_NoiseLayerReplace": NoiseLayerReplaceNode, # Experimental Nodes "ADE_AnimateDiffModelSettingsSimple": AnimateDiffModelSettingsSimple, "ADE_AnimateDiffModelSettings": AnimateDiffModelSettingsAdvanced, @@ -249,10 +221,13 @@ 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_NoiseLayerReplace": "Noise Layer [Replace] 🎭🅐🅓", # Experimental Nodes "ADE_AnimateDiffModelSettingsSimple": "EXP Motion Model Settings (Simple) 🎭🅐🅓", "ADE_AnimateDiffModelSettings": "EXP Motion Model Settings (Advanced) 🎭🅐🅓", diff --git a/animatediff/nodes_extras.py b/animatediff/nodes_extras.py index c105daa..52f60df 100644 --- a/animatediff/nodes_extras.py +++ b/animatediff/nodes_extras.py @@ -39,7 +39,7 @@ class CheckpointLoaderSimpleWithNoiseSelect: }, "optional": { "use_custom_scale_factor": ("BOOLEAN", {"default": False}), - "scale_factor": ("FLOAT", {"default": 0.18215, "min": 0.0, "max": 1.0}) + "scale_factor": ("FLOAT", {"default": 0.18215, "min": 0.0, "max": 1.0, "step": 0.00001}) } } RETURN_TYPES = ("MODEL", "CLIP", "VAE") diff --git a/animatediff/nodes_sample.py b/animatediff/nodes_sample.py index e09c77b..3dd9514 100644 --- a/animatediff/nodes_sample.py +++ b/animatediff/nodes_sample.py @@ -1,16 +1,97 @@ -from .sample_settings import SeedNoiseType +from torch import Tensor + +from .sample_settings import NoiseLayerAdd, NoiseLayerGroup, NoiseLayerReplace, NoiseLayerType, NoiseNormalize, SeedNoiseGeneration, SampleSettings + class SampleSettingsNode: @classmethod def INPUT_TYPES(s): return { "required": { - "seed_noise": (SeedNoiseType.LIST), + "seed_gen": (SeedNoiseGeneration.LIST, ), }, "optional": { "noise_layers": ("NOISE_LAYERS", ), } } - RETURN_TYPES = ("SAMPLE_SETTINGS") + RETURN_TYPES = ("SAMPLE_SETTINGS",) + RETURN_NAMES = ("settings",) + CATEGORY = "Animate Diff 🎭🅐🅓" + FUNCTION = "create_settings" + def create_settings(self, seed_gen: str, noise_layers: NoiseLayerGroup=None): + sampling_settings = SampleSettings(seed_gen=seed_gen, noise_layers=noise_layers) + return (sampling_settings,) + + +class NoiseLayerReplaceNode: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "noise_type": (NoiseLayerType.LIST,), + "seed_gen_override": (SeedNoiseGeneration.LIST_WITH_OVERRIDE,), + "seed_offset": ("INT", {"default": 0}), + }, + "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, 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, 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": { + "noise_type": (NoiseLayerType.LIST,), + "seed_gen_override": (SeedNoiseGeneration.LIST_WITH_OVERRIDE,), + "seed_offset": ("INT", {"default": 0, "min": -999999999999999}), + "weighted_average": ("BOOLEAN", {"default": True}), + "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}), + "normalize": (NoiseNormalize.LIST,), + }, + "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, noise_type: str, seed_gen_override: str, seed_offset: int, + weighted_average: bool, noise_weight: float, balance_multiplier: float, normalize: str, + 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, 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, weighted_average=weighted_average, normalize=normalize) + prev_noise_layers.add_to_start(layer) + return (prev_noise_layers,) diff --git a/animatediff/sample_settings.py b/animatediff/sample_settings.py index 32f66e0..ed7896d 100644 --- a/animatediff/sample_settings.py +++ b/animatediff/sample_settings.py @@ -1,35 +1,117 @@ import torch from torch import Tensor +import comfy.sample + from .logger import logger -class SamplingSettings: - def __init__(self): - self.seed_noise_type: str = SeedNoiseType.DEFAULT - self.noise_layers = NoiseLayerGroup() - - class NoiseLayerType: DEFAULT = "default" CONSTANT = "constant" + LIST = [DEFAULT, CONSTANT] + + +class NoiseApplication: + ADD = "add" + REPLACE = "replace" + + LIST = [ADD, REPLACE] + + +class NoiseNormalize: + DISABLE = "disable" + NORMAL = "normal" + + LIST = [DISABLE, NORMAL] + + +class SampleSettings: + def __init__(self, seed_gen: str=None, noise_layers: 'NoiseLayerGroup'=None, negative_cond_flipflop=False): + 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.negative_cond_flipflop = negative_cond_flipflop + + def prepare_noise(self, seed: int, latents: Tensor, noise: Tensor): + if self.seed_gen == SeedNoiseGeneration.AUTO1111: + noise = SeedNoiseGeneration.create_noise_auto1111(seed, latents) + # TODO: 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) + # 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): - self.mask: Tensor = None - self.layer_type: str = NoiseLayerType.DEFAULT - self.options: dict = {} + def __init__(self, noise_type: str, 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.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) -> Tensor: + if self.seed_override is not None: + seed = self.seed_override + seed += self.seed_offset + return SeedNoiseGeneration.create_noise(existing_seed_gen=existing_seed_gen, seed_gen=self.seed_gen_override, noise_type=self.noise_type, seed=seed, latents=latents) + + def apply_layer_noise(self, new_noise: Tensor, old_noise: Tensor) -> Tensor: + return old_noise + + +class NoiseLayerReplace(NoiseLayer): + def __init__(self, noise_type: str, seed_gen_override: str, seed_offset: int, seed_override: int=None, mask: Tensor=None): + super().__init__(noise_type, seed_gen_override, seed_offset, seed_override, mask) + self.application = NoiseApplication.REPLACE + + def apply_layer_noise(self, new_noise: Tensor, old_noise: Tensor) -> Tensor: + # if no mask, use new_noise + if self.mask is None: + return new_noise + # replicate noise_mask operations + return new_noise + + +class NoiseLayerAdd(NoiseLayer): + def __init__(self, noise_type: str, seed_gen_override: str, seed_offset: int, seed_override: int=None, mask: Tensor=None, + noise_weight=1.0, balance_multiplier=1.0, weighted_average=True, normalize:str = NoiseNormalize.DISABLE): + super().__init__(noise_type, seed_gen_override, seed_offset, seed_override, mask) + self.application = NoiseApplication.ADD + self.noise_weight = noise_weight + self.balance_multiplier = balance_multiplier + self.weighted_average = weighted_average + self.normalize = normalize + + def apply_layer_noise(self, new_noise: Tensor, old_noise: Tensor) -> Tensor: + # if mask is present, apply mask to new_noise + if self.mask is not None: + pass + if self.weighted_average: + final_noise = old_noise * (1.0-(self.noise_weight*self.balance_multiplier)) + new_noise * self.noise_weight + else: + final_noise = old_noise + new_noise * self.noise_weight + # TODO: perform normalization as requested + return final_noise class NoiseLayerGroup: def __init__(self): - self.layers = list[NoiseLayer] = [] + 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] @@ -43,44 +125,36 @@ class NoiseLayerGroup: return cloned -class SeedNoiseType: - DEFAULT = "default" +class SeedNoiseGeneration: + COMFY = "comfy" AUTO1111 = "auto1111" + USE_EXISTING = "use existing" - LIST = [DEFAULT, AUTO1111] + LIST = [COMFY, AUTO1111] + LIST_WITH_OVERRIDE = [USE_EXISTING, COMFY, 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.AUTO1111: - return cls.prepare_noise_auto1111(latents, noise, context_length, seed) - logger.warning(f"Noise type {noise_type} not recognized, proceeding with default noise.") + def create_noise(cls, existing_seed_gen: str, seed_gen: str, noise_type: str, seed: int, latents: Tensor): + # 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) + elif seed_gen == cls.AUTO1111: + return cls.create_noise_auto1111(seed, latents, noise_type) + raise ValueError(f"Noise seed_gen {seed_gen} is not recognized.") + + @staticmethod + def create_noise_comfy(seed: int, latents: Tensor, noise_type: str): + if noise_type == NoiseLayerType.CONSTANT: + return SeedNoiseGeneration.create_noise_constant(seed, latents) + noise = comfy.sample.prepare_noise(latents, seed) 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): + + @staticmethod + def create_noise_auto1111(seed: int, latents: Tensor, noise_type: str): + if noise_type == NoiseLayerType.CONSTANT: + return SeedNoiseGeneration.create_noise_constant(seed, latents) # auto1111 applies growing seeds for a batch length = latents.shape[0] single_shape = (1, latents.shape[1], latents.shape[2], latents.shape[3]) @@ -91,8 +165,14 @@ class SeedNoiseType: 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_noise_constant(seed: int, latents: Tensor): + # constant noise is just the noise for the first latent copied latent-length amount of times + 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) class __OLDNoiseType: diff --git a/animatediff/sampling.py b/animatediff/sampling.py index f956fe0..ca9ad74 100644 --- a/animatediff/sampling.py +++ b/animatediff/sampling.py @@ -15,7 +15,7 @@ import comfy.utils from comfy.controlnet import ControlBase from .context import get_context_scheduler -from .sample_settings import SeedNoiseType +from .sample_settings import SeedNoiseGeneration from .motion_utils import GroupNormAD from .model_utils import ModelTypeSD, wrap_function_to_inject_xformers_bug_info from .model_injection import InjectionParams, ModelPatcherAndInjector, MotionModelPatcher @@ -244,9 +244,12 @@ 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 = SeedNoiseType.prepare_noise(params.noise_type, latents=latents, noise=noise, context_length=params.context_length, seed=seed) + + # if sampling settings provided, time to get wild with it + if model.sample_settings is not None: + # if noise is not disabled, do noise stuff + if not disable_noise: + noise = model.sample_settings.prepare_noise(seed, latents, noise) # apply params to motion model apply_params_to_motion_model(model.motion_model, params) From 18684dfe2a2ee9221c4d9578391ad501b53e70e1 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Wed, 27 Dec 2023 00:45:26 -0600 Subject: [PATCH 06/12] Separate old Noise Layer [Add] node into [Add] and [Add Weighted], support masks, added Empty noise type --- animatediff/nodes.py | 4 +- animatediff/nodes_sample.py | 69 +++++++++++++---- animatediff/sample_settings.py | 135 +++++++++++++++++++++------------ animatediff/sampling.py | 14 +--- 4 files changed, 145 insertions(+), 77 deletions(-) diff --git a/animatediff/nodes.py b/animatediff/nodes.py index f1063a6..fd560c7 100644 --- a/animatediff/nodes.py +++ b/animatediff/nodes.py @@ -12,7 +12,7 @@ from .model_injection import InjectionParams, ModelPatcherAndInjector, MotionMod from .sample_settings import SampleSettings, SeedNoiseGeneration from .sampling import motion_sample_factory -from .nodes_sample import SampleSettingsNode, NoiseLayerAddNode, NoiseLayerReplaceNode +from .nodes_sample import NoiseLayerAddWeightedNode, SampleSettingsNode, NoiseLayerAddNode, NoiseLayerReplaceNode from .nodes_extras import AnimateDiffUnload, EmptyLatentImageLarge, CheckpointLoaderSimpleWithNoiseSelect from .nodes_experimental import AnimateDiffModelSettingsSimple, AnimateDiffModelSettingsAdvanced, AnimateDiffModelSettingsAdvancedAttnStrengths from .nodes_deprecated import AnimateDiffLoader_Deprecated, AnimateDiffLoaderAdvanced_Deprecated, AnimateDiffCombine_Deprecated @@ -205,6 +205,7 @@ NODE_CLASS_MAPPINGS = { "ADE_AnimateDiffModelSettings_Release": AnimateDiffModelSettings, # Noise Layer Nodes "ADE_NoiseLayerAdd": NoiseLayerAddNode, + "ADE_NoiseLayerAddWeighted": NoiseLayerAddWeightedNode, "ADE_NoiseLayerReplace": NoiseLayerReplaceNode, # Experimental Nodes "ADE_AnimateDiffModelSettingsSimple": AnimateDiffModelSettingsSimple, @@ -227,6 +228,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "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] 🎭🅐🅓", # Experimental Nodes "ADE_AnimateDiffModelSettingsSimple": "EXP Motion Model Settings (Simple) 🎭🅐🅓", diff --git a/animatediff/nodes_sample.py b/animatediff/nodes_sample.py index 3dd9514..cd8dde4 100644 --- a/animatediff/nodes_sample.py +++ b/animatediff/nodes_sample.py @@ -1,6 +1,6 @@ from torch import Tensor -from .sample_settings import NoiseLayerAdd, NoiseLayerGroup, NoiseLayerReplace, NoiseLayerType, NoiseNormalize, SeedNoiseGeneration, SampleSettings +from .sample_settings import NoiseLayerAdd, NoiseLayerAddWeighted, NoiseLayerGroup, NoiseLayerReplace, NoiseLayerType, NoiseNormalize, SeedNoiseGeneration, SampleSettings class SampleSettingsNode: @@ -8,10 +8,12 @@ class SampleSettingsNode: def INPUT_TYPES(s): return { "required": { - "seed_gen": (SeedNoiseGeneration.LIST, ), + "batch_offset": ("INT", {"default": 0, "min": 0}), + "noise_type": (NoiseLayerType.LIST,), + "seed_gen": (SeedNoiseGeneration.LIST,), }, "optional": { - "noise_layers": ("NOISE_LAYERS", ), + "noise_layers": ("NOISE_LAYERS",), } } @@ -20,8 +22,8 @@ class SampleSettingsNode: CATEGORY = "Animate Diff 🎭🅐🅓" FUNCTION = "create_settings" - def create_settings(self, seed_gen: str, noise_layers: NoiseLayerGroup=None): - sampling_settings = SampleSettings(seed_gen=seed_gen, noise_layers=noise_layers) + def create_settings(self, batch_offset: int, noise_type: str, seed_gen: str, noise_layers: NoiseLayerGroup=None): + sampling_settings = SampleSettings(batch_offset=batch_offset, noise_type=noise_type, seed_gen=seed_gen, noise_layers=noise_layers) return (sampling_settings,) @@ -30,9 +32,10 @@ class NoiseLayerReplaceNode: 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}), + "seed_offset": ("INT", {"default": 0, "min": -999999999999999}), }, "optional": { "prev_noise_layers": ("NOISE_LAYERS",), @@ -45,14 +48,14 @@ class NoiseLayerReplaceNode: CATEGORY = "Animate Diff 🎭🅐🅓/noise layers" FUNCTION = "create_layers" - def create_layers(self, noise_type: str, seed_gen_override: str, seed_offset: int, + 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, seed_gen_override=seed_gen_override, seed_offset=seed_offset, + 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,) @@ -63,13 +66,11 @@ class NoiseLayerAddNode: 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}), - "weighted_average": ("BOOLEAN", {"default": True}), "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}), - "normalize": (NoiseNormalize.LIST,), }, "optional": { "prev_noise_layers": ("NOISE_LAYERS",), @@ -82,16 +83,54 @@ class NoiseLayerAddNode: CATEGORY = "Animate Diff 🎭🅐🅓/noise layers" FUNCTION = "create_layers" - def create_layers(self, noise_type: str, seed_gen_override: str, seed_offset: int, - weighted_average: bool, noise_weight: float, balance_multiplier: float, normalize: str, + 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, seed_gen_override=seed_gen_override, seed_offset=seed_offset, + 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, balance_multiplier=balance_multiplier, weighted_average=weighted_average, normalize=normalize) + 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,) diff --git a/animatediff/sample_settings.py b/animatediff/sample_settings.py index ed7896d..aeb24fd 100644 --- a/animatediff/sample_settings.py +++ b/animatediff/sample_settings.py @@ -6,18 +6,30 @@ import comfy.sample 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" - LIST = [DEFAULT, CONSTANT] + LIST = [DEFAULT, CONSTANT, EMPTY] class NoiseApplication: ADD = "add" + ADD_WEIGHTED = "add_weighted" REPLACE = "replace" - LIST = [ADD, REPLACE] + LIST = [ADD, ADD_WEIGHTED, REPLACE] class NoiseNormalize: @@ -28,15 +40,19 @@ class NoiseNormalize: class SampleSettings: - def __init__(self, seed_gen: str=None, noise_layers: 'NoiseLayerGroup'=None, negative_cond_flipflop=False): + def __init__(self, batch_offset: int=0, noise_type: str=None, seed_gen: str=None, noise_layers: 'NoiseLayerGroup'=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.negative_cond_flipflop = negative_cond_flipflop def prepare_noise(self, seed: int, latents: Tensor, noise: Tensor): - if self.seed_gen == SeedNoiseGeneration.AUTO1111: - noise = SeedNoiseGeneration.create_noise_auto1111(seed, latents) - # TODO: apply noise layers + # replace initial noise if not batch_offset 0 or Comfy seed_gen or not NoiseType default + if self.batch_offset != 0 or self.noise_type != NoiseLayerType.DEFAULT or self.seed_gen != SeedNoiseGeneration.COMFY: + noise = SeedNoiseGeneration.create_noise(existing_seed_gen=self.seed_gen, seed_gen=self.seed_gen, noise_type=self.noise_type, batch_offset=self.batch_offset, + seed=seed, latents=latents) + # 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) @@ -47,9 +63,10 @@ class SampleSettings: class NoiseLayer: - def __init__(self, noise_type: str, seed_gen_override: str, seed_offset: int, seed_override: int=None, mask: Tensor=None): + 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 @@ -59,45 +76,51 @@ class NoiseLayer: if self.seed_override is not None: seed = self.seed_override seed += self.seed_offset - return SeedNoiseGeneration.create_noise(existing_seed_gen=existing_seed_gen, seed_gen=self.seed_gen_override, noise_type=self.noise_type, seed=seed, latents=latents) + return SeedNoiseGeneration.create_noise(existing_seed_gen=existing_seed_gen, seed_gen=self.seed_gen_override, + noise_type=self.noise_type, batch_offset=self.batch_offset, seed=seed, latents=latents) 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, seed_gen_override: str, seed_offset: int, seed_override: int=None, mask: Tensor=None): - super().__init__(noise_type, seed_gen_override, seed_offset, seed_override, mask) + 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: - # if no mask, use new_noise - if self.mask is None: - return new_noise - # replicate noise_mask operations - return new_noise + 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, seed_gen_override: str, seed_offset: int, seed_override: int=None, mask: Tensor=None, - noise_weight=1.0, balance_multiplier=1.0, weighted_average=True, normalize:str = NoiseNormalize.DISABLE): - super().__init__(noise_type, seed_gen_override, seed_offset, seed_override, mask) - self.application = NoiseApplication.ADD + 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.balance_multiplier = balance_multiplier - self.weighted_average = weighted_average - self.normalize = normalize + self.application = NoiseApplication.ADD def apply_layer_noise(self, new_noise: Tensor, old_noise: Tensor) -> Tensor: - # if mask is present, apply mask to new_noise - if self.mask is not None: - pass - if self.weighted_average: - final_noise = old_noise * (1.0-(self.noise_weight*self.balance_multiplier)) + new_noise * self.noise_weight - else: - final_noise = old_noise + new_noise * self.noise_weight - # TODO: perform normalization as requested - return final_noise + 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: @@ -134,45 +157,59 @@ class SeedNoiseGeneration: LIST_WITH_OVERRIDE = [USE_EXISTING, COMFY, AUTO1111] @classmethod - def create_noise(cls, existing_seed_gen: str, seed_gen: str, noise_type: str, seed: int, latents: Tensor): + def create_noise(cls, existing_seed_gen: str, seed_gen: str, noise_type: str, batch_offset: int, seed: int, latents: Tensor): # 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) + return cls.create_noise_comfy(seed, latents, noise_type, batch_offset) elif seed_gen == cls.AUTO1111: - return cls.create_noise_auto1111(seed, latents, noise_type) + return cls.create_noise_auto1111(seed, latents, noise_type, batch_offset) raise ValueError(f"Noise seed_gen {seed_gen} is not recognized.") @staticmethod - def create_noise_comfy(seed: int, latents: Tensor, noise_type: str): + def create_noise_comfy(seed: int, latents: Tensor, noise_type: str, batch_offset: int): + common_noise = SeedNoiseGeneration._create_common_noise(seed, latents, noise_type, batch_offset) + if common_noise is not None: + return common_noise if noise_type == NoiseLayerType.CONSTANT: - return SeedNoiseGeneration.create_noise_constant(seed, latents) - noise = comfy.sample.prepare_noise(latents, seed) - return noise + 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]) + offset_noises = torch.randn(offset_shape, dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu") + return offset_noises[batch_offset:] @staticmethod - def create_noise_auto1111(seed: int, latents: Tensor, noise_type: str): + def create_noise_auto1111(seed: int, latents: Tensor, noise_type: str, batch_offset: int): + common_noise = SeedNoiseGeneration._create_common_noise(seed, latents, noise_type, batch_offset) + if common_noise is not None: + return common_noise if noise_type == NoiseLayerType.CONSTANT: - return SeedNoiseGeneration.create_noise_constant(seed, latents) + 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) + 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")) return torch.cat(all_noises, dim=0) - + @staticmethod - def create_noise_constant(seed: int, latents: Tensor): - # constant noise is just the noise for the first latent copied latent-length amount of times - 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) + def _create_common_noise(seed: int, latents: Tensor, noise_type: str, batch_offset: int): + if noise_type == NoiseLayerType.EMPTY: + return torch.zeros_like(latents) + return None class __OLDNoiseType: diff --git a/animatediff/sampling.py b/animatediff/sampling.py index ca9ad74..9fd5923 100644 --- a/animatediff/sampling.py +++ b/animatediff/sampling.py @@ -15,7 +15,7 @@ import comfy.utils from comfy.controlnet import ControlBase from .context import get_context_scheduler -from .sample_settings import SeedNoiseGeneration +from .sample_settings import 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 @@ -147,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}.") @@ -244,7 +234,7 @@ 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 sampling settings provided, time to get wild with it if model.sample_settings is not None: # if noise is not disabled, do noise stuff From 290162d5e8741358470672817b7d02e3871564df Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sat, 30 Dec 2023 03:44:33 -0600 Subject: [PATCH 07/12] First real attempt at FreeInit implementation - not yet compared to diffusers implementation, will do so after I get sleep --- animatediff/freeinit.py | 162 +++++++++++++++++++++++++++++++++ animatediff/nodes.py | 10 +- animatediff/nodes_sample.py | 52 ++++++++++- animatediff/sample_settings.py | 90 +++++++++++++++++- animatediff/sampling.py | 80 +++++++++++----- 5 files changed, 366 insertions(+), 28 deletions(-) create mode 100644 animatediff/freeinit.py diff --git a/animatediff/freeinit.py b/animatediff/freeinit.py new file mode 100644 index 0000000..ac9edfa --- /dev/null +++ b/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 + diff --git a/animatediff/nodes.py b/animatediff/nodes.py index fd560c7..6db98ca 100644 --- a/animatediff/nodes.py +++ b/animatediff/nodes.py @@ -12,14 +12,14 @@ from .model_injection import InjectionParams, ModelPatcherAndInjector, MotionMod from .sample_settings import SampleSettings, SeedNoiseGeneration from .sampling import motion_sample_factory -from .nodes_sample import NoiseLayerAddWeightedNode, SampleSettingsNode, NoiseLayerAddNode, NoiseLayerReplaceNode +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: @@ -207,6 +207,9 @@ NODE_CLASS_MAPPINGS = { "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, @@ -230,6 +233,9 @@ NODE_DISPLAY_NAME_MAPPINGS = { "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) 🎭🅐🅓", diff --git a/animatediff/nodes_sample.py b/animatediff/nodes_sample.py index cd8dde4..b522747 100644 --- a/animatediff/nodes_sample.py +++ b/animatediff/nodes_sample.py @@ -1,6 +1,7 @@ from torch import Tensor -from .sample_settings import NoiseLayerAdd, NoiseLayerAddWeighted, NoiseLayerGroup, NoiseLayerReplace, NoiseLayerType, NoiseNormalize, SeedNoiseGeneration, SampleSettings +from .freeinit import FreeInitFilter +from .sample_settings import FreeInitOptions, IterationOptions, NoiseLayerAdd, NoiseLayerAddWeighted, NoiseLayerGroup, NoiseLayerReplace, NoiseLayerType, NoiseNormalize, SeedNoiseGeneration, SampleSettings class SampleSettingsNode: @@ -14,6 +15,7 @@ class SampleSettingsNode: }, "optional": { "noise_layers": ("NOISE_LAYERS",), + "iter_opts": ("ITER_OPTS",), } } @@ -22,8 +24,8 @@ class SampleSettingsNode: CATEGORY = "Animate Diff 🎭🅐🅓" FUNCTION = "create_settings" - def create_settings(self, batch_offset: int, noise_type: str, seed_gen: str, noise_layers: NoiseLayerGroup=None): - sampling_settings = SampleSettings(batch_offset=batch_offset, noise_type=noise_type, seed_gen=seed_gen, noise_layers=noise_layers) + def create_settings(self, batch_offset: int, noise_type: str, seed_gen: str, noise_layers: NoiseLayerGroup=None, iter_opts: IterationOptions=None): + sampling_settings = SampleSettings(batch_offset=batch_offset, noise_type=noise_type, seed_gen=seed_gen, noise_layers=noise_layers, iter_opts=iter_opts) return (sampling_settings,) @@ -134,3 +136,47 @@ class NoiseLayerAddWeightedNode: 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 = ("ITER_OPTS",) + CATEGORY = "Animate Diff 🎭🅐🅓/iter 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}), + } + } + + RETURN_TYPES = ("ITER_OPTS",) + CATEGORY = "Animate Diff 🎭🅐🅓/iter 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): + 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,) diff --git a/animatediff/sample_settings.py b/animatediff/sample_settings.py index aeb24fd..253db7d 100644 --- a/animatediff/sample_settings.py +++ b/animatediff/sample_settings.py @@ -2,7 +2,9 @@ import torch from torch import Tensor import comfy.sample +from model_patcher import ModelPatcher +from .freeinit import FreeInitFilter, freq_mix_3d, get_freq_filter from .logger import logger @@ -40,11 +42,12 @@ class NoiseNormalize: class SampleSettings: - def __init__(self, batch_offset: int=0, noise_type: str=None, seed_gen: str=None, noise_layers: 'NoiseLayerGroup'=None, negative_cond_flipflop=False): + def __init__(self, batch_offset: int=0, noise_type: str=None, seed_gen: str=None, noise_layers: 'NoiseLayerGroup'=None, iter_opts=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.iter_opts = iter_opts if iter_opts else IterationOptions() self.negative_cond_flipflop = negative_cond_flipflop def prepare_noise(self, seed: int, latents: Tensor, noise: Tensor): @@ -212,6 +215,91 @@ class SeedNoiseGeneration: return None +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): + 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): + # 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 + # 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 + # 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 - combines low freq. noise from noised_latents and high freq. noise from z_rand + noise = freq_mix_3d(x=noised_latents, noise=z_rand, LPF=self.freq_filter) + return cached_latents, noise #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 + + class __OLDNoiseType: DEFAULT = "default" REPEATED = "repeated" diff --git a/animatediff/sampling.py b/animatediff/sampling.py index 9fd5923..1ff140f 100644 --- a/animatediff/sampling.py +++ b/animatediff/sampling.py @@ -9,13 +9,13 @@ 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 .sample_settings import SeedNoiseGeneration, prepare_mask_ad +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 @@ -177,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: @@ -191,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): @@ -203,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 @@ -244,12 +246,7 @@ def motion_sample_factory(orig_comfy_sample: Callable) -> Callable: # 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 - + # callback setup original_callback = kwargs.get("callback", None) def ad_callback(step, x0, x, total_steps): if original_callback is not None: @@ -259,11 +256,50 @@ 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.iter_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) + 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 From fda07b5b3630111aff278e1a7af96ab850989ba8 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Mon, 1 Jan 2024 22:21:02 -0600 Subject: [PATCH 08/12] FreeInit high freq noise gen nose based on Sample Settings, refactored some code, added repeated_context noise type --- animatediff/model_injection.py | 2 + animatediff/nodes_sample.py | 13 +++-- animatediff/sample_settings.py | 102 +++++++++++++++++++++++---------- animatediff/sampling.py | 17 +++--- 4 files changed, 93 insertions(+), 41 deletions(-) diff --git a/animatediff/model_injection.py b/animatediff/model_injection.py index 9e6b29e..6a16a1d 100644 --- a/animatediff/model_injection.py +++ b/animatediff/model_injection.py @@ -327,6 +327,8 @@ class InjectionParams: 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 diff --git a/animatediff/nodes_sample.py b/animatediff/nodes_sample.py index b522747..3109cf0 100644 --- a/animatediff/nodes_sample.py +++ b/animatediff/nodes_sample.py @@ -16,6 +16,7 @@ class SampleSettingsNode: "optional": { "noise_layers": ("NOISE_LAYERS",), "iter_opts": ("ITER_OPTS",), + "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}), } } @@ -24,8 +25,8 @@ class SampleSettingsNode: CATEGORY = "Animate Diff 🎭🅐🅓" FUNCTION = "create_settings" - def create_settings(self, batch_offset: int, noise_type: str, seed_gen: str, noise_layers: NoiseLayerGroup=None, iter_opts: IterationOptions=None): - sampling_settings = SampleSettings(batch_offset=batch_offset, noise_type=noise_type, seed_gen=seed_gen, noise_layers=noise_layers, iter_opts=iter_opts) + def create_settings(self, batch_offset: int, noise_type: str, seed_gen: str, noise_layers: NoiseLayerGroup=None, iter_opts: IterationOptions=None, seed_override: int=None,): + sampling_settings = SampleSettings(batch_offset=batch_offset, noise_type=noise_type, seed_gen=seed_gen, noise_layers=noise_layers, iter_opts=iter_opts, seed_override=seed_override) return (sampling_settings,) @@ -148,7 +149,7 @@ class IterationOptionsNode: } RETURN_TYPES = ("ITER_OPTS",) - CATEGORY = "Animate Diff 🎭🅐🅓/iter opts" + CATEGORY = "Animate Diff 🎭🅐🅓/iteration opts" FUNCTION = "create_iter_opts" def create_iter_opts(self, iterations: int): @@ -168,15 +169,17 @@ class FreeInitOptionsNode: "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 = ("ITER_OPTS",) - CATEGORY = "Animate Diff 🎭🅐🅓/iter 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): + 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,) diff --git a/animatediff/sample_settings.py b/animatediff/sample_settings.py index 253db7d..66f15d4 100644 --- a/animatediff/sample_settings.py +++ b/animatediff/sample_settings.py @@ -22,8 +22,9 @@ class NoiseLayerType: DEFAULT = "default" CONSTANT = "constant" EMPTY = "empty" + REPEATED_CONTEXT = "repeated_context" - LIST = [DEFAULT, CONSTANT, EMPTY] + LIST = [DEFAULT, CONSTANT, EMPTY, REPEATED_CONTEXT] class NoiseApplication: @@ -42,23 +43,27 @@ class NoiseNormalize: class SampleSettings: - def __init__(self, batch_offset: int=0, noise_type: str=None, seed_gen: str=None, noise_layers: 'NoiseLayerGroup'=None, iter_opts=None, negative_cond_flipflop=False): + def __init__(self, batch_offset: int=0, noise_type: str=None, seed_gen: str=None, noise_layers: 'NoiseLayerGroup'=None, iter_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.iter_opts = iter_opts if iter_opts else IterationOptions() + self.seed_override = seed_override self.negative_cond_flipflop = negative_cond_flipflop - def prepare_noise(self, seed: int, latents: Tensor, noise: Tensor): + 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 # replace initial noise if not batch_offset 0 or Comfy seed_gen or not NoiseType default - if self.batch_offset != 0 or self.noise_type != NoiseLayerType.DEFAULT or self.seed_gen != SeedNoiseGeneration.COMFY: - noise = SeedNoiseGeneration.create_noise(existing_seed_gen=self.seed_gen, seed_gen=self.seed_gen, noise_type=self.noise_type, batch_offset=self.batch_offset, - seed=seed, latents=latents) + if self.batch_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) + 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 @@ -75,12 +80,12 @@ class NoiseLayer: self.seed_override = seed_override self.mask = mask - def create_layer_noise(self, existing_seed_gen: str, seed: int, latents: Tensor) -> Tensor: + 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 - seed += self.seed_offset - return SeedNoiseGeneration.create_noise(existing_seed_gen=existing_seed_gen, seed_gen=self.seed_gen_override, - noise_type=self.noise_type, batch_offset=self.batch_offset, seed=seed, latents=latents) + 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 @@ -150,7 +155,6 @@ class NoiseLayerGroup: cloned.add(layer) return cloned - class SeedNoiseGeneration: COMFY = "comfy" AUTO1111 = "auto1111" @@ -159,20 +163,22 @@ class SeedNoiseGeneration: LIST = [COMFY, AUTO1111] LIST_WITH_OVERRIDE = [USE_EXISTING, COMFY, AUTO1111] + + @classmethod - def create_noise(cls, existing_seed_gen: str, seed_gen: str, noise_type: str, batch_offset: int, seed: int, latents: Tensor): + 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) + 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) + 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, batch_offset: int): - common_noise = SeedNoiseGeneration._create_common_noise(seed, latents, noise_type, batch_offset) + 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: @@ -184,12 +190,17 @@ class SeedNoiseGeneration: # 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]) - offset_noises = torch.randn(offset_shape, dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu") - return offset_noises[batch_offset:] + 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, batch_offset: int): - common_noise = SeedNoiseGeneration._create_common_noise(seed, latents, noise_type, batch_offset) + 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: @@ -206,13 +217,39 @@ class SeedNoiseGeneration: 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")) - return torch.cat(all_noises, dim=0) + 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_common_noise(seed: int, latents: Tensor, noise_type: str, batch_offset: int): + 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: @@ -238,6 +275,10 @@ class IterationOptions: 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) @@ -254,21 +295,24 @@ class FreeInitOptions(IterationOptions): 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): + 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 + # 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 - # 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) + 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, LPF=self.freq_filter) - return cached_latents, noise #noise + 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 diff --git a/animatediff/sampling.py b/animatediff/sampling.py index 1ff140f..b2f29fa 100644 --- a/animatediff/sampling.py +++ b/animatediff/sampling.py @@ -237,15 +237,16 @@ def motion_sample_factory(orig_comfy_sample: Callable, is_custom: bool=False) -> disable_noise = kwargs.get("disable_noise") or False seed = kwargs["seed"] - # if sampling settings provided, time to get wild with it - if model.sample_settings is not None: - # if noise is not disabled, do noise stuff - if not disable_noise: - noise = model.sample_settings.prepare_noise(seed, latents, noise) - # apply params to motion model apply_params_to_motion_model(model.motion_model, params) + # 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): @@ -289,7 +290,9 @@ def motion_sample_factory(orig_comfy_sample: Callable, is_custom: bool=False) -> 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) + 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() From b22e7cdf48eaedfb7f4952f932ad860d6ad493d2 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Mon, 1 Jan 2024 23:34:11 -0600 Subject: [PATCH 09/12] Renamed iter_opts to iteration_opts --- animatediff/nodes_sample.py | 10 +++++----- animatediff/sample_settings.py | 4 ++-- animatediff/sampling.py | 2 +- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/animatediff/nodes_sample.py b/animatediff/nodes_sample.py index 3109cf0..4d613dc 100644 --- a/animatediff/nodes_sample.py +++ b/animatediff/nodes_sample.py @@ -15,7 +15,7 @@ class SampleSettingsNode: }, "optional": { "noise_layers": ("NOISE_LAYERS",), - "iter_opts": ("ITER_OPTS",), + "iteration_opts": ("ITERATION_OPTS",), "seed_override": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "forceInput": True}), } } @@ -25,8 +25,8 @@ class SampleSettingsNode: CATEGORY = "Animate Diff 🎭🅐🅓" FUNCTION = "create_settings" - def create_settings(self, batch_offset: int, noise_type: str, seed_gen: str, noise_layers: NoiseLayerGroup=None, iter_opts: IterationOptions=None, seed_override: int=None,): - sampling_settings = SampleSettings(batch_offset=batch_offset, noise_type=noise_type, seed_gen=seed_gen, noise_layers=noise_layers, iter_opts=iter_opts, seed_override=seed_override) + def create_settings(self, batch_offset: int, noise_type: str, seed_gen: str, 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, noise_layers=noise_layers, iteration_opts=iteration_opts, seed_override=seed_override) return (sampling_settings,) @@ -148,7 +148,7 @@ class IterationOptionsNode: } } - RETURN_TYPES = ("ITER_OPTS",) + RETURN_TYPES = ("ITERATION_OPTS",) CATEGORY = "Animate Diff 🎭🅐🅓/iteration opts" FUNCTION = "create_iter_opts" @@ -173,7 +173,7 @@ class FreeInitOptionsNode: } } - RETURN_TYPES = ("ITER_OPTS",) + RETURN_TYPES = ("ITERATION_OPTS",) CATEGORY = "Animate Diff 🎭🅐🅓/iteration opts" FUNCTION = "create_iter_opts" diff --git a/animatediff/sample_settings.py b/animatediff/sample_settings.py index 66f15d4..3363b60 100644 --- a/animatediff/sample_settings.py +++ b/animatediff/sample_settings.py @@ -43,12 +43,12 @@ class NoiseNormalize: class SampleSettings: - def __init__(self, batch_offset: int=0, noise_type: str=None, seed_gen: str=None, noise_layers: 'NoiseLayerGroup'=None, iter_opts=None, seed_override:int=None, negative_cond_flipflop=False): + def __init__(self, batch_offset: int=0, noise_type: str=None, seed_gen: str=None, 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.iter_opts = iter_opts if iter_opts else IterationOptions() + self.iteration_opts = iteration_opts if iteration_opts else IterationOptions() self.seed_override = seed_override self.negative_cond_flipflop = negative_cond_flipflop diff --git a/animatediff/sampling.py b/animatediff/sampling.py index b2f29fa..192ec55 100644 --- a/animatediff/sampling.py +++ b/animatediff/sampling.py @@ -259,7 +259,7 @@ def motion_sample_factory(orig_comfy_sample: Callable, is_custom: bool=False) -> iter_opts = IterationOptions() if model.sample_settings is not None: - iter_opts = model.sample_settings.iter_opts + iter_opts = model.sample_settings.iteration_opts iter_opts.initialize(latents) # cache initial noise and latents, if needed if iter_opts.cache_init_latents: From 42447bad5e03a2918e40d6b9c03fd0fc79a78a46 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 2 Jan 2024 02:02:38 -0600 Subject: [PATCH 10/12] Support int lists for seed_override inputs --- animatediff/sample_settings.py | 36 +++++++++++++++++++++++++++------- 1 file changed, 29 insertions(+), 7 deletions(-) diff --git a/animatediff/sample_settings.py b/animatediff/sample_settings.py index 3363b60..fd5645b 100644 --- a/animatediff/sample_settings.py +++ b/animatediff/sample_settings.py @@ -1,3 +1,4 @@ +from collections.abc import Iterable import torch from torch import Tensor @@ -43,22 +44,29 @@ class NoiseNormalize: class SampleSettings: - def __init__(self, batch_offset: int=0, noise_type: str=None, seed_gen: str=None, noise_layers: 'NoiseLayerGroup'=None, iteration_opts=None, seed_override:int=None, negative_cond_flipflop=False): + 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 - # replace initial noise if not batch_offset 0 or Comfy seed_gen or not NoiseType default - if self.batch_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) + # 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 @@ -83,6 +91,9 @@ class NoiseLayer: 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) @@ -163,8 +174,6 @@ class SeedNoiseGeneration: 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 @@ -224,6 +233,19 @@ class SeedNoiseGeneration: 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: From b13afab5635c2e439b58e7eb7e14c39407732ae2 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 2 Jan 2024 02:23:01 -0600 Subject: [PATCH 11/12] Add seed_offset to Sample Settings node --- animatediff/nodes_sample.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/animatediff/nodes_sample.py b/animatediff/nodes_sample.py index 4d613dc..e4e7f61 100644 --- a/animatediff/nodes_sample.py +++ b/animatediff/nodes_sample.py @@ -12,6 +12,7 @@ class SampleSettingsNode: "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",), @@ -25,8 +26,8 @@ class SampleSettingsNode: CATEGORY = "Animate Diff 🎭🅐🅓" FUNCTION = "create_settings" - def create_settings(self, batch_offset: int, noise_type: str, seed_gen: str, 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, noise_layers=noise_layers, iteration_opts=iteration_opts, seed_override=seed_override) + 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,) From 6730c8e24001bcb144b250a0ca8bd7af25283cce Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 2 Jan 2024 02:44:21 -0600 Subject: [PATCH 12/12] Removed no longer used code, fixed an import to have comfy prefix --- animatediff/nodes_sample.py | 2 +- animatediff/sample_settings.py | 62 +--------------------------------- 2 files changed, 2 insertions(+), 62 deletions(-) diff --git a/animatediff/nodes_sample.py b/animatediff/nodes_sample.py index e4e7f61..f3777c3 100644 --- a/animatediff/nodes_sample.py +++ b/animatediff/nodes_sample.py @@ -1,7 +1,7 @@ from torch import Tensor from .freeinit import FreeInitFilter -from .sample_settings import FreeInitOptions, IterationOptions, NoiseLayerAdd, NoiseLayerAddWeighted, NoiseLayerGroup, NoiseLayerReplace, NoiseLayerType, NoiseNormalize, SeedNoiseGeneration, SampleSettings +from .sample_settings import FreeInitOptions, IterationOptions, NoiseLayerAdd, NoiseLayerAddWeighted, NoiseLayerGroup, NoiseLayerReplace, NoiseLayerType, SeedNoiseGeneration, SampleSettings class SampleSettingsNode: diff --git a/animatediff/sample_settings.py b/animatediff/sample_settings.py index fd5645b..bd163e0 100644 --- a/animatediff/sample_settings.py +++ b/animatediff/sample_settings.py @@ -3,7 +3,7 @@ import torch from torch import Tensor import comfy.sample -from model_patcher import ModelPatcher +from comfy.model_patcher import ModelPatcher from .freeinit import FreeInitFilter, freq_mix_3d, get_freq_filter from .logger import logger @@ -364,63 +364,3 @@ class FreeInitOptions(IterationOptions): # 3. noise reinitialization latents = freq_mix_3d(x=latents, noise=z_rand, LPF=self.freq_filter) return latents, noise - - -class __OLDNoiseType: - DEFAULT = "default" - REPEATED = "repeated" - CONSTANT = "constant" - CONSTANT_ADDED = "constant_added" - AUTO1111 = "auto1111" - - LIST = [DEFAULT, REPEATED, CONSTANT, CONSTANT_ADDED, 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.CONSTANT_ADDED: - new_noise = cls.prepare_noise_constant(latents, noise, context_length, seed) - noise_weight = 0.2 - return noise * (1.0-(noise_weight/3)) + new_noise * noise_weight - 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)