upload mask come back

Author: continue-revolution

README.md

@@ -18,6 +18,7 @@ This extension aim for connecting [AUTOMATIC1111 Stable Diffusion WebUI](https:/
 - `2023/05/22`: [v1.4.1](https://github.com/continue-revolution/sd-webui-segment-anything/releases/tag/v1.4.1) [EditAnything](https://github.com/sail-sg/EditAnything) is ready to use! You can generate random segmentation and copy the output to EditAnything ControlNet.
 - `2023/05/29`: [v1.4.2](https://github.com/continue-revolution/sd-webui-segment-anything/releases/tag/v1.4.2) You may now do SAM inference on CPU by checking "Use CPU for SAM". This is for some MAC users who are not able to do SAM inference on GPU. I discourage other users from using this feature because it is significantly slower than CUDA.
 - `2023/06/01`: [v1.5.0](https://github.com/continue-revolution/sd-webui-segment-anything/releases/tag/v1.5.0) You may now choose to use local GroundingDINO to bypass C++ problem. See [FAQ](#faq)-1 for more detail.
+- `2023/06/04`: [v1.5.1](https://github.com/continue-revolution/sd-webui-segment-anything/releases/tag/v1.5.1) `Upload Mask to ControlNet Inpainting` comes bask in response to [ControlNet inpaint improvement](https://github.com/Mikubill/sd-webui-controlnet/discussions/1464). You should see a new tab beside `AutoSAM` after updating the extension. This feature will again be removed once ControlNet extension has its own uploading feature.
 
 ## FAQ
 

scripts/process_params.py

@@ -60,26 +60,38 @@ def get_input_and_mask(self, mask_blur):
 class SAMProcessUnit:
     def __init__(self, args: Tuple, is_img2img=False):
         self.is_img2img = is_img2img
-        sam_inpaint_args = args[:8]
+        self.sam_inpaint_unit = SAMInpaintUnit(args, is_img2img)
+
         args = args[8:]
-        self.sam_inpaint_unit = SAMInpaintUnit(sam_inpaint_args, is_img2img)
-        
-        self.cnet_seg_output_gallery: List[Dict] = None
-        self.cnet_seg_enable_copy: bool = False
-        self.cnet_seg_idx: int = 0
-        self.cnet_seg_gallery_input: int = 0
-        args = self.init_cnet_seg_process(args)
+        self.cnet_seg_output_gallery: List[Dict]    = None
+        self.cnet_seg_enable_copy: bool             = False
+        self.cnet_seg_idx: int                      = 0
+        self.cnet_seg_gallery_input: int            = 0
+        self.init_cnet_seg_process(args)
 
+        args = args[4:]
         self.crop_inpaint_unit = SAMInpaintUnit(args, is_img2img)
+
+        args = args[8:]
+        self.cnet_upload_enable: bool                   = False
+        self.cnet_upload_num: int                       = 0
+        self.cnet_upload_img_inpaint: Image.Image       = None
+        self.cnet_upload_mask_inpaint: Image.Image      = None
+        self.init_cnet_upload_process(args)
 
 
     def init_cnet_seg_process(self, args):
-        csp_args    = args[:4]
-        self.cnet_seg_output_gallery    = csp_args[0]
-        self.cnet_seg_enable_copy       = csp_args[1]
-        self.cnet_seg_idx               = csp_args[2]
-        self.cnet_seg_gallery_input     = csp_args[3]
-        return args[4:]
+        self.cnet_seg_output_gallery    = args[0]
+        self.cnet_seg_enable_copy       = args[1]
+        self.cnet_seg_idx               = args[2]
+        self.cnet_seg_gallery_input     = args[3]
+    
+
+    def init_cnet_upload_process(self, args):
+        self.cnet_upload_enable         = args[0]
+        self.cnet_upload_num            = args[1]
+        self.cnet_upload_img_inpaint    = args[2] 
+        self.cnet_upload_mask_inpaint   = args[3]
 
 
     def set_process_attributes(self, p):
@@ -103,6 +115,10 @@ def set_process_attributes(self, p):
                 cnet_seg_gallery_index += self.cnet_seg_gallery_input
             self.set_p_value(p, 'control_net_input_image', self.cnet_seg_idx, 
                              Image.open(self.cnet_seg_output_gallery[cnet_seg_gallery_index]['name']))
+        
+        if self.cnet_upload_enable and self.cnet_upload_img_inpaint is not None and self.cnet_upload_mask_inpaint is not None:
+            self.set_p_value(p, 'control_net_input_image', self.cnet_upload_num, 
+                            {"image": self.cnet_upload_img_inpaint, "mask": self.cnet_upload_mask_inpaint.convert("L")})
 
 
     def set_p_value(self, p, attr: str, idx: int, v):

scripts/sam.py

@@ -733,6 +733,24 @@ def layout_show(mode):
                                                 crop_category_input, crop_batch_dilation_amt, crop_batch_source_dir, crop_batch_dest_dir, 
                                                 crop_batch_save_image, crop_batch_save_mask, crop_batch_save_image_with_mask, crop_batch_save_background, *auto_sam_config],
                                         outputs=[crop_batch_progress])
+                            
+                            
+                with gr.TabItem(label="Upload Mask to ControlNet Inpainting"):
+                    gr.Markdown("This panel is for those who want to upload mask to ControlNet inpainting. It is not part of the SAM feature. It might be removed someday when ControlNet support uploading image and mask. "
+                                "It serves as a temporarily workaround to overcome the unavailability of image with mask uploading feature in ControlNet extension.")
+                    with gr.Row():
+                        cnet_upload_enable = gr.Checkbox(value=False, label="Enable uploading manually created mask to SAM.")
+                        cnet_upload_num = gr.Radio(value="0", choices=[str(i) for i in range(max_cn_num())], label='ControlNet Inpaint Number', type="index")
+                    with gr.Column(visible=False) as cnet_upload_panel:
+                        cnet_upload_img_inpaint = gr.Image(label="Image for ControlNet Inpaint", show_label=False, source="upload", interactive=True, type="pil")
+                        cnet_upload_mask_inpaint = gr.Image(label="Mask for ControlNet Inpaint", source="upload", interactive=True, type="pil")
+                    cnet_upload_enable.change(
+                        fn=gr_show,
+                        inputs=[cnet_upload_enable],
+                        outputs=[cnet_upload_panel],
+                        show_progress=False)
+                    cnet_upload_process = (cnet_upload_enable, cnet_upload_num, cnet_upload_img_inpaint, cnet_upload_mask_inpaint)
+                    ui_process += cnet_upload_process
 
                 with gr.Row():
                     switch = gr.Button(value="Switch to Inpaint Upload")