From cd18582578d38081af5614d018781bcdfbe95e0b Mon Sep 17 00:00:00 2001 From: Christian Byrne Date: Sun, 4 May 2025 20:26:57 -0700 Subject: [PATCH] Support saving Comfy `VIDEO` type to buffer (#7939) * get output format when saving to buffer * add unit tests for writing to file or stream with correct fmt * handle `to_format=None` * fix formatting --- comfy_api/input_impl/video_types.py | 46 +++++++++- tests-unit/comfy_api_test/input_impl_test.py | 91 ++++++++++++++++++++ 2 files changed, 135 insertions(+), 2 deletions(-) create mode 100644 tests-unit/comfy_api_test/input_impl_test.py diff --git a/comfy_api/input_impl/video_types.py b/comfy_api/input_impl/video_types.py index 12e5783d..d0b0b36d 100644 --- a/comfy_api/input_impl/video_types.py +++ b/comfy_api/input_impl/video_types.py @@ -12,6 +12,46 @@ import torch from comfy_api.input import VideoInput from comfy_api.util import VideoContainer, VideoCodec, VideoComponents + +def container_to_output_format(container_format: str | None) -> str | None: + """ + A container's `format` may be a comma-separated list of formats. + E.g., iso container's `format` may be `mov,mp4,m4a,3gp,3g2,mj2`. + However, writing to a file/stream with `av.open` requires a single format, + or `None` to auto-detect. + """ + if not container_format: + return None # Auto-detect + + if "," not in container_format: + return container_format + + formats = container_format.split(",") + return formats[0] + + +def get_open_write_kwargs( + dest: str | io.BytesIO, container_format: str, to_format: str | None +) -> dict: + """Get kwargs for writing a `VideoFromFile` to a file/stream with `av.open`""" + open_kwargs = { + "mode": "w", + # If isobmff, preserve custom metadata tags (workflow, prompt, extra_pnginfo) + "options": {"movflags": "use_metadata_tags"}, + } + + is_write_to_buffer = isinstance(dest, io.BytesIO) + if is_write_to_buffer: + # Set output format explicitly, since it cannot be inferred from file extension + if to_format == VideoContainer.AUTO: + to_format = container_format.lower() + elif isinstance(to_format, str): + to_format = to_format.lower() + open_kwargs["format"] = container_to_output_format(to_format) + + return open_kwargs + + class VideoFromFile(VideoInput): """ Class representing video input from a file. @@ -89,7 +129,7 @@ class VideoFromFile(VideoInput): def save_to( self, - path: str, + path: str | io.BytesIO, format: VideoContainer = VideoContainer.AUTO, codec: VideoCodec = VideoCodec.AUTO, metadata: Optional[dict] = None @@ -116,7 +156,9 @@ class VideoFromFile(VideoInput): ) streams = container.streams - with av.open(path, mode='w', options={"movflags": "use_metadata_tags"}) as output_container: + + open_kwargs = get_open_write_kwargs(path, container_format, format) + with av.open(path, **open_kwargs) as output_container: # Copy over the original metadata for key, value in container.metadata.items(): if metadata is None or key not in metadata: diff --git a/tests-unit/comfy_api_test/input_impl_test.py b/tests-unit/comfy_api_test/input_impl_test.py new file mode 100644 index 00000000..5fc21a9a --- /dev/null +++ b/tests-unit/comfy_api_test/input_impl_test.py @@ -0,0 +1,91 @@ +import io +from comfy_api.input_impl.video_types import ( + container_to_output_format, + get_open_write_kwargs, +) +from comfy_api.util import VideoContainer + + +def test_container_to_output_format_empty_string(): + """Test that an empty string input returns None. `None` arg allows default auto-detection.""" + assert container_to_output_format("") is None + + +def test_container_to_output_format_none(): + """Test that None input returns None.""" + assert container_to_output_format(None) is None + + +def test_container_to_output_format_comma_separated(): + """Test that a comma-separated list returns a valid singular format from the list.""" + comma_separated_format = "mp4,mov,m4a" + output_format = container_to_output_format(comma_separated_format) + assert output_format in comma_separated_format + + +def test_container_to_output_format_single(): + """Test that a single format string (not comma-separated list) is returned as is.""" + 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.""" + 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)" + + 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 + + +def test_get_open_write_kwargs_base_options_mode(): + """Test basic kwargs for file path: mode and movflags.""" + kwargs = get_open_write_kwargs("output.mp4", "mp4", VideoContainer.AUTO) + assert kwargs["mode"] == "w", "mode should be set to write" + + fail_msg = "movflags should be set to preserve custom metadata tags" + assert "movflags" in kwargs["options"], fail_msg + assert kwargs["options"]["movflags"] == "use_metadata_tags", fail_msg + + +def test_get_open_write_kwargs_bytesio_auto_format(): + """Test kwargs for BytesIO dest with AUTO format.""" + dest = io.BytesIO() + container_fmt = "mov,mp4,m4a" + kwargs = get_open_write_kwargs(dest, container_fmt, VideoContainer.AUTO) + + assert kwargs["mode"] == "w" + assert kwargs["options"]["movflags"] == "use_metadata_tags" + + fail_msg = ( + "Format should be a valid format from the container's format list when AUTO" + ) + assert kwargs["format"] in container_fmt, fail_msg + + +def test_get_open_write_kwargs_bytesio_specific_format(): + """Test kwargs for BytesIO dest with a specific single format.""" + dest = io.BytesIO() + container_fmt = "avi" + to_fmt = VideoContainer.MP4 + kwargs = get_open_write_kwargs(dest, container_fmt, to_fmt) + + assert kwargs["mode"] == "w" + assert kwargs["options"]["movflags"] == "use_metadata_tags" + + fail_msg = "Format should be the specified format (lowercased) when output format is not AUTO" + assert kwargs["format"] == "mp4", fail_msg + + +def test_get_open_write_kwargs_bytesio_specific_format_list(): + """Test kwargs for BytesIO dest with a specific comma-separated format.""" + dest = io.BytesIO() + container_fmt = "avi" + to_fmt = "mov,mp4,m4a" # A format string that is a list + kwargs = get_open_write_kwargs(dest, container_fmt, to_fmt) + + assert kwargs["mode"] == "w" + assert kwargs["options"]["movflags"] == "use_metadata_tags" + + 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