optimize: use gpu for dilation/erosion

fix: mask cache didn't work on regional sampler
fix: erosion didn't work on ultralytics detectors

modules/impact/config.py

@@ -2,9 +2,9 @@ import configparser
 import os
 
 
-version = "V4.29.1"
+version = "V4.29.2"
 
-dependency_version = 16
+dependency_version = 17
 
 my_path = os.path.dirname(__file__)
 old_config_path = os.path.join(my_path, "impact-pack.ini")

modules/impact/core.py

@@ -49,19 +49,19 @@ def erosion_mask(mask, grow_mask_by):
     w = mask.shape[1]
     h = mask.shape[0]
 
-    mask = mask.clone()
-
+    device = comfy.model_management.get_torch_device()
+    mask = mask.clone().to(device)
     mask2 = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(w, h),
-                                            mode="bilinear")
+                                            mode="bilinear").to(device)
     if grow_mask_by == 0:
         mask_erosion = mask2
     else:
-        kernel_tensor = torch.ones((1, 1, grow_mask_by, grow_mask_by))
+        kernel_tensor = torch.ones((1, 1, grow_mask_by, grow_mask_by)).to(device)
         padding = math.ceil((grow_mask_by - 1) / 2)
 
         mask_erosion = torch.clamp(torch.nn.functional.conv2d(mask2.round(), kernel_tensor, padding=padding), 0, 1)
 
-    return mask_erosion[:, :, :w, :h].round()
+    return mask_erosion[:, :, :w, :h].round().cpu()
 
 
 def ksampler_wrapper(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise,
@@ -115,6 +115,7 @@ class REGIONAL_PROMPT:
     def get_mask_erosion(self, factor):
         if self.mask_erosion is None or self.erosion_factor != factor:
             self.mask_erosion = erosion_mask(self.mask, factor)
+            self.erosion_factor = factor
 
         return self.mask_erosion
 

modules/impact/utils.py

@@ -95,12 +95,15 @@ def dilate_mask(mask, dilation_factor, iter=1):
     if len(mask.shape) == 3:
         mask = mask.squeeze(0)
 
+    kernel = np.ones((abs(dilation_factor), abs(dilation_factor)), np.uint8)
+    gpu_mask = cv2.UMat(mask)
+    gpu_kernel = cv2.UMat(kernel)
     if dilation_factor > 0:
-        kernel = np.ones((dilation_factor, dilation_factor), np.uint8)
-        return cv2.dilate(mask, kernel, iter)
+        result = cv2.dilate(gpu_mask, gpu_kernel, iter)
     else:
-        kernel = np.ones((-dilation_factor, -dilation_factor), np.uint8)
-        return cv2.erode(mask, kernel, iter)
+        result = cv2.erode(gpu_mask, gpu_kernel, iter)
+
+    return result.get()
 
 
 def dilate_masks(segmasks, dilation_factor, iter=1):
@@ -108,12 +111,22 @@ def dilate_masks(segmasks, dilation_factor, iter=1):
         return segmasks
 
     dilated_masks = []
-    kernel = np.ones((dilation_factor, dilation_factor), np.uint8)
+    kernel = np.ones((abs(dilation_factor), abs(dilation_factor)), np.uint8)
+
+    gpu_kernel = cv2.UMat(kernel)
+
     for i in range(len(segmasks)):
         cv2_mask = segmasks[i][1]
-        dilated_mask = cv2.dilate(cv2_mask, kernel, iter)
+        gpu_mask = cv2.UMat(cv2_mask)
+
+        if dilation_factor > 0:
+            dilated_mask = cv2.dilate(gpu_mask, gpu_kernel, iter).get()
+        else:
+            dilated_mask = cv2.erode(gpu_mask, gpu_kernel, iter).get()
+
         item = (segmasks[i][0], dilated_mask, segmasks[i][2])
         dilated_masks.append(item)
+
     return dilated_masks