From a31e893fd9392cd74f37117b8ea3b7a7d61b7558 Mon Sep 17 00:00:00 2001
From: ZHOU He <461607353@qq.com>
Date: Tue, 18 Nov 2025 11:21:32 +0800
Subject: [PATCH] Refactor mask generation.

---
 node.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/node.py b/node.py
index b456039..b8bf63a 100644
--- a/node.py
+++ b/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()