Update app.py

app.py

@@ -73,15 +73,14 @@ def encode_image_to_base64(image):
     buffered = BytesIO()
     image.save(buffered, format="PNG")  # Use PNG for compatibility with RGBA
     return base64.b64encode(buffered.getvalue()).decode('utf-8')
-
-# @spaces.GPU
+@spaces.GPU
 def rm_background(image):
     im = load_img(image, output_type="pil")
     im = im.convert("RGB")
     image_size = im.size
     origin = im.copy()
     image = load_img(im)
-    input_images = transform_image(image).unsqueeze(0).to("cpu")
+    input_images = transform_image(image).unsqueeze(0).to("cuda")
     # Prediction
     with torch.no_grad():
         preds = birefnet(input_images)[-1].sigmoid().cpu()