331 lines
14 KiB
Python
331 lines
14 KiB
Python
"""Real CPU-only IMAGE transfers; no ComfyUI/bootstrap/provider imports."""
|
|
import importlib
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from PIL import Image
|
|
|
|
# Load a private package, not Distributed's root __init__.
|
|
ROOT = Path(__file__).resolve().parents[1] / "serverless_bridge"
|
|
PACKAGE = "image_codec_tests_bridge"
|
|
pkg = types.ModuleType(PACKAGE)
|
|
pkg.__path__ = [str(ROOT)]
|
|
sys.modules[PACKAGE] = pkg
|
|
codec = importlib.import_module(PACKAGE + ".codec")
|
|
images = importlib.import_module(PACKAGE + ".image_codec")
|
|
FORMATS = ("legacy_png", "png", "jpeg", "webp", "ffv1", "h264", "h265", "av1")
|
|
LOSSLESS = ("legacy_png", "png", "ffv1")
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def no_gpu(monkeypatch):
|
|
def forbidden(*a, **kw):
|
|
pytest.fail("GPU query or initialization is forbidden")
|
|
for name in ("is_available", "init", "_lazy_init", "device_count"):
|
|
monkeypatch.setattr(torch.cuda, name, forbidden)
|
|
|
|
|
|
def available(fmt):
|
|
if fmt in images._VIDEO_CODECS:
|
|
av = pytest.importorskip("av")
|
|
name = images._VIDEO_CODECS[fmt][0]
|
|
if name not in av.codecs_available:
|
|
pytest.skip("Software encoder unavailable: " + name)
|
|
|
|
|
|
def sample(channels=3, dtype=torch.float32, shape=(3, 65, 67)):
|
|
g = torch.Generator(device="cpu").manual_seed(42)
|
|
return torch.rand((*shape, channels), generator=g, device="cpu", dtype=dtype)
|
|
|
|
|
|
def quantized(value):
|
|
return torch.from_numpy(np.clip(value.float().numpy() * 255, 0, 255).astype(np.uint8)).float() / 255
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
|
|
def test_exact_unchanged(dtype):
|
|
value = sample(4, dtype, (2, 3, 5))
|
|
value.flatten()[0] = -0.0
|
|
expected = codec.encode_payload("IMAGE", value)
|
|
assert images.encode_image(value) == expected
|
|
result = images.decode_image(*expected)
|
|
assert torch.equal(result.view(torch.uint8), value.view(torch.uint8))
|
|
|
|
|
|
@pytest.mark.parametrize("fmt", FORMATS)
|
|
@pytest.mark.parametrize("channels", [1, 3, 4])
|
|
def test_real_roundtrips_shape_order_quantization_alpha(fmt, channels):
|
|
available(fmt)
|
|
value = sample(channels)
|
|
# Large inter-frame difference also verifies frame ordering.
|
|
value[0] *= 0.2
|
|
value[2] = value[2] * 0.2 + 0.8
|
|
if channels == 4:
|
|
value[..., 3] = torch.linspace(-0.1, 1.1, value[..., 3].numel()).reshape(value[..., 3].shape)
|
|
descriptor, payload = images.encode_image(value, transfer_format=fmt)
|
|
assert descriptor["encoding"] == "image-transfer-v1"
|
|
result = codec.decode_payload("IMAGE", descriptor, payload)
|
|
assert result.shape == value.shape
|
|
assert result.device.type == "cpu" and result.is_contiguous()
|
|
color = 3 if channels == 4 else channels
|
|
expected = quantized(value[..., :color])
|
|
if fmt in LOSSLESS:
|
|
assert torch.equal(result[..., :color], expected)
|
|
else:
|
|
# Random RGB noise loses chroma detail in 4:2:0 even at high quality.
|
|
assert (result[..., :color] - expected).abs().mean() < 0.12
|
|
assert result[0, ..., :color].mean() < result[2, ..., :color].mean()
|
|
if channels == 4:
|
|
assert torch.equal(result[..., 3], value[..., 3])
|
|
|
|
|
|
@pytest.mark.parametrize("fmt", ["png", "ffv1", "jpeg", "webp", "h264", "h265", "av1"])
|
|
def test_float64_alpha_is_exact(fmt):
|
|
available(fmt)
|
|
value = sample(4, torch.float64, (1, 65, 67))
|
|
descriptor, payload = images.encode_image(value, transfer_format=fmt)
|
|
out = images.decode_image(descriptor, payload)
|
|
assert out.dtype == torch.float64
|
|
assert torch.equal(out[..., 3].contiguous().view(torch.uint8), value[..., 3].contiguous().view(torch.uint8))
|
|
|
|
|
|
@pytest.mark.parametrize("fmt", ["jpeg", "webp", "h264", "h265", "av1"])
|
|
def test_lossy_quality_changes_final_payload_size(fmt):
|
|
available(fmt)
|
|
value = sample(shape=(2, 96, 96))
|
|
low = images.encode_image(value, transfer_format=fmt, transfer_quality=15)[1]
|
|
high = images.encode_image(value, transfer_format=fmt, transfer_quality=95)[1]
|
|
assert len(high) > len(low)
|
|
|
|
|
|
def test_truncation_clipping_grayscale_noncontiguous():
|
|
value = torch.tensor([-0.1, 0.5, 1.1, 0.12345], device="cpu").reshape(1, 2, 2, 1).transpose(1, 2)
|
|
out = images.decode_image(*images.encode_image(value, transfer_format="png"))
|
|
assert torch.equal(out, quantized(value))
|
|
assert out[0, 0, 1, 0] == 1
|
|
|
|
|
|
@pytest.mark.parametrize("value", [None, [], torch.ones(2, 3, 3), torch.ones(1, 2, 3, 2), torch.ones(1, 2, 3, 0), torch.ones(1, 2, 3, 3, dtype=torch.uint8), torch.ones(1, 2, 3, 3, device="meta"), torch.ones(1, 2, 3, 3).to_sparse()])
|
|
def test_invalid_values(value):
|
|
with pytest.raises(ValueError):
|
|
images.encode_image(value, transfer_format="png")
|
|
|
|
|
|
@pytest.mark.parametrize("bad", [float("nan"), float("inf"), -float("inf")])
|
|
def test_nonfinite_rejected(bad):
|
|
with pytest.raises(ValueError, match="finite"):
|
|
images.encode_image(torch.full((1, 1, 1, 3), bad), transfer_format="png")
|
|
|
|
|
|
@pytest.mark.parametrize("options", [{"transfer_format": "PNG"}, {"transfer_format": []}, {"transfer_quality": True}, {"transfer_quality": 0}, {"transfer_quality": 101}, {"transfer_quality": 1.5}, {"max_bytes": True}, {"max_bytes": -1}])
|
|
def test_strict_options(options):
|
|
with pytest.raises(ValueError):
|
|
images.encode_image(sample(shape=(1, 1, 1)), **options)
|
|
|
|
|
|
@pytest.mark.parametrize("change", [{"extra": 1}, {"format": "exact"}, {"quality": True}, {"layout": "BCHW"}, {"shape": [1, 2**60, 1, 3]}, {"shape": [True, 1, 1, 3]}, {"parts": [0]}, {"parts": [True]}, {"parts": [2**60]}, {"alpha": {"extra": 0}}])
|
|
def test_bad_descriptors(change):
|
|
descriptor, payload = images.encode_image(sample(shape=(1, 3, 5)), transfer_format="png")
|
|
descriptor.update(change)
|
|
with pytest.raises(ValueError):
|
|
images.decode_image(descriptor, payload)
|
|
|
|
|
|
def test_truncated_appended_and_mismatched_parts():
|
|
d, p = images.encode_image(sample(shape=(2, 3, 5)), transfer_format="png")
|
|
for payload in (p[:-1], p + b"x"):
|
|
with pytest.raises(ValueError):
|
|
images.decode_image(d, payload)
|
|
for change in ({"shape": [1, 3, 5, 3]}, {"parts": [len(p)]}, {"shape": [2, 2, 5, 3]}):
|
|
bad = dict(d, **change)
|
|
with pytest.raises(ValueError):
|
|
images.decode_image(bad, p)
|
|
|
|
|
|
@pytest.mark.parametrize("fmt", ["png", "jpeg", "webp", "ffv1", "h264", "h265", "av1"])
|
|
def test_payload_and_decoded_budget(fmt):
|
|
available(fmt)
|
|
value = sample(shape=(1, 65, 67))
|
|
d, p = images.encode_image(value, transfer_format=fmt)
|
|
with pytest.raises(ValueError, match="limit|budget"):
|
|
images.decode_image(d, p, max_bytes=100)
|
|
with pytest.raises(ValueError, match="limit|budget"):
|
|
images.encode_image(value, transfer_format=fmt, max_bytes=100)
|
|
|
|
|
|
def test_oversized_shape_before_materialization(monkeypatch):
|
|
value = torch.ones(1).expand(1, 8192, 8192, 4)
|
|
def forbidden(*a, **kw):
|
|
pytest.fail("Materialized an oversized tensor")
|
|
monkeypatch.setattr(torch.Tensor, "contiguous", forbidden)
|
|
monkeypatch.setattr(torch, "isfinite", forbidden)
|
|
with pytest.raises(ValueError):
|
|
images.encode_image(value, transfer_format="png")
|
|
|
|
|
|
def test_header_dimensions_checked_before_pixels(monkeypatch):
|
|
d, p = images.encode_image(sample(shape=(1, 9, 11)), transfer_format="png")
|
|
d["shape"] = [1, 1, 1, 3]
|
|
def forbidden(*a, **kw):
|
|
pytest.fail("Pixels decoded before header dimensions were checked")
|
|
monkeypatch.setattr(Image.Image, "load", forbidden)
|
|
with pytest.raises(ValueError):
|
|
images.decode_image(d, p)
|
|
|
|
|
|
def test_other_kinds_do_not_dispatch():
|
|
d, p = images.encode_image(sample(shape=(1, 3, 5)), transfer_format="png")
|
|
for kind in ("MASK", "AUDIO", "STRING"):
|
|
with pytest.raises(ValueError):
|
|
codec.decode_payload(kind, d, p)
|
|
|
|
|
|
def test_missing_video_encoder_is_explicit(monkeypatch):
|
|
monkeypatch.setitem(images._VIDEO_CODECS, "h264", ("missing_software_encoder", 51, "yuv420p"))
|
|
with pytest.raises(ValueError, match="encoder|unavailable"):
|
|
images.encode_image(sample(shape=(1, 64, 64)), transfer_format="h264")
|
|
|
|
|
|
def test_bounded_sink_catches_final_flush_bytes():
|
|
sink = images._BoundedIO(4)
|
|
sink.write(b"1234")
|
|
with pytest.raises(ValueError, match="limit"):
|
|
sink.write(b"5")
|
|
assert sink.getvalue() == b"1234"
|
|
|
|
|
|
def test_crf_mapping():
|
|
assert images._quality_to_crf(90, 51) == 13
|
|
assert images._quality_to_crf(90, 63) == 16
|
|
assert images._quality_to_crf(100, 63) == 0
|
|
|
|
|
|
@pytest.mark.parametrize("fmt", FORMATS)
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
|
|
def test_compressed_preserves_source_dtype_and_exact_alpha(fmt, dtype):
|
|
available(fmt)
|
|
value = sample(4, dtype, (1, 17, 19))
|
|
value[..., 3].flatten()[0] = -0.0
|
|
d, p = images.encode_image(value, transfer_format=fmt)
|
|
with torch.device("meta"):
|
|
result = images.decode_image(d, p)
|
|
assert result.device.type == "cpu" and result.dtype == dtype
|
|
assert d["dtype"] == str(dtype).removeprefix("torch.")
|
|
assert torch.equal(result[..., 3].contiguous().view(torch.uint8), value[..., 3].contiguous().view(torch.uint8))
|
|
if fmt in LOSSLESS:
|
|
assert torch.equal(result[..., :3], quantized(value[..., :3]).to(dtype))
|
|
|
|
|
|
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
|
|
def test_video_count_and_header_fail_before_native_decoder(fmt, monkeypatch):
|
|
available(fmt)
|
|
d, p = images.encode_image(sample(shape=(2, 65, 67)), transfer_format=fmt)
|
|
def forbidden(*a, **kw):
|
|
pytest.fail("Native decoder reached before header validation")
|
|
monkeypatch.setattr(images, "_decode_video", forbidden)
|
|
for shape in ([1, 65, 67, 3], [3, 65, 67, 3], [2, 32, 32, 3]):
|
|
bad = dict(d, shape=shape)
|
|
with pytest.raises(ValueError):
|
|
images.decode_image(bad, p)
|
|
for payload in (p[:-1], p[:40], p + b"extra"):
|
|
bad = dict(d, parts=[len(payload)])
|
|
with pytest.raises(ValueError):
|
|
images.decode_image(bad, payload)
|
|
|
|
|
|
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
|
|
def test_video_does_not_use_native_container_probe(fmt, monkeypatch):
|
|
available(fmt)
|
|
d, p = images.encode_image(sample(shape=(2, 65, 67)), transfer_format=fmt)
|
|
av = images._import_av()
|
|
def forbidden(*a, **kw):
|
|
pytest.fail("Unbounded native stream probing is forbidden")
|
|
monkeypatch.setattr(av, "open", forbidden)
|
|
assert images.decode_image(d, p).shape == (2, 65, 67, 3)
|
|
|
|
|
|
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
|
|
def test_real_video_flush_bound(fmt):
|
|
available(fmt)
|
|
frames = [np.zeros((64, 64, 3), dtype=np.uint8)]
|
|
good = images._encode_video(iter(frames), fmt, 90, 64, 64, 100000)
|
|
with pytest.raises(ValueError, match="limit"):
|
|
images._encode_video(iter(frames), fmt, 90, 64, 64, len(good) - 1)
|
|
|
|
|
|
def test_alpha_invalid_dtype_nonfinite_and_extra_fields():
|
|
import copy
|
|
d, p = images.encode_image(sample(4, shape=(1, 3, 5)), transfer_format="png")
|
|
for field, value in (("dtype", "float64"), ("shape", [1, 3, 6, 1]), ("extra", 1)):
|
|
bad = copy.deepcopy(d)
|
|
bad["alpha"]["descriptor"][field] = value
|
|
with pytest.raises(ValueError):
|
|
images.decode_image(bad, p)
|
|
import struct
|
|
bad = p[:-4] + struct.pack("<f", float("nan"))
|
|
with pytest.raises(ValueError, match="finite"):
|
|
images.decode_image(d, bad)
|
|
|
|
|
|
def test_descriptor_structure_bound_and_frame_limit():
|
|
d, p = images.encode_image(sample(shape=(1, 3, 5)), transfer_format="png")
|
|
for bad in (dict(d, quality="x" * 70000), dict(d, parts=[1] * 241), dict(d, shape=[241, 1, 1, 3]), dict(d, dtype="x")):
|
|
with pytest.raises(ValueError):
|
|
images.decode_image(bad, p)
|
|
with pytest.raises(ValueError):
|
|
images.encode_image(torch.zeros(241, 1, 1, 3), transfer_format="png")
|
|
|
|
|
|
@pytest.mark.parametrize("fmt", ["png", "jpeg", "webp"])
|
|
def test_all_still_headers_checked_before_pillow_open(fmt, monkeypatch):
|
|
d, p = images.encode_image(sample(shape=(1, 31, 33)), transfer_format=fmt)
|
|
d["shape"] = [1, 2, 2, 3]
|
|
def forbidden(*a, **kw):
|
|
pytest.fail("Pillow opened a mismatched header")
|
|
monkeypatch.setattr(Image, "open", forbidden)
|
|
with pytest.raises(ValueError, match="dimensions"):
|
|
images.decode_image(d, p)
|
|
|
|
|
|
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "h265", "av1"])
|
|
def test_video_longer_batch_order_and_coded_allocation_bound(fmt):
|
|
available(fmt)
|
|
value = torch.linspace(0, 1, 16).reshape(16, 1, 1, 1).expand(16, 64, 64, 3)
|
|
d, p = images.encode_image(value, transfer_format=fmt)
|
|
header = images._video_header(p, fmt, 64, 64, 16)
|
|
coded = images._coded_frame_bound(p, fmt, header)
|
|
assert 16 <= coded <= images.MAX_FRAMES * 4
|
|
result = images.decode_image(d, p)
|
|
assert torch.all((result.mean(dim=(1, 2, 3)) - value.mean(dim=(1, 2, 3))).abs() < 0.03)
|
|
|
|
|
|
@pytest.mark.parametrize("fmt", ["h264", "h265", "av1"])
|
|
def test_packet_unit_bounds_before_native_decode(fmt):
|
|
available(fmt)
|
|
d, p = images.encode_image(sample(shape=(1, 64, 64)), transfer_format=fmt)
|
|
config, bounds = images._video_header(p, fmt, 64, 64, 1)
|
|
start, end = bounds[0]
|
|
packet = p[start:end]
|
|
with pytest.raises(ValueError):
|
|
images._coded_frame_bound(packet[:-1], fmt, (config, [(0, len(packet) - 1)]))
|
|
# An adversarial single block can contain many coded frames: bound it too.
|
|
enormous = packet * (images.MAX_FRAMES * 4 + 1)
|
|
with pytest.raises(ValueError, match="limit"):
|
|
images._coded_frame_bound(enormous, fmt, (config, [(0, len(enormous))]))
|
|
|
|
|
|
@pytest.mark.parametrize("fmt", ["h264", "h265", "av1"])
|
|
def test_bitstream_resolution_cannot_bypass_native_pixel_limit(fmt, monkeypatch):
|
|
available(fmt)
|
|
d, p = images.encode_image(sample(shape=(1, 256, 256)), transfer_format=fmt)
|
|
header = images._video_header(p, fmt, 256, 256, 1)
|
|
av = images._import_av()
|
|
# A lying container is not enough to disable max_pixels in the decoder.
|
|
# The bitstream still contains 256x256 pixels but the budget permits 64x64.
|
|
with pytest.raises((ValueError, av.error.FFmpegError)):
|
|
list(images._decode_video(p, fmt, 64, 64, 1, header))
|