Update app.py

app.py

@@ -3,7 +3,7 @@ import cv2
 import numpy as np
 import torch
 import gradio as gr
-import spaces  # Added import for spaces
+import spaces  # Required for @spaces.GPU
 
 from PIL import Image, ImageOps
 from transformers import AutoModelForImageSegmentation
@@ -49,24 +49,30 @@ class ImagePreprocessor():
         self.transform_image = transforms.Compose([
             transforms.Resize(resolution),
             transforms.ToTensor(),
-            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+            transforms.Normalize([0.485, 0.456, 0.406],
+                                 [0.229, 0.224, 0.225]),
         ])
 
     def proc(self, image: Image.Image) -> torch.Tensor:
         image = self.transform_image(image)
         return image
 
+# Load the model
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     'zhengpeng7/BiRefNet-matting', trust_remote_code=True)
 birefnet.to(device)
 birefnet.eval()
 
-@spaces.GPU  # Added the @spaces.GPU decorator
-def remove_background(image):
+def remove_background_wrapper(image):
     if image is None:
         raise gr.Error("Please upload an image.")
-
     image_ori = Image.fromarray(image).convert('RGB')
+    # Call the processing function
+    foreground, background, pred_pil, reverse_mask = remove_background(image_ori)
+    return foreground, background, pred_pil, reverse_mask
+
+@spaces.GPU  # Decorate the processing function
+def remove_background(image_ori):
     original_size = image_ori.size
 
     # Preprocess the image
@@ -100,7 +106,7 @@ def remove_background(image):
     return foreground, background, pred_pil, reverse_mask
 
 iface = gr.Interface(
-    fn=remove_background,
+    fn=remove_background_wrapper,
     inputs=gr.Image(type="numpy"),
     outputs=[
         gr.Image(type="pil", label="Foreground"),