Update app.py

app.py

@@ -7,12 +7,18 @@ from transformers import AutoModelForImageSegmentation
 import torch
 from torchvision import transforms
 
+# 检查 CUDA 是否可用
+if torch.cuda.is_available():
+    device = "cuda"
+else:
+    device = "cpu"
+
 torch.set_float32_matmul_precision(["high", "highest"][0])
 
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "briaai/RMBG-2.0", trust_remote_code=True
 )
-birefnet.to("cuda")
+birefnet.to(device)
 transform_image = transforms.Compose(
     [
         transforms.Resize((1024, 1024)),
@@ -37,7 +43,7 @@ def fn(image):
 @spaces.GPU
 def process(image):
     image_size = image.size
-    input_images = transform_image(image).unsqueeze(0).to("cuda")
+    input_images = transform_image(image).unsqueeze(0).to(device)
     # Prediction
     with torch.no_grad():
         preds = birefnet(input_images)[-1].sigmoid().cpu()
@@ -80,4 +86,4 @@ demo = gr.TabbedInterface(
 )
 
 if __name__ == "__main__":
-    demo.launch(show_error=True)
+    demo.launch(share=True, show_error=True)