Update app.py

app.py

@@ -16,7 +16,8 @@ def read_content(file_path: str) -> str:
 
     return content
 
-def predict(dict, prompt="", negative_prompt="", guidance_scale=7.5, steps=20, strength=1.0, scheduler="EulerDiscreteScheduler"):
+def predict(dict, mask_image,mask_use,prompt="", negative_prompt="", guidance_scale=7.5, steps=20, strength=1.0, scheduler="EulerDiscreteScheduler"):
+    print(dict.keys(),mask_image,mask_use)
     if negative_prompt == "":
         negative_prompt = None
     scheduler_class_name = scheduler.split("-")[0]
@@ -31,8 +32,11 @@ def predict(dict, prompt="", negative_prompt="", guidance_scale=7.5, steps=20, s
     pipe.scheduler = scheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler", **add_kwargs)
     
     init_image = dict["image"].convert("RGB").resize((1024, 1024))
-    mask = dict["mask"].convert("RGB").resize((1024, 1024))
     
+    if mask_use:
+        mask=mask_image.convert("L").resize((1024, 1024))
+    else:
+        mask = dict["mask"].convert("RGB").resize((1024, 1024))
     output = pipe(prompt = prompt, negative_prompt=negative_prompt, image=init_image, mask_image=mask, guidance_scale=guidance_scale, num_inference_steps=int(steps), strength=strength)
     
     return output.images[0], gr.update(visible=True)
@@ -80,7 +84,9 @@ with image_blocks as demo:
     gr.HTML(read_content("header.html"))
     with gr.Row():
                 with gr.Column():
-                    image = gr.Image(source='upload', tool='sketch', elem_id="image_upload", type="pil", label="Upload",height=400)
+                    image = gr.Image(source='upload', tool='sketch', elem_id="image_upload", type="pil", label="image_upload",height=400)
+                    mask_image = gr.Image(source='upload', elem_id="mask_upload", type="pil", label="mask_upload",height=400)
+                    mask_use=gr.Checkbox(label="Use Mask")
                     with gr.Row(elem_id="prompt-container", mobile_collapse=False, equal_height=True):
                         with gr.Row():
                             prompt = gr.Textbox(placeholder="Your prompt (what you want in place of what is erased)", show_label=False, elem_id="prompt")
@@ -104,8 +110,8 @@ with image_blocks as demo:
                         share_button = gr.Button("Share to community", elem_id="share-btn",visible=True)
             
 
-    btn.click(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out, share_btn_container], api_name='run')
-    prompt.submit(fn=predict, inputs=[image, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out, share_btn_container])
+    btn.click(fn=predict, inputs=[image,mask_image,mask_use, prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out, share_btn_container], api_name='run')
+    prompt.submit(fn=predict, inputs=[image,mask_image,mask_use , prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out, share_btn_container])
     share_button.click(None, [], [], _js=share_js)
 
     gr.Examples(