update code

app.py

@@ -60,39 +60,33 @@ import os
 import subprocess
 import sys
 
-def auto_set_cuda_home():
-    """
-    Automatically detect and set CUDA_HOME environment variable.
-    """
-    if "CUDA_HOME" not in os.environ:
-        potential_paths = [
-            "/usr/local/cuda",
-        ]
-        # Also scan for /usr/local/cuda-*
-        try:
-            ls_out = subprocess.check_output("ls /usr/local | grep cuda", shell=True).decode().splitlines()
-            for line in ls_out:
-                full_path = os.path.join("/usr/local", line.strip())
-                if os.path.isdir(full_path):
-                    potential_paths.append(full_path)
-        except Exception:
-            pass
-
-        for path in potential_paths:
-            nvcc_path = os.path.join(path, "bin", "nvcc")
-            if os.path.exists(nvcc_path):
-                print(f"[INFO] Detected CUDA at: {path}")
-                os.environ["CUDA_HOME"] = path
-                os.environ["PATH"] = f'{os.path.join(path, "bin")}:' + os.environ.get("PATH", "")
-                os.environ["LD_LIBRARY_PATH"] = f'{os.path.join(path, "lib64")}:' + os.environ.get("LD_LIBRARY_PATH", "")
-                return
-
-        print("[WARNING] CUDA not found. Some plugins may fail to compile.")
-    else:
-        print(f"[INFO] CUDA_HOME is already set to: {os.environ['CUDA_HOME']}")
+
+def install_cuda_toolkit():
+    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} ...")
+    subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
+    subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
+
+    print("[INFO] Installing CUDA Toolkit silently ...")
+    subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
+
+    print("[INFO] Setting CUDA environment variables ...")
+    os.environ["CUDA_HOME"] = "/usr/local/cuda"
+    os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ.get("PATH", ""))
+    os.environ["LD_LIBRARY_PATH"] = "%s/lib64:%s" % (
+        os.environ["CUDA_HOME"],
+        os.environ.get("LD_LIBRARY_PATH", "")
+    )
+
+    # Optional: set architecture list for compilation (Ampere and Ada)
+    os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6;8.9"
+
+    print("[INFO] CUDA 12.1 installation complete. CUDA_HOME set to /usr/local/cuda")
 
 # 🔧 Set CUDA_HOME before anything else
-auto_set_cuda_home()
+install_cuda_toolkit()
 
 # Configure logging settings
 logging.basicConfig(