diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py index 7011d9c13..5c12ff926 100644 --- a/comfy_extras/nodes_images.py +++ b/comfy_extras/nodes_images.py @@ -189,6 +189,8 @@ class ImageAddNoise(IO.ComfyNode): def execute(cls, image, seed, strength) -> IO.NodeOutput: generator = torch.manual_seed(seed) s = torch.clip((image + strength * torch.randn(image.size(), generator=generator, device="cpu").to(image)), min=0.0, max=1.0) + if image.shape[-1] == 4: # alpha stores transparency, not color + s[..., 3] = image[..., 3] return IO.NodeOutput(s) repeat = execute # TODO: remove diff --git a/tests-unit/comfy_extras_test/image_add_noise_alpha_test.py b/tests-unit/comfy_extras_test/image_add_noise_alpha_test.py new file mode 100644 index 000000000..3f71e834c --- /dev/null +++ b/tests-unit/comfy_extras_test/image_add_noise_alpha_test.py @@ -0,0 +1,63 @@ +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_images import ImageAddNoise # noqa: E402 + + +def image(channels, value=0.5, alpha=0.5, size=8): + t = torch.full((1, size, size, channels), value) + if channels == 4: + t[..., 3] = alpha + return t + + +def test_rgb_gets_noise_on_every_channel(): + src = image(3) + + s = ImageAddNoise.execute(src, 0, 0.5).result[0] + + assert s.shape == src.shape + assert not torch.equal(s, src) + + +def test_rgba_keeps_alpha_untouched(): + src = image(4) + + s = ImageAddNoise.execute(src, 0, 0.5).result[0] + + assert s.shape[-1] == 4 + assert torch.equal(s[..., 3], src[..., 3]) + assert not torch.equal(s[..., :3], src[..., :3]) + + +def test_fully_transparent_pixels_stay_transparent(): + src = image(4, alpha=0.0) + + s = ImageAddNoise.execute(src, 0, 1.0).result[0] + + assert torch.all(s[..., 3] == 0.0) + + +def test_only_the_alpha_channel_is_guarded(): + """Colour channels must still get exactly the unguarded noise result.""" + src = image(4) + generator = torch.manual_seed(0) + unguarded = torch.clip(src + 0.5 * torch.randn(src.size(), generator=generator), min=0.0, max=1.0) + + s = ImageAddNoise.execute(src, 0, 0.5).result[0] + + assert torch.equal(s[..., :3], unguarded[..., :3]) + assert torch.equal(s[..., 3], src[..., 3]) + + +def test_does_not_mutate_input(): + src = image(4) + before = src.clone() + + ImageAddNoise.execute(src, 0, 0.5) + + assert torch.equal(src, before)