feat: 💫 export to prores -> export with ffmpeg

melMass · Aug 11, 2023 · a4d99d9 · a4d99d9
1 parent 4fc84d6
commit a4d99d9
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 11 deletions.
diff --git a/node_list.json b/node_list.json
@@ -1,5 +1,6 @@
 {
     "Animation Builder (mtb)": "Convenient way to manage basic animation maths at the core of many of my workflows",
+    "Any To String (mtb)": "Tries to take any input and convert it to a string",
     "Bbox (mtb)": "The bounding box (BBOX) custom type used by other nodes",
     "Bbox From Mask (mtb)": "From a mask extract the bounding box",
     "Blur (mtb)": "Blur an image using a Gaussian filter.",
@@ -9,7 +10,7 @@
     "Crop (mtb)": "Crops an image and an optional mask to a given bounding box\n\n    The bounding box can be given as a tuple of (x, y, width, height) or as a BBOX type\n    The BBOX input takes precedence over the tuple input\n    ",
     "Debug (mtb)": "Experimental node to debug any Comfy values, support for more types and widgets is planned",
     "Deep Bump (mtb)": "Normal & height maps generation from single pictures",
-    "Export To Prores (mtb)": "Export to ProRes 4444 (Experimental)",
+    "Export With Ffmpeg (mtb)": "Export with FFmpeg (Experimental)",
     "Face Swap (mtb)": "Face swap using deepinsight/insightface models",
     "Film Interpolation (mtb)": "Google Research FILM frame interpolation for large motion",
     "Fit Number (mtb)": "Fit the input float using a source and target range",

diff --git a/nodes/io.py b/nodes/io.py
@@ -11,8 +11,8 @@
 from typing import Optional, List
 
 
-class ExportToProres:
-    """Export to ProRes 4444 (Experimental)"""
+class ExportWithFfmpeg:
+    """Export with FFmpeg (Experimental)"""
 
     @classmethod
     def INPUT_TYPES(cls):
@@ -22,6 +22,11 @@ def INPUT_TYPES(cls):
                 # "frames": ("FRAMES",),
                 "fps": ("FLOAT", {"default": 24, "min": 1}),
                 "prefix": ("STRING", {"default": "export"}),
+                "format": (["mov", "mp4", "mkv", "avi"], {"default": "mov"}),
+                "codec": (
+                    ["prores_ks", "libx264", "libx265"],
+                    {"default": "prores_ks"},
+                ),
             }
         }
 
@@ -35,11 +40,15 @@ def export_prores(
         images: torch.Tensor,
         fps: float,
         prefix: str,
+        format: str,
+        codec: str,
     ):
         if images.size(0) == 0:
             return ("",)
         output_dir = Path(folder_paths.get_output_directory())
-        file_id = f"{prefix}_{uuid.uuid4()}.mov"
+        pix_fmt = "rgb48le" if codec == "prores_ks" else "yuv420p"
+        file_ext = format
+        file_id = f"{prefix}_{uuid.uuid4()}.{file_ext}"
 
         log.debug(f"Exporting to {output_dir / file_id}")
 
@@ -64,17 +73,13 @@ def export_prores(
             "-s",
             f"{width}x{height}",
             "-pix_fmt",
-            "rgb48le",
+            pix_fmt,
             "-r",
             str(fps),
             "-i",
             "-",
             "-c:v",
-            "prores_ks",
-            "-profile:v",
-            "4",
-            "-pix_fmt",
-            "yuva444p10le",
+            codec,
             "-r",
             str(fps),
             "-y",
@@ -187,4 +192,4 @@ def save_gif(
         return {"ui": {"gif": results}}
 
 
-__nodes__ = [SaveGif, ExportToProres]
+__nodes__ = [SaveGif, ExportWithFfmpeg]