created app.py

app.py

@@ -0,0 +1,93 @@
+import gradio as gr
+import torch
+import uuid
+from PIL import Image
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
+from typing import Union, List
+from loadimg import load_img  # Your helper to load from URL or file
+
+torch.set_float32_matmul_precision("high")
+
+# Load BiRefNet model
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    "ZhengPeng7/BiRefNet", trust_remote_code=True
+)
+birefnet.to("cuda")
+
+# Image transformation
+transform_image = transforms.Compose([
+    transforms.Resize((1024, 1024)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+])
+
+def process(image: Image.Image) -> Image.Image:
+    image_size = image.size
+    input_tensor = transform_image(image).unsqueeze(0).to("cuda")
+
+    with torch.no_grad():
+        preds = birefnet(input_tensor)[-1].sigmoid().cpu()
+
+    pred = preds[0].squeeze()
+    mask = transforms.ToPILImage()(pred).resize(image_size).convert("L")
+    binary_mask = mask.point(lambda p: 255 if p > 127 else 0)
+
+    white_bg = Image.new("RGB", image_size, (255, 255, 255))
+    result = Image.composite(image, white_bg, binary_mask)
+    return result
+
+def handler(image=None, image_url=None, batch_urls=None) -> Union[str, List[str], None]:
+    results = []
+
+    try:
+        # Single image upload
+        if image is not None:
+            image = image.convert("RGB")
+            processed = process(image)
+            filename = f"output_{uuid.uuid4().hex[:8]}.png"
+            processed.save(filename)
+            return filename
+
+        # Single image from URL
+        if image_url:
+            im = load_img(image_url, output_type="pil").convert("RGB")
+            processed = process(im)
+            filename = f"output_{uuid.uuid4().hex[:8]}.png"
+            processed.save(filename)
+            return filename
+
+        # Batch of URLs
+        if batch_urls:
+            urls = [u.strip() for u in batch_urls.split(",") if u.strip()]
+            for url in urls:
+                try:
+                    im = load_img(url, output_type="pil").convert("RGB")
+                    processed = process(im)
+                    filename = f"output_{uuid.uuid4().hex[:8]}.png"
+                    processed.save(filename)
+                    results.append(filename)
+                except Exception as e:
+                    print(f"Error with {url}: {e}")
+            return results if results else None
+
+    except Exception as e:
+        print("General error:", e)
+
+    return None
+
+# Interface
+demo = gr.Interface(
+    fn=handler,
+    inputs=[
+        gr.Image(label="Upload Image", type="pil", optional=True),
+        gr.Textbox(label="Paste Image URL", optional=True),
+        gr.Textbox(label="Comma-separated Image URLs (Batch)", optional=True),
+    ],
+    outputs=gr.File(label="Output File(s)", file_count="multiple"),
+    title="Background Remover (White Fill)",
+    description="Upload an image, paste a URL, or send a batch of URLs to remove the background and replace it with white.",
+)
+
+if __name__ == "__main__":
+    demo.launch(show_error=True, mcp_server=True)