Update app.py

app.py

@@ -2,40 +2,58 @@ import gradio as gr
 from PIL import Image
 import torch
 import numpy as np
-from transformers import U2NetForImageSegmentation, U2NetProcessor
-
-# Load model and processor
-processor = U2NetProcessor.from_pretrained("akhaliq/U-2-Net")
-model = U2NetForImageSegmentation.from_pretrained("akhaliq/U-2-Net")
-
-def remove_background(image: Image.Image) -> Image.Image:
-    # Convert to RGB
-    image = image.convert("RGB")
-    
-    # Prepare input
-    inputs = processor(images=image, return_tensors="pt")
-    
-    # Run inference
+import os
+from torchvision import transforms
+from u2net import U2NET # Load model class
+
+# Load the model (download once, reuse)
+model_path = "u2net.pth"
+if not os.path.exists(model_path):
+    import requests
+    url = "https://huggingface.co/akhaliq/U-2-Net/resolve/main/u2net.pth"
+    with open(model_path, "wb") as f:
+        f.write(requests.get(url).content)
+
+# Load model to CPU
+net = U2NET(3,1)
+net.load_state_dict(torch.load(model_path, map_location='cpu'))
+net.eval()
+
+# Preprocessing
+transform = transforms.Compose([
+    transforms.Resize((320,320)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406],
+                         [0.229, 0.224, 0.225])
+])
+
+# Post-process output mask
+def normalize_prediction(pred):
+    ma = torch.max(pred)
+    mi = torch.min(pred)
+    return (pred - mi) / (ma - mi)
+
+# Main background remover function
+def remove_bg(input_image):
+    image = input_image.convert("RGB")
+    orig_size = image.size
+    img_tensor = transform(image).unsqueeze(0)
+
     with torch.no_grad():
-        outputs = model(**inputs)
-    
-    # Get predicted mask
-    pred = outputs.predictions[0][0]
-    mask = (pred > 0.5).numpy().astype(np.uint8) * 255
-    mask_image = Image.fromarray(mask).resize(image.size)
-    
-    # Create RGBA image with transparency
-    result = image.copy()
-    result.putalpha(mask_image)
-    return result
-
-# Gradio UI
-interface = gr.Interface(
-    fn=remove_background,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Image(type="pil", label="Image with Transparent Background"),
-    title="🪄 AI Background Remover (remove.bg Clone)",
-    description="Removes background from images using U-2-Net model from Hugging Face.",
-)
+        d1, *_ = net(img_tensor)
+        mask = normalize_prediction(d1[0][0])
+        mask = mask.squeeze().cpu().numpy()
+        mask = Image.fromarray((mask * 255).astype(np.uint8)).resize(orig_size)
 
-interface.launch()
+    # Add alpha channel using mask
+    image.putalpha(mask)
+    return image
+
+# Gradio app
+gr.Interface(
+    fn=remove_bg,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=gr.Image(type="pil", label="Transparent PNG"),
+    title="🪄 Remove.bg Clone - AI Background Remover",
+    description="Upload any image to remove the background using U-2-Net (Hugging Face version)."
+).launch()