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

@@ -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)
 

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)