Update app.py

app.py

@@ -1,7 +1,3 @@
-
-
-
-
 import gradio as gr
 from rembg import remove
 from PIL import Image
@@ -12,13 +8,16 @@ def remove_background(input_image):
     output = remove(input) # remove background
     return output
 
-# define input and output interfaces
-
-input_image = gr.inputs.Image(type="filepath", label="Input Image")
-output_image = gr.inputs.Image(type="filepath", label="Output Image")
-
+def download_image(output_image):
+    img_bytes = output_image.getvalue()
+    return gradio.Interface.download(img_bytes, "output_image.png")
 
+# define input and output interfaces
+input_image = gr.inputs.Image(type="file", label="Input Image")
+output_image = gr.outputs.Image(type="file", label="Output Image")
 
+# define the download button
+download_button = gr.InterfaceFn(download_image, output_image, title="Download")
 
 # create the interface
 iface = gr.Interface(
@@ -26,9 +25,12 @@ iface = gr.Interface(
     inputs=input_image, 
     outputs=output_image,
     title='Image Background Remover',
-    description='Upload an image and it will remove.'
+    description='Upload an image and it will remove.',
+    allow_download=True
 )
 
-# launch the interface
-iface.launch()
+# combine the main interface and the download button
+interface_with_download = gr.Interface([iface, download_button], "grid")
 
+# launch the interface
+interface_with_download.launch()