73 lines
2.4 KiB
Python
73 lines
2.4 KiB
Python
# (c) TenserTensor <tenser.tensor@proton.me> || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
|
|
|
|
|
|
from typing import override
|
|
|
|
from kornia.filters import canny
|
|
|
|
from comfy import model_management
|
|
from comfy_api.latest import io, ui, ComfyExtension
|
|
from .nodes_image import get_image_files, load_image
|
|
|
|
CATEGORY = "TenserTensor/Detector"
|
|
|
|
|
|
def detect_edge(image, low_threshold, high_threshold):
|
|
torch_device = model_management.get_torch_device()
|
|
intermediate_device = model_management.intermediate_device()
|
|
timage = image.to(torch_device).movedim(-1, 1)
|
|
o_images = canny(timage, low_threshold, high_threshold)
|
|
|
|
return o_images[1].to(intermediate_device).expand(-1, 3, -1, -1).movedim(1, -1)
|
|
|
|
|
|
class TT_CannyEdgeDetectorNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="TT_CannyEdgeDetectorNode",
|
|
display_name="TT Canny Edge Detector",
|
|
category=CATEGORY,
|
|
description="",
|
|
inputs=[
|
|
io.Float.Input("low_threshold", default=0.4, min=0.01, max=0.99, step=0.01),
|
|
io.Float.Input("high_threshold", default=0.8, min=0.01, max=0.99, step=0.01),
|
|
io.Combo.Input("reference_image", options=get_image_files(), upload=io.UploadType.image)
|
|
],
|
|
outputs=[
|
|
io.Image.Output("IMAGE")
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs) -> io.NodeOutput:
|
|
timage, _ = load_image(kwargs.get("reference_image"))
|
|
cn_image = detect_edge(timage, kwargs.get("low_threshold"), kwargs.get("high_threshold"))
|
|
|
|
return io.NodeOutput(cn_image, ui=ui.PreviewImage(timage, cls=cls))
|
|
|
|
|
|
# ==============================================================================
|
|
# V3 entrypoint — registers context nodes with ComfyUI
|
|
# ==============================================================================
|
|
|
|
class DetectorNodesExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
TT_CannyEdgeDetectorNode,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> DetectorNodesExtension:
|
|
return DetectorNodesExtension()
|
|
|
|
|
|
# ==============================================================================
|
|
# Re-exports for backward compatibility
|
|
# ==============================================================================
|
|
|
|
__all__ = [
|
|
"TT_CannyEdgeDetectorNode",
|
|
]
|