Support HDR video saving, AV1 codec, mkv and webm. (#15741)

2026-08-24 18:10:27 +08:00 · 2026-08-20 16:29:20 -07:00
parent de6b062fb5
commit dcbcf8c2e1
6 changed files with 772 additions and 73 deletions
--- a/tests-unit/comfy_api_test/input_impl_test.py
+++ b/tests-unit/comfy_api_test/input_impl_test.py
@@ -2,8 +2,9 @@ import io
 from comfy_api.input_impl.video_types import (
    container_to_output_format,
    get_open_write_kwargs,
+    video_encoder_options,
 )
-from comfy_api.util import VideoContainer
+from comfy_api.util import VideoCodec, VideoContainer


 def test_container_to_output_format_empty_string():
@@ -36,7 +37,7 @@ def test_get_open_write_kwargs_filepath_no_format():
    kwargs_specific = get_open_write_kwargs("output.avi", "mp4", "avi")
    fail_msg = "Format should not be set for file paths (Specific)"
    assert "format" not in kwargs_specific, fail_msg
-    assert kwargs_specific["options"]["movflags"] == "use_metadata_tags"
+    assert "options" not in kwargs_specific


 def test_get_open_write_kwargs_base_options_mode():
@@ -90,3 +91,16 @@ def test_get_open_write_kwargs_bytesio_specific_format_list():

    fail_msg = "Format should be a valid format from the specified format list when output format is not AUTO"
    assert kwargs["format"] in to_fmt, fail_msg
+
+
+def test_get_open_write_kwargs_does_not_pass_movflags_to_matroska_or_webm():
+    for format, suffix in ((VideoContainer.MKV, "mkv"), (VideoContainer.WEBM, "webm")):
+        assert "options" not in get_open_write_kwargs(f"output.{suffix}", "mp4", format)
+        assert "options" not in get_open_write_kwargs(io.BytesIO(), "mp4", format)
+
+
+def test_av1_zero_crf_uses_lossless_mode():
+    assert video_encoder_options(VideoCodec.AV1, 0) == {"svtav1-params": "lossless=1"}
+    assert video_encoder_options(VideoCodec.AV1, 30.0) == {"crf": "30.0"}
+    assert video_encoder_options(VideoCodec.H264, 0) == {"crf": "0"}
+    assert video_encoder_options(VideoCodec.AV1, None) == {}
--- a/tests-unit/comfy_api_test/video_types_test.py
+++ b/tests-unit/comfy_api_test/video_types_test.py
@@ -5,11 +5,13 @@ import os
 import sys
 import av
 import io
+import numpy as np
 from fractions import Fraction
 from comfy_api.input_impl.video_types import VideoFromFile, VideoFromComponents
 from comfy_api.util.video_types import VideoComponents, VideoContainer, VideoCodec
 from comfy_api.input.basic_types import AudioInput
 from av.error import InvalidDataError
+from av.video.reformatter import ColorPrimaries, ColorRange, ColorTrc

 EPSILON = 0.0001

@@ -132,6 +134,11 @@ def test_video_from_file_get_dimensions(simple_video_file):
    assert height == 4


+def test_video_color_space_defaults_to_srgb(simple_video_file, video_components):
+    assert VideoFromFile(simple_video_file).get_color_space() == "sRGB"
+    assert VideoFromComponents(video_components).get_color_space() == "sRGB"
+
+
 def test_video_from_file_bytesio_input():
    """VideoFromFile works with BytesIO input"""
    buffer = io.BytesIO()
@@ -258,6 +265,456 @@ def test_save_to_h264_crf_controls_quality(tmp_path):
    assert os.path.getsize(transcoded) < os.path.getsize(high_quality)


+def video_packet_bytes(path):
+    with av.open(path) as container:
+        return [bytes(packet) for packet in container.demux(container.streams.video[0]) if packet.size]
+
+
+def decoded_video_frames(path):
+    with av.open(path) as container:
+        frames = []
+        for frame in container.decode(video=0):
+            bytes_per_sample = max(component.bits for component in frame.format.components)
+            bytes_per_sample = (bytes_per_sample + 7) // 8
+            plane_sizes = (
+                (frame.width * bytes_per_sample, frame.height),
+                ((frame.width // 2) * bytes_per_sample, frame.height // 2),
+                ((frame.width // 2) * bytes_per_sample, frame.height // 2),
+            )
+            frames.append(tuple(
+                b"".join(
+                    bytes(plane)[row * plane.line_size:row * plane.line_size + row_size]
+                    for row in range(rows)
+                )
+                for plane, (row_size, rows) in zip(frame.planes, plane_sizes)
+            ))
+        return frames
+
+
+@pytest.mark.parametrize(
+    "format,suffix,codec",
+    [
+        (VideoContainer.MP4, "mp4", VideoCodec.H264),
+        (VideoContainer.MKV, "mkv", VideoCodec.H264),
+        (VideoContainer.MKV, "mkv", VideoCodec.AV1),
+        (VideoContainer.WEBM, "webm", VideoCodec.AV1),
+    ],
+)
+def test_video_from_components_auto_color_space_matches_srgb(tmp_path, format, suffix, codec):
+    components = VideoComponents(
+        images=torch.rand(3, 64, 64, 3, generator=torch.Generator().manual_seed(23)),
+        frame_rate=Fraction(30),
+    )
+    auto = str(tmp_path / f"auto.{suffix}")
+    srgb = str(tmp_path / f"srgb.{suffix}")
+
+    VideoFromComponents(components).save_to(
+        auto,
+        format=format,
+        codec=codec,
+        crf=0,
+    )
+    VideoFromComponents(components).save_to(
+        srgb,
+        format=format,
+        codec=codec,
+        crf=0,
+        color_space="sRGB",
+    )
+
+    assert decoded_video_frames(auto) == decoded_video_frames(srgb)
+    for path in (auto, srgb):
+        with av.open(path) as container:
+            stream = container.streams.video[0]
+            assert stream.color_primaries == ColorPrimaries.BT709
+            assert stream.color_trc == ColorTrc.IEC61966_2_1
+            assert stream.colorspace == 1
+            assert stream.color_range == ColorRange.MPEG
+
+
+def create_hdr_av1_video(path, transfer, color_range):
+    images = np.random.default_rng(17).integers(0, 65536, (3, 64, 64, 3), dtype=np.uint16)
+    with av.open(path, mode="w") as container:
+        stream = container.add_stream("libsvtav1", rate=30)
+        stream.width = 64
+        stream.height = 64
+        stream.pix_fmt = "yuv420p10le"
+        stream.options = {"svtav1-params": "lossless=1"}
+        stream.color_primaries = ColorPrimaries.BT2020
+        stream.color_trc = transfer
+        stream.colorspace = 9
+        stream.color_range = color_range
+        for image in images:
+            frame = av.VideoFrame.from_ndarray(image, format="rgb48le").reformat(format="yuv420p10le")
+            frame.color_primaries = ColorPrimaries.BT2020
+            frame.color_trc = transfer
+            frame.colorspace = 9
+            frame.color_range = color_range
+            container.mux(stream.encode(frame))
+        container.mux(stream.encode(None))
+
+
+def test_save_to_av1_crf_controls_quality(tmp_path):
+    generator = torch.Generator().manual_seed(11)
+    components = VideoComponents(
+        images=torch.rand(12, 64, 64, 3, generator=generator),
+        frame_rate=Fraction(30),
+    )
+    high_quality = str(tmp_path / "high_quality.mkv")
+    low_quality = str(tmp_path / "low_quality.mkv")
+
+    VideoFromComponents(components).save_to(
+        high_quality,
+        format=VideoContainer.MKV,
+        codec=VideoCodec.AV1,
+        crf=0,
+    )
+    VideoFromComponents(components).save_to(
+        low_quality,
+        format=VideoContainer.MKV,
+        codec=VideoCodec.AV1,
+        crf=63,
+    )
+
+    assert os.path.getsize(high_quality) > os.path.getsize(low_quality)
+
+
+@pytest.mark.parametrize(
+    "format,suffix,codec,video_codec,audio_codec,audio_rate",
+    [
+        (VideoContainer.AUTO, "mp4", VideoCodec.AUTO, "h264", "aac", 44100),
+        (VideoContainer.MP4, "mp4", VideoCodec.H264, "h264", "aac", 44100),
+        (VideoContainer.MP4, "mp4", VideoCodec.AV1, "av1", "aac", 44100),
+        (VideoContainer.MKV, "mkv", VideoCodec.AUTO, "h264", "aac", 44100),
+        (VideoContainer.MKV, "mkv", VideoCodec.H264, "h264", "aac", 44100),
+        (VideoContainer.MKV, "mkv", VideoCodec.AV1, "av1", "aac", 44100),
+        (VideoContainer.WEBM, "webm", VideoCodec.AUTO, "av1", "opus", 48000),
+        (VideoContainer.WEBM, "webm", VideoCodec.AV1, "av1", "opus", 48000),
+    ],
+)
+def test_save_components_container_codec_and_audio_matrix(
+    tmp_path, format, suffix, codec, video_codec, audio_codec, audio_rate
+):
+    components = VideoComponents(
+        images=torch.rand(3, 64, 64, 3),
+        audio=AudioInput({
+            "waveform": torch.rand(1, 2, 4410),
+            "sample_rate": 44100,
+        }),
+        frame_rate=Fraction(30),
+    )
+    path = str(tmp_path / f"components.{suffix}")
+
+    VideoFromComponents(components).save_to(
+        path,
+        format=format,
+        codec=codec,
+        crf=30,
+        metadata={"prompt": {"test": "video"}},
+    )
+
+    with av.open(path) as container:
+        assert container.streams.video[0].codec.canonical_name == video_codec
+        assert container.streams.video[0].format.name == "yuv420p"
+        assert container.streams.audio[0].codec.canonical_name == audio_codec
+        assert container.streams.audio[0].sample_rate == audio_rate
+        prompt = container.metadata.get("PROMPT", container.metadata.get("prompt"))
+        assert prompt == '{"test": "video"}'
+        assert sum(1 for _ in container.decode(video=0)) == 3
+    with av.open(path) as container:
+        assert sum(frame.samples for frame in container.decode(audio=0)) > 0
+
+
+@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_av1_mkv_color_space(tmp_path, color_space, transfer, pix_fmt, primaries, colorspace):
+    components = VideoComponents(
+        images=torch.rand(2, 64, 64, 3),
+        frame_rate=Fraction(30),
+    )
+    path = str(tmp_path / "hdr.mkv")
+    remuxed = str(tmp_path / "remuxed.mkv")
+
+    VideoFromComponents(components).save_to(
+        path,
+        format=VideoContainer.MKV,
+        codec=VideoCodec.AV1,
+        crf=30,
+        color_space=color_space,
+        metadata={"prompt": {"test": "hdr"}},
+    )
+
+    with av.open(path) as container:
+        stream = container.streams.video[0]
+        assert stream.codec.canonical_name == "av1"
+        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 container.metadata["PROMPT"] == '{"test": "hdr"}'
+    assert VideoFromFile(path).get_color_space() == color_space
+
+    source_packets = video_packet_bytes(path)
+    VideoFromFile(path).save_to(
+        remuxed,
+        format=VideoContainer.MKV,
+        codec=VideoCodec.AV1,
+    )
+    with av.open(remuxed) as container:
+        stream = container.streams.video[0]
+        assert stream.codec.canonical_name == "av1"
+        assert stream.format.name == pix_fmt
+        assert stream.color_primaries == primaries
+        assert stream.color_trc == transfer
+        assert container.metadata["PROMPT"] == '{"test": "hdr"}'
+    assert video_packet_bytes(remuxed) == source_packets
+
+
+@pytest.mark.parametrize(
+    "transfer,color_range,color_space",
+    [
+        (ColorTrc.SMPTE2084, ColorRange.MPEG, None),
+        (ColorTrc.SMPTE2084, ColorRange.MPEG, "HDR PQ"),
+        (ColorTrc.ARIB_STD_B67, ColorRange.MPEG, None),
+        (ColorTrc.ARIB_STD_B67, ColorRange.MPEG, "HDR"),
+        (ColorTrc.ARIB_STD_B67, ColorRange.JPEG, "HDR"),
+    ],
+)
+def test_save_to_loaded_hdr_preserves_color(tmp_path, transfer, color_range, color_space):
+    source = str(tmp_path / "source.mkv")
+    remuxed = str(tmp_path / "auto_encoding.mkv")
+    reencoded = str(tmp_path / "auto_color_space.webm")
+    create_hdr_av1_video(source, transfer, color_range)
+
+    video = VideoFromFile(source)
+    video.save_to(remuxed, format=VideoContainer.MKV, codec=VideoCodec.AV1)
+    video.save_to(
+        reencoded,
+        format=VideoContainer.WEBM,
+        codec=VideoCodec.AV1,
+        crf=0,
+        color_space=color_space,
+    )
+
+    assert video_packet_bytes(remuxed) == video_packet_bytes(source)
+    source_frames = decoded_video_frames(source)
+    reencoded_frames = decoded_video_frames(reencoded)
+    assert len(source_frames) == len(reencoded_frames)
+    assert all(np.array_equal(source_frame, output_frame) for source_frame, output_frame in zip(source_frames, reencoded_frames))
+
+    for path in (remuxed, reencoded):
+        with av.open(path) as container:
+            stream = container.streams.video[0]
+            assert stream.codec.canonical_name == "av1"
+            assert stream.format.name == "yuv420p10le"
+            assert stream.color_primaries == ColorPrimaries.BT2020
+            assert stream.color_trc == transfer
+            assert stream.colorspace == 9
+            assert stream.color_range == color_range
+
+
+@pytest.mark.parametrize("color_space", ["sRGB", "HDR PQ"])
+def test_save_to_loaded_hdr_rejects_color_conversion(tmp_path, color_space):
+    source = str(tmp_path / "source.mkv")
+    output = str(tmp_path / "wrong_transfer.webm")
+    create_hdr_av1_video(source, ColorTrc.ARIB_STD_B67, ColorRange.MPEG)
+
+    with pytest.raises(ValueError, match=f"Cannot save HDR video as {color_space} without color conversion"):
+        VideoFromFile(source).save_to(
+            output,
+            format=VideoContainer.WEBM,
+            codec=VideoCodec.AV1,
+            crf=30,
+            color_space=color_space,
+        )
+
+    assert not os.path.exists(output)
+
+
+def test_save_to_av1_webm_transcodes_audio(tmp_path):
+    components = VideoComponents(
+        images=torch.rand(2, 64, 64, 3),
+        audio=AudioInput({
+            "waveform": torch.rand(1, 2, 4410),
+            "sample_rate": 44100,
+        }),
+        frame_rate=Fraction(30),
+    )
+    source = str(tmp_path / "source.mp4")
+    path = str(tmp_path / "output.webm")
+    VideoFromComponents(components).save_to(source, color_space="HDR")
+
+    VideoFromFile(source).save_to(
+        path,
+        format=VideoContainer.WEBM,
+        codec=VideoCodec.AV1,
+        crf=30,
+        color_space="HDR",
+    )
+
+    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.color_primaries == ColorPrimaries.BT2020
+        assert video_stream.color_trc == ColorTrc.ARIB_STD_B67
+        assert video_stream.colorspace == 9
+        assert container.streams.audio[0].codec.name == "opus"
+        assert container.streams.audio[0].sample_rate == 48000
+        assert sum(1 for _ in container.decode(video=0)) == 2
+    with av.open(path) as container:
+        assert sum(frame.samples for frame in container.decode(audio=0)) > 0
+
+
+@pytest.mark.parametrize("source_codec", [VideoCodec.H264, VideoCodec.AV1])
+def test_save_loaded_mkv_to_webm_auto_transcodes_incompatible_streams(tmp_path, source_codec):
+    components = VideoComponents(
+        images=torch.rand(3, 64, 64, 3),
+        audio=AudioInput({
+            "waveform": torch.rand(1, 2, 4410),
+            "sample_rate": 44100,
+        }),
+        frame_rate=Fraction(30),
+    )
+    source = str(tmp_path / "source.mkv")
+    output = str(tmp_path / "output.webm")
+    VideoFromComponents(components).save_to(
+        source,
+        format=VideoContainer.MKV,
+        codec=source_codec,
+        crf=63 if source_codec == VideoCodec.AV1 else 30,
+    )
+
+    VideoFromFile(source).save_to(
+        output,
+        format=VideoContainer.WEBM,
+        codec=VideoCodec.AUTO,
+    )
+
+    with av.open(output) as container:
+        assert container.streams.video[0].codec.canonical_name == "av1"
+        assert container.streams.audio[0].codec.canonical_name == "opus"
+        assert container.streams.audio[0].sample_rate == 48000
+        assert sum(1 for _ in container.decode(video=0)) == 3
+    with av.open(output) as container:
+        assert sum(frame.samples for frame in container.decode(audio=0)) > 0
+
+
+def test_save_loaded_h264_mkv_to_webm_h264_rejected_before_creating_output(tmp_path):
+    components = VideoComponents(
+        images=torch.rand(1, 64, 64, 3),
+        frame_rate=Fraction(30),
+    )
+    source = str(tmp_path / "source.mkv")
+    output = tmp_path / "output.webm"
+    VideoFromComponents(components).save_to(
+        source,
+        format=VideoContainer.MKV,
+        codec=VideoCodec.H264,
+    )
+
+    with pytest.raises(ValueError, match="WebM output requires the AV1 codec"):
+        VideoFromFile(source).save_to(
+            str(output),
+            format=VideoContainer.WEBM,
+            codec=VideoCodec.H264,
+        )
+
+    assert not output.exists()
+
+
+def test_save_loaded_webm_auto_remuxes_compatible_streams(tmp_path):
+    components = VideoComponents(
+        images=torch.rand(3, 64, 64, 3),
+        audio=AudioInput({
+            "waveform": torch.rand(1, 2, 4410),
+            "sample_rate": 44100,
+        }),
+        frame_rate=Fraction(30),
+    )
+    source = str(tmp_path / "source.webm")
+    output = str(tmp_path / "output.webm")
+    VideoFromComponents(components).save_to(
+        source,
+        format=VideoContainer.WEBM,
+        codec=VideoCodec.AV1,
+        crf=63,
+    )
+    source_packets = video_packet_bytes(source)
+
+    VideoFromFile(source).save_to(
+        output,
+        format=VideoContainer.WEBM,
+        codec=VideoCodec.AUTO,
+    )
+
+    assert video_packet_bytes(output) == source_packets
+    with av.open(output) as container:
+        assert container.streams.video[0].codec.canonical_name == "av1"
+        assert container.streams.audio[0].codec.canonical_name == "opus"
+        assert sum(1 for _ in container.decode(video=0)) == 3
+
+
+@pytest.mark.parametrize("format", [VideoContainer.MKV, VideoContainer.WEBM])
+def test_save_to_av1_file_like_output(format):
+    components = VideoComponents(
+        images=torch.rand(1, 64, 64, 3),
+        frame_rate=Fraction(30),
+    )
+    output = io.BytesIO()
+
+    VideoFromComponents(components).save_to(
+        output,
+        format=format,
+        codec=VideoCodec.AV1,
+        crf=63,
+    )
+
+    output.seek(0)
+    with av.open(output) as container:
+        assert container.streams.video[0].codec.canonical_name == "av1"
+        assert sum(1 for _ in container.decode(video=0)) == 1
+
+
+def test_save_to_rejects_h264_webm_before_creating_output(tmp_path):
+    components = VideoComponents(
+        images=torch.rand(1, 64, 64, 3),
+        frame_rate=Fraction(30),
+    )
+    path = tmp_path / "invalid.webm"
+
+    with pytest.raises(ValueError, match="WebM output requires the AV1 codec"):
+        VideoFromComponents(components).save_to(
+            str(path),
+            format=VideoContainer.WEBM,
+            codec=VideoCodec.H264,
+        )
+
+    assert not path.exists()
+
+
+def test_save_to_rejects_unknown_color_space(tmp_path):
+    components = VideoComponents(
+        images=torch.rand(1, 64, 64, 3),
+        frame_rate=Fraction(30),
+    )
+
+    with pytest.raises(ValueError, match="Unsupported video color space: HLG"):
+        VideoFromComponents(components).save_to(
+            str(tmp_path / "invalid.mkv"),
+            format=VideoContainer.MKV,
+            codec=VideoCodec.AV1,
+            color_space="HLG",
+        )
+
+
 def test_save_to_mp4_writes_metadata_before_media(video_components, tmp_path):
    encoded = tmp_path / "encoded.mp4"
    remuxed = tmp_path / "remuxed.mp4"