Fix Quantize Image on images with an alpha channel

Quantize allocated its result buffer from the input shape but only ever
filled three channels, so a 4 channel image raised a RuntimeError. Quantize
the colour channels and carry the original alpha through untouched.

CORE-393

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

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)