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