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

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)]:

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