From fc6c05323dd3a7eaddd5e4d0b249ed2d366be912 Mon Sep 17 00:00:00 2001
From: kijai <jukka.seppanen@poke.fi>
Date: Sat, 16 Sep 2023 19:46:49 +0300
Subject: [PATCH] Add webm node preview

added video preview to the node, with pingpong loop option
---
 animatediff/nodes.py  | 21 +++++++++++----
 web/js/gif_preview.js | 62 ++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 74 insertions(+), 9 deletions(-)

diff --git a/animatediff/nodes.py b/animatediff/nodes.py
index d0ac39b..1fa3823 100644
--- a/animatediff/nodes.py
+++ b/animatediff/nodes.py
@@ -540,6 +540,7 @@ class AnimateDiffCombineVideo:
                 ),
                 "save_image": (["Enabled", "Disabled"],),
                 "filename_prefix": ("STRING", {"default": "AnimateDiff"}),
+                "pingpong": ("BOOLEAN", {"default": False}),
             },
             "hidden": {
                 "prompt": "PROMPT",
@@ -558,6 +559,7 @@ class AnimateDiffCombineVideo:
         frame_rate: int,
         save_image="Enabled",
         filename_prefix="AnimateDiff",
+        pingpong=False,
         prompt=None,
         extra_pnginfo=None,
     ):
@@ -591,7 +593,7 @@ class AnimateDiffCombineVideo:
             pnginfo=metadata,
             compress_level=4,
         )
-
+        
         # save webm
         ffmpeg_path = shutil.which("ffmpeg")
         if ffmpeg_path is None:
@@ -601,20 +603,29 @@ class AnimateDiffCombineVideo:
         dimensions = f"{first_image.width}x{first_image.height}"
         args = [ffmpeg_path, "-v", "panic", "-n", "-f", "rawvideo", "-pix_fmt", "rgb24", "-s",
                 dimensions, "-r", str(frame_rate), "-i", "-", "-pix_fmt", "yuv420p", file_path]
+        # convert images to numpy
+        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))
+            frames.append(img)
+        if pingpong:
+            frames += frames[::-1]
+
         with subprocess.Popen(args, stdin=subprocess.PIPE) as proc:
-            for image in images:
-                proc.stdin.write(np.clip(255.0*image.cpu().numpy(),0,255).astype(np.uint8).tobytes())
+            for frame in frames:
+                proc.stdin.write(frame.tobytes())
         print("Saved webm to", file_path, os.path.exists(file_path))
 
 
         previews = [
             {
-                "filename": file,
+                "filename": file_webm,
                 "subfolder": subfolder,
                 "type": "output" if save_image == "Enabled" else "temp",
             }
         ]
-        return {"ui": {"images": previews}}
+        return {"ui": {"video": previews}}
 
 
 class CheckpointLoaderSimpleWithNoiseSelect:
diff --git a/web/js/gif_preview.js b/web/js/gif_preview.js
index 6aa5ea7..1e83d06 100644
--- a/web/js/gif_preview.js
+++ b/web/js/gif_preview.js
@@ -42,7 +42,7 @@ function offsetDOMWidget(
     }
     return true
   }
-  
+
   export const cleanupNode = (node) => {
     if (!hasWidgets(node)) {
       return
@@ -60,7 +60,43 @@ function offsetDOMWidget(
     }
   }
 
-const DEBUG_IMG = (name, val) => {
+const CreateVideoElement = (name, val) => {
+    const w = {
+      name,
+      type: 'video',
+      value: val,
+      draw: function (ctx, node, widgetWidth, widgetY, height) {
+        const [cw, ch] = this.computeSize(widgetWidth)
+        offsetDOMWidget(this, ctx, node, widgetWidth, widgetY, ch)
+      },
+      computeSize: function (width) {
+        const ratio = this.inputRatio || 1
+        if (width) {
+          return [width, width / ratio + 4]
+        }
+        return [128, 128]
+      },
+      onRemoved: function () {
+        if (this.inputEl) {
+          this.inputEl.remove()
+        }
+      },
+    }
+
+    w.inputEl = document.createElement('video')
+    w.inputEl.src = w.value
+    w.inputEl.setAttribute('type', 'video/webm');
+    w.inputEl.autoplay = true
+    w.inputEl.loop = true
+    w.inputEl.controls = false; // Add controls to the video element
+    w.inputEl.onload = function () {
+      w.inputRatio = w.inputEl.naturalWidth / w.inputEl.naturalHeight
+    }
+    document.body.appendChild(w.inputEl)
+    return w
+  }
+
+  const CreateImgElement = (name, val) => {
     const w = {
       name,
       type: 'image',
@@ -85,6 +121,7 @@ const DEBUG_IMG = (name, val) => {
 
     w.inputEl = document.createElement('img')
     w.inputEl.src = w.value
+    w.inputEl.controls = false; // Add controls to the video element
     w.inputEl.onload = function () {
       w.inputRatio = w.inputEl.naturalWidth / w.inputEl.naturalHeight
     }
@@ -96,6 +133,7 @@ const gif_preview = {
     name: 'gif.preview',
     async beforeRegisterNodeDef(nodeType, nodeData, app) {
         switch (nodeData.name) {
+            case 'ADE_AnimateDiffCombineVideo':
             case 'ADE_AnimateDiffCombine':{
                 const onExecuted = nodeType.prototype.onExecuted
                 nodeType.prototype.onExecuted = function (message) {
@@ -112,9 +150,9 @@ const gif_preview = {
                     }
 
                     let imgURLs = []
+                    let videoURLs = []
                     if (message) {
                     if (message.gif) {
-                        console.log("found gif")
                         imgURLs = imgURLs.concat(
                         message.gif.map((params) => {
                             return api.apiURL(
@@ -123,14 +161,30 @@ const gif_preview = {
                         })
                         )
                     }
+                    if (message.video) {
+                      videoURLs = videoURLs.concat(
+                      message.video.map((params) => {
+                          return api.apiURL(
+                          '/view?' + new URLSearchParams(params).toString()
+                          )
+                      })
+                      )
+                  }
                     let i = 0
                     for (const img of imgURLs) {
                         const w = this.addCustomWidget(
-                        DEBUG_IMG(`${prefix}_${i}`, img)
+                        CreateImgElement(`${prefix}_${i}`, img)
                         )
                         w.parent = this
                         i++
                     }
+                    for (const video of videoURLs) {
+                      const w = this.addCustomWidget(
+                      CreateVideoElement(`${prefix}_${i}`, video)
+                      )
+                      w.parent = this
+                      i++
+                  }
                     }
                     const onRemoved = this.onRemoved
                     this.onRemoved = () => {