Update app.py

app.py

@@ -1,13 +1,40 @@
 import spaces
 import os
 import gradio as gr
+import shutil
+
+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
+
+os.environ["CUDA_HOME"] = find_cuda()
+
+if cuda_path:
+    print(f"CUDA installation found at: {cuda_path}")
+else:
+    print("CUDA installation not found")
+
+
+
 @spaces.GPU(duration=180)
 def install_setup():
     os.system("pip install git+https://github.com/facebookresearch/detectron2.git")
     os.system("git clone https://github.com/Visual-AI/Mr.DETR.git MrDETR && cd MrDETR && rm -f requirements.txt && cd ..")
     # os.system("cp multi_scale_deform_attn.py MrDETR/detrex/layers/ && cd MrDETR && pip install . & cd ..")
     # os.system("cd MrDETR && export CUDA_HOME=$(dirname $(dirname $(which nvcc))) && pip install . & cd ..")
-    os.environ["CUDA_HOME"] = os.environ.get("CUDA_PATH")
     os.system("cd MrDETR && pip install . & cd ..")
     import sys 
     sys.path.append("MrDETR/")