update README & add sliding window example workling

README.md

@@ -11,6 +11,12 @@
 - Community modules: [manshoety/AD_Stabilized_Motion](https://huggingface.co/manshoety/AD_Stabilized_Motion) | [CiaraRowles/TemporalDiff](https://huggingface.co/CiaraRowles/TemporalDiff)
 - AnimateDiff v2 [mm_sd_v15_v2.ckpt](https://huggingface.co/guoyww/animatediff/blob/main/mm_sd_v15_v2.ckpt)
 
+## Update 2023/09/21
+
+#### **Sliding Window** is now available!
+
+Sliding window is trigged automatically when generating more than 16 frames. To adjust the trigger number and other options, use `SlidingWindowOptions` node. See the sample workflow bellow.
+
 ## Nodes
 
 #### AnimateDiffLoader
@@ -20,12 +26,13 @@
 #### AnimateDiffSampler
 
 - Mostly the same with `KSampler`
-- Use `AnimateDiffLoader` to load the motion module
+- `motion_module`: use `AnimateDiffLoader` to load the motion module
 - `inject_method`: should left default
 - `frame_number`: animation length
 - `latent_image`: You can pass an `EmptyLatentImage`
+- `sliding_window_opts`: custom sliding window options
 
-<img width="370" alt="image" src="https://github.com/ArtVentureX/comfyui-animatediff/assets/133728487/f22d6b36-ce36-44cc-80e8-dffe6f77b296">
+<img width="370" alt="image" src="https://github.com/ArtVentureX/comfyui-animatediff/assets/133728487/a352195d-f40c-494d-bd3d-30ee88174b88">
 
 #### AnimateDiffCombine
 
@@ -33,10 +40,28 @@
 - `frame_rate`: number of frame per second
 - `loop_count`: use 0 for infinite loop
 - `save_image`: should GIF be saved to disk
-- `format`: supports `image/gif`, `image/webp` (better compression) or `video/webm` (need `ffmpeg` installed and available in PATH)
+- `format`: supports `image/gif`, `image/webp` (better compression), `video/webm`, `video/h264-mp4`, `video/h265-mp4`. To use video formats, you'll need [ffmpeg](https://ffmpeg.org/download.html) installed and available in **`PATH`**
 
 <img width="370" alt="image" src="https://github.com/ArtVentureX/comfyui-animatediff/assets/133728487/381c5acc-06ef-43da-ada0-3dc76f37a3e4">
 
+#### SlidingWindowOptions
+
+Custom sliding window options
+
+- `context_length`: number of frame per _window_. Use **16** to get the best results. Reduce it if you have low VRAM.
+- `closed_loop`: try to make the GIF a closed loop. Will take longer to render.
+
+<img width="370" alt="image" src="https://github.com/ArtVentureX/comfyui-animatediff/assets/133728487/6679a8dd-bf96-419f-8934-ea2b046dd23c">
+
+#### LoadVideo
+
+Load GIF or video as images. Usefull to load a GIF as ControlNet input.
+
+- `frame_start`: Skip some begining frames and start at `frame_start`
+- `frame_limit`: Only take `frame_limit` frames
+
+<img width="370" alt="image" src="https://github.com/ArtVentureX/comfyui-animatediff/assets/133728487/684176d5-6369-4a27-9f33-e721e0fe1876">
+
 ## Workflows
 
 ### Simple txt2gif
@@ -51,6 +76,17 @@ Samples:
 
 ![animate_diff_02](https://github.com/ArtVentureX/comfyui-animatediff/assets/133728487/c39b26f7-a2af-4dc4-902f-c363e2e6f39a)
 
+### Long duration with sliding window
+
+<img width="1280" alt="image" src="https://github.com/ArtVentureX/comfyui-animatediff/assets/133728487/0f8bfb87-83cb-4119-9777-e3948ec0cb5c">
+
+Workflow: [sliding-window.json](https://github.com/ArtVentureX/comfyui-animatediff/blob/main/workflows/sliding-window.json)
+
+Samples:
+![AnimateDiff_00096_](https://github.com/ArtVentureX/comfyui-animatediff/assets/133728487/e1da7a66-e615-475d-9400-41eff484ad49)
+
+![AnimateDiff_00099_](https://github.com/ArtVentureX/comfyui-animatediff/assets/133728487/4faa7e5e-cdaa-49da-8759-46d779c0e0b6)
+
 ### Latent upscale
 
 Upscale latent output using `LatentUpscale` then do a 2nd pass with `AnimateDiffSampler`.

animatediff/model_utils.py

@@ -1,7 +1,16 @@
 import os
 import hashlib
+from typing import Dict
 
 import folder_paths
+import comfy.model_management as model_management
+from comfy.utils import load_torch_file, calculate_parameters
+
+from .logger import logger
+from .motion_module import MotionWrapper
+
+
+motion_modules: Dict[str, MotionWrapper] = {}
 
 
 folder_paths.folder_names_and_paths["AnimateDiff"] = (
@@ -11,12 +20,6 @@ folder_paths.folder_names_and_paths["AnimateDiff"] = (
     ],
     folder_paths.supported_pt_extensions,
 )
-folder_paths.folder_names_and_paths["video_formats"] = (
-    [
-        os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "video_formats"),
-    ],
-    [".json"]
-)
 
 
 def get_available_models():
@@ -31,3 +34,23 @@ def get_model_hash(file_path):
     with open(file_path, "rb") as f:
         bytes = f.read()  # read entire file as bytes
         return hashlib.sha256(bytes).hexdigest()
+
+
+def load_motion_module(model_name: str):
+    model_path = get_model_path(model_name)
+    model_hash = get_model_hash(model_path)
+    if model_hash not in motion_modules:
+        logger.info(f"Loading motion module {model_name}")
+        mm_state_dict = load_torch_file(model_path)
+        motion_module = MotionWrapper.from_state_dict(mm_state_dict, model_name)
+
+        params = calculate_parameters(mm_state_dict, "")
+        if model_management.should_use_fp16(model_params=params):
+            logger.info(f"Converting motion module to fp16.")
+            motion_module.half()
+        offload_device = model_management.unet_offload_device()
+        motion_module = motion_module.to(offload_device)
+
+        motion_modules[model_hash] = motion_module
+
+    return motion_modules[model_hash]

animatediff/motion_module.py

@@ -41,6 +41,7 @@ class MotionWrapper(nn.Module):
         self.down_blocks = nn.ModuleList([])
         self.up_blocks = nn.ModuleList([])
         self.mid_block = None
+        self.encoding_max_len = encoding_max_len
 
         for c in (320, 640, 1280, 1280):
             self.down_blocks.append(
@@ -56,7 +57,7 @@ class MotionWrapper(nn.Module):
             )
 
     @classmethod
-    def from_pretrained(cls, mm_state_dict: dict[str, Tensor], mm_type: str):
+    def from_state_dict(cls, mm_state_dict: dict[str, Tensor], mm_type: str):
         encoding_max_len = get_encoding_max_len(mm_state_dict)
         is_v2 = has_mid_block(mm_state_dict)
 

animatediff/nodes.py

@@ -3,176 +3,22 @@ import json
 import torch
 import numpy as np
 import hashlib
-from typing import Dict, List
+from typing import List
 from torch import Tensor
-from torch.nn.functional import group_norm
 from PIL import Image, ImageSequence
 from PIL.PngImagePlugin import PngInfo
-from einops import rearrange
 
 import folder_paths
-import comfy.ldm.modules.diffusionmodules.openaimodel as openaimodel
-import comfy.model_management as model_management
-from comfy.model_base import BaseModel
-from comfy.ldm.modules.attention import SpatialTransformer
-from comfy.utils import load_torch_file, calculate_parameters
-from nodes import KSampler
 
-from .logger import logger
-from .motion_module import MotionWrapper, VanillaTemporalModule
-from .model_utils import get_available_models, get_model_path, get_model_hash
+from .model_utils import get_available_models, load_motion_module
 from .utils import pil2tensor
+from .sampler import AnimateDiffSampler, AnimateDiffSlidingWindowOptions
 
 
-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, SpatialTransformer):
-            x = layer(x, context, 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
+SLIDING_CONTEXT_LENGTH = 16
 
-
-def groupnorm_mm_factory(video_length: int):
-    def groupnorm_mm_forward(self, input: Tensor) -> Tensor:
-        # axes_factor normalizes batch based on total conds and unconds passed in batch;
-        # the conds and unconds per batch can change based on VRAM optimizations that may kick in
-        axes_factor = input.size(0) // video_length
-
-        input = rearrange(input, "(b f) c h w -> b c f h w", b=axes_factor)
-        input = group_norm(input, self.num_groups,
-                           self.weight, self.bias, self.eps)
-        input = rearrange(input, "b c f h w -> (b f) c h w", b=axes_factor)
-        return input
-
-    return groupnorm_mm_forward
-
-
-orig_forward_timestep_embed = openaimodel.forward_timestep_embed
-orig_maximum_batch_area = model_management.maximum_batch_area
-orig_groupnorm_forward = torch.nn.GroupNorm.forward
-openaimodel.forward_timestep_embed = forward_timestep_embed
-
-motion_modules: Dict[str, MotionWrapper] = {}
-
-
-def load_motion_module(model_name: str):
-    model_path = get_model_path(model_name)
-    model_hash = get_model_hash(model_path)
-    if model_hash not in motion_modules:
-        logger.info(f"Loading motion module {model_name}")
-        mm_state_dict = load_torch_file(model_path)
-        motion_module = MotionWrapper.from_pretrained(
-            mm_state_dict, model_name)
-
-        params = calculate_parameters(mm_state_dict, "")
-        if model_management.should_use_fp16(model_params=params):
-            logger.info(f"Converting motion module to fp16.")
-            motion_module.half()
-        offload_device = model_management.unet_offload_device()
-        motion_module = motion_module.to(offload_device)
-
-        motion_modules[model_hash] = motion_module
-
-    return motion_modules[model_hash]
-
-
-def inject_motion_module_to_unet_legacy(unet, motion_module: MotionWrapper):
-    for mm_idx, unet_idx in enumerate([1, 2, 4, 5, 7, 8, 10, 11]):
-        mm_idx0, mm_idx1 = mm_idx // 2, mm_idx % 2
-        unet.input_blocks[unet_idx].append(
-            motion_module.down_blocks[mm_idx0].motion_modules[mm_idx1]
-        )
-
-    for unet_idx in range(12):
-        mm_idx0, mm_idx1 = unet_idx // 3, unet_idx % 3
-        if unet_idx % 2 == 2:
-            unet.output_blocks[unet_idx].insert(
-                -1, motion_module.up_blocks[mm_idx0].motion_modules[mm_idx1]
-            )
-        else:
-            unet.output_blocks[unet_idx].append(
-                motion_module.up_blocks[mm_idx0].motion_modules[mm_idx1]
-            )
-    if motion_module.is_v2:
-        unet.middle_block.insert(-1, motion_module.mid_block.motion_modules[0])
-
-    unet.motion_module = motion_module
-
-
-def eject_motion_module_from_unet_legacy(unet):
-    for unet_idx in [1, 2, 4, 5, 7, 8, 10, 11]:
-        unet.input_blocks[unet_idx].pop(-1)
-
-    for unet_idx in range(12):
-        if unet_idx % 2 == 2:
-            unet.output_blocks[unet_idx].pop(-2)
-        else:
-            unet.output_blocks[unet_idx].pop(-1)
-
-    if unet.motion_module.is_v2:
-        unet.middle_block.pop(-2)
-
-    del unet.motion_module
-
-
-def inject_motion_module_to_unet(unet, motion_module: MotionWrapper):
-    for mm_idx, unet_idx in enumerate([1, 2, 4, 5, 7, 8, 10, 11]):
-        mm_idx0, mm_idx1 = mm_idx // 2, mm_idx % 2
-        unet.input_blocks[unet_idx].append(
-            motion_module.down_blocks[mm_idx0].motion_modules[mm_idx1]
-        )
-
-    for unet_idx in range(12):
-        mm_idx0, mm_idx1 = unet_idx // 3, unet_idx % 3
-        if unet_idx % 3 == 2 and unet_idx != 11:
-            unet.output_blocks[unet_idx].insert(
-                -1, motion_module.up_blocks[mm_idx0].motion_modules[mm_idx1]
-            )
-        else:
-            unet.output_blocks[unet_idx].append(
-                motion_module.up_blocks[mm_idx0].motion_modules[mm_idx1]
-            )
-    if motion_module.is_v2:
-        unet.middle_block.insert(-1, motion_module.mid_block.motion_modules[0])
-
-    unet.motion_module = motion_module
-
-
-def eject_motion_module_from_unet(unet):
-    for unet_idx in [1, 2, 4, 5, 7, 8, 10, 11]:
-        unet.input_blocks[unet_idx].pop(-1)
-
-    for unet_idx in range(12):
-        if unet_idx % 3 == 2 and unet_idx != 11:
-            unet.output_blocks[unet_idx].pop(-2)
-        else:
-            unet.output_blocks[unet_idx].pop(-1)
-
-    if unet.motion_module.is_v2:
-        unet.middle_block.pop(-2)
-
-    del unet.motion_module
-
-
-injectors = {
-    "legacy": inject_motion_module_to_unet_legacy,
-    "default": inject_motion_module_to_unet,
-}
-
-ejectors = {
-    "legacy": eject_motion_module_from_unet_legacy,
-    "default": eject_motion_module_from_unet,
-}
+video_formats_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "video_formats")
+video_formats = ["video/" + x[:-5] for x in os.listdir(video_formats_dir)]
 
 
 class AnimateDiffModuleLoader:
@@ -197,134 +43,6 @@ class AnimateDiffModuleLoader:
         return (motion_module,)
 
 
-class AnimateDiffSampler(KSampler):
-    @classmethod
-    def INPUT_TYPES(s):
-        inputs = {
-            "required": {
-                "motion_module": ("MOTION_MODULE",),
-                "inject_method": (["default", "legacy"],),
-                "frame_number": (
-                    "INT",
-                    {"default": 16, "min": 2, "max": 32, "step": 1},
-                ),
-            }
-        }
-        inputs["required"].update(KSampler.INPUT_TYPES()["required"])
-        return inputs
-
-    FUNCTION = "animatediff_sample"
-    CATEGORY = "Animate Diff"
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.prev_beta = None
-        self.prev_linear_start = None
-        self.prev_linear_end = None
-
-    def override_beta_schedule(self, model: BaseModel):
-        logger.info(f"Override beta schedule.")
-        self.prev_beta = model.get_buffer("betas").cpu().clone()
-        self.prev_linear_start = model.linear_start
-        self.prev_linear_end = model.linear_end
-        model.register_schedule(
-            given_betas=None,
-            beta_schedule="sqrt_linear",
-            timesteps=1000,
-            linear_start=0.00085,
-            linear_end=0.012,
-            cosine_s=8e-3,
-        )
-
-    def restore_beta_schedule(self, model: BaseModel):
-        logger.info(f"Restoring beta schedule.")
-        model.register_schedule(
-            given_betas=self.prev_beta,
-            linear_start=self.prev_linear_start,
-            linear_end=self.prev_linear_end,
-        )
-        self.prev_beta = None
-        self.prev_linear_start = None
-        self.prev_linear_end = None
-
-    def inject_motion_module(
-        self, model, motion_module: MotionWrapper, inject_method: str, frame_number: int
-    ):
-        model = model.clone()
-        unet = model.model.diffusion_model
-
-        logger.info(f"Injecting motion module with method {inject_method}.")
-        motion_module.set_video_length(frame_number)
-        injectors[inject_method](unet, motion_module)
-        self.override_beta_schedule(model.model)
-        if not motion_module.is_v2:
-            logger.info(f"Hacking GroupNorm.forward function.")
-            torch.nn.GroupNorm.forward = groupnorm_mm_factory(frame_number)
-
-        return model
-
-    def eject_motion_module(self, model, inject_method):
-        unet = model.model.diffusion_model
-
-        self.restore_beta_schedule(model.model)
-        if not unet.motion_module.is_v2:
-            logger.info(f"Restore GroupNorm.forward function.")
-            torch.nn.GroupNorm.forward = orig_groupnorm_forward
-
-        logger.info(f"Ejecting motion module with method {inject_method}.")
-        ejectors[inject_method](unet)
-
-    def animatediff_sample(
-        self,
-        motion_module,
-        inject_method,
-        frame_number,
-        model,
-        seed,
-        steps,
-        cfg,
-        sampler_name,
-        scheduler,
-        positive,
-        negative,
-        latent_image,
-        denoise=1.0,
-    ):
-        model = self.inject_motion_module(
-            model, motion_module, inject_method, frame_number
-        )
-
-        init_frames = len(latent_image["samples"])
-        samples = latent_image["samples"][:init_frames, :, :, :].clone().cpu()
-
-        if init_frames < frame_number:
-            last_frame = samples[-1].unsqueeze(0)
-            repeated_last_frames = last_frame.repeat(
-                frame_number - init_frames, 1, 1, 1
-            )
-            samples = torch.cat((samples, repeated_last_frames), dim=0)
-
-        latent_image = {"samples": samples}
-
-        try:
-            return super().sample(
-                model,
-                seed,
-                steps,
-                cfg,
-                sampler_name,
-                scheduler,
-                positive,
-                negative,
-                latent_image,
-                denoise=denoise,
-            )
-        except:
-            raise
-        finally:
-            self.eject_motion_module(model, inject_method)
-
-
 class AnimateDiffCombine:
     @classmethod
     def INPUT_TYPES(s):
@@ -336,11 +54,10 @@ class AnimateDiffCombine:
                     {"default": 8, "min": 1, "max": 24, "step": 1},
                 ),
                 "loop_count": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
-                "save_image": ([True, False],),
+                "save_image": ("BOOLEAN", {"default": True}),
                 "filename_prefix": ("STRING", {"default": "animate_diff"}),
-                "format": (["image/gif", "image/webp"] +
-                           ["video/"+x[:-5] for x in folder_paths.get_filename_list("video_formats")],),
-                "pingpong": ([False, True],),
+                "format": (["image/gif", "image/webp"] + video_formats,),
+                "pingpong": ("BOOLEAN", {"default": False}),
             },
             "hidden": {
                 "prompt": "PROMPT",
@@ -373,11 +90,7 @@ class AnimateDiffCombine:
             frames.append(img)
 
         # save image
-        output_dir = (
-            folder_paths.get_output_directory()
-            if save_image
-            else folder_paths.get_temp_directory()
-        )
+        output_dir = folder_paths.get_output_directory() if save_image else folder_paths.get_temp_directory()
         (
             full_output_folder,
             filename,
@@ -426,16 +139,31 @@ class AnimateDiffCombine:
             ffmpeg_path = shutil.which("ffmpeg")
             if ffmpeg_path is None:
                 raise ProcessLookupError("Could not find ffmpeg")
-            video_format_path = folder_paths.get_full_path(
-                "video_formats", format_ext + ".json")
-            with open(video_format_path, 'r') as stream:
+            video_format_path = os.path.join(video_formats_dir, format_ext + ".json")
+            with open(video_format_path, "r") as stream:
                 video_format = json.load(stream)
             file = f"{filename}_{counter:05}_.{video_format['extension']}"
             file_path = os.path.join(full_output_folder, file)
             dimensions = f"{frames[0].width}x{frames[0].height}"
-            args = [ffmpeg_path, "-v", "error", "-f", "rawvideo", "-pix_fmt", "rgb24",
-                    "-s", dimensions, "-r", str(frame_rate), "-i", "-"] \
-                + video_format['main_pass'] + [file_path]
+            args = (
+                [
+                    ffmpeg_path,
+                    "-v",
+                    "error",
+                    "-f",
+                    "rawvideo",
+                    "-pix_fmt",
+                    "rgb24",
+                    "-s",
+                    dimensions,
+                    "-r",
+                    str(frame_rate),
+                    "-i",
+                    "-",
+                ]
+                + video_format["main_pass"]
+                + [file_path]
+            )
 
             env = os.environ
             if "environment" in video_format:
@@ -462,17 +190,16 @@ class LoadVideo:
         if not os.path.exists(input_dir):
             os.makedirs(input_dir, exist_ok=True)
 
-        files = [f"video/{f}" for f in os.listdir(input_dir) if os.path.isfile(
-            os.path.join(input_dir, f))]
+        files = [f"video/{f}" for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
 
         return {
             "required": {
                 "video": (sorted(files), {"video_upload": True}),
             },
             "optional": {
-                "frame_start": ("INT", {"default": 0, "min": 0, "max": 0xffffffff, "step": 1}),
+                "frame_start": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF, "step": 1}),
                 "frame_limit": ("INT", {"default": 16, "min": 1, "max": 10240, "step": 1}),
-            }
+            },
         }
 
     CATEGORY = "Animate Diff/Utils"
@@ -517,7 +244,6 @@ class LoadVideo:
         return frames
 
     def load(self, video: str, frame_start=0, frame_limit=16):
-        print("path", video)
         video_path = folder_paths.get_annotated_filepath(video)
         (_, ext) = os.path.splitext(video_path)
 
@@ -534,7 +260,7 @@ class LoadVideo:
     def IS_CHANGED(s, image, *args, **kwargs):
         image_path = folder_paths.get_annotated_filepath(image)
         m = hashlib.sha256()
-        with open(image_path, 'rb') as f:
+        with open(image_path, "rb") as f:
             m.update(f.read())
         return m.digest().hex()
 
@@ -572,7 +298,7 @@ class ImageChunking:
             "required": {
                 "images": ("IMAGE",),
                 "chunk_size": ("INT", {"default": 16, "min": 1, "max": 1024, "step": 1}),
-                "allow_remainder": ([True, False],),
+                "allow_remainder": ("BOOLEAN", {"default": True}),
             },
         }
 
@@ -584,29 +310,29 @@ class ImageChunking:
     def chunk(self, images: Tensor, chunk_size: int, allow_remainder: bool):
         # Check if tensor is divisible into chunks of chunk_size
         if images.shape[0] % chunk_size != 0 and not allow_remainder:
-            raise ValueError(
-                "Tensor's first dimension is not divisible by chunk size")
+            raise ValueError("Tensor's first dimension is not divisible by chunk size")
 
         # Use torch.chunk to divide the tensor
-        chunk_count = images.shape[0] // chunk_size + \
-            images.shape[0] % chunk_size
+        chunk_count = images.shape[0] // chunk_size + images.shape[0] % chunk_size
 
         print("chunk_count", chunk_count)
         chunks = torch.chunk(images, chunk_count, dim=0)
 
-        return (list(chunks), )
+        return (list(chunks),)
 
 
 NODE_CLASS_MAPPINGS = {
     "AnimateDiffModuleLoader": AnimateDiffModuleLoader,
     "AnimateDiffCombine": AnimateDiffCombine,
     "AnimateDiffSampler": AnimateDiffSampler,
+    "AnimateDiffSlidingWindowOptions": AnimateDiffSlidingWindowOptions,
     "LoadVideo": LoadVideo,
     "ImageSizeAndBatchSize": ImageSizeAndBatchSize,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "AnimateDiffModuleLoader": "Animate Diff Module Loader",
     "AnimateDiffSampler": "Animate Diff Sampler",
+    "AnimateDiffSlidingWindowOptions": "Sliding Window Options",
     "AnimateDiffCombine": "Animate Diff Combine",
     "LoadVideo": "Load Video",
     "ImageSizeAndBatchSize": "Get Image Size + Batch Size",

animatediff/sampler.py

@@ -0,0 +1,316 @@
+import torch
+from torch import Tensor
+from torch.nn.functional import group_norm
+from einops import rearrange
+
+import comfy.ldm.modules.diffusionmodules.openaimodel as openaimodel
+import comfy.model_management as model_management
+from comfy.model_base import BaseModel
+from comfy.ldm.modules.attention import SpatialTransformer
+from nodes import KSampler
+
+from .logger import logger
+from .motion_module import MotionWrapper, VanillaTemporalModule
+from .sliding_schedule import ContextSchedules
+from .sliding_context_sampling import SlidingContext, inject_sampling_function, eject_sampling_function
+
+
+SLIDING_CONTEXT_LENGTH = 16
+
+
+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, SpatialTransformer):
+            x = layer(x, context, 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 groupnorm_mm_factory(video_length: int):
+    def groupnorm_mm_forward(self, input: Tensor) -> Tensor:
+        # axes_factor normalizes batch based on total conds and unconds passed in batch;
+        # the conds and unconds per batch can change based on VRAM optimizations that may kick in
+        axes_factor = input.size(0) // video_length
+
+        input = rearrange(input, "(b f) c h w -> b c f h w", b=axes_factor)
+        input = group_norm(input, self.num_groups, self.weight, self.bias, self.eps)
+        input = rearrange(input, "b c f h w -> (b f) c h w", b=axes_factor)
+        return input
+
+    return groupnorm_mm_forward
+
+
+orig_forward_timestep_embed = openaimodel.forward_timestep_embed
+orig_maximum_batch_area = model_management.maximum_batch_area
+orig_groupnorm_forward = torch.nn.GroupNorm.forward
+
+
+def inject_motion_module_to_unet_legacy(unet, motion_module: MotionWrapper):
+    for mm_idx, unet_idx in enumerate([1, 2, 4, 5, 7, 8, 10, 11]):
+        mm_idx0, mm_idx1 = mm_idx // 2, mm_idx % 2
+        unet.input_blocks[unet_idx].append(motion_module.down_blocks[mm_idx0].motion_modules[mm_idx1])
+
+    for unet_idx in range(12):
+        mm_idx0, mm_idx1 = unet_idx // 3, unet_idx % 3
+        if unet_idx % 2 == 2:
+            unet.output_blocks[unet_idx].insert(-1, motion_module.up_blocks[mm_idx0].motion_modules[mm_idx1])
+        else:
+            unet.output_blocks[unet_idx].append(motion_module.up_blocks[mm_idx0].motion_modules[mm_idx1])
+    if motion_module.is_v2:
+        unet.middle_block.insert(-1, motion_module.mid_block.motion_modules[0])
+
+    unet.motion_module = motion_module
+
+
+def eject_motion_module_from_unet_legacy(unet):
+    for unet_idx in [1, 2, 4, 5, 7, 8, 10, 11]:
+        unet.input_blocks[unet_idx].pop(-1)
+
+    for unet_idx in range(12):
+        if unet_idx % 2 == 2:
+            unet.output_blocks[unet_idx].pop(-2)
+        else:
+            unet.output_blocks[unet_idx].pop(-1)
+
+    if unet.motion_module.is_v2:
+        unet.middle_block.pop(-2)
+
+    del unet.motion_module
+
+
+def inject_motion_module_to_unet(unet, motion_module: MotionWrapper):
+    for mm_idx, unet_idx in enumerate([1, 2, 4, 5, 7, 8, 10, 11]):
+        mm_idx0, mm_idx1 = mm_idx // 2, mm_idx % 2
+        unet.input_blocks[unet_idx].append(motion_module.down_blocks[mm_idx0].motion_modules[mm_idx1])
+
+    for unet_idx in range(12):
+        mm_idx0, mm_idx1 = unet_idx // 3, unet_idx % 3
+        if unet_idx % 3 == 2 and unet_idx != 11:
+            unet.output_blocks[unet_idx].insert(-1, motion_module.up_blocks[mm_idx0].motion_modules[mm_idx1])
+        else:
+            unet.output_blocks[unet_idx].append(motion_module.up_blocks[mm_idx0].motion_modules[mm_idx1])
+    if motion_module.is_v2:
+        unet.middle_block.insert(-1, motion_module.mid_block.motion_modules[0])
+
+    unet.motion_module = motion_module
+
+
+def eject_motion_module_from_unet(unet):
+    for unet_idx in [1, 2, 4, 5, 7, 8, 10, 11]:
+        unet.input_blocks[unet_idx].pop(-1)
+
+    for unet_idx in range(12):
+        if unet_idx % 3 == 2 and unet_idx != 11:
+            unet.output_blocks[unet_idx].pop(-2)
+        else:
+            unet.output_blocks[unet_idx].pop(-1)
+
+    if unet.motion_module.is_v2:
+        unet.middle_block.pop(-2)
+
+    del unet.motion_module
+
+
+injectors = {
+    "legacy": inject_motion_module_to_unet_legacy,
+    "default": inject_motion_module_to_unet,
+}
+
+ejectors = {
+    "legacy": eject_motion_module_from_unet_legacy,
+    "default": eject_motion_module_from_unet,
+}
+
+
+class AnimateDiffSlidingWindowOptions:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "context_length": ("INT", {"default": SLIDING_CONTEXT_LENGTH, "min": 2, "max": 32}),
+                "context_stride": ("INT", {"default": 1, "min": 1, "max": 32}),
+                "context_overlap": ("INT", {"default": 4, "min": 0, "max": 32}),
+                "context_schedule": (ContextSchedules.CONTEXT_SCHEDULE_LIST,),
+                "closed_loop": ("BOOLEAN", {"default": False}),
+            }
+        }
+
+    RETURN_TYPES = ("SLIDING_WINDOW_OPTS",)
+    FUNCTION = "init_options"
+    CATEGORY = "Animate Diff"
+
+    def init_options(self, context_length, context_stride, context_overlap, context_schedule, closed_loop):
+        ctx = SlidingContext(
+            context_length=context_length,
+            context_stride=context_stride,
+            context_overlap=context_overlap,
+            context_schedule=context_schedule,
+            closed_loop=closed_loop,
+        )
+
+        return (ctx,)
+
+
+class AnimateDiffSampler(KSampler):
+    @classmethod
+    def INPUT_TYPES(s):
+        inputs = {
+            "required": {
+                "motion_module": ("MOTION_MODULE",),
+                "inject_method": (["default", "legacy"],),
+                "frame_number": (
+                    "INT",
+                    {"default": 16, "min": 2, "max": 10000, "step": 1},
+                ),
+            }
+        }
+        inputs["required"].update(KSampler.INPUT_TYPES()["required"])
+        inputs["optional"] = {"sliding_window_opts": ("SLIDING_WINDOW_OPTS",)}
+        return inputs
+
+    FUNCTION = "animatediff_sample"
+    CATEGORY = "Animate Diff"
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.prev_beta = None
+        self.prev_linear_start = None
+        self.prev_linear_end = None
+
+    def override_beta_schedule(self, model: BaseModel):
+        self.prev_beta = model.get_buffer("betas").cpu().clone().detach()
+        self.prev_linear_start = model.linear_start
+        self.prev_linear_end = model.linear_end
+        model.register_schedule(
+            given_betas=None,
+            beta_schedule="sqrt_linear",
+            timesteps=1000,
+            linear_start=0.00085,
+            linear_end=0.012,
+            cosine_s=8e-3,
+        )
+
+    def restore_beta_schedule(self, model: BaseModel):
+        model.register_schedule(
+            given_betas=self.prev_beta,
+            linear_start=self.prev_linear_start,
+            linear_end=self.prev_linear_end,
+        )
+        self.prev_beta = None
+        self.prev_linear_start = None
+        self.prev_linear_end = None
+
+    def inject_motion_module(self, model, motion_module: MotionWrapper, inject_method: str, frame_number: int):
+        model = model.clone()
+        unet = model.model.diffusion_model
+
+        logger.info(f"Injecting motion module with method {inject_method}.")
+        motion_module.set_video_length(frame_number)
+        injectors[inject_method](unet, motion_module)
+        self.override_beta_schedule(model.model)
+        openaimodel.forward_timestep_embed = forward_timestep_embed
+        if not motion_module.is_v2:
+            logger.info(f"Hacking GroupNorm.forward function.")
+            torch.nn.GroupNorm.forward = groupnorm_mm_factory(frame_number)
+
+        return model
+
+    def inject_sliding_sampler(self, video_length, sliding_window_opts: SlidingContext = None):
+        ctx = sliding_window_opts.copy() if sliding_window_opts else SlidingContext()
+        ctx.video_length = video_length
+
+        inject_sampling_function(ctx)
+
+    def eject_motion_module(self, model, inject_method):
+        unet = model.model.diffusion_model
+
+        self.restore_beta_schedule(model.model)
+        openaimodel.forward_timestep_embed = orig_forward_timestep_embed
+        if not unet.motion_module.is_v2:
+            logger.info(f"Restore GroupNorm.forward function.")
+            torch.nn.GroupNorm.forward = orig_groupnorm_forward
+
+        logger.info(f"Ejecting motion module with method {inject_method}.")
+        ejectors[inject_method](unet)
+
+    def eject_sliding_sampler(self):
+        eject_sampling_function()
+
+    def animatediff_sample(
+        self,
+        motion_module,
+        inject_method,
+        frame_number,
+        model,
+        seed,
+        steps,
+        cfg,
+        sampler_name,
+        scheduler,
+        positive,
+        negative,
+        latent_image,
+        denoise=1.0,
+        sliding_window_opts: SlidingContext = None,
+        **kwargs,
+    ):
+        # init latents
+        samples = latent_image["samples"]
+        init_frames = len(samples)
+        if init_frames < frame_number:
+            # TODO: apply different noise to each frame
+            last_frame = samples[-1].clone().cpu().unsqueeze(0)
+            repeated_last_frames = last_frame.repeat(frame_number - init_frames, 1, 1, 1)
+            samples = torch.cat((samples, repeated_last_frames), dim=0)
+
+        latent_image = {"samples": samples}
+
+        # validate context_length
+        context_length = sliding_window_opts.context_length if sliding_window_opts else SLIDING_CONTEXT_LENGTH
+        is_sliding = frame_number > context_length
+        video_length = context_length if is_sliding else frame_number
+
+        if video_length > motion_module.encoding_max_len:
+            error = f'{"context_length" if is_sliding else "frame_number"} = {video_length}'
+            raise ValueError(
+                f"AnimateDiff model {motion_module.mm_type} has upper limit of {motion_module.encoding_max_len} frames, but received {error}."
+            )
+
+        # inject motion module
+        model = self.inject_motion_module(model, motion_module, inject_method, video_length)
+
+        # inject sliding sampler
+        if is_sliding:
+            self.inject_sliding_sampler(frame_number, sliding_window_opts=sliding_window_opts)
+
+        try:
+            return super().sample(
+                model,
+                seed,
+                steps,
+                cfg,
+                sampler_name,
+                scheduler,
+                positive,
+                negative,
+                latent_image,
+                denoise=denoise,
+                **kwargs,
+            )
+        except:
+            raise
+        finally:
+            # eject motion module
+            self.eject_motion_module(model, inject_method)
+
+            # eject sliding sampler
+            if is_sliding:
+                self.eject_sliding_sampler()

