Faster blur.

comfyanonymous · Apr 19, 2024 · 133dc33 · 133dc33
1 parent 5d08802
commit 133dc33
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/comfy_extras/nodes_post_processing.py b/comfy_extras/nodes_post_processing.py
@@ -5,6 +5,7 @@
 import math
 
 import comfy.utils
+import comfy.model_management
 
 
 class Blend:
@@ -102,6 +103,7 @@ def blur(self, image: torch.Tensor, blur_radius: int, sigma: float):
         if blur_radius == 0:
             return (image,)
 
+        image = image.to(comfy.model_management.get_torch_device())
         batch_size, height, width, channels = image.shape
 
         kernel_size = blur_radius * 2 + 1
@@ -112,7 +114,7 @@ def blur(self, image: torch.Tensor, blur_radius: int, sigma: float):
         blurred = F.conv2d(padded_image, kernel, padding=kernel_size // 2, groups=channels)[:,:,blur_radius:-blur_radius, blur_radius:-blur_radius]
         blurred = blurred.permute(0, 2, 3, 1)
 
-        return (blurred,)
+        return (blurred.to(comfy.model_management.intermediate_device()),)
 
 class Quantize:
     def __init__(self):