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

ComfyUI/custom_nodes/ComfyUI-BrushNet/brushnet_nodes.py

@@ -493,26 +493,31 @@ class BlendInpaint:
             cut_width = min(cut_width, width - x0)
             cut_height = min(cut_height, height - y0)
 
+            # Scale inpainted image to match the cut size
             scaled_inpaint = F.interpolate(inpaint[i].permute(2, 0, 1).unsqueeze(0), size=(cut_height, cut_width), mode='bilinear', align_corners=False).squeeze(0).permute(1, 2, 0)
 
-            result = original[i].clone()
-            result[y0:y0+cut_height, x0:x0+cut_width] = scaled_inpaint
-
-            # Create a new mask for blending
-            blend_mask = torch.zeros((height, width), device=mask.device, dtype=mask.dtype)
-            blend_mask[y0:y0+cut_height, x0:x0+cut_width] = 1.0
+            # Create a mask for the inpainted region
+            inpaint_mask = torch.zeros((height, width), device=mask.device, dtype=mask.dtype)
+            inpaint_mask[y0:y0+cut_height, x0:x0+cut_width] = F.interpolate(mask[i][None, None, :, :], size=(cut_height, cut_width), mode='nearest').squeeze()
 
-            # Apply Gaussian blur to the blend mask
-            blurred_mask = transform(blend_mask.unsqueeze(0).unsqueeze(0)).squeeze(0).squeeze(0)
+            # Apply Gaussian blur to the inpaint mask
+            blurred_mask = transform(inpaint_mask.unsqueeze(0).unsqueeze(0)).squeeze(0).squeeze(0)
             blurred.append(blurred_mask)
 
-            # Apply the blended mask
-            ret.append(original[i] * (1.0 - blurred_mask[:,:,None]) + result * blurred_mask[:,:,None])
+            # Create the result by blending only the masked area
+            result = original[i].clone()
+            result[y0:y0+cut_height, x0:x0+cut_width] = (
+                original[i][y0:y0+cut_height, x0:x0+cut_width] * (1 - blurred_mask[y0:y0+cut_height, x0:x0+cut_width, None]) +
+                scaled_inpaint * blurred_mask[y0:y0+cut_height, x0:x0+cut_width, None]
+            )
+
+            ret.append(result)
 
         return (torch.stack(ret), torch.stack(blurred))
 
 
 
+
 def scale_mask_and_image(image, mask, width, height):
     h0, w0 = mask.shape
     iy, ix = (mask == 1).nonzero(as_tuple=True)