Update app.py

app.py

@@ -1,19 +1,20 @@
 import gradio as gr
-from PIL import Image
-import torch
 from transformers import pipeline
+import torch
+import numpy as np
+from PIL import Image
+import io
 
 def remove_background(input_image):
     try:
-        # Initialize the segmentation pipeline
-        segmentor = pipeline(
-            "image-segmentation",
-            model="briaai/RMBG-1.4",
-            device="cpu"  # Use CPU for inference
-        )
+        # Initialize the pipeline
+        segmentor = pipeline("image-segmentation", 
+                           model="briaai/RMBG-1.4",
+                           device=-1)  # CPU inference
         
         # Process the image
         result = segmentor(input_image)
+        
         return result['output_image']
     except Exception as e:
         raise gr.Error(f"Error processing image: {str(e)}")
@@ -55,7 +56,6 @@ 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",
@@ -71,7 +71,6 @@ with gr.Blocks(css=css) as demo:
     with gr.Row():
         clear_btn = gr.Button("Clear")
         process_btn = gr.Button("Remove Background", variant="primary")
-        download_btn = gr.Button("Download Result")
 
     # Event handlers
     process_btn.click(
@@ -85,17 +84,4 @@ with gr.Blocks(css=css) as demo:
         outputs=[input_image, output_image]
     )
 
-    # Example images
-    gr.Examples(
-        examples=[
-            ["example1.jpg"],
-            ["example2.jpg"],
-            ["example3.jpg"]
-        ],
-        inputs=input_image,
-        outputs=output_image,
-        fn=remove_background,
-        cache_examples=True
-    )
-
 demo.launch()