From 3fd0681484d05380ed2155c474c46b49b0379fb5 Mon Sep 17 00:00:00 2001 From: Glary-Bot Date: Sat, 15 Aug 2026 00:20:02 +0000 Subject: [PATCH] Fix Detect Edges (Canny) on images with an alpha channel kornia's canny only accepts 1 or 3 channel input, so a 4 channel image raised a RuntimeError. Run edge detection on the colour channels. CORE-393 --- comfy_extras/nodes_canny.py | 2 +- .../comfy_extras_test/canny_alpha_test.py | 58 +++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) create mode 100644 tests-unit/comfy_extras_test/canny_alpha_test.py diff --git a/comfy_extras/nodes_canny.py b/comfy_extras/nodes_canny.py index 462f6fea0..a0205f52d 100644 --- a/comfy_extras/nodes_canny.py +++ b/comfy_extras/nodes_canny.py @@ -30,7 +30,7 @@ class Canny(io.ComfyNode): @classmethod def execute(cls, image, low_threshold, high_threshold) -> io.NodeOutput: - output = canny(image.to(device=comfy.model_management.get_torch_device(), dtype=torch.float32).movedim(-1, 1), low_threshold, high_threshold) + output = canny(image[..., :3].to(device=comfy.model_management.get_torch_device(), dtype=torch.float32).movedim(-1, 1), low_threshold, high_threshold) img_out = output[1].to(device=comfy.model_management.intermediate_device(), dtype=comfy.model_management.intermediate_dtype()).repeat(1, 3, 1, 1).movedim(1, -1) return io.NodeOutput(img_out) diff --git a/tests-unit/comfy_extras_test/canny_alpha_test.py b/tests-unit/comfy_extras_test/canny_alpha_test.py new file mode 100644 index 000000000..3b0a82996 --- /dev/null +++ b/tests-unit/comfy_extras_test/canny_alpha_test.py @@ -0,0 +1,58 @@ +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_canny import Canny # noqa: E402 + + +def edged_image(channels, alpha=0.8, size=16): + """Half black, half white, so there is a real edge down the middle.""" + t = torch.zeros(1, size, size, channels) + t[:, :, size // 2:, :3] = 1.0 + if channels == 4: + t[..., 3] = alpha + return t + + +def test_rgb_detects_the_edge(): + out = Canny.execute(edged_image(3), 0.4, 0.8).result[0] + + assert out.shape[-1] == 3 + assert out.max() > 0.0 + + +def test_rgba_does_not_raise(): + out = Canny.execute(edged_image(4), 0.4, 0.8).result[0] + + assert out.shape[-1] == 3 + assert out.max() > 0.0 + + +def test_rgba_and_rgb_give_the_same_edges(): + """Alpha must not influence edge detection.""" + from_rgb = Canny.execute(edged_image(3), 0.4, 0.8).result[0] + from_rgba = Canny.execute(edged_image(4), 0.4, 0.8).result[0] + + assert torch.equal(from_rgb, from_rgba) + + +def test_alpha_pattern_does_not_change_the_result(): + opaque = edged_image(4, alpha=1.0) + transparent = edged_image(4, alpha=0.0) + + from_opaque = Canny.execute(opaque, 0.4, 0.8).result[0] + from_transparent = Canny.execute(transparent, 0.4, 0.8).result[0] + + assert torch.equal(from_opaque, from_transparent) + + +def test_does_not_mutate_input(): + src = edged_image(4) + before = src.clone() + + Canny.execute(src, 0.4, 0.8) + + assert torch.equal(src, before)