Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import numpy as np
 import cv2
+import torch
 
 # import spaces
 
@@ -23,7 +24,8 @@ choice_mapping = {
 # @spaces.GPU
 def predict(model_choice: str, annotations, image):
     config_file, ckpt_path = choice_mapping[str(model_choice)]
-    sam2_model = build_sam2(config_file, ckpt_path, device="cuda")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    sam2_model = build_sam2(config_file, ckpt_path, device=device)
     predictor = SAM2ImagePredictor(sam2_model)
     predictor.set_image(image)
     coordinates = np.array(