diff --git a/extras/inpaint_mask.py b/extras/inpaint_mask.py
index 9773eb47..08d3c3cc 100644
--- a/extras/inpaint_mask.py
+++ b/extras/inpaint_mask.py
@@ -1,7 +1,7 @@
from rembg import remove, new_session
-def generate_mask_from_image(image, mask_model):
+def generate_mask_from_image(image, mask_model, extras):
if image is None:
return
@@ -11,5 +11,6 @@ def generate_mask_from_image(image, mask_model):
return remove(
image,
session=new_session(mask_model),
- only_mask=True
+ only_mask=True,
+ **extras
)
diff --git a/modules/config.py b/modules/config.py
index 11481bb9..1010fe6a 100644
--- a/modules/config.py
+++ b/modules/config.py
@@ -390,6 +390,12 @@ default_inpaint_mask_model = get_config_item_or_set_default(
validator=lambda x: x in modules.flags.inpaint_mask_models
)
+default_inpaint_mask_cloth_category = get_config_item_or_set_default(
+ key='default_inpaint_mask_cloth_category',
+ default_value='full',
+ validator=lambda x: x in modules.flags.inpaint_mask_cloth_category
+)
+
config_dict["default_loras"] = default_loras = default_loras[:5] + [['None', 1.0] for _ in range(5 - len(default_loras))]
# mapping config to meta parameter
diff --git a/modules/flags.py b/modules/flags.py
index ad2acfee..7279dba5 100644
--- a/modules/flags.py
+++ b/modules/flags.py
@@ -48,9 +48,11 @@ performance_selections = [
output_formats = ['png', 'jpg', 'webp']
inpaint_mask_models = [
- 'u2net', 'u2netp', 'silueta', 'isnet-general-use', 'isnet-anime'
+ 'u2net', 'u2netp', 'u2net_human_seg', 'u2net_cloth_seg', 'silueta', 'isnet-general-use', 'isnet-anime'
]
+inpaint_mask_cloth_category = ['full', 'upper', 'lower']
+
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.)'
diff --git a/webui.py b/webui.py
index c12e5cad..bacd380f 100644
--- a/webui.py
+++ b/webui.py
@@ -197,29 +197,46 @@ with shared.gradio_root:
with gr.TabItem(label='Inpaint or Outpaint') as inpaint_tab:
with gr.Row():
- inpaint_input_image = grh.Image(label='Drag inpaint or outpaint image to here', source='upload', type='numpy', tool='sketch', height=500, brush_color="#FFFFFF", elem_id='inpaint_canvas')
- inpaint_mask_image = grh.Image(label='Mask Upload', source='upload', type='numpy', height=500, visible=False)
+ with gr.Column():
+ inpaint_input_image = grh.Image(label='Drag inpaint or outpaint image to here', source='upload', type='numpy', tool='sketch', height=500, brush_color="#FFFFFF", elem_id='inpaint_canvas')
+ inpaint_mode = gr.Dropdown(choices=modules.flags.inpaint_options, value=modules.flags.inpaint_option_default, label='Method')
+ 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')
+ 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 \U0001F4D4 Document')
+ example_inpaint_prompts.click(lambda x: x[0], inputs=example_inpaint_prompts, outputs=inpaint_additional_prompt, show_progress=False, queue=False)
- with gr.Row():
- 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)
+ with gr.Column(visible=False) as inpaint_mask_generation_col:
+ inpaint_mask_image = grh.Image(label='Mask Upload', source='upload', type='numpy',
+ height=500, visible=False)
+ inpaint_mask_model = gr.Dropdown(label='Mask generation model',
+ choices=flags.inpaint_mask_models,
+ value=modules.config.default_inpaint_mask_model,
+ visible=False)
+ inpaint_mask_cloth_category = gr.Dropdown(label='Cloth category',
+ choices=flags.inpaint_mask_cloth_category,
+ value=modules.config.default_inpaint_mask_cloth_category,
+ 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)
+ def generate_mask(image, mask_model, cloth_category):
+ from extras.inpaint_mask import generate_mask_from_image
+ return generate_mask_from_image(image, mask_model, {"cloth_category": cloth_category})
- generate_mask_button.click(fn=generate_mask, inputs=[inpaint_input_image, inpaint_mask_model],
- outputs=inpaint_mask_image)
+ generate_mask_button.click(fn=generate_mask,
+ inputs=[
+ inpaint_input_image, inpaint_mask_model,
+ inpaint_mask_cloth_category
+ ],
+ 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 \U0001F4D4 Document')
- example_inpaint_prompts.click(lambda x: x[0], inputs=example_inpaint_prompts, outputs=inpaint_additional_prompt, show_progress=False, queue=False)
+ inpaint_mask_model.change(lambda x: gr.update(visible=x == 'u2net_cloth_seg'),
+ inputs=inpaint_mask_model,
+ outputs=inpaint_mask_cloth_category,
+ queue=False, show_progress=False)
with gr.TabItem(label='Describe') as desc_tab:
with gr.Row():
@@ -504,7 +521,7 @@ with shared.gradio_root:
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],
+ inpaint_mask_model, inpaint_mask_generation_col],
queue=False, show_progress=False)
with gr.Tab(label='FreeU'):