Update app.py

app.py

@@ -1,32 +1,24 @@
 import gradio as gr
-from transformers import pipeline
+from rembg import remove
 from PIL import Image
-
-# -------------------------------
-# Load Hugging Face background removal model
-# -------------------------------
-pipe = pipeline(
-    task="image-segmentation",
-    model="briaai/RMBG-1.4",
-    trust_remote_code=True
-)
+import io
 
 # -------------------------------
 # Function to remove background
 # -------------------------------
 def remove_bg(image: Image.Image):
-    # Run model → this returns a PIL mask directly
-    mask = pipe(image)
+    # Convert PIL image to bytes
+    img_byte_arr = io.BytesIO()
+    image.save(img_byte_arr, format='PNG')
+    img_bytes = img_byte_arr.getvalue()
 
-    # Ensure mask is grayscale
-    mask = mask.convert("L")
+    # Remove background using rembg
+    output_bytes = remove(img_bytes)
 
-    # Apply mask to original image
-    image = image.convert("RGBA")
-    mask = mask.resize(image.size)
-    image.putalpha(mask)
+    # Convert bytes back to PIL image
+    output_image = Image.open(io.BytesIO(output_bytes))
 
-    return image
+    return output_image
 
 # -------------------------------
 # Gradio UI
@@ -36,7 +28,7 @@ demo = gr.Interface(
     inputs=gr.Image(type="pil", label="Upload an image"),
     outputs=gr.Image(type="pil", label="Background Removed"),
     title="Background Remover",
-    description="Upload an image to remove its background using Hugging Face RMBG model.",
+    description="Upload an image to remove its background using the rembg library.",
 )
 
 # -------------------------------