Update app.py

app.py

@@ -1,21 +1,11 @@
 import gradio as gr
 from PIL import Image
-from transformers import pipeline
-
-pipe = pipeline("image-segmentation", model="briaai/RMBG-2.0", trust_remote_code=True)
+from rembg import remove
 
 def remove_background(image):
-    result = pipe(image)
-    if isinstance(result, list):
-        result = result[0]
-    output = result.get("image") or result.get("mask")
-    if output is None:
-        return image
-    rgba = image.convert("RGBA")
-    mask = output.convert("L").resize(rgba.size)
-    rgba.putalpha(mask)
-    white_bg = Image.new("RGBA", rgba.size, (255, 255, 255, 255))
-    white_bg.paste(rgba, mask=mask)
+    output = remove(image)
+    white_bg = Image.new("RGBA", output.size, (255, 255, 255, 255))
+    white_bg.paste(output, mask=output.split()[3])
     return white_bg.convert("RGB")
 
 demo = gr.Interface(