Merge pull request #27 from AustinMroz/upstream_video_format

ffmpeg improvements: webm quality, and additional video formats

README.md

@@ -5,20 +5,34 @@
 ## How to Use
 
 1. Clone this repo into `custom_nodes` folder.
-2. Download motion modules from [Google Drive](https://drive.google.com/drive/folders/1EqLC65eR1-W-sGD0Im7fkED6c8GkiNFI) | [HuggingFace](https://huggingface.co/guoyww/animatediff) | [CivitAI](https://civitai.com/models/108836) | [Baidu NetDisk](https://pan.baidu.com/s/18ZpcSM6poBqxWNHtnyMcxg?pwd=et8y). You only need to download one of `mm_sd_v14.ckpt` | `mm_sd_v15.ckpt`. Put the model weights under `comfyui-animatediff/models/`. DO NOT change model filename.
+2. Download motion modules and put them under `comfyui-animatediff/models/`.
+  * Original modules: [Google Drive](https://drive.google.com/drive/folders/1EqLC65eR1-W-sGD0Im7fkED6c8GkiNFI) | [HuggingFace](https://huggingface.co/guoyww/animatediff) | [CivitAI](https://civitai.com/models/108836) | [Baidu NetDisk](https://pan.baidu.com/s/18ZpcSM6poBqxWNHtnyMcxg?pwd=et8y)
+  * 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/15
+## Nodes
 
-- You can now use community models from [manshoety/AD_Stabilized_Motion](https://huggingface.co/manshoety/AD_Stabilized_Motion) or [CiaraRowles/TemporalDiff](https://huggingface.co/CiaraRowles/TemporalDiff)
+#### AnimateDiffLoader
-- Supports AnimateDiff v2 [mm_sd_v15_v2.ckpt](https://huggingface.co/guoyww/animatediff/blob/main/mm_sd_v15_v2.ckpt) model
-- Fix image is grayed out.
-- New node: **AnimateDiffSampler** and **AnimateDiffLoader**
-  - Mostly the same with `KSampler`
-  - Use `AnimateDiffLoader` to load the motion module
-  - `inject_method`: should left default. See [this issue](https://github.com/ArtVentureX/comfyui-animatediff#gif-has-wartermark-after-update-to-the-latest-version) for more details.
-  - `frame_number`: animation length
 
-<img width="506" 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/9d756d01-ea45-4d1c-8e48-56f2725c7ca1">
+
+#### AnimateDiffSampler
+- Mostly the same with `KSampler`
+- Use `AnimateDiffLoader` to load the motion module
+- `inject_method`: should left default
+- `frame_number`: animation length
+- `latent_image`: You can pass an `EmptyLatentImage`
+
+<img width="370" alt="image" src="https://github.com/ArtVentureX/comfyui-animatediff/assets/133728487/f22d6b36-ce36-44cc-80e8-dffe6f77b296">
+
+#### AnimateDiffCombine
+- Combine GIF frames and produce the GIF image
+- `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)
+
+<img width="370" alt="image" src="https://github.com/ArtVentureX/comfyui-animatediff/assets/133728487/381c5acc-06ef-43da-ada0-3dc76f37a3e4">
 
 #### Example Workflow
 
@@ -39,39 +53,14 @@ Workflow file: https://github.com/ArtVentureX/comfyui-animatediff/blob/main/work
 
 ![AnimateDiff_00007_](https://github.com/ArtVentureX/comfyui-animatediff/assets/8894763/e6cd53cb-9878-45da-a58a-a15851882386)
 
-See: https://github.com/continue-revolution/sd-webui-animatediff/issues/38
-
-Main reasons:
-
-- Promt are too long (more than 75 tokens)
-- Resolution are too high
-- Number of frame too high
-
 Work around:
 
 - Shorter your prompt and negative prompt
 - Reduce resolution. AnimateDiff is trained on 512x512 images so it works best with 512x512 output.
-- Shouldn't generate longer than 16 frames. AnimateDiff is trained to output the best results with 16 frames.
+- Disable xformers with `--disable-xformers`
 
-### GIF has Wartermark after update to the latest version
+### GIF has Wartermark (especially when using mm_sd_v15)
 
 See: https://github.com/continue-revolution/sd-webui-animatediff/issues/31
 
-As mentioned in the issue thread, it seems to be due to the training dataset. The new version is the correct implementation and produces smoother GIFs compared to the older version.
+Training data used by the authors of the AnimateDiff paper contained Shutterstock watermarks. Since mm_sd_v15 was finetuned on finer, less drastic movement, the motion module attempts to replicate the transparency of that watermark and does not get blurred away like mm_sd_v14. Try other community finetuned modules.
-
-<table  class="center">
-    <tr>
-    <td>Old revision</td>
-    <td>New revision</td>
-    </tr>
-    <tr>
-    <td><img src="https://github.com/ArtVentureX/comfyui-animatediff/assets/133728487/8f1a6233-875f-4f0c-aa60-ba93e73b7d64" /></td>
-    <td><img src="https://github.com/ArtVentureX/comfyui-animatediff/assets/133728487/a2029eba-f519-437c-a0b5-1f881e099a20" /></td>
-    </tr>
-    <tr>
-    <td><img src="https://github.com/ArtVentureX/comfyui-animatediff/assets/133728487/41ec449f-1955-466c-bd38-6f2a55d654f8" /></td>
-    <td><img src="https://github.com/ArtVentureX/comfyui-animatediff/assets/133728487/766c2891-5d27-4052-99f9-be9862620919" /></td>
-    </tr>
-</table>
-
-I played around with both version and found that the watermark only present in some models, not always. To use the **old (legacy)** method, change `injection_method` to `legacy` in the `AnimateDiffSampler` node.

__init__.py

@@ -5,4 +5,6 @@ from .animatediff.model_utils import get_available_models
 if len(get_available_models()) == 0:
     logger.error("No models available. Please download one and put it in models folder")
 
-__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+WEB_DIRECTORY = "./web"
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]

animatediff/model_utils.py

@@ -11,6 +11,12 @@ 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():

animatediff/motion_module.py

@@ -5,7 +5,6 @@ from torch import Tensor, nn
 import math
 from einops import rearrange, repeat
 
-from comfy.utils import load_torch_file
 from comfy.ldm.modules.attention import FeedForward, CrossAttention
 
 
@@ -57,14 +56,12 @@ class MotionWrapper(nn.Module):
             )
 
     @classmethod
-    def from_pretrained(cls, checkpoint_path: str):
+    def from_pretrained(cls, mm_state_dict: dict[str, Tensor], mm_type: str):
-        mm_state_dict = load_torch_file(checkpoint_path)
-        mm_type = os.path.basename(checkpoint_path)
         encoding_max_len = get_encoding_max_len(mm_state_dict)
         is_v2 = has_mid_block(mm_state_dict)
 
         mm = cls(mm_type, encoding_max_len=encoding_max_len, is_v2=is_v2)
-        mm.load_state_dict(mm_state_dict)
+        mm.load_state_dict(mm_state_dict, strict=False)
         return mm
 
     def set_video_length(self, video_length: int):

animatediff/nodes.py

@@ -14,7 +14,7 @@ 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.cli_args import args as cli_args
+from comfy.utils import load_torch_file, calculate_parameters
 from nodes import KSampler
 
 from .logger import logger
@@ -67,10 +67,15 @@ def load_motion_module(model_name: str):
     model_hash = get_model_hash(model_path)
     if model_hash not in motion_modules:
         logger.info(f"Loading motion module {model_name}")
-        motion_module = MotionWrapper.from_pretrained(model_path)
+        mm_state_dict = load_torch_file(model_path)
-        if not cli_args.force_fp32:
+        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
 
@@ -215,7 +220,7 @@ class AnimateDiffSampler(KSampler):
 
     def override_beta_schedule(self, model: BaseModel):
         logger.info(f"Override beta schedule.")
-        self.prev_beta = model.get_buffer("betas")
+        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(
@@ -245,6 +250,7 @@ class AnimateDiffSampler(KSampler):
         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:
@@ -258,7 +264,7 @@ class AnimateDiffSampler(KSampler):
 
         self.restore_beta_schedule(model.model)
         if not unet.motion_module.is_v2:
-            logger.info(f"Restore GroupNorm32 forward function.")
+            logger.info(f"Restore GroupNorm.forward function.")
             torch.nn.GroupNorm.forward = orig_groupnorm_forward
 
         logger.info(f"Ejecting motion module with method {inject_method}.")
@@ -326,8 +332,11 @@ class AnimateDiffCombine:
                     {"default": 8, "min": 1, "max": 24, "step": 1},
                 ),
                 "loop_count": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
-                "save_image": (["Enabled", "Disabled"],),
+                "save_image": ([True, False],),
-                "filename_prefix": ("STRING", {"default": "AnimateDiff"}),
+                "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],),
             },
             "hidden": {
                 "prompt": "PROMPT",
@@ -335,7 +344,7 @@ class AnimateDiffCombine:
             },
         }
 
-    RETURN_TYPES = ()
+    RETURN_TYPES = ("GIF",)
     OUTPUT_NODE = True
     CATEGORY = "Animate Diff"
     FUNCTION = "generate_gif"
@@ -345,22 +354,24 @@ class AnimateDiffCombine:
         images,
         frame_rate: int,
         loop_count: int,
-        save_image="Enabled",
+        save_image=True,
         filename_prefix="AnimateDiff",
+        format="image/gif",
+        pingpong=False,
         prompt=None,
         extra_pnginfo=None,
     ):
         # convert images to numpy
-        pil_images: List[Image.Image] = []
+        frames: List[Image.Image] = []
         for image in images:
             img = 255.0 * image.cpu().numpy()
             img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))
-            pil_images.append(img)
+            frames.append(img)
 
         # save image
         output_dir = (
             folder_paths.get_output_directory()
-            if save_image == "Enabled"
+            if save_image
             else folder_paths.get_temp_directory()
         )
         (
@@ -381,34 +392,62 @@ class AnimateDiffCombine:
         # save first frame as png to keep metadata
         file = f"{filename}_{counter:05}_.png"
         file_path = os.path.join(full_output_folder, file)
-        pil_images[0].save(
+        frames[0].save(
             file_path,
             pnginfo=metadata,
             compress_level=4,
         )
+        if pingpong:
+            frames = frames + frames[-2:0:-1]
 
-        # save gif
+        format_type, format_ext = format.split("/")
-        file = f"{filename}_{counter:05}_.gif"
+
+        if format_type == "image":
+            file = f"{filename}_{counter:05}_.{format_ext}"
             file_path = os.path.join(full_output_folder, file)
-        pil_images[0].save(
+            frames[0].save(
                 file_path,
+                format=format_ext.upper(),
                 save_all=True,
-            append_images=pil_images[1:],
+                append_images=frames[1:],
                 duration=round(1000 / frame_rate),
                 loop=loop_count,
                 compress_level=4,
             )
+        else:
+            # save webm
+            import shutil
+            import subprocess
 
-        print("Saved gif to", file_path, os.path.exists(file_path))
+            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 = 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]
+
+            env=os.environ
+            if "environment" in video_format:
+                env.update(video_format["environment"])
+            with subprocess.Popen(args, stdin=subprocess.PIPE, env=env) as proc:
+                for frame in frames:
+                    proc.stdin.write(frame.tobytes())
 
         previews = [
             {
                 "filename": file,
                 "subfolder": subfolder,
-                "type": "output" if save_image == "Enabled" else "temp",
+                "type": "output" if save_image else "temp",
+                "format": format,
             }
         ]
-        return {"ui": {"images": previews}}
+        return {"ui": {"gifs": previews}}
 
 
 NODE_CLASS_MAPPINGS = {

video_formats/av1-webm.json

@@ -0,0 +1,10 @@
+{
+    "main_pass":
+    [
+        "-n", "-c:v", "libsvtav1",
+        "-pix_fmt", "yuv420p10le",
+        "-crf", "23"
+    ],
+     "extension": "webm",
+     "environment": {"SVT_LOG": "1"}
+}

video_formats/h264-mp4.json

@@ -0,0 +1,9 @@
+{
+    "main_pass":
+    [
+        "-n", "-c:v", "libx264",
+        "-pix_fmt", "yuv420p",
+        "-crf", "19"
+    ],
+     "extension": "mp4"
+}

video_formats/h265-mp4.json

@@ -0,0 +1,11 @@
+{
+    "main_pass":
+    [
+        "-n", "-c:v", "libx265",
+        "-pix_fmt", "yuv420p10le",
+        "-preset", "medium",
+        "-crf", "22",
+        "-x265-params", "log-level=quiet"
+    ],
+     "extension": "mp4"
+}

video_formats/webm.json

@@ -0,0 +1,9 @@
+{
+    "main_pass":
+    [
+        "-n",
+        "-pix_fmt", "yuv420p",
+        "-crf", "23"
+    ],
+     "extension": "webm"
+}

web/js/gif_preview.js

@@ -0,0 +1,146 @@
+import { app } from '../../../scripts/app.js'
+import { api } from '../../../scripts/api.js'
+
+function offsetDOMWidget(
+  widget,
+  ctx,
+  node,
+  widgetWidth,
+  widgetY,
+  height
+) {
+  const margin = 10
+  const elRect = ctx.canvas.getBoundingClientRect()
+  const transform = new DOMMatrix()
+    .scaleSelf(
+      elRect.width / ctx.canvas.width,
+      elRect.height / ctx.canvas.height
+    )
+    .multiplySelf(ctx.getTransform())
+    .translateSelf(0, widgetY + margin)
+
+  const scale = new DOMMatrix().scaleSelf(transform.a, transform.d)
+  Object.assign(widget.inputEl.style, {
+    transformOrigin: '0 0',
+    transform: scale,
+    left: `${transform.e}px`,
+    top: `${transform.d + transform.f}px`,
+    width: `${widgetWidth}px`,
+    height: `${(height || widget.parent?.inputHeight || 32) - margin}px`,
+    position: 'absolute',
+    background: !node.color ? '' : node.color,
+    color: !node.color ? '' : 'white',
+    zIndex: 5, //app.graph._nodes.indexOf(node),
+  })
+}
+
+export const hasWidgets = (node) => {
+  if (!node.widgets || !node.widgets?.[Symbol.iterator]) {
+    return false
+  }
+  return true
+}
+
+export const cleanupNode = (node) => {
+  if (!hasWidgets(node)) {
+    return
+  }
+
+  for (const w of node.widgets) {
+    if (w.canvas) {
+      w.canvas.remove()
+    }
+    if (w.inputEl) {
+      w.inputEl.remove()
+    }
+    // calls the widget remove callback
+    w.onRemoved?.()
+  }
+}
+
+const CreatePreviewElement = (name, val, format) => {
+  const [type] = format.split('/')
+
+  const w = {
+    name,
+    type,
+    value: val,
+    draw: function (ctx, node, widgetWidth, widgetY, height) {
+      const [cw, ch] = this.computeSize(widgetWidth)
+      offsetDOMWidget(this, ctx, node, widgetWidth, widgetY, ch)
+    },
+    computeSize: function (_) {
+      const ratio = this.inputRatio || 1
+      const width = Math.max(220, this.parent.size[0])
+      return [width, (width / ratio + 10)]
+    },
+    onRemoved: function () {
+      if (this.inputEl) {
+        this.inputEl.remove()
+      }
+    },
+  }
+
+  w.inputEl = document.createElement(type === 'video' ? 'video' : 'img')
+  w.inputEl.src = w.value
+  if (type === 'video') {
+    w.inputEl.setAttribute('type', 'video/webm');
+    w.inputEl.autoplay = true
+    w.inputEl.loop = true
+    w.inputEl.controls = false;
+  }
+  w.inputEl.onload = function () {
+    w.inputRatio = w.inputEl.naturalWidth / w.inputEl.naturalHeight
+  }
+  document.body.appendChild(w.inputEl)
+  return w
+}
+
+const gif_preview = {
+  name: 'AnimateDiff.gif_preview',
+  async beforeRegisterNodeDef(nodeType, nodeData, app) {
+    switch (nodeData.name) {
+      case 'AnimateDiffCombine': {
+        const onExecuted = nodeType.prototype.onExecuted
+        nodeType.prototype.onExecuted = function (message) {
+          const prefix = 'ad_gif_preview_'
+          const r = onExecuted ? onExecuted.apply(this, message) : undefined
+
+          if (this.widgets) {
+            const pos = this.widgets.findIndex((w) => w.name === `${prefix}_0`)
+            if (pos !== -1) {
+              for (let i = pos; i < this.widgets.length; i++) {
+                this.widgets[i].onRemoved?.()
+              }
+              this.widgets.length = pos
+            }
+            if (message?.gifs) {
+              message.gifs.forEach((params, i) => {
+                const previewUrl = api.apiURL(
+                  '/view?' + new URLSearchParams(params).toString()
+                )
+                const w = this.addCustomWidget(
+                  CreatePreviewElement(`${prefix}_${i}`, previewUrl, params.format || 'image/gif')
+                )
+                w.parent = this
+              })
+            }
+
+            const onRemoved = this.onRemoved
+            this.onRemoved = () => {
+              cleanupNode(this)
+              return onRemoved?.()
+            }
+          }
+
+          // keep width and update height
+          this.setSize([this.size[0], this.computeSize([this.size[0], this.size[1]])[1]])
+          return r
+        }
+        break
+      }
+    }
+  }
+}
+
+app.registerExtension(gif_preview)