137 lines
4.3 KiB
Python
137 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Auto-install the correct llama-cpp-python wheel for the current system.
|
|
|
|
Detects CUDA / Metal / CPU and installs from the abetlen prebuilt wheel index.
|
|
Run with the target Python interpreter, e.g.:
|
|
|
|
/path/to/ComfyUI/venv/bin/python install_llama_cpp.py
|
|
|
|
Use --dry-run to see what would be installed without running pip.
|
|
"""
|
|
|
|
import argparse
|
|
import platform
|
|
import subprocess
|
|
import sys
|
|
|
|
|
|
def detect_cuda_version() -> str | None:
|
|
"""Return the installed CUDA runtime version as '12.8', or None."""
|
|
try:
|
|
import torch
|
|
if torch.cuda.is_available() and torch.version.cuda:
|
|
return torch.version.cuda
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
["nvcc", "--version"],
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
check=False,
|
|
)
|
|
if result.returncode == 0:
|
|
for line in result.stdout.splitlines():
|
|
if "release" in line.lower():
|
|
parts = line.split("release")
|
|
ver = parts[-1].strip().split(",")[0].strip()
|
|
return ver
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
return None
|
|
|
|
|
|
def _cuda_wheel_tag(cuda_version: str) -> str:
|
|
"""Map a CUDA version string to the closest supported wheel tag.
|
|
|
|
llama-cpp-python provides wheels for cu122, cu123, cu124, cu125,
|
|
cu126, cu128, cu130, cu132. We pick the exact tag if it exists,
|
|
otherwise the next higher tag. If the version is below the lowest
|
|
supported tag we fall back to CPU to avoid runtime incompatibility.
|
|
"""
|
|
ver = cuda_version.split(".")
|
|
try:
|
|
major = int(ver[0])
|
|
minor = int(ver[1].split("_")[0]) if len(ver) > 1 else 0
|
|
except Exception:
|
|
return "cpu"
|
|
|
|
if major < 12:
|
|
return "cpu"
|
|
|
|
installed = major * 100 + minor
|
|
tags = [122, 123, 124, 125, 126, 128, 130, 132]
|
|
|
|
# Exact match first, then the smallest tag that is >= installed.
|
|
for tag in tags:
|
|
if tag >= installed:
|
|
return f"cu{tag}"
|
|
|
|
# Newer than all available tags: use the newest (forward compatibility).
|
|
return f"cu{tags[-1]}"
|
|
|
|
|
|
def detect_backend() -> tuple[str, str]:
|
|
"""Return (backend_tag, reason) for the wheel index."""
|
|
system = platform.system().lower()
|
|
machine = platform.machine().lower()
|
|
|
|
if system == "darwin" and machine in ("arm64", "aarch64"):
|
|
return "metal", "macOS Apple Silicon"
|
|
|
|
cuda_version = detect_cuda_version()
|
|
if cuda_version:
|
|
tag = _cuda_wheel_tag(cuda_version)
|
|
return tag, f"CUDA {cuda_version}"
|
|
|
|
return "cpu", "no CUDA detected"
|
|
|
|
|
|
def build_pip_command(backend: str, llama_cpp_version: str | None = None) -> list[str]:
|
|
index_url = f"https://abetlen.github.io/llama-cpp-python/whl/{backend}"
|
|
cmd = [sys.executable, "-m", "pip", "install", "--upgrade", "--force-reinstall"]
|
|
if backend != "cpu":
|
|
cmd += ["--no-cache-dir"]
|
|
if llama_cpp_version:
|
|
spec = f"llama-cpp-python=={llama_cpp_version}"
|
|
else:
|
|
spec = "llama-cpp-python"
|
|
cmd += [spec, "--extra-index-url", index_url]
|
|
return cmd
|
|
|
|
|
|
def format_install_command(llama_cpp_version: str | None = None) -> str:
|
|
"""Return a shell command string the user can run to install the right wheel."""
|
|
backend, reason = detect_backend()
|
|
cmd = build_pip_command(backend, llama_cpp_version)
|
|
return " ".join(cmd)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Install llama-cpp-python with the right backend")
|
|
parser.add_argument("--dry-run", action="store_true", help="Print command without running it")
|
|
parser.add_argument("--backend", type=str, default=None, help="Override backend (e.g. cu128, cpu, metal)")
|
|
parser.add_argument("--version", type=str, default=None, help="Pin llama-cpp-python version")
|
|
args = parser.parse_args()
|
|
|
|
backend = args.backend or detect_backend()[0]
|
|
reason = "user override" if args.backend else detect_backend()[1]
|
|
|
|
print(f"[LD-PromptNodes] Detected backend: {backend} ({reason})")
|
|
cmd = build_pip_command(backend, args.version)
|
|
print(f"[LD-PromptNodes] Running:\n {' '.join(cmd)}")
|
|
|
|
if args.dry_run:
|
|
print("[LD-PromptNodes] Dry run - not executing.")
|
|
return 0
|
|
|
|
result = subprocess.run(cmd)
|
|
return result.returncode
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|