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