From d70b0bc43c2d1d0a2212e9493926ddf231c80a41 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 29 Aug 2023 17:58:40 -0400 Subject: [PATCH] Use the GPU for the canny preprocessor when available. --- comfy_extras/nodes_canny.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/comfy_extras/nodes_canny.py b/comfy_extras/nodes_canny.py index d7c3f132..94d453f2 100644 --- a/comfy_extras/nodes_canny.py +++ b/comfy_extras/nodes_canny.py @@ -3,7 +3,7 @@ import math import torch import torch.nn.functional as F - +import comfy.model_management def get_canny_nms_kernel(device=None, dtype=None): """Utility function that returns 3x3 kernels for the Canny Non-maximal suppression.""" @@ -290,8 +290,8 @@ class Canny: CATEGORY = "image/preprocessors" def detect_edge(self, image, low_threshold, high_threshold): - output = canny(image.movedim(-1, 1), low_threshold, high_threshold) - img_out = output[1].repeat(1, 3, 1, 1).movedim(1, -1) + output = canny(image.to(comfy.model_management.get_torch_device()).movedim(-1, 1), low_threshold, high_threshold) + img_out = output[1].cpu().repeat(1, 3, 1, 1).movedim(1, -1) return (img_out,) NODE_CLASS_MAPPINGS = {