Add colorspace option and change bit_depth to a combo on CreateVideo. (#15810)

comfy_extras/nodes_video.py

@@ -72,16 +72,6 @@ class SaveWEBM(io.ComfyNode):
 
         return io.NodeOutput(images, ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
 
-def _save_video_color_space_input():
-    return io.Combo.Input(
-        "color_space",
-        options=["auto", "sRGB", "HDR", "HDR PQ"],
-        default="auto",
-        display_name="color space",
-        tooltip="Auto uses sRGB for videos created from images and preserves recognized colors on loaded videos. sRGB writes SDR BT.709/sRGB. HDR writes 10-bit BT.2020/HLG; HDR PQ writes BT.2020/PQ. Other input pixels must already use the selected color space.",
-    )
-
-
 def _save_video_codec_input(supported_codecs: list[str], *, optional=False, hidden=False):
     codec_options = []
     if "auto" in supported_codecs:
@@ -100,7 +90,6 @@ def _save_video_codec_input(supported_codecs: list[str], *, optional=False, hidd
                                 "re-encode",
                                 [
                                     io.Float.Input("crf", default=23.0, min=0.0, max=51.0, step=1.0, tooltip="Lower values produce higher quality and larger files."),
-                                    _save_video_color_space_input(),
                                 ],
                             ),
                         ],
@@ -124,7 +113,6 @@ def _save_video_codec_input(supported_codecs: list[str], *, optional=False, hidd
                                 "re-encode",
                                 [
                                     io.Float.Input("crf", default=30.0, min=0.0, max=63.0, step=1.0, tooltip="Lower values produce higher quality and larger files."),
-                                    _save_video_color_space_input(),
                                 ],
                             ),
                         ],
@@ -164,7 +152,7 @@ class SaveVideo(io.ComfyNode):
                         io.DynamicCombo.Option("mkv", [_save_video_codec_input(["auto", "h264", "av1"])]),
                         io.DynamicCombo.Option("webm", [_save_video_codec_input(["auto", "av1"])]),
                     ],
-                    tooltip="The output container. Auto preserves the source container when possible; MP4, MKV, and WebM select a specific container.",
+                    tooltip="The output container. Auto uses MP4 for Auto/H.264 and WebM for AV1. MP4, MKV, and WebM select a specific container.",
                 ),
                 _save_video_codec_input(["auto", "h264", "av1"], optional=True, hidden=True),
             ],
@@ -183,10 +171,9 @@ class SaveVideo(io.ComfyNode):
         if codec is None:
             codec = {"codec": "auto"}
         codec_name = codec["codec"]
+        if format_name == "auto":
+            format_name = "webm" if codec_name == "av1" else "mp4"
         encoding = codec.get("encoding") or {}
-        color_space = encoding.get("color_space")
-        if color_space == "auto":
-            color_space = None
         width, height = video.get_dimensions()
         full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
             filename_prefix,
@@ -210,7 +197,6 @@ class SaveVideo(io.ComfyNode):
             codec=Types.VideoCodec(codec_name),
             metadata=saved_metadata,
             crf=encoding.get("crf"),
-            color_space=color_space,
         )
 
         return io.NodeOutput(video, ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
@@ -230,16 +216,19 @@ class CreateVideo(io.ComfyNode):
                 io.Image.Input("images", tooltip="The images to create a video from."),
                 io.Float.Input("fps", default=30.0, min=1.0, max=120.0, step=1.0),
                 io.Audio.Input("audio", optional=True, tooltip="The audio to add to the video."),
-                io.Int.Input(
+                io.Combo.Input(
                     "bit_depth",
-                    min=8,
-                    max=10,
-                    default=8,
-                    step=2,
-                    tooltip="Bit depth of the created video. 10-bit keeps smoother gradients with less"
-                    " banding, but some players and downstream nodes may not support it.",
+                    options=["auto", 8, 10],
+                    default="auto",
+                    tooltip="Auto uses 8-bit for sRGB and 10-bit for HDR. Explicit 8-bit and 10-bit choices are independent of colorspace.",
                     optional=True,
-                    display_mode=io.NumberDisplay.number,
+                ),
+                io.Combo.Input(
+                    "color_space",
+                    options=["sRGB", "HDR", "HDR PQ"],
+                    default="sRGB",
+                    optional=True,
+                    tooltip="Colorspace of the input images. HDR selects BT.2020/HLG and HDR PQ selects BT.2020/PQ.",
                 ),
             ],
             outputs=[
@@ -249,12 +238,15 @@ class CreateVideo(io.ComfyNode):
 
     @classmethod
     def execute(
-        cls, images: Input.Image, fps: float, audio: Optional[Input.Audio] = None, bit_depth: int = 8,
+        cls, images: Input.Image, fps: float, audio: Optional[Input.Audio] = None, bit_depth: int | str = "auto", color_space: str = "sRGB",
     ) -> io.NodeOutput:
+        if bit_depth == "auto":
+            bit_depth = 10 if color_space in ("HDR", "HDR PQ") else 8
         return io.NodeOutput(
             InputImpl.VideoFromComponents(
                 Types.VideoComponents(images=images, audio=audio, frame_rate=Fraction(fps)),
                 bit_depth=bit_depth,
+                color_space=color_space,
             )
         )
 
@@ -274,7 +266,7 @@ class GetVideoComponents(io.ComfyNode):
                 io.Image.Output(display_name="images"),
                 io.Audio.Output(display_name="audio"),
                 io.Float.Output(display_name="fps"),
-                io.Int.Output(display_name="bit_depth"),
+                io.Combo.Output(display_name="bit_depth"),
                 io.Combo.Output(display_name="color_space"),
             ],
         )