Update app.py

app.py

@@ -33,7 +33,7 @@ def resize_image(image):
 # Background removal process
 def process(image, progress=gr.Progress()):
     if image is None:
-        return None, None, gr.update(visible=False)
+        return None, gr.update(visible=False)
     
     progress(0, desc="Starting processing...")
     
@@ -84,8 +84,8 @@ def process(image, progress=gr.Progress()):
     output_array = np.array(new_im.convert("RGBA"))
     
     progress(1.0, desc="Done!")
-    # Return the image, filepath and update button visibility
-    return output_array, filepath, gr.update(visible=True)
+    # Return the processed image and update download component visibility
+    return output_array, gr.update(visible=True, value=filepath)
 
 # Gradio interface setup
 title = "Background Removal"
@@ -94,24 +94,28 @@ description = r"""Background removal model developed by <a href='https://BRIA.AI
 examples = [['./input.jpg']]
 
 # Create the Gradio interface
-with gr.Blocks() as demo:
+with gr.Blocks(css="""
+    .download-btn {
+        margin-top: 10px !important;
+    }
+    .result-container {
+        display: flex;
+        flex-direction: column;
+    }
+""") as demo:
     gr.Markdown(f"# {title}")
     gr.Markdown(description)
     
-    # Store the filepath
-    filepath = gr.State()
-    
     with gr.Row():
         with gr.Column(scale=1):
             input_image = gr.Image(type="numpy", label="Upload Image")
         
-        with gr.Column(scale=1):
+        with gr.Column(scale=1, elem_classes="result-container"):
             output_image = gr.Image(type="numpy", label="Result")
-            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)
+            # Use File component as a download button
+            download_file = gr.File(label="Download Processed Image", 
+                                  elem_classes="download-btn",
+                                  visible=False)
     
     # Set up example images
     gr.Examples(examples, inputs=input_image)
@@ -120,34 +124,9 @@ with gr.Blocks() as demo:
     input_image.change(
         fn=process,
         inputs=input_image,
-        outputs=[output_image, filepath, download_btn],
+        outputs=[output_image, download_file],
         show_progress="full"
     )
-    
-    # Handle download button click
-    download_btn.click(
-        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__":
     demo.launch(share=False)