Update app.py

app.py

@@ -60,16 +60,16 @@ birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet-ma
 birefnet.to(device)
 birefnet.eval()
 
+@spaces.GPU
 def predict(image):
     if image is None:
         raise gr.Error("Please upload an image.")
 
-    image_ori = Image.fromarray(image)
-    image = image_ori.convert('RGB')
+    image_ori = Image.fromarray(image).convert('RGB')
     
     # Preprocess the image
     image_preprocessor = ImagePreprocessor(resolution=(1024, 1024))
-    image_proc = image_preprocessor.proc(image)
+    image_proc = image_preprocessor.proc(image_ori)
     image_proc = image_proc.unsqueeze(0)
 
     # Prediction
@@ -79,8 +79,8 @@ def predict(image):
 
     # Show Results
     pred_pil = transforms.ToPILImage()(pred)
-    image_masked = refine_foreground(image, pred_pil)
-    image_masked.putalpha(pred_pil.resize(image.size))
+    image_masked = refine_foreground(image_ori, pred_pil)
+    image_masked.putalpha(pred_pil.resize(image_ori.size))
 
     torch.cuda.empty_cache()
 
@@ -94,8 +94,6 @@ iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="numpy"),
     outputs=gr.Image(type="filepath"),
-    title="BiRefNet Matting",
-    description="Upload an image to perform matting using BiRefNet."
 )
 
 if __name__ == "__main__":