diff --git a/comfy_extras/nodes_video.py b/comfy_extras/nodes_video.py index c9a239989..72fc387db 100644 --- a/comfy_extras/nodes_video.py +++ b/comfy_extras/nodes_video.py @@ -72,6 +72,16 @@ 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: @@ -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, ) diff --git a/tests-unit/comfy_api_test/video_types_test.py b/tests-unit/comfy_api_test/video_types_test.py index f90913e06..6ac882229 100644 --- a/tests-unit/comfy_api_test/video_types_test.py +++ b/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", [