Update app.py

app.py

@@ -1,52 +1,42 @@
 import gradio as gr
-from transformers import pipeline
-import torch
-import numpy as np
 from PIL import Image
-import io
-import base64
+import torch
+from transformers import pipeline
 
 def remove_background(input_image):
     try:
-        # Initialize the pipeline
-        segmentor = pipeline("image-segmentation", 
-                           model="briaai/RMBG-1.4",
-                           device=-1)  # CPU inference
+        # Initialize the segmentation pipeline
+        segmentor = pipeline(
+            "image-segmentation",
+            model="briaai/RMBG-1.4",
+            device="cpu"  # Use CPU for inference
+        )
         
         # Process the image
         result = segmentor(input_image)
-        
-        # Return both original and processed images
         return result['output_image']
     except Exception as e:
         raise gr.Error(f"Error processing image: {str(e)}")
 
-# Create the Gradio interface
+# Create the Gradio interface with custom styling
 css = """
 .gradio-container {
     font-family: 'Segoe UI', sans-serif;
     background: linear-gradient(135deg, #1a1a1a 0%, #2d2d2d 100%);
 }
+.gr-image {
+    border-radius: 12px;
+    border: 2px solid rgba(255, 215, 0, 0.3);
+}
 .gr-button {
     background: linear-gradient(45deg, #FFD700, #FFA500);
     border: none;
     color: black;
 }
 .gr-button:hover {
-    background: linear-gradient(45deg, #FFA500, #FFD700);
     transform: translateY(-2px);
     box-shadow: 0 4px 12px rgba(255, 215, 0, 0.3);
 }
-.gr-form {
-    background: rgba(255, 255, 255, 0.1);
-    border-radius: 16px;
-    padding: 20px;
-    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.2);
-}
-.gr-image {
-    border-radius: 12px;
-    border: 2px solid rgba(255, 215, 0, 0.3);
-}
 """
 
 with gr.Blocks(css=css) as demo:
@@ -54,7 +44,7 @@ with gr.Blocks(css=css) as demo:
         """
         <div style="text-align: center; max-width: 800px; margin: 0 auto; padding: 20px;">
             <h1 style="font-size: 2.5rem; margin-bottom: 1rem; background: linear-gradient(45deg, #FFD700, #FFA500); -webkit-background-clip: text; -webkit-text-fill-color: transparent;">
-                Background Removal Tool
+                AI Background Remover
             </h1>
             <p style="color: #cccccc; font-size: 1.1rem; margin-bottom: 2rem;">
                 Powered by RMBG V1.4 model from BRIA AI
@@ -65,33 +55,34 @@ with gr.Blocks(css=css) as demo:
     
     with gr.Row():
         with gr.Column():
+            # Correct implementation of input image
             input_image = gr.Image(
                 label="Upload Image",
                 type="pil",
-                tool="upload",
+                sources=["upload", "clipboard"]
             )
             
-            with gr.Row():
-                clear_btn = gr.Button("Clear")
-                submit_btn = gr.Button("Remove Background", variant="primary")
-
         with gr.Column():
             output_image = gr.Image(
                 label="Result",
-                type="pil",
+                type="pil"
             )
-            download_btn = gr.Button("Download Result")
+    
+    with gr.Row():
+        clear_btn = gr.Button("Clear")
+        process_btn = gr.Button("Remove Background", variant="primary")
+        download_btn = gr.Button("Download Result")
 
     # Event handlers
-    submit_btn.click(
+    process_btn.click(
         fn=remove_background,
         inputs=[input_image],
-        outputs=[output_image],
+        outputs=[output_image]
     )
     
     clear_btn.click(
         lambda: (None, None),
-        outputs=[input_image, output_image],
+        outputs=[input_image, output_image]
     )
 
     # Example images
@@ -99,19 +90,12 @@ with gr.Blocks(css=css) as demo:
         examples=[
             ["example1.jpg"],
             ["example2.jpg"],
-            ["example3.jpg"],
+            ["example3.jpg"]
         ],
         inputs=input_image,
         outputs=output_image,
         fn=remove_background,
-        cache_examples=True,
-    )
-
-    # Download functionality
-    download_btn.click(
-        lambda x: x,
-        inputs=[output_image],
-        outputs=[output_image],
+        cache_examples=True
     )
 
 demo.launch()