fix mask (#1914)

Mikubill · Aug 13, 2023 · c542a8f · c542a8f
1 parent 50b5436
commit c542a8f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 9 deletions.
diff --git a/scripts/controlnet.py b/scripts/controlnet.py
@@ -628,7 +628,10 @@ def choose_input_image(
             else:
                 input_image = HWC3(image['image'])
 
-            have_mask = 'mask' in image and not ((image['mask'][:, :, 0] == 0).all() or (image['mask'][:, :, 0] == 255).all())
+            have_mask = 'mask' in image and not (
+                (image['mask'][:, :, 0] <= 5).all() or 
+                (image['mask'][:, :, 0] >= 250).all()
+            )
 
             if 'inpaint' in unit.module:
                 logger.info("using inpaint as input")
@@ -639,7 +642,7 @@ def choose_input_image(
                     alpha = np.zeros_like(color)[:, :, 0:1]
                 input_image = np.concatenate([color, alpha], axis=2)
             else:
-                if have_mask:
+                if have_mask and not shared.opts.data.get("controlnet_ignore_noninpaint_mask", False):
                     logger.info("using mask as input")
                     input_image = HWC3(image['mask'][:, :, 0])
                     unit.module = 'none'  # Always use black bg and white line
@@ -1056,6 +1059,9 @@ def on_ui_settings():
         False, "Disable control type selection", gr.Checkbox, {"interactive": True}, section=section))
     shared.opts.add_option("controlnet_disable_openpose_edit", shared.OptionInfo(
         False, "Disable openpose edit", gr.Checkbox, {"interactive": True}, section=section))
+    shared.opts.add_option("controlnet_ignore_noninpaint_mask", shared.OptionInfo(
+        False, "Ignore mask on ControlNet input image if control type is not inpaint", 
+        gr.Checkbox, {"interactive": True}, section=section))
 
 
 batch_hijack.instance.do_hijack()

diff --git a/scripts/controlnet_ui/controlnet_ui_group.py b/scripts/controlnet_ui/controlnet_ui_group.py
@@ -562,17 +562,17 @@ def run_annotator(image, module, pres, pthr_a, pthr_b, t2i_w, t2i_h, pp, rm):
                     *self.openpose_editor.update(''),
                 )
 
-            img = HWC3(image["image"])
-            if not (
-                (image["mask"][:, :, 0] == 0).all()
-                or (image["mask"][:, :, 0] == 255).all()
-            ):
-                img = HWC3(image["mask"][:, :, 0])
-
+            img = HWC3(image["image"])            
+            has_mask = not (
+                (image["mask"][:, :, 0] <= 5).all()
+                or (image["mask"][:, :, 0] >= 250).all()
+            )
             if "inpaint" in module:
                 color = HWC3(image["image"])
                 alpha = image["mask"][:, :, 0:1]
                 img = np.concatenate([color, alpha], axis=2)
+            elif has_mask and not shared.opts.data.get("controlnet_ignore_noninpaint_mask", False):
+                img = HWC3(image["mask"][:, :, 0])
 
             module = global_state.get_module_basename(module)
             preprocessor = self.preprocessors[module]