Update app.py

app.py

@@ -1,9 +1,34 @@
 from transformers import pipeline, SamModel, SamProcessor
 import torch
 import numpy as np
+import spaces
 import gradio as gr
 from PIL import Image
 
+def find_cuda():
+    # Check if CUDA_HOME or CUDA_PATH environment variables are set
+    cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
+
+    if cuda_home and os.path.exists(cuda_home):
+        return cuda_home
+
+    # Search for the nvcc executable in the system's PATH
+    nvcc_path = shutil.which('nvcc')
+
+    if nvcc_path:
+        # Remove the 'bin/nvcc' part to get the CUDA installation path
+        cuda_path = os.path.dirname(os.path.dirname(nvcc_path))
+        return cuda_path
+
+    return None
+
+cuda_path = find_cuda()
+
+if cuda_path:
+    print(f"CUDA installation found at: {cuda_path}")
+else:
+    print("CUDA installation not found")
+    
 # check if cuda is available
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -19,6 +44,7 @@ def apply_mask(image, mask, color):
         image[:, :, c] = np.where(mask, color[c], image[:, :, c])
     return image
 
+@spaces.GPU
 def query(image, texts, threshold):
     texts = texts.split(",")
     predictions = detector(