update code

app.py

@@ -56,6 +56,43 @@ import shutil
 
 # Suppress warnings (especially for PyTorch)
 warnings.filterwarnings("ignore")
+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']}")
+
+# 🔧 Set CUDA_HOME before anything else
+auto_set_cuda_home()
 
 # Configure logging settings
 logging.basicConfig(