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

This commit is contained in:
comfyanonymous
2026-08-20 16:29:20 -07:00
committed by GitHub
parent de6b062fb5
commit dcbcf8c2e1
6 changed files with 772 additions and 73 deletions

View File

@@ -30,15 +30,24 @@ class VideoInput(ABC):
metadata: Optional[dict] = None,
bit_depth: int | None = None,
crf: float | None = None,
color_space: str | None = None,
):
"""
Abstract method to save the video input to a file.
bit_depth selects the encoded bit depth; None keeps the video's native depth.
crf selects the H.264 constant rate factor; None uses the encoder default.
crf selects the H.264 or AV1 constant rate factor; None uses the encoder default.
color_space="sRGB" writes SDR BT.709/sRGB video. "HDR" writes 10-bit BT.2020/HLG video;
"HDR PQ" selects BT.2020/PQ.
Tensor-created videos default to sRGB when color_space is None. Loaded videos keep matching recognized native color
properties; other input pixels must already use the selected color space.
"""
pass
def get_color_space(self) -> str:
"""Return the video's color space as sRGB, HDR, HDR PQ, or auto when unspecified."""
return "auto"
@abstractmethod
def as_trimmed(
self,

View File

@@ -1,6 +1,6 @@
from av.container import InputContainer
from av.subtitles.stream import SubtitleStream
from av.video.reformatter import ColorRange
from av.video.reformatter import ColorPrimaries, ColorRange, ColorTrc
from fractions import Fraction
from typing import Optional
from .._input import AudioInput, VideoInput
@@ -16,6 +16,38 @@ from .._util import VideoContainer, VideoCodec, VideoComponents
import logging
VIDEO_ENCODERS = {
VideoCodec.H264: "h264",
VideoCodec.AV1: "libsvtav1",
}
VIDEO_CONTAINER_FORMATS = {
VideoContainer.MP4: "mp4",
VideoContainer.MKV: "matroska",
VideoContainer.WEBM: "webm",
}
WEBM_STREAM_CODECS = {
"video": {"av1", "vp8", "vp9"},
"audio": {"opus", "vorbis"},
"subtitle": {"webvtt"},
}
BT2020_NCL = 9
BT709_NCL = 1
HDR_COLOR_TRANSFERS = {
"HDR": ColorTrc.ARIB_STD_B67,
"HDR PQ": ColorTrc.SMPTE2084,
}
VIDEO_COLOR_TRANSFERS = {
"sRGB": ColorTrc.IEC61966_2_1,
**HDR_COLOR_TRANSFERS,
}
VIDEO_TRANSFER_COLOR_SPACES = {
ColorTrc.BT709: "sRGB",
ColorTrc.IEC61966_2_1: "sRGB",
ColorTrc.ARIB_STD_B67: "HDR",
ColorTrc.SMPTE2084: "HDR PQ",
}
def container_to_output_format(container_format: str | None) -> str | None:
"""
A container's `format` may be a comma-separated list of formats.
@@ -37,22 +69,24 @@ def get_open_write_kwargs(
) -> dict:
"""Get kwargs for writing a `VideoFromFile` to a file/stream with `av.open`"""
is_write_to_buffer = isinstance(dest, io.BytesIO)
is_mp4_file = not is_write_to_buffer and os.path.splitext(dest)[1].lower() == ".mp4"
movflags = "use_metadata_tags+faststart" if is_mp4_file else "use_metadata_tags"
open_kwargs = {
"mode": "w",
# If isobmff, preserve custom metadata tags (workflow, prompt, extra_pnginfo)
"options": {"movflags": movflags},
}
open_kwargs = {"mode": "w"}
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, VideoContainer):
to_format = VIDEO_CONTAINER_FORMATS[to_format]
elif isinstance(to_format, str):
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(".")
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"
open_kwargs["options"] = {"movflags": movflags}
return open_kwargs
@@ -100,19 +134,66 @@ def write_output_metadata(container: InputContainer, output, metadata: dict | No
output.metadata[key] = value if isinstance(value, str) else json.dumps(value)
def mp4_output_open_kwargs(path: str | io.BytesIO, format: VideoContainer, codec: VideoCodec) -> dict:
if format != VideoContainer.AUTO and format != VideoContainer.MP4:
raise ValueError("Only MP4 format is supported for now")
if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
raise ValueError("Only H264 codec is supported for now")
def video_output_config(path: str | io.BytesIO, format: VideoContainer, codec: VideoCodec) -> tuple[dict, VideoContainer, VideoCodec]:
if isinstance(format, str):
format = VideoContainer(format)
if isinstance(codec, str):
codec = VideoCodec(codec)
if format == VideoContainer.AUTO:
extension = os.path.splitext(os.fspath(path))[1].lower() if isinstance(path, (str, os.PathLike)) else ""
format = {
".mkv": VideoContainer.MKV,
".webm": VideoContainer.WEBM,
}.get(extension, VideoContainer.MP4)
if codec == VideoCodec.AUTO:
codec = VideoCodec.AV1 if format == VideoContainer.WEBM else VideoCodec.H264
if format == VideoContainer.WEBM and codec != VideoCodec.AV1:
raise ValueError("WebM output requires the AV1 codec")
# FFmpeg's faststart pass reopens the output by filename, so it cannot be used with file-like objects.
movflags = "use_metadata_tags+faststart" if isinstance(path, (str, os.PathLike)) else "use_metadata_tags"
open_kwargs = {"mode": "w", "options": {"movflags": movflags}}
if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
open_kwargs["format"] = format.value
elif isinstance(path, io.BytesIO):
open_kwargs["format"] = "mp4" # no file extension to infer the format from
return open_kwargs
open_kwargs = {"mode": "w", "format": VIDEO_CONTAINER_FORMATS[format]}
if format == VideoContainer.MP4:
movflags = "use_metadata_tags+faststart" if isinstance(path, (str, os.PathLike)) else "use_metadata_tags"
open_kwargs["options"] = {"movflags": movflags}
return open_kwargs, format, codec
def set_video_color_properties(target, color_space):
is_hdr = color_space in HDR_COLOR_TRANSFERS
target.color_primaries = ColorPrimaries.BT2020 if is_hdr else ColorPrimaries.BT709
target.color_trc = VIDEO_COLOR_TRANSFERS[color_space]
target.colorspace = BT2020_NCL if is_hdr else BT709_NCL
target.color_range = ColorRange.MPEG
def copy_color_properties(source, target):
target.color_primaries = source.color_primaries
target.color_trc = source.color_trc
target.colorspace = source.colorspace
target.color_range = source.color_range
def video_stream_color_space(stream) -> str | None:
if stream is None:
return None
return VIDEO_TRANSFER_COLOR_SPACES.get(stream.color_trc)
def video_encoder_options(codec: VideoCodec, crf: float | None) -> dict[str, str]:
if crf is None:
return {}
if codec == VideoCodec.AV1 and crf == 0:
return {"svtav1-params": "lossless=1"}
return {"crf": str(crf)}
def webm_streams_compatible(streams) -> bool:
for stream in streams:
allowed_codecs = WEBM_STREAM_CODECS.get(stream.type)
if allowed_codecs is not None and stream.codec_context is not None and stream.codec.canonical_name not in allowed_codecs:
return False
return True
class VideoFromFile(VideoInput):
@@ -167,6 +248,13 @@ class VideoFromFile(VideoInput):
video_stream = container.streams.video[0] if len(container.streams.video) > 0 else None
return video_stream_bit_depth(video_stream)
def get_color_space(self) -> str:
if isinstance(self.__file, io.BytesIO):
self.__file.seek(0)
with av.open(self.__file, mode="r") as container:
video_stream = container.streams.video[0] if len(container.streams.video) > 0 else None
return video_stream_color_space(video_stream) or "sRGB"
def get_duration(self) -> float:
"""
Returns the duration of the video in seconds.
@@ -465,16 +553,28 @@ class VideoFromFile(VideoInput):
metadata: Optional[dict] = None,
bit_depth: int | None = None,
crf: float | None = None,
color_space: str | None = None,
):
if color_space is not None and color_space not in VIDEO_COLOR_TRANSFERS:
raise ValueError(f"Unsupported video color space: {color_space}")
_, output_format, _ = video_output_config(path, format, codec)
if isinstance(self.__file, io.BytesIO):
self.__file.seek(0) # Reset the BytesIO object to the beginning
with av.open(self.__file, mode='r') as container:
container_format = container.format.name
video_stream = container.streams.video[0] if len(container.streams.video) > 0 else None
video_encoding = video_stream.codec.name if video_stream is not None else None
video_encoding = video_stream.codec.canonical_name if video_stream is not None else None
source_bit_depth = video_stream_bit_depth(video_stream)
source_color_space = video_stream_color_space(video_stream)
if source_color_space is not None and color_space is not None and source_color_space != color_space:
raise ValueError(
f"Cannot save {source_color_space} video as {color_space} without color conversion; "
f"use auto or {source_color_space}"
)
reuse_streams = True
if format != VideoContainer.AUTO and format not in container_format.split(","):
if format != VideoContainer.AUTO and VIDEO_CONTAINER_FORMATS[VideoContainer(format)] not in container_format.split(","):
reuse_streams = False
if output_format == VideoContainer.WEBM and not webm_streams_compatible(container.streams):
reuse_streams = False
if codec != VideoCodec.AUTO and codec != video_encoding and video_encoding is not None:
reuse_streams = False
@@ -482,13 +582,15 @@ class VideoFromFile(VideoInput):
reuse_streams = False
if crf is not None:
reuse_streams = False
if color_space is not None:
reuse_streams = False
if self.__start_time or self.__duration:
reuse_streams = False
if not reuse_streams:
if bit_depth is None:
bit_depth = source_bit_depth
return self._save_transcoded(container, path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth, crf=crf)
return self._save_transcoded(container, path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth, crf=crf, color_space=color_space)
streams = container.streams
@@ -522,9 +624,10 @@ class VideoFromFile(VideoInput):
metadata: dict | None,
bit_depth: int,
crf: float | None = None,
color_space: str | None = None,
):
"""Re-encode to H.264/AAC one frame at a time; peak memory does not scale with video length."""
open_kwargs = mp4_output_open_kwargs(path, format, codec)
"""Re-encode one frame at a time; peak memory does not scale with video length."""
open_kwargs, output_format, output_codec = video_output_config(path, format, codec)
video_stream = self._get_first_video_stream(container)
start_time, duration = self.get_active_trim_window()
start_pts = int(start_time / video_stream.time_base)
@@ -539,6 +642,10 @@ class VideoFromFile(VideoInput):
container.seek(start_pts, stream=video_stream)
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
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"
rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
@@ -558,6 +665,8 @@ class VideoFromFile(VideoInput):
logging.warning("Audio stream parameters could not be determined; ignoring audio.")
audio_stream = None
if audio_stream is not None:
if output_format == VideoContainer.WEBM:
sample_rate = 48000
audio_time_base = Fraction(1, sample_rate)
layout = {1: "mono", 2: "stereo", 6: "5.1"}.get(channels, "stereo")
resampler = av.audio.resampler.AudioResampler(format="fltp", layout=layout, rate=sample_rate)
@@ -655,24 +764,28 @@ class VideoFromFile(VideoInput):
else:
out_width, out_height = frame.width, frame.height
if out_width % 2 or out_height % 2:
raise ValueError(f"H.264 output requires even dimensions, got {out_width}x{out_height}")
raise ValueError(f"{output_codec.value.upper()} output requires even dimensions, got {out_width}x{out_height}")
source_size = (frame.width, frame.height)
output = av.open(path, **open_kwargs)
# Add metadata before writing any streams
write_output_metadata(container, output, metadata)
out_video = output.add_stream("h264", rate=rate)
out_video = output.add_stream(VIDEO_ENCODERS[output_codec], rate=rate)
# no B-frames: reordering makes mp4 sample durations follow decode order,
# so irregular-VFR spans and trim windows land wrong
out_video.codec_context.max_b_frames = 0
out_video.width = out_width
out_video.height = out_height
out_video.pix_fmt = pix_fmt
if crf is not None:
out_video.options = {"crf": str(crf)}
out_video.options = video_encoder_options(output_codec, crf)
if preserve_source_color:
copy_color_properties(video_stream, out_video.codec_context)
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
out_video.codec_context.time_base = video_stream.time_base
if audio_stream is not None:
out_audio = output.add_stream("aac", rate=sample_rate, layout=layout)
audio_codec = "libopus" if output_format == VideoContainer.WEBM else "aac"
out_audio = output.add_stream(audio_codec, rate=sample_rate, layout=layout)
if (frame.width, frame.height) != source_size:
# encoding would silently rescale the new geometry into the old one
raise ValueError(
@@ -697,11 +810,15 @@ class VideoFromFile(VideoInput):
rotation_filter = (g_src, g_sink)
rotation_filter[0].push(frame)
frame = rotation_filter[1].pull()
if frame.color_range == ColorRange.JPEG:
if frame.color_range == ColorRange.JPEG and not preserve_source_color:
# compress full-range sources (yuvj/MJPEG) to limited 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)
elif color_space is not None:
set_video_color_properties(frame, color_space)
frame_output_end = None
if frame.pts is not None:
if video_pts_offset is None:
@@ -830,6 +947,9 @@ class VideoFromComponents(VideoInput):
def get_bit_depth(self) -> int:
return self.__bit_depth
def get_color_space(self) -> str:
return "sRGB"
def save_to(
self,
path: str,
@@ -838,12 +958,19 @@ class VideoFromComponents(VideoInput):
metadata: Optional[dict] = None,
bit_depth: int | None = None,
crf: float | None = None,
color_space: str | None = None,
):
"""Save the video to a file path or BytesIO buffer."""
open_kwargs = mp4_output_open_kwargs(path, format, codec)
if color_space is None:
color_space = "sRGB"
if color_space is not None and color_space not in VIDEO_COLOR_TRANSFERS:
raise ValueError(f"Unsupported video color space: {color_space}")
open_kwargs, output_format, output_codec = video_output_config(path, format, codec)
# None means "use the depth this video was created with" (CreateVideo's choice).
if bit_depth is None:
bit_depth = self.__bit_depth
if color_space in HDR_COLOR_TRANSFERS:
bit_depth = max(bit_depth, 10)
is_10bit = bit_depth >= 10
with av.open(path, **open_kwargs) as output:
# Add metadata before writing any streams
@@ -854,22 +981,28 @@ class VideoFromComponents(VideoInput):
frame_rate = Fraction(round(self.__components.frame_rate * 1000), 1000)
# Create a video stream
pix_fmt = "yuv420p10le" if is_10bit else "yuv420p"
video_stream = output.add_stream('h264', rate=frame_rate)
video_stream = output.add_stream(VIDEO_ENCODERS[output_codec], rate=frame_rate)
video_stream.width = self.__components.images.shape[2]
video_stream.height = self.__components.images.shape[1]
video_stream.pix_fmt = pix_fmt
if crf is not None:
video_stream.options = {"crf": str(crf)}
video_stream.options = video_encoder_options(output_codec, crf)
if color_space is not None:
set_video_color_properties(video_stream.codec_context, color_space)
# Create an audio stream
audio_sample_rate = 1
audio_resampler = None
audio_stream: Optional[av.AudioStream] = None
if self.__components.audio:
audio_sample_rate = int(self.__components.audio['sample_rate'])
source_audio_sample_rate = int(self.__components.audio['sample_rate'])
audio_sample_rate = 48000 if output_format == VideoContainer.WEBM else source_audio_sample_rate
waveform = self.__components.audio['waveform']
waveform = waveform[0, :, :math.ceil((audio_sample_rate / frame_rate) * self.__components.images.shape[0])]
waveform = waveform[0, :, :math.ceil((source_audio_sample_rate / frame_rate) * self.__components.images.shape[0])]
layout = {1: 'mono', 2: 'stereo', 6: '5.1'}.get(waveform.shape[0], 'stereo')
audio_stream = output.add_stream('aac', rate=audio_sample_rate, layout=layout)
audio_codec = "libopus" if output_format == VideoContainer.WEBM else "aac"
audio_stream = output.add_stream(audio_codec, rate=audio_sample_rate, layout=layout)
if audio_sample_rate != source_audio_sample_rate:
audio_resampler = av.audio.resampler.AudioResampler(format="fltp", layout=layout, rate=audio_sample_rate)
# Encode video
for i, frame in enumerate(self.__components.images):
@@ -880,7 +1013,14 @@ class VideoFromComponents(VideoInput):
else:
img = (frame * 255).clamp(0, 255).byte().cpu().numpy() # shape: (H, W, 3)
frame = av.VideoFrame.from_ndarray(img, format='rgb24')
frame = frame.reformat(format=pix_fmt)
dst_colorspace = None
if color_space == "sRGB":
dst_colorspace = BT709_NCL
elif color_space in HDR_COLOR_TRANSFERS:
dst_colorspace = BT2020_NCL
frame = frame.reformat(format=pix_fmt, dst_colorspace=dst_colorspace)
if color_space is not None:
set_video_color_properties(frame, color_space)
packet = video_stream.encode(frame)
output.mux(packet)
@@ -890,9 +1030,14 @@ class VideoFromComponents(VideoInput):
if audio_stream and self.__components.audio:
frame = av.AudioFrame.from_ndarray(waveform.float().cpu().contiguous().numpy(), format='fltp', layout=layout)
frame.sample_rate = audio_sample_rate
frame.sample_rate = source_audio_sample_rate
frame.pts = 0
output.mux(audio_stream.encode(frame))
frames = [frame] if audio_resampler is None else audio_resampler.resample(frame)
for frame in frames:
output.mux(audio_stream.encode(frame))
if audio_resampler is not None:
for frame in audio_resampler.resample(None):
output.mux(audio_stream.encode(frame))
# Flush encoder
output.mux(audio_stream.encode(None))

View File

@@ -7,6 +7,7 @@ from .._input import ImageInput, AudioInput, MaskInput
class VideoCodec(str, Enum):
AUTO = "auto"
H264 = "h264"
AV1 = "av1"
@classmethod
def as_input(cls) -> list[str]:
@@ -18,6 +19,8 @@ class VideoCodec(str, Enum):
class VideoContainer(str, Enum):
AUTO = "auto"
MP4 = "mp4"
MKV = "mkv"
WEBM = "webm"
@classmethod
def as_input(cls) -> list[str]:
@@ -35,6 +38,10 @@ class VideoContainer(str, Enum):
value = cls(value)
if value == VideoContainer.MP4 or value == VideoContainer.AUTO:
return "mp4"
if value == VideoContainer.MKV:
return "mkv"
if value == VideoContainer.WEBM:
return "webm"
return ""
@dataclass

View File

@@ -72,6 +72,70 @@ class SaveWEBM(io.ComfyNode):
return io.NodeOutput(images, ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
def _save_video_codec_input(supported_codecs: list[str], *, optional=False, hidden=False):
codec_options = []
if "auto" in supported_codecs:
codec_options.append(io.DynamicCombo.Option("auto", []))
if "h264" in supported_codecs:
codec_options.append(
io.DynamicCombo.Option(
"h264",
[
io.DynamicCombo.Input(
"encoding",
display_name="encoding mode",
options=[
io.DynamicCombo.Option("auto", []),
io.DynamicCombo.Option(
"re-encode",
[io.Float.Input("crf", default=23.0, min=0.0, max=51.0, step=1.0, tooltip="Lower values produce higher quality and larger files.")],
),
],
optional=True,
tooltip="Automatic preserves compatible H.264 streams. Re-encode applies a custom CRF.",
),
],
)
)
if "av1" in supported_codecs:
codec_options.append(
io.DynamicCombo.Option(
"av1",
[
io.DynamicCombo.Input(
"encoding",
display_name="encoding mode",
options=[
io.DynamicCombo.Option("auto", []),
io.DynamicCombo.Option(
"re-encode",
[
io.Float.Input("crf", default=30.0, min=0.0, max=63.0, step=1.0, tooltip="Lower values produce higher quality and larger files."),
io.Combo.Input(
"color_space",
options=["auto", "sRGB", "HDR", "HDR PQ"],
default="auto",
display_name="color space",
tooltip="Auto uses sRGB for videos created from images and preserves recognized colors on loaded videos. sRGB writes SDR BT.709/sRGB. HDR writes 10-bit BT.2020/HLG; HDR PQ writes BT.2020/PQ. Other input pixels must already use the selected color space.",
),
],
),
],
optional=True,
tooltip="Automatic preserves compatible AV1 streams. Re-encode applies custom encoding options.",
),
],
)
)
return io.DynamicCombo.Input(
"codec",
options=codec_options,
optional=optional,
tooltip="The output video codec. Auto preserves a compatible source stream. H.264 re-encoding supports SDR; AV1 re-encoding supports SDR, HDR (HLG), and HDR PQ.",
extra_dict={"hidden": True} if hidden else None,
)
class SaveVideo(io.ComfyNode):
@classmethod
def define_schema(cls):
@@ -85,42 +149,37 @@ class SaveVideo(io.ComfyNode):
inputs=[
io.Video.Input("video", tooltip="The video to save."),
io.String.Input("filename_prefix", default="video/ComfyUI", tooltip="The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."),
io.Combo.Input("format", options=Types.VideoContainer.as_input(), default="auto", tooltip="The format to save the video as."),
io.DynamicCombo.Input(
"codec",
"format",
options=[
io.DynamicCombo.Option("auto", []),
io.DynamicCombo.Option(
"h264",
[
io.DynamicCombo.Input(
"encoding",
display_name="encoding mode",
options=[
io.DynamicCombo.Option("auto", []),
io.DynamicCombo.Option(
"re-encode",
[io.Float.Input("crf", default=23.0, min=0.0, max=51.0, step=1.0, tooltip="Lower values produce higher quality and larger files.")],
),
],
optional=True,
tooltip="Automatic preserves compatible H.264 streams. Re-encode applies a custom CRF.",
),
],
),
io.DynamicCombo.Option("auto", [_save_video_codec_input(["auto", "h264", "av1"])]),
io.DynamicCombo.Option("mp4", [_save_video_codec_input(["auto", "h264", "av1"])]),
io.DynamicCombo.Option("mkv", [_save_video_codec_input(["auto", "h264", "av1"])]),
io.DynamicCombo.Option("webm", [_save_video_codec_input(["auto", "av1"])]),
],
tooltip="The codec to use for the video.",
tooltip="The output container. Auto preserves the source container when possible; MP4, MKV, and WebM select a specific container.",
),
_save_video_codec_input(["auto", "h264", "av1"], optional=True, hidden=True),
],
hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
is_output_node=True,
outputs=[io.Video.Output("video")],
outputs=[io.Video.Output("video", tooltip="The input video, unchanged.")],
)
@classmethod
def execute(cls, video: Input.Video, filename_prefix, format: str, codec: io.DynamicCombo.Type) -> io.NodeOutput:
def execute(cls, video: Input.Video, filename_prefix, format: io.DynamicCombo.Type | str, codec: io.DynamicCombo.Type | 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"}
codec_name = codec["codec"]
encoding = codec.get("encoding") or {}
color_space = encoding.get("color_space")
if color_space == "auto":
color_space = None
width, height = video.get_dimensions()
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
filename_prefix,
@@ -137,13 +196,14 @@ class SaveVideo(io.ComfyNode):
metadata["prompt"] = cls.hidden.prompt
if len(metadata) > 0:
saved_metadata = metadata
file = f"{filename}_{counter:05}_.{Types.VideoContainer.get_extension(format)}"
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),
codec=codec_name,
format=Types.VideoContainer(format_name),
codec=Types.VideoCodec(codec_name),
metadata=saved_metadata,
crf=encoding.get("crf"),
color_space=color_space,
)
return io.NodeOutput(video, ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
@@ -199,7 +259,7 @@ class GetVideoComponents(io.ComfyNode):
search_aliases=["extract frames", "split video", "video to images", "demux"],
display_name="Get Video Components",
category="video",
description="Extracts all components from a video: frames, audio, framerate, and bit depth.",
description="Extracts video frames, audio, frame rate, bit depth, and color space.",
inputs=[
io.Video.Input("video", tooltip="The video to extract components from."),
],
@@ -208,13 +268,20 @@ class GetVideoComponents(io.ComfyNode):
io.Audio.Output(display_name="audio"),
io.Float.Output(display_name="fps"),
io.Int.Output(display_name="bit_depth"),
io.Combo.Output(display_name="color_space"),
],
)
@classmethod
def execute(cls, video: Input.Video) -> io.NodeOutput:
components = video.get_components()
return io.NodeOutput(components.images, components.audio, float(components.frame_rate), video.get_bit_depth())
return io.NodeOutput(
components.images,
components.audio,
float(components.frame_rate),
video.get_bit_depth(),
video.get_color_space(),
)
class LoadVideo(io.ComfyNode):

View File

@@ -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) == {}

View File

@@ -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"