Upload app.py with huggingface_hub

app.py

@@ -23,16 +23,28 @@ if not os.path.exists(MODEL_PATH):
     log.info("Download complete")
 
 def _setup_cuda_paths():
+    nvidia_pkgs = ["nvidia", "nvidia.cuda_runtime", "nvidia.cublas"]
     candidates = set()
+    for mod_name in nvidia_pkgs:
+        try:
+            mod = __import__(mod_name, fromlist=["__path__"])
+            pkg_path = mod.__path__[0]
+            candidates.add(os.path.join(os.path.dirname(pkg_path), mod_name.split(".")[-1], "lib"))
+        except Exception:
+            pass
     for p in sys.path:
         candidates.add(os.path.join(p, "nvidia", "cuda_runtime", "lib"))
+        candidates.add(os.path.join(p, "nvidia", "cublas", "lib"))
     candidates.add("/usr/local/lib/python3.12/site-packages/nvidia/cuda_runtime/lib")
+    candidates.add("/usr/local/lib/python3.12/site-packages/nvidia/cublas/lib")
     current = os.environ.get("LD_LIBRARY_PATH", "")
+    merged = []
     for p in candidates:
-        if os.path.exists(os.path.join(p, "libcudart.so.12")) and p not in current:
-            os.environ["LD_LIBRARY_PATH"] = p + ":" + current
-            log.info(f"CUDA libs found at {p}")
-            return
+        if os.path.exists(p) and p not in merged:
+            merged.append(p)
+    for p in merged:
+        os.environ["LD_LIBRARY_PATH"] = p + ":" + os.environ.get("LD_LIBRARY_PATH", "")
+    log.info(f"CUDA lib paths: {merged}")
 
 @spaces.GPU
 def cuda_test():