From 414c5d3bb850e3a7416b66f4fae613d549a611f5 Mon Sep 17 00:00:00 2001
From: Austin Mroz <AustinMroz@utexas.edu>
Date: Sun, 17 Sep 2023 19:45:35 -0500
Subject: [PATCH] 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.
---
 animatediff/model_utils.py  |  6 ++++++
 animatediff/nodes.py        | 23 ++++++++++++++++-------
 video_formats/av1-webm.json | 10 ++++++++++
 video_formats/h264-mp4.json |  9 +++++++++
 video_formats/h265-mp4.json | 11 +++++++++++
 video_formats/webm.json     |  9 +++++++++
 6 files changed, 61 insertions(+), 7 deletions(-)
 create mode 100644 video_formats/av1-webm.json
 create mode 100644 video_formats/h264-mp4.json
 create mode 100644 video_formats/h265-mp4.json
 create mode 100644 video_formats/webm.json

diff --git a/animatediff/model_utils.py b/animatediff/model_utils.py
index c9fe398..7a6320a 100644
--- a/animatediff/model_utils.py
+++ b/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():
diff --git a/animatediff/nodes.py b/animatediff/nodes.py
index ffe10b6..35cd9fd 100644
--- a/animatediff/nodes.py
+++ b/animatediff/nodes.py
@@ -334,7 +334,8 @@ class AnimateDiffCombine:
                 "loop_count": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
                 "save_image": ([True, False],),
                 "filename_prefix": ("STRING", {"default": "animate_diff"}),
-                "format": (["image/gif", "image/webp", "video/webm"],),
+                "format": (["image/gif", "image/webp"] +
+                           ["video/"+x[:-5] for x in folder_paths.get_filename_list("video_formats")],),
                 "pingpong": ([False, True],),
             },
             "hidden": {
@@ -399,12 +400,11 @@ class AnimateDiffCombine:
         if pingpong:
             frames = frames + frames[-2:0:-1]
 
-        # save gif
         format_type, format_ext = format.split("/")
-        file = f"{filename}_{counter:05}_.{format_ext}"
-        file_path = os.path.join(full_output_folder, file)
 
         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(),
@@ -422,11 +422,20 @@ 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 = 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", "panic", "-n", "-f", "rawvideo", "-pix_fmt", "rgb24", "-s",
-                    dimensions, "-r", str(frame_rate), "-i", "-", "-pix_fmt", "yuv420p", 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]
 
-            with subprocess.Popen(args, stdin=subprocess.PIPE) as proc:
+            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())
 
diff --git a/video_formats/av1-webm.json b/video_formats/av1-webm.json
new file mode 100644
index 0000000..137ad18
--- /dev/null
+++ b/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"}
+}
diff --git a/video_formats/h264-mp4.json b/video_formats/h264-mp4.json
new file mode 100644
index 0000000..6b50b12
--- /dev/null
+++ b/video_formats/h264-mp4.json
@@ -0,0 +1,9 @@
+{
+    "main_pass":
+    [
+        "-n", "-c:v", "libx264",
+        "-pix_fmt", "yuv420p",
+        "-crf", "19"
+    ],
+     "extension": "mp4"
+}
diff --git a/video_formats/h265-mp4.json b/video_formats/h265-mp4.json
new file mode 100644
index 0000000..a8b677b
--- /dev/null
+++ b/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"
+}
diff --git a/video_formats/webm.json b/video_formats/webm.json
new file mode 100644
index 0000000..1551e2c
--- /dev/null
+++ b/video_formats/webm.json
@@ -0,0 +1,9 @@
+{
+    "main_pass":
+    [
+        "-n",
+        "-pix_fmt", "yuv420p",
+        "-crf", "23"
+    ],
+     "extension": "webm"
+}