Upload app_hg.py with huggingface_hub

app_hg.py

@@ -28,7 +28,6 @@ import warnings
 import argparse
 import gradio as gr
 from glob import glob
-print(glob("/usr/local/cuda*"))
 import shutil
 import torch
 import numpy as np
@@ -52,6 +51,32 @@ parser.add_argument("--save_memory", default=True) # , action="store_true")
 parser.add_argument("--device", default="cuda:0", type=str)
 args = parser.parse_args()
 
+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
+
+cuda_path = find_cuda()
+
+if cuda_path:
+    print(f"CUDA installation found at: {cuda_path}")
+else:
+    print("CUDA installation not found")
+
+
+
 def download_models():
     # Create weights directory if it doesn't exist
     os.makedirs("weights", exist_ok=True)