test

app.py

@@ -88,20 +88,28 @@ 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(duration=120)
 def upsample_features(image, model_option):
     with torch.no_grad():
         from setuptools import setup, find_packages
         from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
 
-        cuda_available = torch.cuda.is_available()
-        cuda_version = torch.version.cuda
-
-        if cuda_available:
-            print(f"CUDA is available. CUDA Version: {cuda_version}")
-            print(f"CUDA Device Name: {torch.cuda.get_device_name(0)}")
-        else:
-            print("CUDA is not available.")
+        cuda_home = find_cuda_home()
+        if cuda_home:
+            print(cuda_home)
+            os.environ['CUDA_HOME'] = cuda_home
 
         setup(
             name='featup',