Add HDR color space options to h264 codec in Save Video node. (#15764)

Author: comfyanonymous

comfy_extras/nodes_video.py

@@ -72,6 +72,16 @@ def execute(cls, images, codec, fps, filename_prefix, crf) -> io.NodeOutput:
 
         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:
@@ -88,11 +98,14 @@ def _save_video_codec_input(supported_codecs: list[str], *, optional=False, hidd
                             io.DynamicCombo.Option("auto", []),
                             io.DynamicCombo.Option(
                                 "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.")],
+                                [
+                                    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(),
+                                ],
                             ),
                         ],
                         optional=True,
-                        tooltip="Automatic preserves compatible H.264 streams. Re-encode applies a custom CRF.",
+                        tooltip="Automatic preserves compatible H.264 streams. Re-encode applies custom encoding options.",
                     ),
                 ],
             )
@@ -111,13 +124,7 @@ 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."),
-                                    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.",
-                                    ),
+                                    _save_video_color_space_input(),
                                 ],
                             ),
                         ],
@@ -131,7 +138,7 @@ def _save_video_codec_input(supported_codecs: list[str], *, optional=False, hidd
         "codec",
         options=codec_options,
         optional=optional,
-        tooltip="The output video codec. Auto preserves a compatible source stream. H.264 re-encoding supports SDR; AV1 re-encoding supports SDR, HDR (HLG), and HDR PQ.",
+        tooltip="The output video codec. Auto preserves a compatible source stream. H.264 and AV1 re-encoding support SDR, HDR (HLG), and HDR PQ.",
         extra_dict={"hidden": True} if hidden else None,
     )
 

tests-unit/comfy_api_test/video_types_test.py

@@ -477,6 +477,48 @@ def test_save_to_av1_mkv_color_space(tmp_path, color_space, transfer, pix_fmt, p
     assert video_packet_bytes(remuxed) == source_packets
 
 
+@pytest.mark.parametrize(
+    "format,suffix",
+    [
+        (VideoContainer.MP4, "mp4"),
+        (VideoContainer.MKV, "mkv"),
+    ],
+)
+@pytest.mark.parametrize(
+    "color_space,transfer,pix_fmt,primaries,colorspace",
+    [
+        ("sRGB", ColorTrc.IEC61966_2_1, "yuv420p", ColorPrimaries.BT709, 1),
+        ("HDR", ColorTrc.ARIB_STD_B67, "yuv420p10le", ColorPrimaries.BT2020, 9),
+        ("HDR PQ", ColorTrc.SMPTE2084, "yuv420p10le", ColorPrimaries.BT2020, 9),
+    ],
+)
+def test_save_to_h264_color_space(tmp_path, format, suffix, color_space, transfer, pix_fmt, primaries, colorspace):
+    components = VideoComponents(
+        images=torch.rand(2, 64, 64, 3),
+        frame_rate=Fraction(30),
+    )
+    path = str(tmp_path / f"h264.{suffix}")
+
+    VideoFromComponents(components).save_to(
+        path,
+        format=format,
+        codec=VideoCodec.H264,
+        crf=23,
+        color_space=color_space,
+    )
+
+    with av.open(path) as container:
+        stream = container.streams.video[0]
+        assert stream.codec.canonical_name == "h264"
+        assert stream.format.name == pix_fmt
+        assert stream.color_primaries == primaries
+        assert stream.color_trc == transfer
+        assert stream.colorspace == colorspace
+        assert stream.color_range == ColorRange.MPEG
+        assert sum(1 for _ in container.decode(video=0)) == 2
+    assert VideoFromFile(path).get_color_space() == color_space
+
+
 @pytest.mark.parametrize(
     "transfer,color_range,color_space",
     [