Don't adjust the alpha channel in the image color adjustment nodes

Normalize Image Colors, Adjust Brightness and Adjust Contrast applied their
maths across every channel, so RGBA input came back with its transparency
rescaled. Fully opaque images happen to survive all three transforms, but
soft alpha edges do not.

CORE-392

comfy_extras/nodes_dataset.py

@@ -1077,7 +1077,10 @@ class NormalizeImagesNode(ImageProcessingNode):
 
     @classmethod
     def _process(cls, image, mean, std):
-        return (image - mean) / std
+        out = (image - mean) / std
+        if image.shape[-1] == 4:  # alpha stores transparency, not color
+            out[..., 3] = image[..., 3]
+        return out
 
 
 class AdjustBrightnessNode(ImageProcessingNode):
@@ -1099,7 +1102,10 @@ class AdjustBrightnessNode(ImageProcessingNode):
 
     @classmethod
     def _process(cls, image, factor):
-        return (image * factor).clamp(0.0, 1.0)
+        out = (image * factor).clamp(0.0, 1.0)
+        if image.shape[-1] == 4:  # alpha stores transparency, not color
+            out[..., 3] = image[..., 3]
+        return out
 
 
 class AdjustContrastNode(ImageProcessingNode):
@@ -1121,7 +1127,10 @@ class AdjustContrastNode(ImageProcessingNode):
 
     @classmethod
     def _process(cls, image, factor):
-        return ((image - 0.5) * factor + 0.5).clamp(0.0, 1.0)
+        out = ((image - 0.5) * factor + 0.5).clamp(0.0, 1.0)
+        if image.shape[-1] == 4:  # alpha stores transparency, not color
+            out[..., 3] = image[..., 3]
+        return out
 
 
 class ShuffleDatasetNode(ImageProcessingNode):

tests-unit/comfy_extras_test/dataset_color_nodes_alpha_test.py

@@ -0,0 +1,78 @@
+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_dataset import (  # noqa: E402
+    AdjustBrightnessNode,
+    AdjustContrastNode,
+    NormalizeImagesNode,
+)
+
+# (node, kwargs) for each colour adjustment that must leave alpha alone.
+CASES = [
+    (NormalizeImagesNode, {"mean": 0.5, "std": 0.5}),
+    (AdjustBrightnessNode, {"factor": 1.5}),
+    (AdjustContrastNode, {"factor": 1.5}),
+]
+IDS = [node.__name__ for node, _ in CASES]
+
+
+def image(channels, value=0.25, alpha=0.8, size=4):
+    # alpha 0.8 on purpose: 0.0, 0.5 and 1.0 are fixed points of one or more of
+    # these transforms, so they would hide the bug.
+    t = torch.full((1, size, size, channels), value)
+    if channels == 4:
+        t[..., 3] = alpha
+    return t
+
+
+@pytest.mark.parametrize(("node", "kwargs"), CASES, ids=IDS)
+def test_rgb_is_adjusted_on_every_channel(node, kwargs):
+    src = image(3)
+
+    out = node._process(src, **kwargs)
+
+    assert out.shape == src.shape
+    assert not torch.equal(out, src)
+
+
+@pytest.mark.parametrize(("node", "kwargs"), CASES, ids=IDS)
+def test_rgba_keeps_alpha_untouched(node, kwargs):
+    src = image(4)
+
+    out = node._process(src, **kwargs)
+
+    assert out.shape[-1] == 4
+    assert torch.equal(out[..., 3], src[..., 3])
+
+
+@pytest.mark.parametrize(("node", "kwargs"), CASES, ids=IDS)
+def test_colour_channels_still_change(node, kwargs):
+    src = image(4)
+
+    out = node._process(src, **kwargs)
+
+    assert not torch.equal(out[..., :3], src[..., :3])
+
+
+@pytest.mark.parametrize(("node", "kwargs"), CASES, ids=IDS)
+def test_transparent_pixels_stay_transparent(node, kwargs):
+    src = image(4, alpha=0.0)
+
+    out = node._process(src, **kwargs)
+
+    assert torch.all(out[..., 3] == 0.0)
+
+
+@pytest.mark.parametrize(("node", "kwargs"), CASES, ids=IDS)
+def test_does_not_mutate_input(node, kwargs):
+    src = image(4)
+    before = src.clone()
+
+    node._process(src, **kwargs)
+
+    assert torch.equal(src, before)