From 82fc5fc441f3b463dc6be35fdd3badfbf08842fd Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 9 Jan 2024 09:00:07 -0600 Subject: [PATCH] Added FreeNoise noise_type, cleaned up forward_timestep_embed_factory --- animatediff/model_injection.py | 1 + animatediff/sample_settings.py | 57 ++++++++++++++++++++++++---- animatediff/sampling.py | 68 +++++++++++++--------------------- 3 files changed, 76 insertions(+), 50 deletions(-) diff --git a/animatediff/model_injection.py b/animatediff/model_injection.py index 686446d..ad11947 100644 --- a/animatediff/model_injection.py +++ b/animatediff/model_injection.py @@ -510,6 +510,7 @@ class InjectionParams: def set_noise_extra_args(self, noise_extra_args: dict): noise_extra_args["context_length"] = self.context_length + noise_extra_args["context_overlap"] = self.context_overlap 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/sample_settings.py b/animatediff/sample_settings.py index 8816321..b675c42 100644 --- a/animatediff/sample_settings.py +++ b/animatediff/sample_settings.py @@ -25,8 +25,9 @@ class NoiseLayerType: CONSTANT = "constant" EMPTY = "empty" REPEATED_CONTEXT = "repeated_context" + FREENOISE = "FreeNoise" - LIST = [DEFAULT, CONSTANT, EMPTY, REPEATED_CONTEXT] + LIST = [DEFAULT, CONSTANT, EMPTY, REPEATED_CONTEXT, FREENOISE] class NoiseApplication: @@ -209,7 +210,7 @@ class SeedNoiseGeneration: 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) + derivative_noise = SeedNoiseGeneration._create_derivative_noise(final_noise, noise_type=noise_type, seed=seed, extra_args=extra_args) if derivative_noise is not None: return derivative_noise return final_noise @@ -235,7 +236,7 @@ class SeedNoiseGeneration: all_noises.append(torch.randn(single_shape, dtype=latents.dtype, layout=latents.layout, generator=generator, device="cpu")) final_noise = torch.cat(all_noises, dim=0) # convert to derivative noise type, if needed - derivative_noise = SeedNoiseGeneration._create_derivative_noise(final_noise, noise_type=noise_type, extra_args=extra_args) + derivative_noise = SeedNoiseGeneration._create_derivative_noise(final_noise, noise_type=noise_type, seed=seed, extra_args=extra_args) if derivative_noise is not None: return derivative_noise return final_noise @@ -260,14 +261,14 @@ class SeedNoiseGeneration: return None @staticmethod - def _create_derivative_noise(noise: Tensor, noise_type: str, extra_args: dict): + def _create_derivative_noise(noise: Tensor, noise_type: str, seed: int, 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) + return derivative_func(noise=noise, seed=seed, extra_args=extra_args) @staticmethod - def _convert_to_repeated_context(noise: Tensor, extra_args: dict): + def _convert_to_repeated_context(noise: Tensor, extra_args: dict, **kwargs): # if no context_length, return unmodified noise context_length: int = extra_args["context_length"] if context_length is None: @@ -277,8 +278,50 @@ class SeedNoiseGeneration: cat_count = (length // context_length) + 1 return torch.cat([noise] * cat_count, dim=0)[:length] + @staticmethod + def _convert_to_freenoise(noise: Tensor, extra_args: dict, seed: int, **kwargs): + # if no context_length, return unmodified noise + context_length: int = extra_args["context_length"] + context_overlap: int = extra_args["context_overlap"] + video_length: int = noise.shape[0] + if context_length is None: + return noise + delta = context_length - context_overlap + generator = torch.manual_seed(seed) -DERIVATIVE_NOISE_FUNC_MAP = {NoiseLayerType.REPEATED_CONTEXT: SeedNoiseGeneration._convert_to_repeated_context,} + for start_idx in range(0, video_length-context_length, delta): + # start_idx corresponds to the beginning of a context window + # goal: place shuffled in the delta region right after the end of the context window + # if space after context window is not enough to place the noise, adjust and finish + place_idx = start_idx + context_length + # if place_idx is outside the valid indexes, we are already finished + if place_idx >= video_length: + break + end_idx = place_idx - 1 + # if there is not enough room to copy delta amount of indexes, copy limited amount and finish + if end_idx + delta >= video_length: + final_delta = video_length - place_idx + # generate list of indexes in final delta region + list_idx = torch.Tensor(list(range(start_idx,start_idx+final_delta))).to(torch.long) + # shuffle list + list_idx = list_idx[torch.randperm(final_delta, generator=generator)] + # apply shuffled indexes + noise[place_idx:place_idx+final_delta] = noise[list_idx] + break + # otherwise, do normal behavior + # generate list of indexes in delta region + list_idx = torch.Tensor(list(range(start_idx,start_idx+delta))).to(torch.long) + # shuffle list + list_idx = list_idx[torch.randperm(delta, generator=generator)] + # apply shuffled indexes + noise[place_idx:place_idx+delta] = noise[list_idx] + return noise + + +DERIVATIVE_NOISE_FUNC_MAP = { + NoiseLayerType.REPEATED_CONTEXT: SeedNoiseGeneration._convert_to_repeated_context, + NoiseLayerType.FREENOISE: SeedNoiseGeneration._convert_to_freenoise, + } class IterationOptions: diff --git a/animatediff/sampling.py b/animatediff/sampling.py index 3731c9a..93ffa46 100644 --- a/animatediff/sampling.py +++ b/animatediff/sampling.py @@ -78,49 +78,31 @@ ADGS = AnimateDiffHelper_GlobalState() # refer to forward_timestep_embed in comfy/ldm/modules/diffusionmodules/openaimodel.py def forward_timestep_embed_factory() -> Callable: - if hasattr(attention, "SpatialVideoTransformer"): - def forward_timestep_embed(ts, x, emb, context=None, transformer_options={}, output_shape=None, time_context=None, num_video_frames=None, image_only_indicator=None): - for layer in ts: - if isinstance(layer, openaimodel.VideoResBlock): - x = layer(x, emb, num_video_frames, image_only_indicator) - elif isinstance(layer, openaimodel.TimestepBlock): - x = layer(x, emb) - elif isinstance(layer, VanillaTemporalModule): - x = layer(x, context) - elif isinstance(layer, attention.SpatialVideoTransformer): - x = layer(x, context, time_context, num_video_frames, image_only_indicator, transformer_options) - if "transformer_index" in transformer_options: - transformer_options["transformer_index"] += 1 - if "current_index" in transformer_options: # keep this for backward compat, for now - transformer_options["current_index"] += 1 - elif isinstance(layer, attention.SpatialTransformer): - x = layer(x, context, transformer_options) - if "transformer_index" in transformer_options: - transformer_options["transformer_index"] += 1 - if "current_index" in transformer_options: # keep this for backward compat, for now - transformer_options["current_index"] += 1 - elif isinstance(layer, openaimodel.Upsample): - x = layer(x, output_shape=output_shape) - else: - x = layer(x) - return x - # keep old version for backwards compatibility (TODO: remove at end of 2023) - else: - def forward_timestep_embed(ts, x, emb, context=None, transformer_options={}, output_shape=None): - for layer in ts: - if isinstance(layer, openaimodel.TimestepBlock): - x = layer(x, emb) - elif isinstance(layer, VanillaTemporalModule): - x = layer(x, context) - elif isinstance(layer, attention.SpatialTransformer): - x = layer(x, context, transformer_options) - if "current_index" in transformer_options: - transformer_options["current_index"] += 1 - elif isinstance(layer, openaimodel.Upsample): - x = layer(x, output_shape=output_shape) - else: - x = layer(x) - return x + def forward_timestep_embed(ts, x, emb, context=None, transformer_options={}, output_shape=None, time_context=None, num_video_frames=None, image_only_indicator=None): + for layer in ts: + if isinstance(layer, openaimodel.VideoResBlock): + x = layer(x, emb, num_video_frames, image_only_indicator) + elif isinstance(layer, openaimodel.TimestepBlock): + x = layer(x, emb) + elif isinstance(layer, VanillaTemporalModule): + x = layer(x, context) + elif isinstance(layer, attention.SpatialVideoTransformer): + x = layer(x, context, time_context, num_video_frames, image_only_indicator, transformer_options) + if "transformer_index" in transformer_options: + transformer_options["transformer_index"] += 1 + if "current_index" in transformer_options: # keep this for backward compat, for now + transformer_options["current_index"] += 1 + elif isinstance(layer, attention.SpatialTransformer): + x = layer(x, context, transformer_options) + if "transformer_index" in transformer_options: + transformer_options["transformer_index"] += 1 + if "current_index" in transformer_options: # keep this for backward compat, for now + transformer_options["current_index"] += 1 + elif isinstance(layer, openaimodel.Upsample): + x = layer(x, output_shape=output_shape) + else: + x = layer(x) + return x return forward_timestep_embed