From 0d4d9222c6922e123e81c9fea104275df33b497d Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Wed, 19 Feb 2025 07:11:49 -0500
Subject: [PATCH] Add early experimental SaveWEBM node to save .webm files.

The frontend part isn't done yet so there is no video preview on the node
or dragging the webm on the interface to load the workflow yet.

This uses a new dependency: PyAV.
---
 comfy_extras/nodes_video.py | 75 +++++++++++++++++++++++++++++++++++++
 nodes.py                    |  1 +
 requirements.txt            |  1 +
 3 files changed, 77 insertions(+)
 create mode 100644 comfy_extras/nodes_video.py

diff --git a/comfy_extras/nodes_video.py b/comfy_extras/nodes_video.py
new file mode 100644
index 00000000..f3922e03
--- /dev/null
+++ b/comfy_extras/nodes_video.py
@@ -0,0 +1,75 @@
+import os
+import av
+import torch
+import folder_paths
+import json
+from fractions import Fraction
+
+
+class SaveWEBM:
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "output"
+        self.prefix_append = ""
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"images": ("IMAGE", ),
+                     "filename_prefix": ("STRING", {"default": "ComfyUI"}),
+                     "codec": (["vp9", "av1"],),
+                     "fps": ("FLOAT", {"default": 24.0, "min": 0.01, "max": 1000.0, "step": 0.01}),
+                     "crf": ("FLOAT", {"default": 32.0, "min": 0, "max": 63.0, "step": 1, "tooltip": "Higher crf means lower quality with a smaller file size, lower crf means higher quality higher filesize."}),
+                     },
+                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+                }
+
+    RETURN_TYPES = ()
+    FUNCTION = "save_images"
+
+    OUTPUT_NODE = True
+
+    CATEGORY = "image/video"
+
+    EXPERIMENTAL = True
+
+    def save_images(self, images, codec, fps, filename_prefix, crf, prompt=None, extra_pnginfo=None):
+        filename_prefix += self.prefix_append
+        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
+
+        file = f"{filename}_{counter:05}_.webm"
+        container = av.open(os.path.join(full_output_folder, file), mode="w")
+
+        if prompt is not None:
+            container.metadata["prompt"] = json.dumps(prompt)
+
+        if extra_pnginfo is not None:
+            for x in extra_pnginfo:
+                container.metadata[x] = json.dumps(extra_pnginfo[x])
+
+        codec_map = {"vp9": "libvpx-vp9", "av1": "libaom-av1"}
+        stream = container.add_stream(codec_map[codec], rate=Fraction(round(fps * 1000), 1000))
+        stream.width = images.shape[-2]
+        stream.height = images.shape[-3]
+        stream.pix_fmt = "yuv420p"
+        stream.bit_rate = 0
+        stream.options = {'crf': str(crf)}
+
+        for frame in images:
+            frame = av.VideoFrame.from_ndarray(torch.clamp(frame[..., :3] * 255, min=0, max=255).to(device=torch.device("cpu"), dtype=torch.uint8).numpy(), format="rgb24")
+            for packet in stream.encode(frame):
+                container.mux(packet)
+        container.close()
+
+        results = [{
+            "filename": file,
+            "subfolder": subfolder,
+            "type": self.type
+        }]
+
+        return {"ui": {"images": results, "animated": (True,)}}  # TODO: frontend side
+
+
+NODE_CLASS_MAPPINGS = {
+    "SaveWEBM": SaveWEBM,
+}
diff --git a/nodes.py b/nodes.py
index b39adc65..4e68af79 100644
--- a/nodes.py
+++ b/nodes.py
@@ -2265,6 +2265,7 @@ def init_builtin_extra_nodes():
         "nodes_hooks.py",
         "nodes_load_3d.py",
         "nodes_cosmos.py",
+        "nodes_video.py",
         "nodes_lumina2.py",
     ]
 
diff --git a/requirements.txt b/requirements.txt
index 3bc945a1..3229fe81 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -19,3 +19,4 @@ psutil
 kornia>=0.7.1
 spandrel
 soundfile
+av