Thanks to visit codestin.com
Credit goes to github.com

Skip to content

[pull] master from comfyanonymous:master #39

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
May 5, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 50 additions & 3 deletions comfy_api/input_impl/video_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,46 @@
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.
Expand Down Expand Up @@ -89,7 +129,7 @@ def get_components(self) -> VideoComponents:

def save_to(
self,
path: str,
path: str | io.BytesIO,
format: VideoContainer = VideoContainer.AUTO,
codec: VideoCodec = VideoCodec.AUTO,
metadata: Optional[dict] = None
Expand All @@ -116,7 +156,9 @@ def save_to(
)

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:
Expand Down Expand Up @@ -211,7 +253,12 @@ def save_to(
start = i * samples_per_frame
end = start + samples_per_frame
# TODO(Feature) - Add support for stereo audio
chunk = self.__components.audio['waveform'][0, 0, start:end].unsqueeze(0).numpy()
chunk = (
self.__components.audio["waveform"][0, 0, start:end]
.unsqueeze(0)
.contiguous()
.numpy()
)
audio_frame = av.AudioFrame.from_ndarray(chunk, format='fltp', layout='mono')
audio_frame.sample_rate = audio_sample_rate
audio_frame.pts = i * samples_per_frame
Expand Down
91 changes: 91 additions & 0 deletions tests-unit/comfy_api_test/input_impl_test.py
Original file line number Diff line number Diff line change
@@ -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
Loading