diff --git a/frontend/frontend.py b/frontend/frontend.py index 185886507..7c838f881 100644 --- a/frontend/frontend.py +++ b/frontend/frontend.py @@ -157,13 +157,19 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x,imgproc=lambda img2img_btn_editor = gr.Button("Generate", variant="primary", elem_id="img2img_edit_btn") with gr.Row().style(equal_height=False): with gr.Column(): - gr.Markdown('#### Img2Img Input') - img2img_image_editor = gr.Image(value=sample_img2img, source="upload", interactive=True, - type="pil", tool="select", elem_id="img2img_editor", - image_mode="RGBA") - img2img_image_mask = gr.Image(value=sample_img2img, source="upload", interactive=True, - type="pil", tool="sketch", visible=False, - elem_id="img2img_mask") + with gr.Tabs(): + with gr.TabItem("Img2Img Input"): + #gr.Markdown('#### Img2Img Input') + img2img_image_editor = gr.Image(value=sample_img2img, source="upload", interactive=True, + type="pil", tool="select", elem_id="img2img_editor", + image_mode="RGBA") + img2img_image_mask = gr.Image(value=sample_img2img, source="upload", interactive=True, + type="pil", tool="sketch", visible=False, + elem_id="img2img_mask") + + with gr.TabItem("Img2Img Mask Input"): + img2img_mask_input = gr.Image(label="Mask",source="upload", interactive=False, + type="pil", visible=True) with gr.Tabs(): with gr.TabItem("Editor Options"): @@ -172,14 +178,14 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x,imgproc=lambda value="Crop", elem_id='edit_mode_select') img2img_mask = gr.Radio(choices=["Keep masked area", "Regenerate only masked area"], label="Mask Mode", type="index", - value=img2img_mask_modes[img2img_defaults['mask_mode']], visible=False) + value=img2img_mask_modes[img2img_defaults['mask_mode']], visible=False) img2img_mask_blur_strength = gr.Slider(minimum=1, maximum=10, step=1, label="How much blurry should the mask be? (to avoid hard edges)", value=3, visible=False) img2img_resize = gr.Radio(label="Resize mode", - choices=["Just resize"], + choices=["Just resize", "Crop and resize", "Resize and fill"], type="index", value=img2img_resize_modes[img2img_defaults['resize_mode']]) @@ -258,7 +264,7 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x,imgproc=lambda uifn.change_image_editor_mode, [img2img_image_editor_mode, img2img_image_editor, img2img_resize, img2img_width, img2img_height], [img2img_image_editor, img2img_image_mask, img2img_btn_editor, img2img_btn_mask, - img2img_painterro_btn, img2img_mask, img2img_mask_blur_strength] + img2img_painterro_btn, img2img_mask, img2img_mask_blur_strength, img2img_mask_input] ) img2img_image_editor.edit( @@ -303,7 +309,7 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x,imgproc=lambda img2img_mask_blur_strength, img2img_steps, img2img_sampling, img2img_toggles, img2img_realesrgan_model_name, img2img_batch_count, img2img_cfg, img2img_denoising, img2img_seed, img2img_height, img2img_width, img2img_resize, - img2img_embeddings] + img2img_embeddings, img2img_mask_input] img2img_outputs = [output_img2img_gallery, output_img2img_seed, output_img2img_params, output_img2img_stats] # If a JobManager was passed in then wrap the Generate functions @@ -314,8 +320,24 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x,imgproc=lambda outputs=img2img_outputs, ) + def generate(*args): + args_list = list(args) + init_info_mask = args_list[3] + # Get the mask input and remove it from the list + mask_input = args_list[18] + del args_list[18] + + # If an external mask is set, use it + if mask_input: + init_info_mask['mask'] = mask_input + + args_list[3] = init_info_mask + + # Return the result of img2img + return img2img_func(*args_list) + img2img_btn_mask.click( - img2img_func, + generate, img2img_inputs, img2img_outputs ) diff --git a/frontend/ui_functions.py b/frontend/ui_functions.py index 421366ea1..a85d15432 100644 --- a/frontend/ui_functions.py +++ b/frontend/ui_functions.py @@ -8,12 +8,31 @@ def change_image_editor_mode(choice, cropped_image, resize_mode, width, height): if choice == "Mask": - return [gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)] - return [gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)] + return [gr.Image.update(visible=False), + gr.Image.update(visible=True), + gr.Button.update("Generate", variant="primary", visible=False), + gr.Button.update("Generate", variant="primary", visible=True), + gr.Button.update("Advanced Editor", visible=False), + gr.Radio.update(choices=["Keep masked area", "Regenerate only masked area"], + label="Mask Mode", + value="Regenerate only masked area", visible=True), + gr.Slider.update(minimum=1, maximum=10, step=1, label="How much blurry should the mask be? (to avoid hard edges)", value=3, visible=True), + gr.Image.update(interactive=True)] + else: + return [gr.Image.update(visible=True), + gr.Image.update(visible=False), + gr.Button.update("Generate", variant="primary", visible=True), + gr.Button.update("Generate", variant="primary", visible=False), + gr.Button.update("Advanced Editor", visible=True), + gr.Radio.update(choices=["Keep masked area", "Regenerate only masked area"], + label="Mask Mode", + value="Regenerate only masked area", visible=False), + gr.Slider.update(minimum=1, maximum=10, step=1, label="How much blurry should the mask be? (to avoid hard edges)", value=3, visible=False), + gr.Image.update(interactive=False)] def update_image_mask(cropped_image, resize_mode, width, height): resized_cropped_image = resize_image(resize_mode, cropped_image, width, height) if cropped_image else None - return gr.update(value=resized_cropped_image) + return gr.Image.update(value=resized_cropped_image) def toggle_options_gfpgan(selection): if 0 in selection: diff --git a/webui.py b/webui.py index 663021feb..f57759040 100644 --- a/webui.py +++ b/webui.py @@ -1273,6 +1273,7 @@ def img2img(prompt: str, image_editor_mode: str, init_info: any, init_info_mask: init_img = init_info_mask["image"] init_img = init_img.convert("RGB") init_img = resize_image(resize_mode, init_img, width, height) + image = image.convert("RGB") init_mask = init_info_mask["mask"] init_mask = resize_image(resize_mode, init_mask, width, height) keep_mask = mask_mode == 0