Update app.py

app.py

@@ -4,22 +4,41 @@ import torch
 import numpy as np
 from PIL import Image
 import io
+import base64
 
 def remove_background(input_image):
     try:
-        # Initialize the pipeline
+        # Initialize the pipeline with trust_remote_code=True
         segmentor = pipeline("image-segmentation", 
                            model="briaai/RMBG-1.4",
-                           device=-1)  # CPU inference
+                           device=-1,
+                           trust_remote_code=True)
         
         # Process the image
         result = segmentor(input_image)
         
-        return result['output_image']
+        # Convert the output image to PNG format
+        if isinstance(result['output_image'], Image.Image):
+            output = result['output_image']
+        else:
+            output = Image.fromarray(result['output_image'])
+            
+        return output, output
+        
     except Exception as e:
         raise gr.Error(f"Error processing image: {str(e)}")
 
-# Create the Gradio interface with custom styling
+def save_image(image):
+    if image is None:
+        return None
+    try:
+        # Save image to bytes
+        buffer = io.BytesIO()
+        image.save(buffer, format="PNG")
+        return buffer.getvalue()
+    except Exception as e:
+        raise gr.Error(f"Error saving image: {str(e)}")
+
 css = """
 .gradio-container {
     font-family: 'Segoe UI', sans-serif;
@@ -28,16 +47,43 @@ css = """
 .gr-image {
     border-radius: 12px;
     border: 2px solid rgba(255, 215, 0, 0.3);
+    transition: transform 0.3s ease;
+}
+.gr-image:hover {
+    transform: scale(1.02);
 }
 .gr-button {
     background: linear-gradient(45deg, #FFD700, #FFA500);
     border: none;
     color: black;
+    padding: 10px 20px;
+    border-radius: 8px;
+    font-weight: bold;
+    transition: all 0.3s ease;
 }
 .gr-button:hover {
     transform: translateY(-2px);
     box-shadow: 0 4px 12px rgba(255, 215, 0, 0.3);
 }
+.container {
+    max-width: 1200px;
+    margin: 0 auto;
+    padding: 2rem;
+}
+.status-msg {
+    text-align: center;
+    margin: 1rem 0;
+    padding: 0.5rem;
+    border-radius: 8px;
+}
+.success {
+    background: rgba(0, 255, 0, 0.1);
+    color: #00ff00;
+}
+.error {
+    background: rgba(255, 0, 0, 0.1);
+    color: #ff0000;
+}
 """
 
 with gr.Blocks(css=css) as demo:
@@ -59,7 +105,8 @@ with gr.Blocks(css=css) as demo:
             input_image = gr.Image(
                 label="Upload Image",
                 type="pil",
-                sources=["upload", "clipboard"]
+                sources=["upload", "clipboard"],
+                tool="select"
             )
             
         with gr.Column():
@@ -67,21 +114,53 @@ with gr.Blocks(css=css) as demo:
                 label="Result",
                 type="pil"
             )
+            download_btn = gr.Button("Download Result", visible=False)
     
     with gr.Row():
         clear_btn = gr.Button("Clear")
         process_btn = gr.Button("Remove Background", variant="primary")
+        
+    status = gr.Textbox(label="Status", visible=False)
 
     # Event handlers
+    def process_and_show(image):
+        if image is None:
+            raise gr.Error("Please upload an image first")
+        try:
+            output, download = remove_background(image)
+            download_btn.visible = True
+            status.visible = True
+            status.value = "✅ Background removed successfully!"
+            return output, download, True, status
+        except Exception as e:
+            status.visible = True
+            status.value = f"❌ Error: {str(e)}"
+            return None, None, False, status
+
+    def clear_all():
+        return {
+            input_image: None,
+            output_image: None,
+            download_btn: gr.Button(visible=False),
+            status: gr.Textbox(visible=False)
+        }
+
     process_btn.click(
-        fn=remove_background,
+        fn=process_and_show,
         inputs=[input_image],
-        outputs=[output_image]
+        outputs=[output_image, output_image, download_btn, status]
     )
     
     clear_btn.click(
-        lambda: (None, None),
-        outputs=[input_image, output_image]
+        fn=clear_all,
+        outputs=[input_image, output_image, download_btn, status]
+    )
+
+    # Download functionality
+    download_btn.click(
+        fn=save_image,
+        inputs=[output_image],
+        outputs=[gr.File(label="Download")]
     )
 
-demo.launch()
+demo.queue().launch()