Spaces:
Running
Running
| """Make the CUDA build of llama-cpp-python load on Windows. | |
| The CUDA wheel ships `ggml-cuda.dll`, which depends on the CUDA runtime | |
| (`cudart64_12.dll`) and cuBLAS (`cublas64_12.dll`, `cublasLt64_12.dll`). Those | |
| are NOT bundled; we install them as pip wheels (`nvidia-cuda-runtime-cu12`, | |
| `nvidia-cublas-cu12`), which drop the DLLs under `site-packages/nvidia/*/bin`. | |
| llama-cpp-python loads its shared library with `winmode=RTLD_GLOBAL`, and under | |
| that flag Windows ignores both `PATH` and `os.add_dll_directory` when resolving | |
| dependencies — so simply registering the nvidia dirs isn't enough. The reliable | |
| fix is to *pre-load* the DLLs ourselves with the default loader (which does | |
| honor `add_dll_directory`) in dependency order. Once a module is mapped into the | |
| process, llama-cpp-python's later load reuses the in-memory copy instead of | |
| searching the disk again. | |
| No-op on non-Windows and when the CUDA/nvidia wheels aren't present (CPU build). | |
| """ | |
| import ctypes | |
| import glob | |
| import os | |
| import sys | |
| _done = False | |
| # llama.dll last — it depends on all the others. ggml-cuda.dll is the one that | |
| # pulls in the CUDA runtime; the rest are loaded so their inter-deps resolve. | |
| _PRELOAD_ORDER = [ | |
| "ggml-base.dll", | |
| "ggml.dll", | |
| "ggml-cpu.dll", | |
| "ggml-cuda.dll", | |
| "llama.dll", | |
| ] | |
| def _site_packages_dirs(): | |
| seen = set() | |
| for p in sys.path: | |
| if p.endswith("site-packages") and p not in seen: | |
| seen.add(p) | |
| yield p | |
| # Fallback: the repo-local venv, in case this runs outside the venv. | |
| local = os.path.join(os.path.dirname(__file__), ".venv", "Lib", "site-packages") | |
| if os.path.isdir(local) and local not in seen: | |
| yield local | |
| def ensure() -> None: | |
| global _done | |
| if _done: | |
| return | |
| _done = True | |
| if sys.platform != "win32": | |
| return | |
| lib_dir = None | |
| for sp in _site_packages_dirs(): | |
| nvidia_root = os.path.join(sp, "nvidia") | |
| if os.path.isdir(nvidia_root): | |
| for bin_dir in glob.glob(os.path.join(nvidia_root, "*", "bin")): | |
| try: | |
| os.add_dll_directory(bin_dir) | |
| except OSError: | |
| pass | |
| candidate = os.path.join(sp, "llama_cpp", "lib") | |
| if os.path.isdir(candidate): | |
| lib_dir = candidate | |
| if not lib_dir or not os.path.exists(os.path.join(lib_dir, "ggml-cuda.dll")): | |
| return # CPU build (no CUDA dll) — nothing to pre-load | |
| try: | |
| os.add_dll_directory(lib_dir) | |
| except OSError: | |
| pass | |
| for name in _PRELOAD_ORDER: | |
| path = os.path.join(lib_dir, name) | |
| if os.path.exists(path): | |
| try: | |
| ctypes.CDLL(path) | |
| except OSError: | |
| # Leave it to llama-cpp-python to surface a precise error. | |
| pass | |