mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-17 14:59:12 +08:00
66 lines
1.5 KiB
Python
66 lines
1.5 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from comfy.cli_args import args as cli_args
|
|
|
|
if not torch.cuda.is_available():
|
|
cli_args.cpu = True
|
|
|
|
from comfy_extras.nodes_post_processing import Quantize # noqa: E402
|
|
|
|
DITHERS = ["none", "floyd-steinberg", "bayer-2", "bayer-4", "bayer-8", "bayer-16"]
|
|
|
|
|
|
def image(channels, alpha=0.8, size=8):
|
|
torch.manual_seed(0)
|
|
t = torch.rand(1, size, size, channels)
|
|
if channels == 4:
|
|
t[..., 3] = alpha
|
|
return t
|
|
|
|
|
|
@pytest.mark.parametrize("dither", DITHERS)
|
|
def test_rgb_still_quantizes(dither):
|
|
src = image(3)
|
|
|
|
out = Quantize.execute(src, 4, dither).result[0]
|
|
|
|
assert out.shape == src.shape
|
|
assert len(torch.unique(out.reshape(-1, 3), dim=0)) <= 4
|
|
|
|
|
|
@pytest.mark.parametrize("dither", DITHERS)
|
|
def test_rgba_does_not_raise_and_keeps_alpha(dither):
|
|
src = image(4)
|
|
|
|
out = Quantize.execute(src, 4, dither).result[0]
|
|
|
|
assert out.shape == src.shape
|
|
assert torch.equal(out[..., 3], src[..., 3])
|
|
|
|
|
|
def test_rgba_colour_channels_are_quantized():
|
|
src = image(4)
|
|
|
|
out = Quantize.execute(src, 4, "none").result[0]
|
|
|
|
assert len(torch.unique(out[..., :3].reshape(-1, 3), dim=0)) <= 4
|
|
|
|
|
|
def test_varying_alpha_is_preserved_per_pixel():
|
|
src = image(4)
|
|
src[0, :, :, 3] = torch.linspace(0.0, 1.0, src.shape[2])
|
|
|
|
out = Quantize.execute(src, 8, "none").result[0]
|
|
|
|
assert torch.equal(out[..., 3], src[..., 3])
|
|
|
|
|
|
def test_does_not_mutate_input():
|
|
src = image(4)
|
|
before = src.clone()
|
|
|
|
Quantize.execute(src, 4, "none")
|
|
|
|
assert torch.equal(src, before)
|