Update app.py

app.py

@@ -6,7 +6,7 @@ from PIL import Image, ImageOps
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 pipeline = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", vae=vae, torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to("cuda")
 
-def generate(image_editor):
+def generate(image_editor, prompt):
     image = image_editor['background'].convert('RGB')
     mask = Image.new("RGBA", image_editor["layers"][0].size, "WHITE") 
     mask.paste(image_editor["layers"][0], (0, 0), image_editor["layers"][0])
@@ -15,7 +15,9 @@ def generate(image_editor):
     image.thumbnail((1024, 1024))
     mask.thumbnail((1024, 1024))
 
-    return image_editor, image, mask
+    final_image = pipeline(prompt=prompt, image=image, mask_image=mask).images[0]
+
+    return image_editor, image, mask, final_image
 
 with gr.Blocks() as demo:
     gr.Markdown("""
@@ -26,6 +28,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             sketch_pad = gr.ImageMask(type='pil', label='Inpaint')
+            prompt = gr.Textbox()
             generate_button = gr.Button("Generate")
         with gr.Column():
             version_gallery = gr.Gallery(label="Versions")
@@ -41,6 +44,6 @@ with gr.Blocks() as demo:
         out2 = gr.Image()
         out3 = gr.Image()
         
-    generate_button.click(fn=generate, inputs=sketch_pad, outputs=[sketch_pad, out1, out2])
+    generate_button.click(fn=generate, inputs=[sketch_pad,prompt], outputs=[sketch_pad, out1, out2, out3])
 
 demo.launch()