Update app.py

app.py

@@ -6,32 +6,54 @@ from PIL import Image
 import io
 import base64
 import requests
+import os
 from torchvision.transforms import Compose, Resize, ToTensor, Normalize
 
 # Download pre-trained DIS (IS-Net) weights
 def download_weights():
-    url = "https://github.com/xuebinqin/DIS/releases/download/v1.0/isnet-general-use.pth"
-    response = requests.get(url)
-    with open("isnet-general-use.pth", "wb") as f:
-        f.write(response.content)
+    weights_path = "isnet-general-use.pth"
+    if not os.path.exists(weights_path):
+        url = "https://github.com/xuebinqin/DIS/releases/download/v1.0/isnet-general-use.pth"
+        try:
+            response = requests.get(url, stream=True)
+            response.raise_for_status()
+            with open(weights_path, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+        except Exception as e:
+            raise Exception(f"Failed to download weights: {str(e)}")
+    return weights_path
 
-# DIS (IS-Net) model definition (simplified)
+# DIS (IS-Net) model architecture (simplified from https://github.com/xuebinqin/DIS)
 class ISNet(torch.nn.Module):
     def __init__(self):
         super(ISNet, self).__init__()
-        # Placeholder for model architecture (simplified)
-        # In practice, use the full IS-Net architecture from https://github.com/xuebinqin/DIS
-        self.conv = torch.nn.Conv2d(3, 1, kernel_size=3, padding=1)
-        # Load actual weights
-        self.load_state_dict(torch.load("isnet-general-use.pth", map_location="cpu"))
+        # Simplified architecture (for demonstration; replace with full IS-Net)
+        # Full architecture: https://github.com/xuebinqin/DIS/blob/main/ISNet.py
+        self.conv1 = torch.nn.Conv2d(3, 64, kernel_size=3, padding=1)
+        self.pool = torch.nn.MaxPool2d(2, 2)
+        self.conv2 = torch.nn.Conv2d(64, 128, kernel_size=3, padding=1)
+        self.upconv = torch.nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
+        self.conv3 = torch.nn.Conv2d(64, 1, kernel_size=3, padding=1)
 
     def forward(self, x):
-        # Simplified forward pass (replace with actual IS-Net forward)
-        return torch.sigmoid(self.conv(x))
+        # Simplified forward pass (replace with full IS-Net forward)
+        x = torch.relu(self.conv1(x))
+        x = self.pool(x)
+        x = torch.relu(self.conv2(x))
+        x = self.upconv(x)
+        x = torch.sigmoid(self.conv3(x))
+        return x
 
 # Initialize model
-download_weights()
-model = ISNet().eval()
+try:
+    weights_path = download_weights()
+    model = ISNet()
+    state_dict = torch.load(weights_path, map_location="cpu", weights_only=True)
+    model.load_state_dict(state_dict)
+    model.eval()
+except Exception as e:
+    raise Exception(f"Model initialization failed: {str(e)}")
 
 def remove_background(image):
     """
@@ -40,6 +62,10 @@ def remove_background(image):
     Output: Base64-encoded PNG with transparent background
     """
     try:
+        # Ensure image is RGB
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+        
         # Preprocess image
         transform = Compose([
             Resize((1024, 1024)),
@@ -77,8 +103,8 @@ iface = gr.Interface(
     fn=remove_background,
     inputs=gr.Image(type="pil", label="Upload Image"),
     outputs=gr.Image(type="pil", label="Image with Background Removed"),
-    title="Background Removal with DIS (IS-Net)",
-    description="Upload an image to remove its background using the open-source DIS (IS-Net) model.",
+    title="DIS Background Removal",
+    description="Remove backgrounds from any image using the open-source DIS (IS-Net) model.",
     allow_flagging="never"
 )