Update app.py

app.py

@@ -84,15 +84,9 @@ def process(image, progress=gr.Progress()):
     output_array = np.array(new_im.convert("RGBA"))
     
     progress(1.0, desc="Done!")
-    # Return image for display and path for download
+    # Return the image, filepath and update button visibility
     return output_array, filepath, gr.update(visible=True)
 
-# Function to handle the download button click
-def download_image(filepath):
-    if filepath and os.path.exists(filepath):
-        return filepath
-    return None
-
 # Gradio interface setup
 title = "Background Removal"
 description = r"""Background removal model developed by <a href='https://BRIA.AI' target='_blank'><b>BRIA.AI</b></a>, trained on a carefully selected dataset and is available as an open-source model for non-commercial use.<br> For testing, upload your image and wait. Read more at model card <a href='https://huggingface.co/briaai/RMBG-1.4' target='_blank'><b>briaai/RMBG-1.4</b></a>. To purchase a commercial license, simply click <a href='https://go.bria.ai/3ZCBTLH' target='_blank'><b>Here</b></a>. <br>"""
@@ -100,29 +94,12 @@ description = r"""Background removal model developed by <a href='https://BRIA.AI
 examples = [['./input.jpg']]
 
 # Create the Gradio interface
-with gr.Blocks(css="""
-    .download-btn {
-        background-color: #4CAF50;
-        border: none;
-        color: white;
-        padding: 10px 24px;
-        text-align: center;
-        text-decoration: none;
-        display: inline-block;
-        font-size: 16px;
-        margin: 4px 2px;
-        cursor: pointer;
-        border-radius: 4px;
-    }
-    .download-btn:hover {
-        background-color: #45a049;
-    }
-""") as demo:
+with gr.Blocks() as demo:
     gr.Markdown(f"# {title}")
     gr.Markdown(description)
     
-    # Store the processed image path
-    image_path = gr.State(None)
+    # Store the filepath
+    filepath = gr.State()
     
     with gr.Row():
         with gr.Column(scale=1):
@@ -130,10 +107,11 @@ with gr.Blocks(css="""
         
         with gr.Column(scale=1):
             output_image = gr.Image(type="numpy", label="Result")
-            download_btn = gr.Button("Download Image", elem_id="download_button", 
-                                     variant="primary", visible=False, 
-                                     elem_classes="download-btn")
-            download_output = gr.File(visible=False)
+            with gr.Row():
+                download_btn = gr.Button("Download Image", variant="primary", visible=False)
+    
+    # Hidden download component that will be triggered
+    download_file = gr.File(interactive=False, visible=False)
     
     # Set up example images
     gr.Examples(examples, inputs=input_image)
@@ -142,15 +120,33 @@ with gr.Blocks(css="""
     input_image.change(
         fn=process,
         inputs=input_image,
-        outputs=[output_image, image_path, download_btn],
+        outputs=[output_image, filepath, download_btn],
         show_progress="full"
     )
     
     # Handle download button click
     download_btn.click(
-        fn=download_image,
-        inputs=[image_path],
-        outputs=[download_output]
+        fn=lambda x: x,  # Just pass through the filepath
+        inputs=[filepath],
+        outputs=[download_file],  # Output to the download component
+        _js="""
+        function(path) {
+            // Create a temporary anchor element
+            const link = document.createElement('a');
+            link.href = '/file=' + path;
+            link.download = path.split('/').pop();
+            
+            // Append to the document
+            document.body.appendChild(link);
+            
+            // Trigger click and remove
+            link.click();
+            document.body.removeChild(link);
+            
+            // Return the path for the hidden component
+            return path;
+        }
+        """
     )
 
 if __name__ == "__main__":