diff --git a/mask.py b/mask.py index dcc85a2..3998c3a 100644 --- a/mask.py +++ b/mask.py @@ -21,6 +21,7 @@ def INPUT_TYPES(s): "required": { "mask": ("MASK",), "amount": ("INT", { "default": 6, "min": 0, "max": 256, "step": 1, }), + "allow_growth": ("BOOLEAN", { "default": True }), "device": (["auto", "cpu", "gpu"],), } } @@ -29,7 +30,7 @@ def INPUT_TYPES(s): FUNCTION = "execute" CATEGORY = "essentials/mask" - def execute(self, mask, amount, device): + def execute(self, mask, amount, allow_growth, device): if amount == 0: return (mask,) @@ -38,6 +39,8 @@ def execute(self, mask, amount, device): elif "cpu" == device: mask = mask.to('cpu') + mask_bin = (mask > 0) + if amount % 2 == 0: amount+= 1 @@ -46,10 +49,13 @@ def execute(self, mask, amount, device): mask = T.functional.gaussian_blur(mask.unsqueeze(1), amount).squeeze(1) + if not allow_growth: + mask = mask * mask_bin + if "gpu" == device or "cpu" == device: mask = mask.to(comfy.model_management.intermediate_device()) - return(mask,) + return (mask, ) class MaskFlip: @classmethod