animatediff/sliding_context_sampling.py

@@ -0,0 +1,487 @@
+import torch
+from torch import Tensor
+import math
+
+import comfy.utils
+import comfy.sample
+import comfy.samplers as comfy_samplers
+import comfy.model_management as model_management
+from comfy.controlnet import ControlBase
+from comfy.model_patcher import ModelPatcher
+
+from .logger import logger
+from .sliding_schedule import get_context_scheduler, ContextSchedules
+
+
+orig_comfy_sample = comfy.sample.sample
+orig_sampling_function = comfy_samplers.sampling_function
+
+
+class SlidingContext:
+    def __init__(
+        self,
+        context_length=16,
+        context_stride=1,
+        context_overlap=4,
+        context_schedule=ContextSchedules.UNIFORM,
+        closed_loop=False,
+        video_length=0,
+        current_step=0,
+        total_steps=0,
+    ):
+        self.context_length = context_length
+        self.context_stride = context_stride
+        self.context_overlap = context_overlap
+        self.context_schedule = context_schedule
+        self.closed_loop = closed_loop
+        self.video_length = video_length
+        self.current_step = current_step
+        self.total_steps = total_steps
+
+    def copy(self):
+        return SlidingContext(
+            context_length=self.context_length,
+            context_stride=self.context_stride,
+            context_overlap=self.context_overlap,
+            context_schedule=self.context_schedule,
+            closed_loop=self.closed_loop,
+            video_length=self.video_length,
+            current_step=self.current_step,
+            total_steps=self.total_steps,
+        )
+
+
+def __sliding_sample_factory(ctx: SlidingContext):
+    logger.info(f"Injecting sliding context sampling function.")
+    logger.info(f"Video length: {ctx.video_length}")
+    logger.info(f"Context length: {ctx.context_length}")
+    logger.info(f"Context schedule: {ctx.context_schedule}")
+
+    context_scheduler = get_context_scheduler(ctx.context_schedule)
+
+    def sample(model: ModelPatcher, *args, **kwargs):
+        orig_callback = kwargs.pop("callback", None)
+        start_step = kwargs.get("start_step") or 0
+
+        # adjust progressbar to account for context frames
+        def callback(step, x0, x, total_steps):
+            if orig_callback:
+                orig_callback(step, x0, x, total_steps)
+
+            ctx.current_step = start_step + step + 1
+
+        try:
+            return orig_comfy_sample(model, *args, **kwargs, callback=callback)
+        except RuntimeError as e:
+            if str(e).startswith("CUDA error: invalid configuration argument"):
+                raise RuntimeError(
+                    f"An xformers bug was encountered in AnimateDiff - to run your workflow, \
+                                    disable xformers in ComfyUI using '--disable-xformers' startup argument."
+                )
+            raise
+
+    def sampling_function(
+        model_function, x, timestep, uncond, cond, cond_scale, cond_concat=None, model_options={}, seed=None
+    ):
+        def get_area_and_mult(cond, x_in, cond_concat_in, timestep_in):
+            area = (x_in.shape[2], x_in.shape[3], 0, 0)
+            strength = 1.0
+            if "timestep_start" in cond[1]:
+                timestep_start = cond[1]["timestep_start"]
+                if timestep_in[0] > timestep_start:
+                    return None
+            if "timestep_end" in cond[1]:
+                timestep_end = cond[1]["timestep_end"]
+                if timestep_in[0] < timestep_end:
+                    return None
+            if "area" in cond[1]:
+                area = cond[1]["area"]
+            if "strength" in cond[1]:
+                strength = cond[1]["strength"]
+
+            adm_cond = None
+            if "adm_encoded" in cond[1]:
+                adm_cond = cond[1]["adm_encoded"]
+
+            input_x = x_in[:, :, area[2] : area[0] + area[2], area[3] : area[1] + area[3]]
+            if "mask" in cond[1]:
+                # Scale the mask to the size of the input
+                # The mask should have been resized as we began the sampling process
+                mask_strength = 1.0
+                if "mask_strength" in cond[1]:
+                    mask_strength = cond[1]["mask_strength"]
+                mask = cond[1]["mask"]
+                assert mask.shape[1] == x_in.shape[2]
+                assert mask.shape[2] == x_in.shape[3]
+                mask = mask[:, area[2] : area[0] + area[2], area[3] : area[1] + area[3]] * mask_strength
+                mask = mask.unsqueeze(1).repeat(input_x.shape[0] // mask.shape[0], input_x.shape[1], 1, 1)
+            else:
+                mask = torch.ones_like(input_x)
+            mult = mask * strength
+
+            if "mask" not in cond[1]:
+                rr = 8
+                if area[2] != 0:
+                    for t in range(rr):
+                        mult[:, :, t : 1 + t, :] *= (1.0 / rr) * (t + 1)
+                if (area[0] + area[2]) < x_in.shape[2]:
+                    for t in range(rr):
+                        mult[:, :, area[0] - 1 - t : area[0] - t, :] *= (1.0 / rr) * (t + 1)
+                if area[3] != 0:
+                    for t in range(rr):
+                        mult[:, :, :, t : 1 + t] *= (1.0 / rr) * (t + 1)
+                if (area[1] + area[3]) < x_in.shape[3]:
+                    for t in range(rr):
+                        mult[:, :, :, area[1] - 1 - t : area[1] - t] *= (1.0 / rr) * (t + 1)
+
+            conditionning = {}
+            conditionning["c_crossattn"] = cond[0]
+            if cond_concat_in is not None and len(cond_concat_in) > 0:
+                cropped = []
+                for x in cond_concat_in:
+                    cr = x[:, :, area[2] : area[0] + area[2], area[3] : area[1] + area[3]]
+                    cropped.append(cr)
+                conditionning["c_concat"] = torch.cat(cropped, dim=1)
+
+            if adm_cond is not None:
+                conditionning["c_adm"] = adm_cond
+
+            control = None
+            if "control" in cond[1]:
+                control = cond[1]["control"]
+
+            patches = None
+            if "gligen" in cond[1]:
+                gligen = cond[1]["gligen"]
+                patches = {}
+                gligen_type = gligen[0]
+                gligen_model = gligen[1]
+                if gligen_type == "position":
+                    gligen_patch = gligen_model.model.set_position(input_x.shape, gligen[2], input_x.device)
+                else:
+                    gligen_patch = gligen_model.model.set_empty(input_x.shape, input_x.device)
+
+                patches["middle_patch"] = [gligen_patch]
+
+            return (input_x, mult, conditionning, area, control, patches)
+
+        def cond_equal_size(c1, c2):
+            if c1 is c2:
+                return True
+            if c1.keys() != c2.keys():
+                return False
+            if "c_crossattn" in c1:
+                s1 = c1["c_crossattn"].shape
+                s2 = c2["c_crossattn"].shape
+                if s1 != s2:
+                    if s1[0] != s2[0] or s1[2] != s2[2]:  # these 2 cases should not happen
+                        return False
+
+                    mult_min = comfy_samplers.lcm(s1[1], s2[1])
+                    diff = mult_min // min(s1[1], s2[1])
+                    if (
+                        diff > 4
+                    ):  # arbitrary limit on the padding because it's probably going to impact performance negatively if it's too much
+                        return False
+            if "c_concat" in c1:
+                if c1["c_concat"].shape != c2["c_concat"].shape:
+                    return False
+            if "c_adm" in c1:
+                if c1["c_adm"].shape != c2["c_adm"].shape:
+                    return False
+            return True
+
+        def can_concat_cond(c1, c2):
+            if c1[0].shape != c2[0].shape:
+                return False
+
+            # control
+            if (c1[4] is None) != (c2[4] is None):
+                return False
+            if c1[4] is not None:
+                if c1[4] is not c2[4]:
+                    return False
+
+            # patches
+            if (c1[5] is None) != (c2[5] is None):
+                return False
+            if c1[5] is not None:
+                if c1[5] is not c2[5]:
+                    return False
+
+            return cond_equal_size(c1[2], c2[2])
+
+        def cond_cat(c_list):
+            c_crossattn = []
+            c_concat = []
+            c_adm = []
+            crossattn_max_len = 0
+            for x in c_list:
+                if "c_crossattn" in x:
+                    c = x["c_crossattn"]
+                    if crossattn_max_len == 0:
+                        crossattn_max_len = c.shape[1]
+                    else:
+                        crossattn_max_len = comfy_samplers.lcm(crossattn_max_len, c.shape[1])
+                    c_crossattn.append(c)
+                if "c_concat" in x:
+                    c_concat.append(x["c_concat"])
+                if "c_adm" in x:
+                    c_adm.append(x["c_adm"])
+            out = {}
+            c_crossattn_out = []
+            for c in c_crossattn:
+                if c.shape[1] < crossattn_max_len:
+                    c = c.repeat(1, crossattn_max_len // c.shape[1], 1)  # padding with repeat doesn't change result
+                c_crossattn_out.append(c)
+
+            if len(c_crossattn_out) > 0:
+                out["c_crossattn"] = torch.cat(c_crossattn_out)
+            if len(c_concat) > 0:
+                out["c_concat"] = torch.cat(c_concat)
+            if len(c_adm) > 0:
+                out["c_adm"] = torch.cat(c_adm)
+            return out
+
+        def calc_cond_uncond_batch(
+            model_function, cond, uncond, x_in, timestep, max_total_area, cond_concat_in, model_options
+        ):
+            out_cond = torch.zeros_like(x_in)
+            out_count = torch.ones_like(x_in) / 100000.0
+
+            out_uncond = torch.zeros_like(x_in)
+            out_uncond_count = torch.ones_like(x_in) / 100000.0
+
+            COND = 0
+            UNCOND = 1
+
+            to_run = []
+            for x in cond:
+                p = get_area_and_mult(x, x_in, cond_concat_in, timestep)
+                if p is None:
+                    continue
+
+                to_run += [(p, COND)]
+            if uncond is not None:
+                for x in uncond:
+                    p = get_area_and_mult(x, x_in, cond_concat_in, timestep)
+                    if p is None:
+                        continue
+
+                    to_run += [(p, UNCOND)]
+
+            while len(to_run) > 0:
+                first = to_run[0]
+                first_shape = first[0][0].shape
+                to_batch_temp = []
+                for x in range(len(to_run)):
+                    if can_concat_cond(to_run[x][0], first[0]):
+                        to_batch_temp += [x]
+
+                to_batch_temp.reverse()
+                to_batch = to_batch_temp[:1]
+
+                for i in range(1, len(to_batch_temp) + 1):
+                    batch_amount = to_batch_temp[: len(to_batch_temp) // i]
+                    if len(batch_amount) * first_shape[0] * first_shape[2] * first_shape[3] < max_total_area:
+                        to_batch = batch_amount
+                        break
+
+                input_x = []
+                mult = []
+                c = []
+                cond_or_uncond = []
+                area = []
+                control = None
+                patches = None
+                for x in to_batch:
+                    o = to_run.pop(x)
+                    p = o[0]
+                    input_x += [p[0]]
+                    mult += [p[1]]
+                    c += [p[2]]
+                    area += [p[3]]
+                    cond_or_uncond += [o[1]]
+                    control = p[4]
+                    patches = p[5]
+
+                batch_chunks = len(cond_or_uncond)
+                input_x = torch.cat(input_x)
+                c = cond_cat(c)
+                timestep_ = torch.cat([timestep] * batch_chunks)
+
+                if control is not None:
+                    c["control"] = control.get_control(input_x, timestep_, c, len(cond_or_uncond))
+
+                transformer_options = {}
+                if "transformer_options" in model_options:
+                    transformer_options = model_options["transformer_options"].copy()
+
+                if patches is not None:
+                    if "patches" in transformer_options:
+                        cur_patches = transformer_options["patches"].copy()
+                        for p in patches:
+                            if p in cur_patches:
+                                cur_patches[p] = cur_patches[p] + patches[p]
+                            else:
+                                cur_patches[p] = patches[p]
+                    else:
+                        transformer_options["patches"] = patches
+
+                transformer_options["cond_or_uncond"] = cond_or_uncond[:]
+                c["transformer_options"] = transformer_options
+
+                if "model_function_wrapper" in model_options:
+                    output = model_options["model_function_wrapper"](
+                        model_function,
+                        {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": cond_or_uncond},
+                    ).chunk(batch_chunks)
+                else:
+                    output = model_function(input_x, timestep_, **c).chunk(batch_chunks)
+                del input_x
+
+                for o in range(batch_chunks):
+                    if cond_or_uncond[o] == COND:
+                        out_cond[:, :, area[o][2] : area[o][0] + area[o][2], area[o][3] : area[o][1] + area[o][3]] += (
+                            output[o] * mult[o]
+                        )
+                        out_count[
+                            :, :, area[o][2] : area[o][0] + area[o][2], area[o][3] : area[o][1] + area[o][3]
+                        ] += mult[o]
+                    else:
+                        out_uncond[
+                            :, :, area[o][2] : area[o][0] + area[o][2], area[o][3] : area[o][1] + area[o][3]
+                        ] += (output[o] * mult[o])
+                        out_uncond_count[
+                            :, :, area[o][2] : area[o][0] + area[o][2], area[o][3] : area[o][1] + area[o][3]
+                        ] += mult[o]
+                del mult
+
+            out_cond /= out_count
+            del out_count
+            out_uncond /= out_uncond_count
+            del out_uncond_count
+
+            return out_cond, out_uncond
+
+        # sliding_calc_cond_uncond_batch inspired by ashen's initial hack for 16-frame sliding context:
+        # https://github.com/comfyanonymous/ComfyUI/compare/master...ashen-sensored:ComfyUI:master
+        def sliding_calc_cond_uncond_batch(
+            model_function, cond, uncond, x_in, timestep, max_total_area, cond_concat_in, model_options
+        ):
+            # figure out how input is split
+            axes_factor = x.size(0) // ctx.video_length
+
+            # prepare final cond, uncond, and out_count
+            cond_final = torch.zeros_like(x)
+            uncond_final = torch.zeros_like(x)
+            out_count_final = torch.zeros((x.shape[0], 1, 1, 1), device=x.device)
+
+            def prepare_control_objects(control: ControlBase, full_idxs: list[int]):
+                if control.previous_controlnet is not None:
+                    prepare_control_objects(control.previous_controlnet, full_idxs)
+                control.sub_idxs = full_idxs
+                control.full_latent_length = ctx.video_length
+                control.context_length = ctx.context_length
+
+            def get_resized_cond(cond_in, full_idxs) -> list:
+                # reuse or resize cond items to match context requirements
+                resized_cond = []
+                # cond object is a list containing a list - outer list is irrelevant, so just loop through it
+                for actual_cond in cond_in:
+                    resized_actual_cond = []
+                    # now we are in the inner list - index 0 is tensor, index 1 is dictionary
+                    for cond_idx, cond_item in enumerate(actual_cond):
+                        if isinstance(cond_item, Tensor):
+                            # check that tensor is the expected length - x.size(0)
+                            if cond_item.size(0) == x.size(0):
+                                pass
+                                # if so, it's subsetting time - tell controls the expected indeces so they can handle them
+                                actual_cond_item = cond_item[full_idxs]
+                                resized_actual_cond.append(actual_cond_item)
+                            else:
+                                resized_actual_cond.append(cond_item)
+                        elif isinstance(cond_item, dict):
+                            # when in dictionary, look for control
+                            if "control" in cond_item:
+                                control_item = cond_item["control"]
+                                if hasattr(control_item, "sub_idxs"):
+                                    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."
+                                    )
+                            resized_actual_cond.append(cond_item)
+                        else:
+                            resized_actual_cond.append(cond_item)
+                    resized_cond.append(resized_actual_cond)
+                return resized_cond
+
+            # perform calc_cond_uncond_batch per context window
+            for ctx_idxs in context_scheduler(
+                ctx.current_step,
+                ctx.total_steps,
+                ctx.video_length,
+                ctx.context_length,
+                ctx.context_stride,
+                ctx.context_overlap,
+                ctx.closed_loop,
+            ):
+                # account for all portions of input frames
+                full_idxs = []
+                for n in range(axes_factor):
+                    for ind in ctx_idxs:
+                        full_idxs.append((ctx.video_length * n) + ind)
+                # get subsections of x, timestep, cond, uncond, cond_concat
+                sub_x = x[full_idxs]
+                sub_timestep = timestep[full_idxs]
+                sub_cond = get_resized_cond(cond, full_idxs) if cond is not None else None
+                sub_uncond = get_resized_cond(uncond, full_idxs) if uncond is not None else None
+                sub_cond_concat = get_resized_cond(cond_concat, full_idxs) if cond_concat is not None else None
+
+                sub_cond_out, sub_uncond_out = calc_cond_uncond_batch(
+                    model_function,
+                    sub_cond,
+                    sub_uncond,
+                    sub_x,
+                    sub_timestep,
+                    max_total_area,
+                    sub_cond_concat,
+                    model_options,
+                )
+
+                cond_final[full_idxs] += sub_cond_out
+                uncond_final[full_idxs] += sub_uncond_out
+                out_count_final[full_idxs] += 1  # increment which indeces were used
+
+            # normalize cond and uncond via division by context usage counts
+            cond_final /= out_count_final
+            uncond_final /= out_count_final
+            return cond_final, uncond_final
+
+        max_total_area = model_management.maximum_batch_area()
+        if math.isclose(cond_scale, 1.0):
+            uncond = None
+
+        cond, uncond = sliding_calc_cond_uncond_batch(
+            model_function, cond, uncond, x, timestep, max_total_area, cond_concat, model_options
+        )
+
+        if "sampler_cfg_function" in model_options:
+            args = {"cond": cond, "uncond": uncond, "cond_scale": cond_scale, "timestep": timestep}
+            return model_options["sampler_cfg_function"](args)
+        else:
+            return uncond + (cond - uncond) * cond_scale
+
+    return (sample, sampling_function)
+
+
+def inject_sampling_function(ctx: SlidingContext):
+    (sample, sampling_function) = __sliding_sample_factory(ctx)
+    comfy.sample.sample = sample
+    comfy_samplers.sampling_function = sampling_function
+
+
+def eject_sampling_function():
+    comfy.sample.sample = orig_comfy_sample
+    comfy_samplers.sampling_function = orig_sampling_function

animatediff/sliding_schedule.py

@@ -0,0 +1,155 @@
+# from https://github.com/neggles/animatediff-cli/blob/main/src/animatediff/pipelines/context.py
+from typing import Callable, Optional
+
+import numpy as np
+
+
+class ContextSchedules:
+    UNIFORM = "uniform"
+    UNIFORM_CONSTANT = "uniform_constant"
+    UNIFORM_V2 = "uniform v2"
+
+    CONTEXT_SCHEDULE_LIST = [UNIFORM, UNIFORM_V2]
+
+
+# Returns fraction that has denominator that is a power of 2
+def ordered_halving(val, print_final=False):
+    # get binary value, padded with 0s for 64 bits
+    bin_str = f"{val:064b}"
+    # flip binary value, padding included
+    bin_flip = bin_str[::-1]
+    # convert binary to int
+    as_int = int(bin_flip, 2)
+    # divide by 1 << 64, equivalent to 2**64, or 18446744073709551616,
+    # or b10000000000000000000000000000000000000000000000000000000000000000 (1 with 64 zero's)
+    final = as_int / (1 << 64)
+    if print_final:
+        print(f"$$$$ final: {final}")
+    return final
+
+
+# Generator that returns lists of latent indeces to diffuse on
+def uniform(
+    step: int = ...,
+    num_steps: Optional[int] = None,
+    num_frames: int = ...,
+    context_size: Optional[int] = None,
+    context_stride: int = 3,
+    context_overlap: int = 4,
+    closed_loop: bool = True,
+    print_final: bool = False,
+):
+    if num_frames <= context_size:
+        yield list(range(num_frames))
+        return
+
+    context_stride = min(context_stride, int(np.ceil(np.log2(num_frames / context_size))) + 1)
+
+    for context_step in 1 << np.arange(context_stride):
+        pad = int(round(num_frames * ordered_halving(step, print_final)))
+        for j in range(
+            int(ordered_halving(step) * context_step) + pad,
+            num_frames + pad + (0 if closed_loop else -context_overlap),
+            (context_size * context_step - context_overlap),
+        ):
+            yield [e % num_frames for e in range(j, j + context_size * context_step, context_step)]
+
+
+def uniform_v2(
+    step: int = ...,
+    num_steps: Optional[int] = None,
+    num_frames: int = ...,
+    context_size: Optional[int] = None,
+    context_stride: int = 3,
+    context_overlap: int = 4,
+    closed_loop: bool = True,
+    print_final: bool = False,
+):
+    if num_frames <= context_size:
+        yield list(range(num_frames))
+        return
+
+    context_stride = min(context_stride, int(np.ceil(np.log2(num_frames / context_size))) + 1)
+
+    pad = int(round(num_frames * ordered_halving(step, print_final)))
+    for context_step in 1 << np.arange(context_stride):
+        j_initial = int(ordered_halving(step) * context_step) + pad
+        for j in range(
+            j_initial,
+            num_frames + pad - context_overlap,
+            (context_size * context_step - context_overlap),
+        ):
+            if context_size * context_step > num_frames:
+                # On the final context_step,
+                # ensure no frame appears in the window twice
+                yield [e % num_frames for e in range(j, j + num_frames, context_step)]
+                continue
+            j = j % num_frames
+            if j > (j + context_size * context_step) % num_frames and not closed_loop:
+                yield [e for e in range(j, num_frames, context_step)]
+                j_stop = (j + context_size * context_step) % num_frames
+                # When  ((num_frames % (context_size - context_overlap)+context_overlap) % context_size != 0,
+                # This can cause 'superflous' runs where all frames in
+                # a context window have already been processed during
+                # the first context window of this stride and step.
+                # While the following commented if should prevent this,
+                # I believe leaving it in is more correct as it maintains
+                # the total conditional passes per frame over a large total steps
+                # if j_stop > context_overlap:
+                yield [e for e in range(0, j_stop, context_step)]
+                continue
+            yield [e % num_frames for e in range(j, j + context_size * context_step, context_step)]
+
+
+def uniform_constant(
+    step: int = ...,
+    num_steps: Optional[int] = None,
+    num_frames: int = ...,
+    context_size: Optional[int] = None,
+    context_stride: int = 3,
+    context_overlap: int = 4,
+    closed_loop: bool = True,
+    print_final: bool = False,
+):
+    if num_frames <= context_size:
+        yield list(range(num_frames))
+        return
+
+    context_stride = min(context_stride, int(np.ceil(np.log2(num_frames / context_size))) + 1)
+
+    # want to avoid loops that connect end to beginning
+
+    for context_step in 1 << np.arange(context_stride):
+        pad = int(round(num_frames * ordered_halving(step, print_final)))
+        for j in range(
+            int(ordered_halving(step) * context_step) + pad,
+            num_frames + pad + (0 if closed_loop else -context_overlap),
+            (context_size * context_step - context_overlap),
+        ):
+            skip_this_window = False
+            prev_val = -1
+            to_yield = []
+            for e in range(j, j + context_size * context_step, context_step):
+                e = e % num_frames
+                # if not a closed loop and loops back on itself, should be skipped
+                if not closed_loop and e < prev_val:
+                    skip_this_window = True
+                    break
+                to_yield.append(e)
+                prev_val = e
+            if skip_this_window:
+                continue
+            # yield if not skipped
+            yield to_yield
+
+
+def get_context_scheduler(name: str) -> Callable:
+    match name:
+        case ContextSchedules.UNIFORM:
+            return uniform
+        case ContextSchedules.UNIFORM_CONSTANT:
+            return uniform_constant
+        case ContextSchedules.UNIFORM_V2:
+            return uniform_v2
+        case _:
+            raise ValueError(f"Unknown context_overlap policy {name}")

workflows/sliding-window.json

@@ -0,0 +1,502 @@
+{
+  "last_node_id": 21,
+  "last_link_id": 36,
+  "nodes": [
+    {
+      "id": 6,
+      "type": "CLIPTextEncode",
+      "pos": [
+        415,
+        186
+      ],
+      "size": {
+        "0": 422.84503173828125,
+        "1": 164.31304931640625
+      },
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "clip",
+          "type": "CLIP",
+          "link": 3
+        }
+      ],
+      "outputs": [
+        {
+          "name": "CONDITIONING",
+          "type": "CONDITIONING",
+          "links": [
+            29
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "CLIPTextEncode"
+      },
+      "widgets_values": [
+        "masterpiece, best quality, 1girl, solo, cherry blossoms, hanami, pink flower, white flower, spring season, wisteria, petals, flower, plum blossoms, outdoors, falling petals, white hair, black eyes"
+      ]
+    },
+    {
+      "id": 8,
+      "type": "VAEDecode",
+      "pos": [
+        1253,
+        191
+      ],
+      "size": {
+        "0": 210,
+        "1": 46
+      },
+      "flags": {},
+      "order": 8,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "samples",
+          "type": "LATENT",
+          "link": 28
+        },
+        {
+          "name": "vae",
+          "type": "VAE",
+          "link": 20
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            19
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VAEDecode"
+      }
+    },
+    {
+      "id": 12,
+      "type": "AnimateDiffCombine",
+      "pos": [
+        1254,
+        290
+      ],
+      "size": {
+        "0": 315,
+        "1": 342
+      },
+      "flags": {},
+      "order": 9,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 19
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "AnimateDiffCombine"
+      },
+      "widgets_values": [
+        8,
+        0,
+        false,
+        "AnimateDiff",
+        "image/gif",
+        false
+      ]
+    },
+    {
+      "id": 16,
+      "type": "AnimateDiffModuleLoader",
+      "pos": [
+        27,
+        345
+      ],
+      "size": {
+        "0": 315,
+        "1": 58
+      },
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "outputs": [
+        {
+          "name": "MOTION_MODULE",
+          "type": "MOTION_MODULE",
+          "links": [
+            24
+          ],
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "AnimateDiffModuleLoader"
+      },
+      "widgets_values": [
+        "mm-Stabilized_mid.pth"
+      ]
+    },
+    {
+      "id": 4,
+      "type": "CheckpointLoaderSimple",
+      "pos": [
+        26,
+        474
+      ],
+      "size": {
+        "0": 315,
+        "1": 98
+      },
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "outputs": [
+        {
+          "name": "MODEL",
+          "type": "MODEL",
+          "links": [
+            25
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "CLIP",
+          "type": "CLIP",
+          "links": [
+            3,
+            5
+          ],
+          "slot_index": 1
+        },
+        {
+          "name": "VAE",
+          "type": "VAE",
+          "links": [],
+          "slot_index": 2
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "CheckpointLoaderSimple"
+      },
+      "widgets_values": [
+        "AnimeLike25D_v11.safetensors"
+      ]
+    },
+    {
+      "id": 13,
+      "type": "VAELoader",
+      "pos": [
+        28,
+        223
+      ],
+      "size": {
+        "0": 315,
+        "1": 58
+      },
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "outputs": [
+        {
+          "name": "VAE",
+          "type": "VAE",
+          "links": [
+            20
+          ],
+          "shape": 3,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "VAELoader"
+      },
+      "widgets_values": [
+        "klF8Anime2.ckpt"
+      ]
+    },
+    {
+      "id": 7,
+      "type": "CLIPTextEncode",
+      "pos": [
+        413,
+        389
+      ],
+      "size": {
+        "0": 425.27801513671875,
+        "1": 180.6060791015625
+      },
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "clip",
+          "type": "CLIP",
+          "link": 5
+        }
+      ],
+      "outputs": [
+        {
+          "name": "CONDITIONING",
+          "type": "CONDITIONING",
+          "links": [
+            30
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "CLIPTextEncode"
+      },
+      "widgets_values": [
+        "embedding:easynegative, embedding:badhandv4, "
+      ]
+    },
+    {
+      "id": 20,
+      "type": "EmptyLatentImage",
+      "pos": [
+        522,
+        621
+      ],
+      "size": {
+        "0": 315,
+        "1": 106
+      },
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "outputs": [
+        {
+          "name": "LATENT",
+          "type": "LATENT",
+          "links": [
+            35
+          ],
+          "shape": 3,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "EmptyLatentImage"
+      },
+      "widgets_values": [
+        512,
+        512,
+        1
+      ]
+    },
+    {
+      "id": 21,
+      "type": "AnimateDiffSlidingWindowOptions",
+      "pos": [
+        517,
+        -34
+      ],
+      "size": {
+        "0": 315,
+        "1": 154
+      },
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "outputs": [
+        {
+          "name": "SLIDING_WINDOW_OPTS",
+          "type": "SLIDING_WINDOW_OPTS",
+          "links": [
+            36
+          ],
+          "shape": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "AnimateDiffSlidingWindowOptions"
+      },
+      "widgets_values": [
+        16,
+        1,
+        4,
+        "uniform",
+        true
+      ]
+    },
+    {
+      "id": 15,
+      "type": "AnimateDiffSampler",
+      "pos": [
+        882,
+        192
+      ],
+      "size": {
+        "0": 315,
+        "1": 350
+      },
+      "flags": {},
+      "order": 7,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "motion_module",
+          "type": "MOTION_MODULE",
+          "link": 24,
+          "slot_index": 0
+        },
+        {
+          "name": "model",
+          "type": "MODEL",
+          "link": 25,
+          "slot_index": 1
+        },
+        {
+          "name": "positive",
+          "type": "CONDITIONING",
+          "link": 29
+        },
+        {
+          "name": "negative",
+          "type": "CONDITIONING",
+          "link": 30
+        },
+        {
+          "name": "latent_image",
+          "type": "LATENT",
+          "link": 35
+        },
+        {
+          "name": "sliding_window_opts",
+          "type": "SLIDING_WINDOW_OPTS",
+          "link": 36,
+          "slot_index": 5
+        }
+      ],
+      "outputs": [
+        {
+          "name": "LATENT",
+          "type": "LATENT",
+          "links": [
+            28
+          ],
+          "shape": 3,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "AnimateDiffSampler"
+      },
+      "widgets_values": [
+        "default",
+        40,
+        345029849956677,
+        "fixed",
+        20,
+        8,
+        "euler",
+        "normal",
+        0.8
+      ]
+    }
+  ],
+  "links": [
+    [
+      3,
+      4,
+      1,
+      6,
+      0,
+      "CLIP"
+    ],
+    [
+      5,
+      4,
+      1,
+      7,
+      0,
+      "CLIP"
+    ],
+    [
+      19,
+      8,
+      0,
+      12,
+      0,
+      "IMAGE"
+    ],
+    [
+      20,
+      13,
+      0,
+      8,
+      1,
+      "VAE"
+    ],
+    [
+      24,
+      16,
+      0,
+      15,
+      0,
+      "MOTION_MODULE"
+    ],
+    [
+      25,
+      4,
+      0,
+      15,
+      1,
+      "MODEL"
+    ],
+    [
+      28,
+      15,
+      0,
+      8,
+      0,
+      "LATENT"
+    ],
+    [
+      29,
+      6,
+      0,
+      15,
+      2,
+      "CONDITIONING"
+    ],
+    [
+      30,
+      7,
+      0,
+      15,
+      3,
+      "CONDITIONING"
+    ],
+    [
+      35,
+      20,
+      0,
+      15,
+      4,
+      "LATENT"
+    ],
+    [
+      36,
+      21,
+      0,
+      15,
+      5,
+      "SLIDING_WINDOW_OPTS"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {},
+  "version": 0.4
+}