Update app.py

app.py

@@ -2,55 +2,46 @@ import gradio as gr
 import torch
 import numpy as np
 from PIL import Image
+from transformers import AutoModelForImageSegmentation
 from torchvision import transforms
 
-# BiRefNet import
-from birefnet import BiRefNet
+# Load RMBG-1.4 (BiRefNet)
+model = AutoModelForImageSegmentation.from_pretrained(
+    "briaai/RMBG-1.4",
+    trust_remote_code=True
+)
+model.eval()
 
-# Device
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load model
-model = BiRefNet.from_pretrained("ZhengPeng7/BiRefNet")
 model.to(device)
-model.eval()
 
-# Image transform
 transform = transforms.Compose([
-    transforms.Resize((1024, 1024)),
-    transforms.ToTensor(),
+    transforms.ToTensor()
 ])
 
-def remove_bg(image):
-    if image is None:
-        return None
+def remove_bg(image: Image.Image):
+    if image.mode != "RGB":
+        image = image.convert("RGB")
 
-    original_size = image.size
-
-    img = transform(image).unsqueeze(0).to(device)
+    inp = transform(image).unsqueeze(0).to(device)
 
     with torch.no_grad():
-        pred = model(img)[-1]
-        mask = torch.sigmoid(pred).cpu().squeeze().numpy()
-
-    mask = Image.fromarray((mask * 255).astype(np.uint8)).resize(original_size)
-
-    image = image.convert("RGBA")
-    image_np = np.array(image)
-    mask_np = np.array(mask)
+        pred = model(inp)[0][0]
 
-    image_np[:, :, 3] = mask_np
+    mask = pred.sigmoid().cpu().numpy()
+    mask = (mask > 0.5).astype(np.uint8) * 255
+    mask = Image.fromarray(mask).resize(image.size)
 
-    return Image.fromarray(image_np)
+    output = image.copy()
+    output.putalpha(mask)
+    return output
 
-# Gradio UI
 demo = gr.Interface(
     fn=remove_bg,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Image(type="pil", label="Background Removed"),
-    title="BiRefNet AI – Background Remover",
-    description="Powered by BiRefNet (Open Source AI Model)"
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="RMBG-1.4 Background Remover",
+    description="BiRefNet based background removal"
 )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()