Update ComfyUI/custom_nodes/ComfyUI-BrushNet/brushnet_nodes.py

ComfyUI/custom_nodes/ComfyUI-BrushNet/brushnet_nodes.py

@@ -6,7 +6,7 @@ import torch
 import torchvision.transforms as T
 import torch.nn.functional as F
 from accelerate import init_empty_weights, load_checkpoint_and_dispatch
-
+import math
 import comfy
 import folder_paths
 
@@ -523,7 +523,9 @@ class BlendInpaint:
 
 
 
-def scale_mask_and_image(image, mask, width, height, side_margin):
+def scale_mask_and_image(image, mask, side_margin, min_side, max_side):
+    min_area = min_side * min_side
+    max_area = max_side * max_side
     h0, w0 = mask.shape
     iy, ix = (mask == 1).nonzero(as_tuple=True)
     
@@ -536,15 +538,14 @@ def scale_mask_and_image(image, mask, width, height, side_margin):
         x_c, y_c = (x_min + x_max) / 2.0, (y_min + y_max) / 2.0
         mask_width, mask_height = x_max - x_min + 1, y_max - y_min + 1
 
-    aspect_ratio = width / height
     mask_aspect_ratio = mask_width / mask_height
 
-    if mask_aspect_ratio > aspect_ratio:
+    if mask_aspect_ratio > 1:
         new_mask_width = mask_width
-        new_mask_height = mask_width / aspect_ratio
+        new_mask_height = mask_width
     else:
         new_mask_height = mask_height
-        new_mask_width = mask_height * aspect_ratio
+        new_mask_width = mask_height
 
     margin = side_margin/100.0
     cut_width = int(new_mask_width * (1 + 2 * margin))
@@ -553,21 +554,32 @@ def scale_mask_and_image(image, mask, width, height, side_margin):
     x0 = max(0, min(w0 - cut_width, int(x_c - cut_width / 2)))
     y0 = max(0, min(h0 - cut_height, int(y_c - cut_height / 2)))
 
-    # Adjust cut dimensions if they exceed image dimensions
     cut_width = min(cut_width, w0 - x0)
     cut_height = min(cut_height, h0 - y0)
 
     cut_image = image[y0:y0+cut_height, x0:x0+cut_width]
     cut_mask = mask[y0:y0+cut_height, x0:x0+cut_width]
 
-    if cut_width >= width and cut_height >= height:
-        # For large masks, return without scaling
-        return cut_image, cut_mask, (x0, y0, cut_width, cut_height)
-    else:
-        # For small masks, scale up to the specified size
-        scaled_image = F.interpolate(cut_image.permute(2, 0, 1).unsqueeze(0), size=(height, width), mode='bilinear', align_corners=False).squeeze(0).permute(1, 2, 0)
-        scaled_mask = F.interpolate(cut_mask.unsqueeze(0).unsqueeze(0).float(), size=(height, width), mode='nearest').squeeze(0).squeeze(0)
+
+    current_area = cut_width * cut_height
+    print(f"current_area: {current_area} min_area: {min_area} max_area: {max_area}")
+
+    if current_area > max_area or current_area <= min_area:
+        if current_area > max_area:
+            print("current_area > max_area")
+            scale_factor = math.sqrt(max_area / current_area)       
+        elif current_area <= min_area:
+            print("current_area <= min_area")
+            scale_factor = math.sqrt(min_area / current_area)  
+        new_width = int(cut_width * scale_factor)
+        new_height = int(cut_height * scale_factor)
+        scaled_image = F.interpolate(cut_image.permute(2, 0, 1).unsqueeze(0), size=(new_height, new_width), mode='bilinear', align_corners=False).squeeze(0).permute(1, 2, 0)
+        scaled_mask = F.interpolate(cut_mask.unsqueeze(0).unsqueeze(0).float(), size=(new_height, new_width), mode='nearest').squeeze(0).squeeze(0)
         return scaled_image, scaled_mask, (x0, y0, cut_width, cut_height)
+    else:
+        print("original size mask")
+        return cut_image, cut_mask, (x0, y0, cut_width, cut_height)
+   
     
 class CutForInpaint:
 
@@ -577,7 +589,9 @@ class CutForInpaint:
                     {    
                         "image": ("IMAGE",),
                         "mask": ("MASK",),
-                        "side_margin_percent": ("INT", {"default": 10, "min": 0, "max": 1000})
+                        "side_margin_percent": ("INT", {"default": 10, "min": 0, "max": 1000}),
+                        "min_side": ("INT", {"default": 512, "min": 128, "max": 4096}),
+                        "max_side": ("INT", {"default": 1536, "min": 128, "max": 4096})
                      },
                 }
 
@@ -587,12 +601,12 @@ class CutForInpaint:
 
     FUNCTION = "cut_for_inpaint"
 
-    def cut_for_inpaint(self, image: torch.Tensor, mask: torch.Tensor, side_margin_percent: int):
+    def cut_for_inpaint(self, image: torch.Tensor, mask: torch.Tensor, side_margin_percent: int, min_side: int, max_side: int):
         ret = []
         msk = []
         org = []
         for i in range(image.shape[0]):
-            cut_image, cut_mask, (x0, y0, cut_width, cut_height) = scale_mask_and_image(image[i], mask[i], 512, 512, side_margin_percent)
+            cut_image, cut_mask, (x0, y0, cut_width, cut_height) = scale_mask_and_image(image[i], mask[i], side_margin_percent, min_side, max_side)
             ret.append(cut_image)
             msk.append(cut_mask)
             org.append(torch.IntTensor([x0, y0, cut_width, cut_height]))