diff --git a/comfy_extras/nodes_post_processing.py b/comfy_extras/nodes_post_processing.py index 763b8a52f..3dc3d5018 100644 --- a/comfy_extras/nodes_post_processing.py +++ b/comfy_extras/nodes_post_processing.py @@ -47,6 +47,8 @@ class Blend(io.ComfyNode): blended_image = cls.blend_mode(image1, image2, blend_mode) blended_image = image1 * (1 - blend_factor) + blended_image * blend_factor blended_image = torch.clamp(blended_image, 0, 1) + if image1.shape[-1] == 4: # alpha stores transparency, not color + blended_image[..., 3] = image1[..., 3] return io.NodeOutput(blended_image) @classmethod diff --git a/tests-unit/comfy_extras_test/image_blend_alpha_test.py b/tests-unit/comfy_extras_test/image_blend_alpha_test.py new file mode 100644 index 000000000..6178f82a5 --- /dev/null +++ b/tests-unit/comfy_extras_test/image_blend_alpha_test.py @@ -0,0 +1,84 @@ +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_post_processing import Blend # noqa: E402 + +MODES = ["normal", "multiply", "screen", "overlay", "soft_light", "difference"] + + +def image(value, alpha=None, size=4): + channels = 3 if alpha is None else 4 + t = torch.full((1, size, size, channels), value) + if alpha is not None: + t[..., 3] = alpha + return t + + +@pytest.mark.parametrize("mode", MODES) +def test_rgb_blend_is_unchanged(mode): + """3 channel images must keep going through the untouched code path.""" + image1, image2 = image(0.8), image(0.3) + + out = Blend.execute(image1, image2, 0.5, mode).result[0] + + assert out.shape == image1.shape + assert torch.all(out >= 0.0) and torch.all(out <= 1.0) + + +@pytest.mark.parametrize("mode", MODES) +def test_rgba_keeps_image1_alpha(mode): + image1, image2 = image(0.8, alpha=0.6), image(0.3, alpha=0.1) + + out = Blend.execute(image1, image2, 0.5, mode).result[0] + + assert out.shape[-1] == 4 + assert torch.equal(out[..., 3], image1[..., 3]) + + +@pytest.mark.parametrize("mode", MODES) +def test_two_opaque_images_stay_opaque(mode): + """Regression: difference mode used to compute 1.0 - 1.0 and erase the image.""" + image1, image2 = image(0.8, alpha=1.0), image(0.3, alpha=1.0) + + out = Blend.execute(image1, image2, 1.0, mode).result[0] + + assert torch.all(out[..., 3] == 1.0) + + +@pytest.mark.parametrize("mode", MODES) +def test_rgb_channels_still_blend_on_rgba(mode): + """Preserving alpha must not stop the colour channels from blending. + + Mid-tones on purpose: pure white over pure black is a fixed point for + several of the modes, so it would not prove anything. + """ + image1, image2 = image(0.6, alpha=1.0), image(0.25, alpha=1.0) + + out = Blend.execute(image1, image2, 1.0, mode).result[0] + + assert not torch.equal(out[..., :3], image1[..., :3]) + + +def test_does_not_mutate_inputs(): + image1, image2 = image(0.8, alpha=0.6), image(0.3, alpha=0.1) + before1, before2 = image1.clone(), image2.clone() + + Blend.execute(image1, image2, 0.5, "difference") + + assert torch.equal(image1, before1) + assert torch.equal(image2, before2) + + +def test_mismatched_channel_counts_still_supported(): + """image_alpha_fix (CORE-103) pads the RGB input; that must keep working.""" + rgba, rgb = image(0.8, alpha=0.5), image(0.3) + + out = Blend.execute(rgba, rgb, 0.5, "normal").result[0] + + assert out.shape[-1] == 4 + assert torch.equal(out[..., 3], rgba[..., 3])