mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-06-14 20:09:24 +08:00
fe54b5e955
Create Video gets a bit_depth option (8-bit/10-bit); the selected depth is carried by the video and applied when it gets encoded. Save Video and Video Slice now keep the source bit depth instead of always quantizing to 8-bit, so 10-bit videos stay 10-bit. 10-bit uses h264 with the yuv420p10le pixel format,so there's no new codec or container. Signed-off-by: bigcat88 <bigcat88@icloud.com>
94 lines
3.8 KiB
Python
94 lines
3.8 KiB
Python
import pytest
|
|
import torch
|
|
import av
|
|
import numpy as np
|
|
from fractions import Fraction
|
|
from comfy_api.latest._input_impl.video_types import VideoFromFile, VideoFromComponents
|
|
from comfy_api.latest._util.video_types import VideoComponents
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def gradient_components():
|
|
"""Narrow horizontal ramp (0.25..0.30) that needs more than 8 bits to stay smooth"""
|
|
width, height, frames = 64, 64, 3
|
|
ramp = torch.linspace(0.25, 0.30, width).view(1, 1, width, 1).expand(frames, height, width, 3)
|
|
return VideoComponents(images=ramp.contiguous(), frame_rate=Fraction(30))
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def src8(gradient_components, tmp_path_factory):
|
|
"""8-bit h264 mp4 (Create Video default)"""
|
|
path = str(tmp_path_factory.mktemp("video") / "src8.mp4")
|
|
VideoFromComponents(gradient_components).save_to(path)
|
|
return path
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def src10(gradient_components, tmp_path_factory):
|
|
"""10-bit h264 mp4 (Create Video with bit_depth=10)"""
|
|
path = str(tmp_path_factory.mktemp("video") / "src10.mp4")
|
|
VideoFromComponents(gradient_components, bit_depth=10).save_to(path)
|
|
return path
|
|
|
|
|
|
def probe(path):
|
|
"""(codec, pix_fmt, bit_depth) of the first video stream"""
|
|
with av.open(path) as container:
|
|
stream = container.streams.video[0]
|
|
return (stream.codec.name, stream.format.name, max(c.bits for c in stream.format.components))
|
|
|
|
|
|
def decoded_levels(path):
|
|
"""Unique tonal levels in the first decoded frame (banding measure)"""
|
|
with av.open(path) as container:
|
|
frame = next(container.decode(container.streams.video[0]))
|
|
return len(np.unique(frame.to_ndarray(format="gbrpf32le")[..., 0]))
|
|
|
|
|
|
def video_packet_bytes(path):
|
|
"""Raw video packet payloads; identical to the source's only for a true remux"""
|
|
with av.open(path) as container:
|
|
return [bytes(p) for p in container.demux(container.streams.video[0]) if p.size]
|
|
|
|
|
|
def test_create_video_bit_depth(src8, src10):
|
|
"""Create Video's bit_depth picks the encoded depth (default 8-bit); 10-bit reduces banding"""
|
|
assert probe(src8) == ("h264", "yuv420p", 8)
|
|
assert probe(src10) == ("h264", "yuv420p10le", 10)
|
|
assert decoded_levels(src10) > 2 * decoded_levels(src8)
|
|
|
|
|
|
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)]:
|
|
path = str(tmp_path / f"{name}.mp4")
|
|
VideoFromFile(src).save_to(path)
|
|
assert probe(path) == probe(src)
|
|
assert video_packet_bytes(path) == video_packet_bytes(src)
|
|
|
|
|
|
def test_save_explicit_depth_reencodes(src8, src10, tmp_path):
|
|
"""An explicit bit_depth different from the source forces a re-encode to that depth"""
|
|
down = str(tmp_path / "down8.mp4")
|
|
VideoFromFile(src10).save_to(down, bit_depth=8)
|
|
assert probe(down) == ("h264", "yuv420p", 8)
|
|
|
|
up = str(tmp_path / "up10.mp4")
|
|
VideoFromFile(src8).save_to(up, bit_depth=10)
|
|
assert probe(up) == ("h264", "yuv420p10le", 10)
|
|
|
|
|
|
def test_trim_keeps_source_depth(src10, tmp_path):
|
|
"""Video Slice re-encodes (trim) but preserves the source's 10-bit depth"""
|
|
path = str(tmp_path / "trim.mp4")
|
|
VideoFromFile(src10).as_trimmed(start_time=0, duration=1 / 30, strict_duration=False).save_to(path)
|
|
assert probe(path) == ("h264", "yuv420p10le", 10)
|
|
|
|
|
|
def test_get_bit_depth(gradient_components, src8, src10):
|
|
"""get_bit_depth reports a video's depth (backs the Get Video Components output)"""
|
|
assert VideoFromFile(src8).get_bit_depth() == 8
|
|
assert VideoFromFile(src10).get_bit_depth() == 10
|
|
assert VideoFromComponents(gradient_components, bit_depth=10).get_bit_depth() == 10
|
|
assert VideoFromComponents(gradient_components).get_bit_depth() == 8
|