Update app.py

app.py

@@ -1,42 +1,17 @@
 import gradio as gr
-import torch
-import numpy as np
-import requests
+from rembg import remove
 from PIL import Image
-from torchvision import transforms
+import io
 
-# Load the U²-Net model
-model_url = "https://huggingface.co/zhanghang1989/ResNet101/resolve/main/u2net.pth"
-model_path = "u2net.pth"
+def remove_bg(image):
+    image = remove(image)
+    return image
 
-# Download model if not present
-if not os.path.exists(model_path):
-    with open(model_path, "wb") as f:
-        f.write(requests.get(model_url).content)
-
-def remove_background(image):
-    transform = transforms.Compose([
-        transforms.Resize((320, 320)),  # Smaller size for speed
-        transforms.ToTensor()
-    ])
-    
-    input_tensor = transform(image).unsqueeze(0)
-    
-    # Process with model (fake example, replace with actual U²-Net processing)
-    mask = np.random.randint(0, 256, (320, 320), dtype=np.uint8)
-    mask = np.stack([mask] * 3, axis=-1)
-
-    image = np.array(image.resize((320, 320)))
-    result = np.dstack((image, mask[:, :, 0]))  # Add transparency
-
-    return Image.fromarray(result)
-
-# Gradio Interface
 iface = gr.Interface(
-    fn=remove_background,
+    fn=remove_bg,
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="pil"),
-    title="Fast AI Background Remover"
+    title="Instant Background Remover"
 )
 
 iface.launch()