diff --git a/comfy_extras/nodes_post_processing.py b/comfy_extras/nodes_post_processing.py index 763b8a52f..9b66ced0f 100644 --- a/comfy_extras/nodes_post_processing.py +++ b/comfy_extras/nodes_post_processing.py @@ -156,11 +156,12 @@ class Quantize(io.ComfyNode): @classmethod def execute(cls, image: torch.Tensor, colors: int, dither: str) -> io.NodeOutput: - batch_size, height, width, _ = image.shape - result = torch.zeros_like(image) + rgb = image[..., :3] + batch_size, height, width, _ = rgb.shape + result = torch.zeros_like(rgb) for b in range(batch_size): - im = Image.fromarray((image[b] * 255).to(torch.uint8).numpy(), mode='RGB') + im = Image.fromarray((rgb[b] * 255).to(torch.uint8).numpy(), mode='RGB') pal_im = im.quantize(colors=colors) # Required as described in https://github.com/python-pillow/Pillow/issues/5836 @@ -175,6 +176,8 @@ class Quantize(io.ComfyNode): quantized_array = torch.tensor(np.array(quantized_image.convert("RGB"))).float() / 255 result[b] = quantized_array + if image.shape[-1] == 4: + result = torch.cat((result, image[..., 3:]), dim=-1) return io.NodeOutput(result) class Sharpen(io.ComfyNode): diff --git a/tests-unit/comfy_extras_test/image_quantize_alpha_test.py b/tests-unit/comfy_extras_test/image_quantize_alpha_test.py new file mode 100644 index 000000000..27802dd98 --- /dev/null +++ b/tests-unit/comfy_extras_test/image_quantize_alpha_test.py @@ -0,0 +1,65 @@ +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-8"] + + +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)