Harden video save compatibility and regressions

Author: christian-byrne

comfy_api/latest/_input_impl/video_types.py

@@ -71,8 +71,7 @@ def get_open_write_kwargs(
     is_write_to_buffer = isinstance(dest, io.BytesIO)
     open_kwargs = {"mode": "w"}
 
-    if is_write_to_buffer:
-        # Set output format explicitly, since it cannot be inferred from file extension
+    if is_write_to_buffer or to_format != VideoContainer.AUTO:
         if to_format == VideoContainer.AUTO:
             to_format = container_format.lower()
         elif isinstance(to_format, VideoContainer):
@@ -81,7 +80,7 @@ def get_open_write_kwargs(
             to_format = to_format.lower()
         open_kwargs["format"] = container_to_output_format(to_format)
 
-    output_format = open_kwargs["format"] if is_write_to_buffer else os.path.splitext(dest)[1].lower().lstrip(".")
+    output_format = open_kwargs["format"] if "format" in open_kwargs else os.path.splitext(dest)[1].lower().lstrip(".")
     if output_format in ("mov", "mp4"):
         # Preserve custom metadata tags (workflow, prompt, extra_pnginfo) in isobmff.
         movflags = "use_metadata_tags" if is_write_to_buffer else "use_metadata_tags+faststart"
@@ -617,10 +616,8 @@ def _save_remuxed(
     ) -> bool:
         streams = container.streams
         with av.open(path, **open_kwargs) as output_container:
-            # Add metadata before writing any streams
             write_output_metadata(container, output_container, metadata)
 
-            # Add streams to the new container. Streams with no codec context cannot be used as an output template.
             stream_map = {}
             for stream in streams:
                 if isinstance(stream, (av.VideoStream, av.AudioStream, SubtitleStream)):
@@ -647,7 +644,6 @@ def _save_remuxed(
                         return False
                     stream_map[stream] = out_stream
 
-            # Write packets to the new container
             for packet in container.demux():
                 if packet.stream in stream_map and packet.dts is not None:
                     packet.stream = stream_map[packet.stream]
@@ -683,6 +679,7 @@ def _save_transcoded(
         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
+        normalize_color_range = video_stream.color_range == ColorRange.JPEG and source_color_space not in HDR_COLOR_TRANSFERS
         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"
@@ -712,15 +709,10 @@ def _save_transcoded(
             if duration:
                 duration_cap = math.ceil(duration * sample_rate)
 
-        # Subtitles are remuxed untouched: there is no subtitle encoder binding, so a stream the
-        # output container cannot store as-is is dropped with a warning naming it, exactly like
-        # the remux path does. Streams FFmpeg has no decoder for cannot template a new stream.
         subtitle_streams = [s for s in container.streams.subtitles if s.codec_context is not None]
         streams = [video_stream] if audio_stream is None else [video_stream, audio_stream]
         streams += subtitle_streams
         subtitle_map = {}
-        # Subtitle packets that arrive before the first kept video frame: the output is not open
-        # yet and the pts rebase offset is not known, so they wait here rather than being lost.
         pending_subtitles = []
         pts_step = max(1, int(round((1 / rate) / video_stream.time_base)))
         video_done = False
@@ -780,15 +772,12 @@ def drain_audio(final=False):
             return cap
 
         def mux_subtitle(packet):
-            """Remux one subtitle packet, rebased onto the trimmed timeline the video was rebased to."""
             out_stream = subtitle_map.get(packet.stream)
             if out_stream is None or packet.dts is None or packet.pts is None or packet.time_base is None:
                 return
             start = float(packet.pts * packet.time_base)
             if start < start_time or (duration and start >= start_time + duration):
                 return
-            # the video's own rebase offset, so subtitles stay in sync with it rather than
-            # with the requested start (a seek lands on the preceding keyframe)
             offset_ticks = video_pts_offset if video_pts_offset is not None else start_pts
             shift = int(round(float(offset_ticks * video_stream.time_base) / packet.time_base))
             packet.pts -= shift
@@ -799,7 +788,6 @@ def mux_subtitle(packet):
             output.mux(packet)
 
         def flush_subtitles():
-            # only once the first video frame fixed the rebase offset
             if output is None or not pending_subtitles or last_video_pts is None:
                 return
             while pending_subtitles:
@@ -853,6 +841,8 @@ def flush_subtitles():
                             out_video.options = video_encoder_options(output_codec, crf)
                             if preserve_source_color:
                                 copy_color_properties(video_stream, out_video.codec_context)
+                                if normalize_color_range:
+                                    out_video.codec_context.color_range = ColorRange.MPEG
                             elif color_space is not None:
                                 set_video_color_properties(out_video.codec_context, color_space)
                             # source pts pass through (rebased to 0), so variable frame rate survives
@@ -897,13 +887,14 @@ def flush_subtitles():
                                 rotation_filter = (g_src, g_sink)
                             rotation_filter[0].push(frame)
                             frame = rotation_filter[1].pull()
-                        if frame.color_range == ColorRange.JPEG and not preserve_source_color:
-                            # compress full-range sources (yuvj/MJPEG) to limited range
+                        if frame.color_range == ColorRange.JPEG and normalize_color_range:
                             frame = frame.reformat(format=pix_fmt, src_color_range="JPEG", dst_color_range="MPEG")
                         else:
                             frame = frame.reformat(format=pix_fmt)
                         if preserve_source_color:
                             copy_color_properties(video_stream, frame)
+                            if normalize_color_range:
+                                frame.color_range = ColorRange.MPEG
                         elif color_space is not None:
                             set_video_color_properties(frame, color_space)
                         frame_output_end = None

comfy_extras/nodes_video.py

@@ -167,14 +167,16 @@ def define_schema(cls):
         )
 
     @classmethod
-    def execute(cls, video: Input.Video, filename_prefix, format: io.DynamicCombo.Type | str, codec: io.DynamicCombo.Type | None = None) -> io.NodeOutput:
+    def execute(cls, video: Input.Video, filename_prefix, format: io.DynamicCombo.Type | str, codec: io.DynamicCombo.Type | str | None = None) -> io.NodeOutput:
         if isinstance(format, dict):
             format_name = format["format"]
             codec = format.get("codec") or codec
         else:
             format_name = format
         if codec is None:
             codec = {"codec": "auto"}
+        elif isinstance(codec, str):
+            codec = {"codec": codec}
         codec_name = codec["codec"]
         encoding = codec.get("encoding") or {}
         color_space = encoding.get("color_space")
@@ -197,14 +199,15 @@ def execute(cls, video: Input.Video, filename_prefix, format: io.DynamicCombo.Ty
             if len(metadata) > 0:
                 saved_metadata = metadata
         file = f"{filename}_{counter:05}_.{Types.VideoContainer.get_extension(format_name)}"
-        video.save_to(
-            os.path.join(full_output_folder, file),
-            format=Types.VideoContainer(format_name),
-            codec=Types.VideoCodec(codec_name),
-            metadata=saved_metadata,
-            crf=encoding.get("crf"),
-            color_space=color_space,
-        )
+        save_options = {
+            "format": Types.VideoContainer(format_name),
+            "codec": Types.VideoCodec(codec_name),
+            "metadata": saved_metadata,
+            "crf": encoding.get("crf"),
+        }
+        if color_space is not None:
+            save_options["color_space"] = color_space
+        video.save_to(os.path.join(full_output_folder, file), **save_options)
 
         return io.NodeOutput(video, ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
 

tests-unit/comfy_api_test/input_impl_test.py

@@ -29,14 +29,12 @@ def test_container_to_output_format_single():
     assert container_to_output_format("mp4") == "mp4"
 
 
-def test_get_open_write_kwargs_filepath_no_format():
-    """Test that 'format' kwarg is NOT set when dest is a file path."""
+def test_get_open_write_kwargs_filepath_format():
     kwargs_auto = get_open_write_kwargs("output.mp4", "mp4", VideoContainer.AUTO)
-    assert "format" not in kwargs_auto, "Format should not be set for file paths (AUTO)"
+    assert "format" not in kwargs_auto
 
     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["format"] == "avi"
     assert "options" not in kwargs_specific
 
 

tests-unit/comfy_api_test/video_types_test.py

@@ -355,6 +355,48 @@ def create_hdr_av1_video(path, transfer, color_range):
         container.mux(stream.encode(None))
 
 
+def create_full_range_srgb_video(path):
+    images = np.random.default_rng(29).integers(0, 256, (3, 64, 64, 3), dtype=np.uint8)
+    with av.open(path, mode="w") as container:
+        stream = container.add_stream("mjpeg", rate=30)
+        stream.width = 64
+        stream.height = 64
+        stream.pix_fmt = "yuvj420p"
+        stream.color_primaries = ColorPrimaries.BT709
+        stream.color_trc = ColorTrc.BT709
+        stream.colorspace = 1
+        stream.color_range = ColorRange.JPEG
+        for image in images:
+            frame = av.VideoFrame.from_ndarray(image, format="rgb24").reformat(format="yuvj420p")
+            frame.color_primaries = ColorPrimaries.BT709
+            frame.color_trc = ColorTrc.BT709
+            frame.colorspace = 1
+            frame.color_range = ColorRange.JPEG
+            container.mux(stream.encode(frame))
+        container.mux(stream.encode(None))
+
+
+def test_save_loaded_srgb_converts_full_range_to_limited(tmp_path):
+    source = str(tmp_path / "source.mov")
+    output = str(tmp_path / "output.mp4")
+    create_full_range_srgb_video(source)
+    with av.open(source) as container:
+        source_stream = container.streams.video[0]
+        source_color = (source_stream.color_primaries, source_stream.color_trc)
+
+    VideoFromFile(source).save_to(
+        output,
+        format=VideoContainer.MP4,
+        codec=VideoCodec.H264,
+        crf=0,
+    )
+
+    with av.open(output) as container:
+        stream = container.streams.video[0]
+        assert (stream.color_primaries, stream.color_trc) == source_color
+        assert stream.color_range == ColorRange.MPEG
+
+
 def test_save_to_av1_crf_controls_quality(tmp_path):
     generator = torch.Generator().manual_seed(11)
     components = VideoComponents(
@@ -824,6 +866,16 @@ def test_save_to_auto_still_remuxes_a_compatible_codec(tmp_path):
         assert decoded_counts(container)[0] == 6
 
 
+def test_save_to_explicit_format_overrides_destination_suffix(simple_video_file, tmp_path):
+    destination = str(tmp_path / "saved.mkv")
+
+    VideoFromFile(simple_video_file).save_to(destination, format=VideoContainer.MP4)
+
+    with av.open(destination) as container:
+        assert "mp4" in container.format.name.split(",")
+        assert decoded_counts(container)[0] == 3
+
+
 def test_save_to_buffer_transcodes_codec_the_container_cannot_store(tmp_path):
     source = create_matroska_source(
         tmp_path, video_codec="mpeg4", audio_codec="pcm_u8", container_format="mov"
@@ -1348,13 +1400,11 @@ def test_save_to_transcode_bakes_rotation():
 
 
 def mov_text_payload(text: str) -> bytes:
-    """A mov_text sample is a 16-bit big-endian length followed by the UTF-8 text."""
     encoded = text.encode("utf-8")
     return len(encoded).to_bytes(2, "big") + encoded
 
 
 def create_subtitled_source(subtitle_codec="mov_text", container_format="mp4", frames=90, fps=30):
-    """mpeg4 video (so save_to must transcode) alongside a subtitle track carrying three cues."""
     buffer = io.BytesIO()
     with av.open(buffer, mode="w", format=container_format) as container:
         video_stream = container.add_stream("mpeg4", rate=fps)
@@ -1380,7 +1430,6 @@ def create_subtitled_source(subtitle_codec="mov_text", container_format="mp4", f
 
 
 def subtitle_cues(buffer):
-    """(seconds, text) for every non-empty cue; the mp4 muxer pads gaps with 2-byte empties."""
     buffer.seek(0)
     with av.open(buffer) as container:
         if not container.streams.subtitles:
@@ -1393,7 +1442,6 @@ def subtitle_cues(buffer):
 
 
 def test_save_to_transcode_keeps_subtitles_the_container_can_store():
-    """Transcoding video must not silently drop a subtitle track the output can hold."""
     output = io.BytesIO()
     VideoFromFile(create_subtitled_source()).save_to(
         output, format=VideoContainer.MP4, codec=VideoCodec.H264
@@ -1407,7 +1455,6 @@ def test_save_to_transcode_keeps_subtitles_the_container_can_store():
 
 
 def test_save_to_transcode_trims_subtitles_with_the_video():
-    """Kept cues rebase onto the trimmed timeline; cues outside the window are dropped."""
     output = io.BytesIO()
     VideoFromFile(create_subtitled_source(), start_time=1, duration=1).save_to(
         output, format=VideoContainer.MP4, codec=VideoCodec.H264
@@ -1417,8 +1464,6 @@ def test_save_to_transcode_trims_subtitles_with_the_video():
 
 
 def test_save_to_transcode_drops_unstorable_subtitles_with_a_warning(caplog):
-    """There is no subtitle encoder binding, so subrip cannot become mov_text: drop it, but
-    name the stream instead of letting the track vanish silently."""
     output = io.BytesIO()
     with caplog.at_level(logging.WARNING):
         VideoFromFile(create_subtitled_source(subtitle_codec="subrip", container_format="matroska")).save_to(
@@ -1437,7 +1482,6 @@ def test_save_to_transcode_drops_unstorable_subtitles_with_a_warning(caplog):
 
 
 def test_save_to_remux_fallback_keeps_subtitles():
-    """The audio-triggered fallback into the transcode path must keep subtitles too."""
     buffer = io.BytesIO()
     with av.open(buffer, mode="w", format="mov") as container:
         video_stream = container.add_stream("mpeg4", rate=30)

tests-unit/comfy_extras_test/nodes_video_test.py

@@ -0,0 +1,50 @@
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+from comfy_api.latest import Types
+from comfy_extras import nodes_video
+
+
+@pytest.fixture
+def save_video(monkeypatch, tmp_path):
+    monkeypatch.setattr(nodes_video.folder_paths, "get_output_directory", lambda: str(tmp_path))
+    monkeypatch.setattr(
+        nodes_video.folder_paths,
+        "get_save_image_path",
+        lambda *args: (str(tmp_path), "output", 1, "", args[0]),
+    )
+    monkeypatch.setattr(nodes_video.SaveVideo, "hidden", SimpleNamespace(prompt=None, extra_pnginfo=None), raising=False)
+    video = Mock()
+    video.get_dimensions.return_value = (64, 64)
+    return video
+
+
+def test_save_video_accepts_legacy_codec_string(save_video):
+    nodes_video.SaveVideo.execute(save_video, "video", "mp4", "h264")
+
+    kwargs = save_video.save_to.call_args.kwargs
+    assert kwargs["format"] == Types.VideoContainer.MP4
+    assert kwargs["codec"] == Types.VideoCodec.H264
+    assert "color_space" not in kwargs
+
+
+def test_save_video_forwards_nested_encoding_options(save_video):
+    nodes_video.SaveVideo.execute(
+        save_video,
+        "video",
+        {
+            "format": "webm",
+            "codec": {
+                "codec": "av1",
+                "encoding": {"crf": 30, "color_space": "HDR"},
+            },
+        },
+    )
+
+    kwargs = save_video.save_to.call_args.kwargs
+    assert kwargs["format"] == Types.VideoContainer.WEBM
+    assert kwargs["codec"] == Types.VideoCodec.AV1
+    assert kwargs["crf"] == 30
+    assert kwargs["color_space"] == "HDR"