Merge pull request #27 from AustinMroz/upstream_video_format

ffmpeg improvements: webm quality, and additional video formats
Add additional video formats and config system
2023-09-18 14:31:52 +07:00 · 2023-09-17 19:45:35 -05:00 · 2023-09-17 11:43:51 +07:00 · 2023-09-17 11:41:55 +07:00 · 2023-09-17 11:08:55 +07:00 · 2023-09-16 17:46:24 +07:00
10 changed files with 290 additions and 72 deletions
@@ -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)
- 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
+#### AnimateDiffLoader

-<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.
-
-<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.
+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.
@@ -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"]
@@ -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():
@@ -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):
-        mm_state_dict = load_torch_file(checkpoint_path)
-        mm_type = os.path.basename(checkpoint_path)
+    def from_pretrained(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)

        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):
@@ -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)
-        if not cli_args.force_fp32:
+        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

@@ -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"],),
-                "filename_prefix": ("STRING", {"default": "AnimateDiff"}),
+                "save_image": ([True, False],),
+                "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
-        file = f"{filename}_{counter:05}_.gif"
-        file_path = os.path.join(full_output_folder, file)
-        pil_images[0].save(
-            file_path,
-            save_all=True,
-            append_images=pil_images[1:],
-            duration=round(1000 / frame_rate),
-            loop=loop_count,
-            compress_level=4,
-        )
+        format_type, format_ext = format.split("/")

-        print("Saved gif to", file_path, os.path.exists(file_path))
+        if format_type == "image":
+            file = f"{filename}_{counter:05}_.{format_ext}"
+            file_path = os.path.join(full_output_folder, file)
+            frames[0].save(
+                file_path,
+                format=format_ext.upper(),
+                save_all=True,
+                append_images=frames[1:],
+                duration=round(1000 / frame_rate),
+                loop=loop_count,
+                compress_level=4,
+            )
+        else:
+            # save webm
+            import shutil
+            import subprocess
+
+            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 = {
@@ -0,0 +1,10 @@
+{
+    "main_pass":
+    [
+        "-n", "-c:v", "libsvtav1",
+        "-pix_fmt", "yuv420p10le",
+        "-crf", "23"
+    ],
+     "extension": "webm",
+     "environment": {"SVT_LOG": "1"}
+}
@@ -0,0 +1,9 @@
+{
+    "main_pass":
+    [
+        "-n", "-c:v", "libx264",
+        "-pix_fmt", "yuv420p",
+        "-crf", "19"
+    ],
+     "extension": "mp4"
+}
@@ -0,0 +1,11 @@
+{
+    "main_pass":
+    [
+        "-n", "-c:v", "libx265",
+        "-pix_fmt", "yuv420p10le",
+        "-preset", "medium",
+        "-crf", "22",
+        "-x265-params", "log-level=quiet"
+    ],
+     "extension": "mp4"
+}
@@ -0,0 +1,9 @@
+{
+    "main_pass":
+    [
+        "-n",
+        "-pix_fmt", "yuv420p",
+        "-crf", "23"
+    ],
+     "extension": "webm"
+}
@@ -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)
Author	SHA1	Message	Date
ArtVenture	4e881671aa	Merge pull request #27 from AustinMroz/upstream_video_format ffmpeg improvements: webm quality, and additional video formats	2023-09-18 14:31:52 +07:00
Austin Mroz	414c5d3bb8	Add additional video formats and config system This ports the video format code written for the upstream changes to the ffmpeg implementation. It improves the quality of webm outputs and adds support for additional codecs (h264, h265, av1) It also improves the logging by passing errors and more selectively blocking the logging of encoders. While h265 has been included, most browsers will be unable to display the resulting video.	2023-09-17 19:45:35 -05:00
ArtVenture	78e04fcdc6	Merge pull request #25 from ArtVentureX/feat/gif_preview Improve GIF preview and support video output	2023-09-17 11:43:51 +07:00
Tung Nguyen	60d14a9840	update README	2023-09-17 11:41:55 +07:00
Tung Nguyen	427cf04893	improve gif preview	2023-09-17 11:08:55 +07:00
Tung Nguyen	87815b7aae	add gif preview & support pingping gif	2023-09-16 17:46:24 +07:00
Tung Nguyen	9ae375fbd8	fix: cannot change frame_number	2023-09-16 17:16:16 +07:00
ArtVenture	d4f5328a47	Merge pull request #23 from ArtVentureX/code-refactor code refactor	2023-09-16 06:28:31 +07:00