Update app.py

app.py

@@ -1,13 +1,14 @@
 # app.py
-# A more user-friendly version with a working custom download button.
+# A user-friendly version with a custom download button, inspired by the watermark remover app.
 
 import gradio as gr
 from rembg import remove
 import traceback
 from PIL import Image
-import tempfile # ✅ NEW: Import the 'tempfile' library
+import tempfile # Used to create temporary files for downloading
 
 # --- Core Function ---
+# This function processes the image. It remains the same.
 def process_image(input_image):
     """
     Removes the background from an image.
@@ -17,24 +18,25 @@ def process_image(input_image):
         raise gr.Error("Please upload an image or select one of the examples.")
     
     try:
+        # Use the default rembg model for stability
         return remove(input_image)
     except Exception as e:
         print(traceback.format_exc())
         raise gr.Error(f"Failed to process image: {e}")
 
 # --- Download Function ---
-# ✅ FIX: This function now correctly saves the image to a temporary
-# file and returns the file path to trigger the download.
-def download_output_image(image):
+# This function correctly prepares the image for download.
+def prepare_for_download(image):
     """
     Saves the processed PIL image to a temporary file and returns the path.
+    This is the correct way to trigger a download in Gradio.
     """
     if image is None:
         # Prevent an error if the download button is clicked before an image is processed.
         return None
     try:
-        # Create a temporary file with a '.png' extension. 'delete=False' is important
-        # so the file is not deleted before Gradio can serve it.
+        # Create a temporary file with a '.png' extension.
+        # 'delete=False' is important so the file is not deleted before Gradio serves it.
         with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
             image.save(tmp.name)
             # Return the path to the temporary file
@@ -44,9 +46,10 @@ def download_output_image(image):
         return None
 
 # --- Gradio App with Modern UI ---
+# We inject CSS to hide the default utility buttons on the output image.
 with gr.Blocks(
     theme=gr.themes.Soft(primary_hue="violet"),
-    css=".meta-buttons { display: none !important; }"
+    css="#output_image .meta-buttons { display: none !important; }"
 ) as demo:
     
     gr.Markdown(
@@ -55,6 +58,7 @@ with gr.Blocks(
         """
     )
     
+    # Main two-column layout
     with gr.Row(variant="panel"):
         
         # --- Input Column ---
@@ -66,7 +70,7 @@ with gr.Blocks(
                 height=400
             )
             
-            
+           
             with gr.Row():
                 clear_btn = gr.ClearButton(
                     value="Clear",
@@ -77,40 +81,37 @@ with gr.Blocks(
 
         # --- Output Column ---
         with gr.Column(scale=1, min_width=300):
+            # We give the output image an ID so we can target it with CSS
             image_output = gr.Image(
+                elem_id="output_image",
                 type="pil",
                 label="Result",
                 height=400,
-                value=None,
-                show_download_button=False,
-                interactive=False 
-            )
-            
-            download_button = gr.Button(
-                "Download Image",
-                icon="https://cdn.iconscout.com/icon/premium/png-256-thumb/download-button-2132338-1794935.png",
-                variant="primary",
+                interactive=False # This disables the upload overlay on the result
             )
             
-            # This hidden file component is used as the target for the download action.
-            download_file_output = gr.File(visible=False)
+            # ✅ NEW: Use a gr.File component for the download, which is more reliable.
+            # This will appear as a download link/button when populated.
+            download_output = gr.File(label="Download Result", visible=False)
 
     # --- Event Listeners ---
+    def on_submit(input_img):
+        # This function runs when the user clicks submit.
+        # It processes the image and then prepares it for download.
+        processed_img = process_image(input_img)
+        download_path = prepare_for_download(processed_img)
+        # It returns the processed image to the preview box,
+        # and the downloadable file to the (now visible) download component.
+        return processed_img, gr.update(value=download_path, visible=True)
+
     submit_btn.click(
-        fn=process_image,
+        fn=on_submit,
         inputs=image_input,
-        outputs=image_output,
-        api_name="predict" 
-    )
-
-    # The download button click event remains the same, but the function it calls is now fixed.
-    download_button.click(
-        fn=download_output_image,
-        inputs=[image_output],
-        outputs=[download_file_output]
+        outputs=[image_output, download_output]
     )
 
-    clear_btn.add(components=[image_output])
+    # When the user clicks "Clear", also hide the download button.
+    clear_btn.click(lambda: gr.update(visible=False), outputs=[download_output])
 
 # --- Launch the App ---
 if __name__ == "__main__":