test

app.py

@@ -88,16 +88,12 @@ model_option = gr.Radio(options, value="dino16",
                         label='Choose a backbone to upsample')
 
 
-def find_cuda_home():
+def find_nvcc():
     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}")
+        result = subprocess.check_output(['find', '/', '-name', 'nvcc', '2>/dev/null'], stderr=subprocess.STDOUT, text=True, shell=True)
+        return result.strip()
+    except subprocess.CalledProcessError as e:
+        print(f"Error occurred: {e}")
         return None
 
 @spaces.GPU(duration=120)
@@ -107,10 +103,11 @@ def upsample_features(image, model_option):
         from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
 
         print(subprocess.check_output(['ls', '/usr/local/']).decode())
-        cuda_home = find_cuda_home()
-        if cuda_home:
-            print(cuda_home)
-            os.environ['CUDA_HOME'] = cuda_home
+        nvcc_path = find_nvcc()
+        if nvcc_path:
+            print(f"CUDA 'nvcc' found at: {nvcc_path}")
+        else:
+            print("CUDA 'nvcc' not found.")
 
         setup(
             name='featup',