Added FreeNoise noise_type, cleaned up forward_timestep_embed_factory

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

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:

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