Merge branch 'feature/add-inpaint-mask-generation'

# Conflicts:
#	modules/config.py
#	requirements_versions.txt

extras/inpaint_mask.py

@@ -0,0 +1,15 @@
+from rembg import remove, new_session
+
+
+def generate_mask_from_image(image, mask_model):
+    if image is None:
+        return
+
+    if 'image' in image:
+        image = image['image']
+
+    return remove(
+        image,
+        session=new_session(mask_model),
+        only_mask=True
+    )

modules/config.py

@@ -384,6 +384,11 @@ default_black_out_nsfw = get_config_item_or_set_default(
     default_value=False,
     validator=lambda x: isinstance(x, bool)
 )
+default_inpaint_mask_model = get_config_item_or_set_default(
+    key='default_inpaint_mask_model',
+    default_value='isnet-general-use',
+    validator=lambda x: isinstance(x, str)
+)
 
 config_dict["default_loras"] = default_loras = default_loras[:5] + [['None', 1.0] for _ in range(5 - len(default_loras))]
 

modules/flags.py

@@ -47,6 +47,10 @@ performance_selections = [
 
 output_formats = ['png', 'jpg', 'webp']
 
+inpaint_mask_models = [
+    'u2net', 'u2netp', 'u2net_human_seg', 'u2net_cloth_seg', 'silueta', 'isnet-general-use', 'isnet-anime', 'sam'
+]
+
 inpaint_option_default = 'Inpaint or Outpaint (default)'
 inpaint_option_detail = 'Improve Detail (face, hand, eyes, etc.)'
 inpaint_option_modify = 'Modify Content (add objects, change background, etc.)'

readme.md

@@ -27,6 +27,7 @@ Included adjustments:
 * ✨ https://github.com/lllyasviel/Fooocus/pull/1938 - automatically describe image on uov image upload if prompt is empty
 * ✨ https://github.com/lllyasviel/Fooocus/pull/1940 - meta data handling, schemes: Fooocus (json) and A1111 (plain text). Compatible with Civitai.
 * ✨ https://github.com/lllyasviel/Fooocus/pull/1979 - prevent outdated history log link after midnight
+* ✨ https://github.com/lllyasviel/Fooocus/pull/2032 - add inpaint mask generation functionality using rembg
 
 ✨ = new feature<br>
 🐛 = bugfix<br>

requirements_versions.txt

@@ -17,4 +17,5 @@ diffusers==0.25.1
 httpx==0.24.1
 onnxruntime==1.16.3
 timm==0.9.2
-translators==5.8.9
+translators==5.8.9
+rembg==2.0.53

webui.py

@@ -204,6 +204,19 @@ with shared.gradio_root:
                             inpaint_additional_prompt = gr.Textbox(placeholder="Describe what you want to inpaint.", elem_id='inpaint_additional_prompt', label='Inpaint Additional Prompt', visible=False)
                             outpaint_selections = gr.CheckboxGroup(choices=['Left', 'Right', 'Top', 'Bottom'], value=[], label='Outpaint Direction')
                             inpaint_mode = gr.Dropdown(choices=modules.flags.inpaint_options, value=modules.flags.inpaint_option_default, label='Method')
+                        with gr.Row(visible=False) as inpaint_mask_generation_row:
+                            inpaint_mask_model = gr.Dropdown(label='Mask generation model',
+                                                             choices=flags.inpaint_mask_models,
+                                                             value=modules.config.default_inpaint_mask_model, visible=False)
+                            generate_mask_button = gr.Button(value='Generate mask from image', visible=False)
+
+                            def generate_mask(image, mask_model):
+                                from extras.inpaint_mask import generate_mask_from_image
+                                return generate_mask_from_image(image, mask_model)
+
+                            generate_mask_button.click(fn=generate_mask, inputs=[inpaint_input_image, inpaint_mask_model],
+                                                       outputs=inpaint_mask_image)
+
                         example_inpaint_prompts = gr.Dataset(samples=modules.config.example_inpaint_prompts, label='Additional Prompt Quick List', components=[inpaint_additional_prompt], visible=False)
                         gr.HTML('* Powered by Fooocus Inpaint Engine <a href="https://github.com/lllyasviel/Fooocus/discussions/414" target="_blank">\U0001F4D4 Document</a>')
                         example_inpaint_prompts.click(lambda x: x[0], inputs=example_inpaint_prompts, outputs=inpaint_additional_prompt, show_progress=False, queue=False)
@@ -488,9 +501,11 @@ with shared.gradio_root:
                                          inpaint_strength, inpaint_respective_field,
                                          inpaint_mask_upload_checkbox, invert_mask_checkbox, inpaint_erode_or_dilate]
 
-                        inpaint_mask_upload_checkbox.change(lambda x: gr.update(visible=x),
-                                                           inputs=inpaint_mask_upload_checkbox,
-                                                           outputs=inpaint_mask_image, queue=False, show_progress=False)
+                        inpaint_mask_upload_checkbox.change(lambda x: [gr.update(visible=x)] * 4,
+                                                            inputs=inpaint_mask_upload_checkbox,
+                                                            outputs=[inpaint_mask_image, generate_mask_button,
+                                                                     inpaint_mask_model, inpaint_mask_generation_row],
+                                                            queue=False, show_progress=False)
 
                     with gr.Tab(label='FreeU'):
                         freeu_enabled = gr.Checkbox(label='Enabled', value=False)