Update app.py

app.py

@@ -8,12 +8,18 @@ from torchvision import transforms
 import uuid
 import os
 
+# Automatically select device based on availability
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Optional: Print which device is being used
+print(f"Using device: {device}")
+
 torch.set_float32_matmul_precision(["high", "highest"][0])
 
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
-birefnet.to("cuda")
+birefnet.to(device)
 
 transform_image = transforms.Compose(
     [
@@ -26,7 +32,7 @@ transform_image = transforms.Compose(
 @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()