diff --git a/nodes.py b/nodes.py index 1a3dd3f48..25a6dc2c3 100644 --- a/nodes.py +++ b/nodes.py @@ -1947,6 +1947,8 @@ class ImageInvert: def invert(self, image): s = 1.0 - image + if image.shape[-1] == 4: # alpha stores transparency, not color + s[..., 3] = image[..., 3] return (s,) class ImageBatch: diff --git a/tests-unit/comfy_test/image_invert_test.py b/tests-unit/comfy_test/image_invert_test.py new file mode 100644 index 000000000..aa39a9126 --- /dev/null +++ b/tests-unit/comfy_test/image_invert_test.py @@ -0,0 +1,46 @@ +import torch + +from comfy.cli_args import args as cli_args + +if not torch.cuda.is_available(): + cli_args.cpu = True + +import nodes # noqa: E402 + + +def test_invert_rgb_inverts_every_channel(): + image = torch.rand(1, 4, 4, 3) + + s, = nodes.ImageInvert().invert(image) + + assert s.shape == image.shape + assert torch.allclose(s, 1.0 - image) + + +def test_invert_rgba_leaves_alpha_untouched(): + image = torch.rand(2, 4, 4, 4) + + s, = nodes.ImageInvert().invert(image) + + assert s.shape == image.shape + assert torch.allclose(s[..., :3], 1.0 - image[..., :3]) + assert torch.equal(s[..., 3], image[..., 3]) + + +def test_invert_rgba_keeps_transparent_pixels_transparent(): + image = torch.zeros(1, 1, 2, 4) + image[..., 3] = torch.tensor([1.0, 0.0]) + + s, = nodes.ImageInvert().invert(image) + + assert s[0, 0, 0, 3] == 1.0 + assert s[0, 0, 1, 3] == 0.0 + + +def test_invert_does_not_mutate_input(): + image = torch.rand(1, 4, 4, 4) + original = image.clone() + + nodes.ImageInvert().invert(image) + + assert torch.equal(image, original)