Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from gradio_imageslider import ImageSlider
 from loadimg import load_img
-import spaces
+# import spaces
 from transformers import AutoModelForImageSegmentation
 import torch
 from torchvision import transforms
@@ -11,7 +11,8 @@ torch.set_float32_matmul_precision(["high", "highest"][0])
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
-birefnet.to("cuda")
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+birefnet.to(device)
 transform_image = transforms.Compose(
     [
         transforms.Resize((1024, 1024)),
@@ -21,14 +22,14 @@ transform_image = transforms.Compose(
 )
 
 
-@spaces.GPU
+# @spaces.GPU
 def fn(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("cuda")
+    input_images = transform_image(image).unsqueeze(0).to(birefnet)
     # Prediction
     with torch.no_grad():
         preds = birefnet(input_images)[-1].sigmoid().cpu()