Upload app.py with huggingface_hub

app.py

@@ -1,17 +1,11 @@
 import os
+import sys
 import logging
 import traceback
 import spaces
 import gradio as gr
 from huggingface_hub import hf_hub_download
 
-# Add PyTorch's bundled CUDA libs to library path BEFORE importing llama_cpp
-import torch
-_torch_lib = os.path.join(os.path.dirname(torch.__file__), "lib")
-_env_ld = os.environ.get("LD_LIBRARY_PATH", "")
-if _torch_lib not in _env_ld:
-    os.environ["LD_LIBRARY_PATH"] = _torch_lib + ":" + _env_ld
-
 logging.basicConfig(level=logging.INFO)
 log = logging.getLogger(__name__)
 
@@ -28,8 +22,28 @@ if not os.path.exists(MODEL_PATH):
     )
     log.info("Download complete")
 
+def _setup_cuda_paths():
+    paths = [
+        "/usr/local/lib/python3.12/site-packages/nvidia/cuda_runtime/lib",
+    ]
+    # Try to find nvidia packages dynamically
+    for mod_name in ["nvidia.cuda_runtime"]:
+        try:
+            mod = __import__(mod_name, fromlist=["__file__"])
+            lib_dir = os.path.join(os.path.dirname(mod.__file__), "lib")
+            if os.path.exists(lib_dir):
+                paths.append(lib_dir)
+        except ImportError:
+            pass
+    current = os.environ.get("LD_LIBRARY_PATH", "")
+    for p in paths:
+        if os.path.exists(p) and p not in current:
+            os.environ["LD_LIBRARY_PATH"] = p + ":" + current
+            log.info(f"Added {p} to LD_LIBRARY_PATH")
+
 @spaces.GPU
 def cuda_test():
+    import torch
     return {
         "cuda_available": torch.cuda.is_available(),
         "device_count": torch.cuda.device_count(),
@@ -39,6 +53,8 @@ def cuda_test():
 @spaces.GPU
 def generate(messages, max_tokens=1024, temperature=0.7, top_p=0.95):
     try:
+        _setup_cuda_paths()
+        log.info(f"LD_LIBRARY_PATH={os.environ.get('LD_LIBRARY_PATH', '')[:200]}")
         log.info("Importing llama_cpp...")
         from llama_cpp import Llama as _Llama
         log.info("Import OK, loading model...")