zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-08-26 02:42:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add colorspace option and change bit_depth to a combo on CreateVideo. (#15810)

Browse Source
This commit is contained in:
comfyanonymous
2026-08-22 18:07:29 -07:00
committed by GitHub
parent 924743af08
commit 9db05e0e1f
5 changed files with 117 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
comfy_api/latest/_input/video_types.py
View File

@@ -37,8 +37,8 @@ class VideoInput(ABC):
bit_depth selects the encoded bit depth; None keeps the video's native depth.
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.
color_space="sRGB" selects SDR BT.709/sRGB, "HDR" selects BT.2020/HLG, and "HDR PQ"
selects BT.2020/PQ. Bit depth is selected independently.
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.
"""

13
comfy_api/latest/_input_impl/video_types.py
View File

@@ -644,8 +644,6 @@ class VideoFromFile(VideoInput):
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)
@@ -932,10 +930,13 @@ class VideoFromComponents(VideoInput):
Class representing video input from tensors.
"""
def __init__(self, components: VideoComponents, bit_depth: int = 8):
def __init__(self, components: VideoComponents, bit_depth: int = 8, color_space: str = "sRGB"):
if color_space not in VIDEO_COLOR_TRANSFERS:
raise ValueError(f"Unsupported video color space: {color_space}")
self.__components = components
# Tensor components have no inherent bit depth; this is the depth used when encoding.
self.__bit_depth = bit_depth
self.__color_space = color_space
def get_components(self) -> VideoComponents:
return VideoComponents(
@@ -948,7 +949,7 @@ class VideoFromComponents(VideoInput):
return self.__bit_depth
def get_color_space(self) -> str:
return "sRGB"
return self.__color_space
def save_to(
self,
@@ -962,15 +963,13 @@ class VideoFromComponents(VideoInput):
):
"""Save the video to a file path or BytesIO buffer."""
if color_space is None:
color_space = "sRGB"
color_space = self.__color_space
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

46
comfy_extras/nodes_video.py
View File

@@ -72,16 +72,6 @@ class SaveWEBM(io.ComfyNode):
return io.NodeOutput(images, ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
def _save_video_color_space_input():
return 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.",
)
def _save_video_codec_input(supported_codecs: list[str], *, optional=False, hidden=False):
codec_options = []
if "auto" in supported_codecs:
@@ -100,7 +90,6 @@ def _save_video_codec_input(supported_codecs: list[str], *, optional=False, hidd
"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."),
_save_video_color_space_input(),
],
),
],
@@ -124,7 +113,6 @@ def _save_video_codec_input(supported_codecs: list[str], *, optional=False, hidd
"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."),
_save_video_color_space_input(),
],
),
],
@@ -164,7 +152,7 @@ class SaveVideo(io.ComfyNode):
io.DynamicCombo.Option("mkv", [_save_video_codec_input(["auto", "h264", "av1"])]),
io.DynamicCombo.Option("webm", [_save_video_codec_input(["auto", "av1"])]),
],
tooltip="The output container. Auto preserves the source container when possible; MP4, MKV, and WebM select a specific container.",
tooltip="The output container. Auto uses MP4 for Auto/H.264 and WebM for AV1. MP4, MKV, and WebM select a specific container.",
),
_save_video_codec_input(["auto", "h264", "av1"], optional=True, hidden=True),
],
@@ -183,10 +171,9 @@ class SaveVideo(io.ComfyNode):
if codec is None:
codec = {"codec": "auto"}
codec_name = codec["codec"]
if format_name == "auto":
format_name = "webm" if codec_name == "av1" else "mp4"
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,
@@ -210,7 +197,6 @@ class SaveVideo(io.ComfyNode):
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)]))
@@ -230,16 +216,19 @@ class CreateVideo(io.ComfyNode):
io.Image.Input("images", tooltip="The images to create a video from."),
io.Float.Input("fps", default=30.0, min=1.0, max=120.0, step=1.0),
io.Audio.Input("audio", optional=True, tooltip="The audio to add to the video."),
io.Int.Input(
io.Combo.Input(
"bit_depth",
min=8,
max=10,
default=8,
step=2,
tooltip="Bit depth of the created video. 10-bit keeps smoother gradients with less"
" banding, but some players and downstream nodes may not support it.",
options=["auto", 8, 10],
default="auto",
tooltip="Auto uses 8-bit for sRGB and 10-bit for HDR. Explicit 8-bit and 10-bit choices are independent of colorspace.",
optional=True,
display_mode=io.NumberDisplay.number,
),
io.Combo.Input(
"color_space",
options=["sRGB", "HDR", "HDR PQ"],
default="sRGB",
optional=True,
tooltip="Colorspace of the input images. HDR selects BT.2020/HLG and HDR PQ selects BT.2020/PQ.",
),
],
outputs=[
@@ -249,12 +238,15 @@ class CreateVideo(io.ComfyNode):
@classmethod
def execute(
cls, images: Input.Image, fps: float, audio: Optional[Input.Audio] = None, bit_depth: int = 8,
cls, images: Input.Image, fps: float, audio: Optional[Input.Audio] = None, bit_depth: int | str = "auto", color_space: str = "sRGB",
) -> io.NodeOutput:
if bit_depth == "auto":
bit_depth = 10 if color_space in ("HDR", "HDR PQ") else 8
return io.NodeOutput(
InputImpl.VideoFromComponents(
Types.VideoComponents(images=images, audio=audio, frame_rate=Fraction(fps)),
bit_depth=bit_depth,
color_space=color_space,
)
)
@@ -274,7 +266,7 @@ class GetVideoComponents(io.ComfyNode):
io.Image.Output(display_name="images"),
io.Audio.Output(display_name="audio"),
io.Float.Output(display_name="fps"),
io.Int.Output(display_name="bit_depth"),
io.Combo.Output(display_name="bit_depth"),
io.Combo.Output(display_name="color_space"),
],
)

