update code

app.py

@@ -62,7 +62,7 @@ import sys
 
 os.environ["MEDIAPIPE_DISABLE_GPU"] = "1"  # Disable GPU for MediaPipe
 def install_cuda_toolkit():
-    CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run"
+    CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run"
     CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
 
     print(f"[INFO] Downloading CUDA Toolkit from {CUDA_TOOLKIT_URL} ...")
@@ -305,8 +305,8 @@ def prepare_working_dir(dir, style):
 def launch_pretrained():
 
     from huggingface_hub import snapshot_download
-    os.system("pip uninstall torch -y")
-    os.system("pip uninstall torchvision -y")
+    os.system("pip uninstall torch")
+    os.system("pip uninstall torchvision")
     os.system("pip install https://download.pytorch.org/whl/cu121/torch-2.4.1%2Bcu121-cp310-cp310-linux_x86_64.whl")
     os.system("pip install https://download.pytorch.org/whl/cu121/torchvision-0.19.1%2Bcu121-cp310-cp310-linux_x86_64.whl")
     snapshot_download(
@@ -957,7 +957,7 @@ if __name__ == '__main__':
     if _get_output(["nvcc", "--version"]) is None:
         logging.info("Installing CUDA toolkit...")
         install_cuda_toolkit()
-        logging.info("install CUDA: %s" % _get_output(["nvcc", "--version"]))
+        logging.info("installCUDA: %s" % _get_output(["nvcc", "--version"]))
 
     else:
         logging.info("Detected CUDA: %s" % _get_output(["nvcc", "--version"]))

recon/torch_utils/custom_ops.py

@@ -19,7 +19,9 @@ import uuid
 import torch
 import torch.utils.cpp_extension
 from torch.utils.file_baton import FileBaton
-
+os.environ["CUDA_HOME"] = "/usr/local/cuda"
+os.environ["PATH"] = f"{os.environ['CUDA_HOME']}/bin:" + os.environ.get("PATH", "")
+os.environ["LD_LIBRARY_PATH"] = f"{os.environ['CUDA_HOME']}/lib64:" + os.environ.get("LD_LIBRARY_PATH", "")
 #----------------------------------------------------------------------------
 # Global options.