From 9db05e0e1f035d1902ffc256fe7a336e549ced34 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Sat, 22 Aug 2026 18:07:29 -0700 Subject: [PATCH] Add colorspace option and change bit_depth to a combo on CreateVideo. (#15810) --- comfy_api/latest/_input/video_types.py | 4 +- comfy_api/latest/_input_impl/video_types.py | 13 ++-- comfy_extras/nodes_video.py | 46 ++++++-------- .../comfy_api_test/video_bit_depth_test.py | 62 +++++++++++++++++++ tests-unit/comfy_api_test/video_types_test.py | 42 ++++++++----- 5 files changed, 117 insertions(+), 50 deletions(-) diff --git a/comfy_api/latest/_input/video_types.py b/comfy_api/latest/_input/video_types.py index eb99518f7..2fbdbb341 100644 --- a/comfy_api/latest/_input/video_types.py +++ b/comfy_api/latest/_input/video_types.py @@ -37,8 +37,8 @@ class VideoInput(ABC): bit_depth selects the encoded bit depth; None keeps the video's native depth. crf selects the H.264 or AV1 constant rate factor; None uses the encoder default. - color_space="sRGB" writes SDR BT.709/sRGB video. "HDR" writes 10-bit BT.2020/HLG video; - "HDR PQ" selects BT.2020/PQ. + color_space="sRGB" selects SDR BT.709/sRGB, "HDR" selects BT.2020/HLG, and "HDR PQ" + selects BT.2020/PQ. Bit depth is selected independently. Tensor-created videos default to sRGB when color_space is None. Loaded videos keep matching recognized native color properties; other input pixels must already use the selected color space. """ diff --git a/comfy_api/latest/_input_impl/video_types.py b/comfy_api/latest/_input_impl/video_types.py index d486cd838..34c6143ef 100644 --- a/comfy_api/latest/_input_impl/video_types.py +++ b/comfy_api/latest/_input_impl/video_types.py @@ -644,8 +644,6 @@ class VideoFromFile(VideoInput): audio_stream = last_decodable_audio_stream(container) source_color_space = video_stream_color_space(video_stream) preserve_source_color = source_color_space is not None - if color_space in HDR_COLOR_TRANSFERS or source_color_space in HDR_COLOR_TRANSFERS: - bit_depth = max(bit_depth, 10) pix_fmt = "yuv420p10le" if bit_depth >= 10 else "yuv420p" rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1) @@ -932,10 +930,13 @@ class VideoFromComponents(VideoInput): Class representing video input from tensors. """ - def __init__(self, components: VideoComponents, bit_depth: int = 8): + def __init__(self, components: VideoComponents, bit_depth: int = 8, color_space: str = "sRGB"): + if color_space not in VIDEO_COLOR_TRANSFERS: + raise ValueError(f"Unsupported video color space: {color_space}") self.__components = components # Tensor components have no inherent bit depth; this is the depth used when encoding. self.__bit_depth = bit_depth + self.__color_space = color_space def get_components(self) -> VideoComponents: return VideoComponents( @@ -948,7 +949,7 @@ class VideoFromComponents(VideoInput): return self.__bit_depth def get_color_space(self) -> str: - return "sRGB" + return self.__color_space def save_to( self, @@ -962,15 +963,13 @@ class VideoFromComponents(VideoInput): ): """Save the video to a file path or BytesIO buffer.""" if color_space is None: - color_space = "sRGB" + color_space = self.__color_space if color_space is not None and color_space not in VIDEO_COLOR_TRANSFERS: raise ValueError(f"Unsupported video color space: {color_space}") open_kwargs, output_format, output_codec = video_output_config(path, format, codec) # None means "use the depth this video was created with" (CreateVideo's choice). if bit_depth is None: bit_depth = self.__bit_depth - if color_space in HDR_COLOR_TRANSFERS: - bit_depth = max(bit_depth, 10) is_10bit = bit_depth >= 10 with av.open(path, **open_kwargs) as output: # Add metadata before writing any streams diff --git a/comfy_extras/nodes_video.py b/comfy_extras/nodes_video.py index 72fc387db..58f58aaf4 100644 --- a/comfy_extras/nodes_video.py +++ b/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"), ], ) diff --git a/tests-unit/comfy_api_test/video_bit_depth_test.py b/tests-unit/comfy_api_test/video_bit_depth_test.py index 6c7bc9163..715a09129 100644 --- a/tests-unit/comfy_api_test/video_bit_depth_test.py +++ b/tests-unit/comfy_api_test/video_bit_depth_test.py @@ -3,8 +3,10 @@ import torch import av import numpy as np from fractions import Fraction +from types import SimpleNamespace from comfy_api.latest._input_impl.video_types import VideoFromFile, VideoFromComponents from comfy_api.latest._util.video_types import VideoComponents +from comfy_extras.nodes_video import CreateVideo, SaveVideo @pytest.fixture(scope="module") @@ -58,6 +60,66 @@ def test_create_video_bit_depth(src8, src10): assert decoded_levels(src10) > 2 * decoded_levels(src8) +@pytest.mark.parametrize( + "bit_depth,color_space,expected_bit_depth", + [ + ("auto", "sRGB", 8), + ("auto", "HDR", 10), + ("auto", "HDR PQ", 10), + (8, "HDR", 8), + (10, "sRGB", 10), + ], +) +def test_create_video_node_bit_depth(gradient_components, bit_depth, color_space, expected_bit_depth): + video = CreateVideo.execute( + gradient_components.images, + float(gradient_components.frame_rate), + bit_depth=bit_depth, + color_space=color_space, + ).args[0] + assert video.get_bit_depth() == expected_bit_depth + assert video.get_color_space() == color_space + + +def test_create_video_node_bit_depth_options(): + bit_depth_input = next(input for input in CreateVideo.define_schema().inputs if input.id == "bit_depth") + assert bit_depth_input.options == ["auto", 8, 10] + assert bit_depth_input.default == "auto" + + +@pytest.mark.parametrize( + "codec,expected_suffix,expected_codec", + [ + ("auto", "mp4", "h264"), + ("h264", "mp4", "h264"), + ("av1", "webm", "av1"), + ], +) +def test_save_video_auto_format(gradient_components, tmp_path, monkeypatch, codec, expected_suffix, expected_codec): + monkeypatch.setattr(SaveVideo, "hidden", SimpleNamespace(prompt=None, extra_pnginfo=None)) + monkeypatch.setattr("comfy_extras.nodes_video.folder_paths.get_output_directory", lambda: str(tmp_path)) + monkeypatch.setattr( + "comfy_extras.nodes_video.folder_paths.get_save_image_path", + lambda *args: (str(tmp_path), "auto", 1, "", "auto"), + ) + video = VideoFromComponents(gradient_components) + + SaveVideo.execute( + video, + "auto", + {"format": "auto", "codec": {"codec": codec}}, + ) + + path = tmp_path / f"auto_00001_.{expected_suffix}" + with av.open(path) as container: + assert container.streams.video[0].codec.canonical_name == expected_codec + + +def test_save_video_has_no_color_space_input(): + schema = SaveVideo.define_schema() + assert all("color_space" not in str(input.as_dict()) for input in schema.inputs) + + def test_save_auto_keeps_source_depth(src8, src10, tmp_path): """Save Video (no bit_depth = auto) stream-copies the source, preserving its depth byte-for-byte""" for name, src in [("p8", src8), ("p10", src10)]: diff --git a/tests-unit/comfy_api_test/video_types_test.py b/tests-unit/comfy_api_test/video_types_test.py index 6ac882229..fe7b8de2a 100644 --- a/tests-unit/comfy_api_test/video_types_test.py +++ b/tests-unit/comfy_api_test/video_types_test.py @@ -139,6 +139,19 @@ def test_video_color_space_defaults_to_srgb(simple_video_file, video_components) assert VideoFromComponents(video_components).get_color_space() == "sRGB" +@pytest.mark.parametrize("color_space", ["sRGB", "HDR", "HDR PQ"]) +@pytest.mark.parametrize("bit_depth", [8, 10]) +def test_video_from_components_color_space(video_components, color_space, bit_depth): + video = VideoFromComponents(video_components, bit_depth=bit_depth, color_space=color_space) + assert video.get_color_space() == color_space + assert video.get_bit_depth() == bit_depth + + +def test_video_from_components_rejects_invalid_color_space(video_components): + with pytest.raises(ValueError, match="Unsupported video color space"): + VideoFromComponents(video_components, color_space="Display P3") + + def test_video_from_file_bytesio_input(): """VideoFromFile works with BytesIO input""" buffer = io.BytesIO() @@ -426,14 +439,15 @@ def test_save_components_container_codec_and_audio_matrix( @pytest.mark.parametrize( - "color_space,transfer,pix_fmt,primaries,colorspace", + "color_space,transfer,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), + ("sRGB", ColorTrc.IEC61966_2_1, ColorPrimaries.BT709, 1), + ("HDR", ColorTrc.ARIB_STD_B67, ColorPrimaries.BT2020, 9), + ("HDR PQ", ColorTrc.SMPTE2084, ColorPrimaries.BT2020, 9), ], ) -def test_save_to_av1_mkv_color_space(tmp_path, color_space, transfer, pix_fmt, primaries, colorspace): +@pytest.mark.parametrize("bit_depth,pix_fmt", [(8, "yuv420p"), (10, "yuv420p10le")]) +def test_save_to_av1_mkv_color_space(tmp_path, color_space, transfer, primaries, colorspace, bit_depth, pix_fmt): components = VideoComponents( images=torch.rand(2, 64, 64, 3), frame_rate=Fraction(30), @@ -441,12 +455,11 @@ def test_save_to_av1_mkv_color_space(tmp_path, color_space, transfer, pix_fmt, p path = str(tmp_path / "hdr.mkv") remuxed = str(tmp_path / "remuxed.mkv") - VideoFromComponents(components).save_to( + VideoFromComponents(components, bit_depth=bit_depth, color_space=color_space).save_to( path, format=VideoContainer.MKV, codec=VideoCodec.AV1, crf=30, - color_space=color_space, metadata={"prompt": {"test": "hdr"}}, ) @@ -485,21 +498,22 @@ def test_save_to_av1_mkv_color_space(tmp_path, color_space, transfer, pix_fmt, p ], ) @pytest.mark.parametrize( - "color_space,transfer,pix_fmt,primaries,colorspace", + "color_space,transfer,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), + ("sRGB", ColorTrc.IEC61966_2_1, ColorPrimaries.BT709, 1), + ("HDR", ColorTrc.ARIB_STD_B67, ColorPrimaries.BT2020, 9), + ("HDR PQ", ColorTrc.SMPTE2084, ColorPrimaries.BT2020, 9), ], ) -def test_save_to_h264_color_space(tmp_path, format, suffix, color_space, transfer, pix_fmt, primaries, colorspace): +@pytest.mark.parametrize("bit_depth,pix_fmt", [(8, "yuv420p"), (10, "yuv420p10le")]) +def test_save_to_h264_color_space(tmp_path, format, suffix, color_space, transfer, primaries, colorspace, bit_depth, pix_fmt): components = VideoComponents( images=torch.rand(2, 64, 64, 3), frame_rate=Fraction(30), ) path = str(tmp_path / f"h264.{suffix}") - VideoFromComponents(components).save_to( + VideoFromComponents(components, bit_depth=bit_depth).save_to( path, format=format, codec=VideoCodec.H264, @@ -604,7 +618,7 @@ def test_save_to_av1_webm_transcodes_audio(tmp_path): with av.open(path) as container: video_stream = container.streams.video[0] assert video_stream.codec.canonical_name == "av1" - assert video_stream.format.name == "yuv420p10le" + assert video_stream.format.name == "yuv420p" assert video_stream.color_primaries == ColorPrimaries.BT2020 assert video_stream.color_trc == ColorTrc.ARIB_STD_B67 assert video_stream.colorspace == 9