62
tests-unit/comfy_api_test/video_bit_depth_test.py
View File

@@ -3,8 +3,10 @@ import torch
import av
import numpy as np
from fractions import Fraction
from types import SimpleNamespace
from comfy_api.latest._input_impl.video_types import VideoFromFile, VideoFromComponents
from comfy_api.latest._util.video_types import VideoComponents
from comfy_extras.nodes_video import CreateVideo, SaveVideo
@pytest.fixture(scope="module")
@@ -58,6 +60,66 @@ def test_create_video_bit_depth(src8, src10):
assert decoded_levels(src10) > 2 * decoded_levels(src8)
@pytest.mark.parametrize(
"bit_depth,color_space,expected_bit_depth",
[
("auto", "sRGB", 8),
("auto", "HDR", 10),
("auto", "HDR PQ", 10),
(8, "HDR", 8),
(10, "sRGB", 10),
],
)
def test_create_video_node_bit_depth(gradient_components, bit_depth, color_space, expected_bit_depth):
video = CreateVideo.execute(
gradient_components.images,
float(gradient_components.frame_rate),
bit_depth=bit_depth,
color_space=color_space,
).args[0]
assert video.get_bit_depth() == expected_bit_depth
assert video.get_color_space() == color_space
def test_create_video_node_bit_depth_options():
bit_depth_input = next(input for input in CreateVideo.define_schema().inputs if input.id == "bit_depth")
assert bit_depth_input.options == ["auto", 8, 10]
assert bit_depth_input.default == "auto"
@pytest.mark.parametrize(
"codec,expected_suffix,expected_codec",
[
("auto", "mp4", "h264"),
("h264", "mp4", "h264"),
("av1", "webm", "av1"),
],
)
def test_save_video_auto_format(gradient_components, tmp_path, monkeypatch, codec, expected_suffix, expected_codec):
monkeypatch.setattr(SaveVideo, "hidden", SimpleNamespace(prompt=None, extra_pnginfo=None))
monkeypatch.setattr("comfy_extras.nodes_video.folder_paths.get_output_directory", lambda: str(tmp_path))
monkeypatch.setattr(
"comfy_extras.nodes_video.folder_paths.get_save_image_path",
lambda *args: (str(tmp_path), "auto", 1, "", "auto"),
)
video = VideoFromComponents(gradient_components)
SaveVideo.execute(
video,
"auto",
{"format": "auto", "codec": {"codec": codec}},
)
path = tmp_path / f"auto_00001_.{expected_suffix}"
with av.open(path) as container:
assert container.streams.video[0].codec.canonical_name == expected_codec
def test_save_video_has_no_color_space_input():
schema = SaveVideo.define_schema()
assert all("color_space" not in str(input.as_dict()) for input in schema.inputs)
def test_save_auto_keeps_source_depth(src8, src10, tmp_path):
"""Save Video (no bit_depth = auto) stream-copies the source, preserving its depth byte-for-byte"""
for name, src in [("p8", src8), ("p10", src10)]:

