added the MaskCoverFourCorners node.

Author: zhangp365

README.md

@@ -51,6 +51,9 @@ Generates a mask from the face model of the Reactor face swap node output. The m
 ## MaskAutoSelector
 Check the three input masks. If any are available, return the first. If none are available, raise an exception.
 
+## MaskCoverFourCorners
+Generates a mask by covering the selected corners with circular edges. This mask can be used as an attention mask to remove watermarks from the corners.
+
 ## CheckpointLoaderSimpleWithSwitch
 Enhanced the official LoadCheckpoint node by integrating three switches. Each switch controls whether a specific component is loaded. When a switch is turned off, the corresponding component will not be loaded. if you use the extra vae and close the model's vae loading, that will save memory.
 

nodes.py

@@ -583,7 +583,72 @@ def get_max_distance(self, keypoints):
                     max_distance = max(max_distance, distance)
 
         return max_distance
+
+
+class MaskCoverFourCorners:
+
+    @classmethod
+    def INPUT_TYPES(self):
+
+        return {
+            "required": {                
+                "width": ("INT", {"default": 1024, "min": 0, "max": 8096, "step": 8}),
+                "height": ("INT", {"default": 1024, "min": 0, "max": 8096, "step": 8}),
+                "radius": ("INT", {"default": 100, "min": 0, "max": 8096, "step": 5}),
+                "draw_top_left": ("BOOLEAN", {"default": False}),
+                "draw_top_right": ("BOOLEAN", {"default": False}),
+                "draw_bottom_right": ("BOOLEAN", {"default": True}),
+                "draw_bottom_left": ("BOOLEAN", {"default": False}),
+            },
+            "optional": {
+                "size_as": ("IMAGE",),
+            }
+        }
+
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = 'mask_get'
+    CATEGORY = 'utils/mask'
+
+    def mask_get(self, size_as=None, width=1024, height=1024, radius=100, 
+                            draw_top_left=False, draw_top_right=False, draw_bottom_right=True, draw_bottom_left=False):
+        if size_as is not None:
+            height, width = size_as.shape[-3:-1]
+        result = self.create_mask_with_arcs(width, height, radius,draw_top_left, draw_top_right,draw_bottom_right,draw_bottom_left)
+
+        result = torch.unsqueeze(torch.tensor(np.clip(result/255, 0, 1)), 0)
+        return (result,)
+    
+    def create_mask_with_arcs(self, width=None, height=None, radius=50, 
+                            draw_top_left=True, draw_top_right=True, draw_bottom_right=True, draw_bottom_left=True):
+        """
+        Creates a mask with circular arcs at the corners.
+        
+        :param width: Width of the mask.
+        :param height: Height of the mask.
+        :param radius: Radius of the circular arcs.
+        :param draw_top_left: Boolean indicating whether to draw an arc at the top-left corner.
+        :param draw_top_right: Boolean indicating whether to draw an arc at the top-right corner.
+        :param draw_bottom_right: Boolean indicating whether to draw an arc at the bottom-right corner.
+        :param draw_bottom_left: Boolean indicating whether to draw an arc at the bottom-left corner.
+        :return: Mask image with arcs drawn.
+        """       
+
+        # Create a white mask
+        mask = np.ones((height, width), dtype=np.uint8) * 255
+        
+        # Draw arcs on the mask, filling them with black color
+        if draw_top_left:  # Top-left corner
+            cv2.ellipse(mask, (0, 0), (radius, radius), 0, 0, 90, 0, -1)  # Fill with black
+        if draw_top_right:  # Top-right corner
+            cv2.ellipse(mask, (width, 0), (radius, radius), 0, 90, 180, 0, -1)  # Fill with black
+        if draw_bottom_right:  # Bottom-right corner
+            cv2.ellipse(mask, (width, height), (radius, radius), 0, 180, 270, 0, -1)  # Fill with black
+        if draw_bottom_left:  # Bottom-left corner
+            cv2.ellipse(mask, (0, height), (radius, radius), 0, 0, -90, 0, -1)  # Fill with black
+        
+        return mask
 
+
 class MaskAutoSelector:
     @classmethod
     def INPUT_TYPES(self):
@@ -1068,6 +1133,7 @@ def forward(self, image, upscale_model, tile_size=512, threshold_of_xl_area=0.9)
     "MaskFastGrow": MaskFastGrow,
     "MaskAutoSelector": MaskAutoSelector,
     "MaskFromFaceModel": MaskFromFaceModel,
+    "MaskCoverFourCorners": MaskCoverFourCorners,
 
     #loader
     "CheckpointLoaderSimpleWithSwitch": CheckpointLoaderSimpleWithSwitch,
@@ -1102,6 +1168,7 @@ def forward(self, image, upscale_model, tile_size=512, threshold_of_xl_area=0.9)
     "MaskFastGrow": "Mask Grow Fast",
     "MaskAutoSelector": "Mask Auto Selector",
     "MaskFromFaceModel": "Mask from FaceModel",
+    "MaskCoverFourCorners": "Mask Cover Four Corners",
 
     # Loader
     "CheckpointLoaderSimpleWithSwitch": "Load Checkpoint with Switch",

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-utils-nodes"
-description = "Nodes:LoadImageWithSwitch, ImageBatchOneOrMore, ModifyTextGender, GenderControlOutput, ImageCompositeMaskedWithSwitch, ColorCorrectOfUtils, SplitMask, MaskFastGrow, CheckpointLoaderSimpleWithSwitch, ImageResizeTo8x, MatchImageRatioToPreset, UpscaleImageWithModelIfNeed, MaskFromFaceModel etc."
-version = "1.1.1"
+description = "Nodes:LoadImageWithSwitch, ImageBatchOneOrMore, ModifyTextGender, GenderControlOutput, ImageCompositeMaskedWithSwitch, ColorCorrectOfUtils, SplitMask, MaskFastGrow, CheckpointLoaderSimpleWithSwitch, ImageResizeTo8x, MatchImageRatioToPreset, UpscaleImageWithModelIfNeed, MaskFromFaceModel, MaskCoverFourCorners etc."
+version = "1.1.2"
 license = { file = "LICENSE" }
 dependencies = []