diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index 71e5ee368..81365717f 100644 --- a/comfy_extras/nodes_dataset.py +++ b/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): diff --git a/tests-unit/comfy_extras_test/dataset_color_nodes_alpha_test.py b/tests-unit/comfy_extras_test/dataset_color_nodes_alpha_test.py new file mode 100644 index 000000000..70b63abd6 --- /dev/null +++ b/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)