Speed up Sharpen node.

guill · Apr 21, 2024 · 4ee9aad · 4ee9aad
1 parent 644a3ae
commit 4ee9aad
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/comfy_extras/nodes_post_processing.py b/comfy_extras/nodes_post_processing.py
@@ -227,6 +227,7 @@ def sharpen(self, image: torch.Tensor, sharpen_radius: int, sigma:float, alpha:
             return (image,)
 
         batch_size, height, width, channels = image.shape
+        image = image.to(comfy.model_management.get_torch_device())
 
         kernel_size = sharpen_radius * 2 + 1
         kernel = gaussian_kernel(kernel_size, sigma, device=image.device) * -(alpha*10)
@@ -241,7 +242,7 @@ def sharpen(self, image: torch.Tensor, sharpen_radius: int, sigma:float, alpha:
 
         result = torch.clamp(sharpened, 0, 1)
 
-        return (result,)
+        return (result.to(comfy.model_management.intermediate_device()),)
 
 class ImageScaleToTotalPixels:
     upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]