|
1 | 1 | from kornia.filters import canny |
| 2 | +from typing_extensions import override |
| 3 | + |
2 | 4 | import comfy.model_management |
| 5 | +from comfy_api.latest import ComfyExtension, io |
3 | 6 |
|
4 | 7 |
|
5 | | -class Canny: |
| 8 | +class Canny(io.ComfyNode): |
6 | 9 | @classmethod |
7 | | - def INPUT_TYPES(s): |
8 | | - return {"required": {"image": ("IMAGE",), |
9 | | - "low_threshold": ("FLOAT", {"default": 0.4, "min": 0.01, "max": 0.99, "step": 0.01}), |
10 | | - "high_threshold": ("FLOAT", {"default": 0.8, "min": 0.01, "max": 0.99, "step": 0.01}) |
11 | | - }} |
12 | | - |
13 | | - RETURN_TYPES = ("IMAGE",) |
14 | | - FUNCTION = "detect_edge" |
| 10 | + def define_schema(cls): |
| 11 | + return io.Schema( |
| 12 | + node_id="Canny", |
| 13 | + category="image/preprocessors", |
| 14 | + inputs=[ |
| 15 | + io.Image.Input("image"), |
| 16 | + io.Float.Input("low_threshold", default=0.4, min=0.01, max=0.99, step=0.01), |
| 17 | + io.Float.Input("high_threshold", default=0.8, min=0.01, max=0.99, step=0.01), |
| 18 | + ], |
| 19 | + outputs=[io.Image.Output()], |
| 20 | + ) |
15 | 21 |
|
16 | | - CATEGORY = "image/preprocessors" |
| 22 | + @classmethod |
| 23 | + def detect_edge(cls, image, low_threshold, high_threshold): |
| 24 | + # Deprecated: use the V3 schema's `execute` method instead of this. |
| 25 | + return cls.execute(image, low_threshold, high_threshold) |
17 | 26 |
|
18 | | - def detect_edge(self, image, low_threshold, high_threshold): |
| 27 | + @classmethod |
| 28 | + def execute(cls, image, low_threshold, high_threshold) -> io.NodeOutput: |
19 | 29 | output = canny(image.to(comfy.model_management.get_torch_device()).movedim(-1, 1), low_threshold, high_threshold) |
20 | 30 | img_out = output[1].to(comfy.model_management.intermediate_device()).repeat(1, 3, 1, 1).movedim(1, -1) |
21 | | - return (img_out,) |
| 31 | + return io.NodeOutput(img_out) |
| 32 | + |
| 33 | + |
| 34 | +class CannyExtension(ComfyExtension): |
| 35 | + @override |
| 36 | + async def get_node_list(self) -> list[type[io.ComfyNode]]: |
| 37 | + return [Canny] |
| 38 | + |
22 | 39 |
|
23 | | -NODE_CLASS_MAPPINGS = { |
24 | | - "Canny": Canny, |
25 | | -} |
| 40 | +async def comfy_entrypoint() -> CannyExtension: |
| 41 | + return CannyExtension() |
0 commit comments