test

app.py

@@ -88,19 +88,42 @@ image_input = gr.Image(label="Choose an image to featurize",
 model_option = gr.Radio(options, value="dino16",
                         label='Choose a backbone to upsample')
 
+
 @spaces.GPU
 def upsample_features(image, model_option):
 
-    # Find all CUDA directories that match /usr/local/cuda*
-    cuda_dirs = glob.glob('/usr/local/cuda*')
-
-    if not cuda_dirs:
+    # Check common CUDA installation paths and set the correct one
+    cuda_paths = [
+        '/usr/local/cuda',
+        '/usr/local/cuda-10.0',
+        '/usr/local/cuda-10.1',
+        '/usr/local/cuda-10.2',
+        '/usr/local/cuda-11.0',
+        '/usr/local/cuda-11.1',
+        '/usr/local/cuda-11.2',
+        '/usr/local/cuda-11.3',
+        '/usr/local/cuda-11.4',
+        '/usr/local/cuda-11.5',
+        '/usr/local/cuda-11.6',
+        '/usr/local/cuda-11.7',
+        '/usr/local/cuda-11.8',
+        '/usr/local/cuda-12.0',
+        '/usr/local/cuda-12.1',
+        '/usr/lib/cuda',
+        '/opt/cuda'
+    ]
+
+    cuda_home = None
+    for path in cuda_paths:
+        if os.path.exists(path):
+            cuda_home = path
+            break
+
+    if cuda_home is None:
         raise EnvironmentError(
-            'No CUDA installation found. Please install CUDA or set CUDA_HOME manually.')
+            'CUDA installation not found. Please install CUDA or set CUDA_HOME manually.')
 
-    # Assume the highest version of CUDA is the one to use
-    cuda_dirs.sort()
-    cuda_home = cuda_dirs[-1]
+    print(cuda_home)
 
     # Set the CUDA_HOME environment variable
     os.environ['CUDA_HOME'] = cuda_home