From 6258103152dab3d59e6ccbfd1f90333e7c5a9e28 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Tue, 30 Apr 2024 23:01:54 -0400
Subject: [PATCH] Replace mask upload with effective region mask (#2830)

---
 scripts/controlnet_ui/controlnet_ui_group.py | 35 ++++----------------
 1 file changed, 7 insertions(+), 28 deletions(-)

diff --git a/scripts/controlnet_ui/controlnet_ui_group.py b/scripts/controlnet_ui/controlnet_ui_group.py
index 2bfc18e22..a155a9dbd 100644
--- a/scripts/controlnet_ui/controlnet_ui_group.py
+++ b/scripts/controlnet_ui/controlnet_ui_group.py
@@ -135,7 +135,6 @@ def __init__(
         ] = [],
         use_preview_as_input: bool = False,
         generated_image: Optional[np.ndarray] = None,
-        mask_image: Optional[np.ndarray] = None,
         enabled: bool = True,
         module: Optional[str] = None,
         model: Optional[str] = None,
@@ -150,11 +149,6 @@ def __init__(
         else:
             input_image = image
 
-        # Prefer uploaded mask_image over hand-drawn mask.
-        if input_image is not None and mask_image is not None:
-            assert isinstance(input_image, dict)
-            input_image["mask"] = mask_image
-
         if merge_gallery_files and input_mode == InputMode.MERGE:
             input_image = [
                 {"image": read_image(file["name"])} for file in merge_gallery_files
@@ -244,7 +238,7 @@ def __init__(
         self.generated_image_group = None
         self.generated_image = None
         self.mask_image_group = None
-        self.mask_image = None
+        self.effective_region_mask = None
         self.batch_tab = None
         self.batch_image_dir = None
         self.merge_tab = None
@@ -296,7 +290,6 @@ def __init__(
 
         # API-only fields
         self.ipadapter_input = gr.State(None)
-        self.effective_region_mask = gr.Image(value=None, visible=False)
 
         ControlNetUiGroup.all_ui_groups.append(self)
 
@@ -369,11 +362,11 @@ def render(self, tabname: str, elem_id_tabname: str) -> None:
                         with gr.Group(
                             visible=False, elem_classes=["cnet-mask-image-group"]
                         ) as self.mask_image_group:
-                            self.mask_image = gr.Image(
+                            self.effective_region_mask = gr.Image(
                                 value=None,
-                                label="Upload Mask",
+                                label="Effective Region Mask",
                                 elem_id=f"{elem_id_tabname}_{tabname}_mask_image",
-                                elem_classes=["cnet-mask-image"],
+                                elem_classes=["cnet-effective-region-mask-image"],
                                 interactive=True,
                             )
 
@@ -481,11 +474,10 @@ def render(self, tabname: str, elem_id_tabname: str) -> None:
                 visible=not self.is_img2img,
             )
             self.mask_upload = gr.Checkbox(
-                label="Mask Upload",
+                label="Effective Region Mask",
                 value=False,
                 elem_classes=["cnet-mask-upload"],
                 elem_id=f"{elem_id_tabname}_{tabname}_controlnet_mask_upload_checkbox",
-                visible=not self.is_img2img,
             )
             self.use_preview_as_input = gr.Checkbox(
                 label="Preview as Input",
@@ -661,7 +653,6 @@ def render(self, tabname: str, elem_id_tabname: str) -> None:
             self.merge_gallery,
             self.use_preview_as_input,
             self.generated_image,
-            self.mask_image,
             # End of Non-persistent fields.
             self.enabled,
             self.module,
@@ -681,6 +672,7 @@ def render(self, tabname: str, elem_id_tabname: str) -> None:
             self.hr_option,
             self.save_detected_map,
             self.advanced_weighting,
+            self.effective_region_mask,
         )
 
         unit = gr.State(self.default_unit)
@@ -1122,23 +1114,10 @@ def register_shift_upload_mask(self):
                 else (gr.update(visible=True), gr.update())
             ),
             inputs=[self.mask_upload],
-            outputs=[self.mask_image_group, self.mask_image],
+            outputs=[self.mask_image_group, self.effective_region_mask],
             show_progress=False,
         )
 
-        if self.upload_independent_img_in_img2img is not None:
-            self.upload_independent_img_in_img2img.change(
-                fn=lambda checked: (
-                    # Uncheck `upload_mask` when not using independent input.
-                    gr.update(visible=False, value=False)
-                    if not checked
-                    else gr.update(visible=True)
-                ),
-                inputs=[self.upload_independent_img_in_img2img],
-                outputs=[self.mask_upload],
-                show_progress=False,
-            )
-
     def register_sync_batch_dir(self):
         def determine_batch_dir(batch_dir, fallback_dir, fallback_fallback_dir):
             if batch_dir: