Refactor mask generation.

node.py

@@ -446,11 +446,12 @@ class MaskToPatchMask:
             "required": {
                 "mask": ("MASK",),
                 "patch_size": ("INT",{"default": 16}),
+                "threshold": ("FLOAT",{"default": 1.0}),
             }
         }
 
     @classmethod
-    def set_mask(cls, mask: torch.Tensor, patch_size):
+    def set_mask(cls, mask: torch.Tensor, patch_size, threshold):
         # Get the dimensions of the input mask
         B, H, W = mask.shape
         
@@ -463,7 +464,7 @@ class MaskToPatchMask:
         
         # Check if any pixel in each 16x16 block is non-zero
         # Take max over the 16x16 blocks, then convert to boolean
-        patch_mask = mask_reshaped.max(dim=2)[0].max(dim=3)[0] > 0
+        patch_mask = mask_reshaped.mean(dim=2).mean(dim=3) >= threshold
         
         # Convert to boolean tensor
         patch_mask = patch_mask.bool()