Update app.py

app.py

@@ -1,18 +1,36 @@
 import gradio as gr
 from rembg import remove
+from rembg.session_factory import new_session
 from PIL import Image
 import io, base64, os
 
 PORT = int(os.environ.get("PORT", 7860))
 
+# ⚡ Preload U²-Net Lite
+session = new_session("u2netp")
+
 def remove_bg_auto_download(image):
-    output = remove(image)
+    if image is None:
+        return ""
 
+    # Convert PIL → bytes
     buf = io.BytesIO()
-    output.save(buf, format="PNG")
-    buf.seek(0)
+    image.save(buf, format="PNG")
+    input_bytes = buf.getvalue()
+
+    # 🔥 REAL background removal
+    output_bytes = remove(input_bytes, session=session)
+
+    # Ensure PIL image w/ alpha
+    output_img = Image.open(io.BytesIO(output_bytes)).convert("RGBA")
+
+    # Back to bytes
+    out_buf = io.BytesIO()
+    output_img.save(out_buf, format="PNG")
+    out_buf.seek(0)
+
+    b64 = base64.b64encode(out_buf.read()).decode()
 
-    b64 = base64.b64encode(buf.read()).decode()
     return f"""
     <script>
       const a = document.createElement('a');
@@ -25,13 +43,13 @@ def remove_bg_auto_download(image):
     """
 
 with gr.Blocks() as demo:
-    gr.Markdown("## Drag & Drop → Auto background removal & download")
-    image = gr.Image(type="pil", label="Drop image")
+    gr.Markdown("## ⚡ Fast Background Removal (U²-Net Lite)")
+    img = gr.Image(type="pil", label="Drag & drop image")
     html = gr.HTML()
 
-    image.change(
+    img.change(
         fn=remove_bg_auto_download,
-        inputs=image,
+        inputs=img,
         outputs=html
     )