test

app.py

@@ -87,11 +87,24 @@ image_input = gr.Image(label="Choose an image to featurize",
 model_option = gr.Radio(options, value="dino16",
                         label='Choose a backbone to upsample')
 
+def find_cuda_home():
+    try:
+        output = subprocess.check_output(['nvcc', '--version']).decode()
+        print(output)
+        for line in output.split('\n'):
+            if 'Cuda compilation tools' in line:
+                version = line.split()[-1]
+                return f"/usr/local/cuda-{version.split('.')[0]}.{version.split('.')[1]}"
+    except Exception as e:
+        print(f"Error finding CUDA_HOME: {e}")
+        return None
+
 @spaces.GPU
 def upsample_features(image, model_option):
-    from torch.utils.cpp_extension import CUDA_HOME
+    cuda_home = find_cuda_home()
+    print(cuda_home)
 
-    print(f"CUDA Home: {CUDA_HOME}")
+    os.environ['CUDA_HOME'] = cuda_home
     os.environ['PATH'] = '/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin'
     os.environ['LD_LIBRARY_PATH'] = '/usr/local/nvidia/lib:/usr/local/nvidia/lib64'