Add blur-type parameter (box/linear/gaussian) for Expand/ShrinkMask nodes

Author: Acly

nodes.py

@@ -20,6 +20,8 @@
 
 from . import mat
 from .util import (
+    BlurKernel,
+    mask_blur,
     gaussian_blur,
     binary_erosion,
     binary_dilation,
@@ -126,6 +128,7 @@ def INPUT_TYPES(s):
 
     def load(self, head: str, patch: str):
         head_file = folder_paths.get_full_path("inpaint", head)
+        assert head_file is not None, f"Inpaint head file not found in inpaint folder: {head}"
         inpaint_head_model = InpaintHead()
         sd = torch.load(head_file, map_location="cpu", weights_only=True)
         inpaint_head_model.load_state_dict(sd)
@@ -486,19 +489,20 @@ def INPUT_TYPES(cls):
                 "mask": ("MASK",),
                 "grow": ("INT", {"default": 16, "min": 0, "max": 8096, "step": 1}),
                 "blur": ("INT", {"default": 7, "min": 0, "max": 8096, "step": 1}),
+                "blur_type": (["box", "linear", "gaussian"], {"default": "gaussian"}),
             }
         }
 
     RETURN_TYPES = ("MASK",)
     CATEGORY = "inpaint"
     FUNCTION = "expand"
 
-    def expand(self, mask: Tensor, grow: int, blur: int):
+    def expand(self, mask: Tensor, grow: int, blur: int, blur_type: str):
         mask = mask_unsqueeze(mask)
         if grow > 0:
             mask = binary_dilation(mask, grow)
         if blur > 0:
-            mask = gaussian_blur(mask, make_odd(blur))
+            mask = mask_blur(mask, make_odd(blur), BlurKernel[blur_type])
         return (mask.squeeze(1),)
 
 
@@ -510,17 +514,18 @@ def INPUT_TYPES(cls):
                 "mask": ("MASK",),
                 "shrink": ("INT", {"default": 1, "min": 0, "max": 8096, "step": 1}),
                 "blur": ("INT", {"default": 0, "min": 0, "max": 8096, "step": 1}),
+                "blur_type": (["box", "linear", "gaussian"], {"default": "gaussian"}),
             }
         }
 
     RETURN_TYPES = ("MASK",)
     CATEGORY = "inpaint"
     FUNCTION = "shrink"
 
-    def shrink(self, mask: Tensor, shrink: int, blur: int):
+    def shrink(self, mask: Tensor, shrink: int, blur: int, blur_type: str):
         mask = mask_unsqueeze(mask)
         if shrink > 0:
             mask = binary_erosion(mask, shrink)
         if blur > 0:
-            mask = gaussian_blur(mask, make_odd(blur))
+            mask = mask_blur(mask, make_odd(blur), BlurKernel[blur_type])
         return (mask.squeeze(1),)

util.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from enum import Enum
 import torch
 import torch.nn.functional as F
 import numpy as np
@@ -20,7 +21,7 @@ def to_torch(image: Tensor, mask: Tensor | None = None):
     image = image.permute(0, 3, 1, 2)  # BHWC -> BCHW
     if mask is not None:
         mask = mask_unsqueeze(mask)
-    if image.shape[2:] != mask.shape[2:]:
+    if mask is not None and image.shape[2:] != mask.shape[2:]:
         raise ValueError(
             f"Image and mask must be the same size. {image.shape[2:]} != {mask.shape[2:]}"
         )
@@ -80,11 +81,11 @@ def undo_resize_square(image: Tensor, original_size: tuple[int, int, int]):
     return image[:, :, 0 : prev_size - pad_h, 0 : prev_size - pad_w]
 
 
-def gaussian_blur(image: Tensor, radius: int, sigma: float = 0):
+def gaussian_blur(image: Tensor, size: int, sigma: float = 0):
     c = image.shape[-3]
     if sigma <= 0:
-        sigma = 0.3 * (radius - 1) + 0.8
-    return kornia.filters.gaussian_blur2d(image, (radius, radius), (sigma, sigma))
+        sigma = 0.3 * (size - 1) + 0.8
+    return kornia.filters.gaussian_blur2d(image, (size, size), (sigma, sigma))
 
 
 def binary_erosion(mask: Tensor, radius: int):
@@ -102,6 +103,31 @@ def binary_dilation(mask: Tensor, radius: int):
     return mask
 
 
+class BlurKernel(Enum):
+    box = "box"
+    linear = "linear"
+    gaussian = "gaussian"
+
+
+def mask_blur(mask: Tensor, size: int, method=BlurKernel.gaussian):
+    if method is BlurKernel.gaussian:
+        return gaussian_blur(mask, size)
+
+    match method:
+        case BlurKernel.box:
+            kernel = torch.ones(1, size, device=mask.device) / size
+        case BlurKernel.linear:
+            kernel = torch.linspace(1 / size, (size - 1) / size, size, device=mask.device)
+            kernel = torch.cat([kernel, torch.ones(1, device=mask.device), kernel.flip(0)])
+            kernel = kernel / kernel.sum()
+            kernel = kernel.unsqueeze(0)
+        case _:
+            raise ValueError(f"Unknown blur kernel: {method}")
+
+    mask = kornia.filters.filter2d_separable(mask, kernel, kernel, border_type="reflect")
+    return mask
+
+
 def make_odd(x):
     if x > 0 and x % 2 == 0:
         return x + 1