42
tests-unit/comfy_api_test/video_types_test.py
View File

@@ -139,6 +139,19 @@ def test_video_color_space_defaults_to_srgb(simple_video_file, video_components)
assert VideoFromComponents(video_components).get_color_space() == "sRGB"
@pytest.mark.parametrize("color_space", ["sRGB", "HDR", "HDR PQ"])
@pytest.mark.parametrize("bit_depth", [8, 10])
def test_video_from_components_color_space(video_components, color_space, bit_depth):
video = VideoFromComponents(video_components, bit_depth=bit_depth, color_space=color_space)
assert video.get_color_space() == color_space
assert video.get_bit_depth() == bit_depth
def test_video_from_components_rejects_invalid_color_space(video_components):
with pytest.raises(ValueError, match="Unsupported video color space"):
VideoFromComponents(video_components, color_space="Display P3")
def test_video_from_file_bytesio_input():
"""VideoFromFile works with BytesIO input"""
buffer = io.BytesIO()
@@ -426,14 +439,15 @@ def test_save_components_container_codec_and_audio_matrix(
@pytest.mark.parametrize(
"color_space,transfer,pix_fmt,primaries,colorspace",
"color_space,transfer,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),
("sRGB", ColorTrc.IEC61966_2_1, ColorPrimaries.BT709, 1),
("HDR", ColorTrc.ARIB_STD_B67, ColorPrimaries.BT2020, 9),
("HDR PQ", ColorTrc.SMPTE2084, ColorPrimaries.BT2020, 9),
],
)
def test_save_to_av1_mkv_color_space(tmp_path, color_space, transfer, pix_fmt, primaries, colorspace):
@pytest.mark.parametrize("bit_depth,pix_fmt", [(8, "yuv420p"), (10, "yuv420p10le")])
def test_save_to_av1_mkv_color_space(tmp_path, color_space, transfer, primaries, colorspace, bit_depth, pix_fmt):
components = VideoComponents(
images=torch.rand(2, 64, 64, 3),
frame_rate=Fraction(30),
@@ -441,12 +455,11 @@ def test_save_to_av1_mkv_color_space(tmp_path, color_space, transfer, pix_fmt, p
path = str(tmp_path / "hdr.mkv")
remuxed = str(tmp_path / "remuxed.mkv")
VideoFromComponents(components).save_to(
VideoFromComponents(components, bit_depth=bit_depth, color_space=color_space).save_to(
path,
format=VideoContainer.MKV,
codec=VideoCodec.AV1,
crf=30,
color_space=color_space,
metadata={"prompt": {"test": "hdr"}},
)
@@ -485,21 +498,22 @@ def test_save_to_av1_mkv_color_space(tmp_path, color_space, transfer, pix_fmt, p
],
)
@pytest.mark.parametrize(
"color_space,transfer,pix_fmt,primaries,colorspace",
"color_space,transfer,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),
("sRGB", ColorTrc.IEC61966_2_1, ColorPrimaries.BT709, 1),
("HDR", ColorTrc.ARIB_STD_B67, ColorPrimaries.BT2020, 9),
("HDR PQ", ColorTrc.SMPTE2084, ColorPrimaries.BT2020, 9),
],
)
def test_save_to_h264_color_space(tmp_path, format, suffix, color_space, transfer, pix_fmt, primaries, colorspace):
@pytest.mark.parametrize("bit_depth,pix_fmt", [(8, "yuv420p"), (10, "yuv420p10le")])
def test_save_to_h264_color_space(tmp_path, format, suffix, color_space, transfer, primaries, colorspace, bit_depth, pix_fmt):
components = VideoComponents(
images=torch.rand(2, 64, 64, 3),
frame_rate=Fraction(30),
)
path = str(tmp_path / f"h264.{suffix}")
VideoFromComponents(components).save_to(
VideoFromComponents(components, bit_depth=bit_depth).save_to(
path,
format=format,
codec=VideoCodec.H264,
@@ -604,7 +618,7 @@ def test_save_to_av1_webm_transcodes_audio(tmp_path):
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.format.name == "yuv420p"
assert video_stream.color_primaries == ColorPrimaries.BT2020
assert video_stream.color_trc == ColorTrc.ARIB_STD_B67
assert video_stream.colorspace == 9