106 lines
3.3 KiB
Python
106 lines
3.3 KiB
Python
"""
|
|
LD ComfyUI Prompt Nodes
|
|
Native ComfyUI nodes for SlopPrompt and Booru tag prompt generation.
|
|
"""
|
|
|
|
import importlib.util
|
|
import os
|
|
import sys
|
|
|
|
NODE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
sys.path.insert(0, NODE_DIR)
|
|
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
WEB_DIRECTORY = "./web"
|
|
|
|
|
|
def _maybe_auto_install_llama_cpp():
|
|
"""Install a CUDA-enabled llama-cpp-python wheel on first start if needed.
|
|
|
|
Only runs when:
|
|
- torch is available and CUDA is present
|
|
- llama-cpp-python is importable but reports no GPU offload support
|
|
- llama-cpp-python is not installed at all
|
|
"""
|
|
try:
|
|
import torch
|
|
if not torch.cuda.is_available():
|
|
return
|
|
except Exception:
|
|
return
|
|
|
|
has_llama = False
|
|
supports_offload = False
|
|
try:
|
|
import llama_cpp
|
|
has_llama = True
|
|
supports_offload = bool(llama_cpp.llama_supports_gpu_offload())
|
|
except Exception:
|
|
pass
|
|
|
|
if has_llama and supports_offload:
|
|
return
|
|
|
|
try:
|
|
from install_llama_cpp import build_pip_command, detect_backend
|
|
except Exception as e:
|
|
print(f"[LD-PromptNodes] Auto-install helper unavailable: {e}")
|
|
return
|
|
|
|
backend, reason = detect_backend()
|
|
if backend == "cpu":
|
|
print("[LD-PromptNodes] No CUDA/Metal detected; skipping llama-cpp-python auto-install.")
|
|
return
|
|
|
|
cmd = build_pip_command(backend)
|
|
print(f"[LD-PromptNodes] Auto-installing llama-cpp-python ({backend}, {reason})...")
|
|
print(f"[LD-PromptNodes] Command: {' '.join(cmd)}")
|
|
try:
|
|
import subprocess
|
|
result = subprocess.run(cmd)
|
|
if result.returncode == 0:
|
|
print("[LD-PromptNodes] llama-cpp-python auto-install completed. Restart ComfyUI to use it.")
|
|
else:
|
|
print(f"[LD-PromptNodes] llama-cpp-python auto-install failed with exit code {result.returncode}.")
|
|
except Exception as e:
|
|
print(f"[LD-PromptNodes] llama-cpp-python auto-install error: {e}")
|
|
|
|
|
|
_maybe_auto_install_llama_cpp()
|
|
|
|
|
|
def _load_module(file_path: str):
|
|
module_name = os.path.basename(file_path)[:-3]
|
|
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[module_name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def _load_directory(directory: str, skip_names: set[str] = None):
|
|
if not os.path.isdir(directory):
|
|
return
|
|
skip_names = skip_names or set()
|
|
for file in sorted(os.listdir(directory)):
|
|
if not file.endswith(".py") or file.startswith("_"):
|
|
continue
|
|
module_name = file[:-3]
|
|
if module_name in skip_names:
|
|
continue
|
|
file_path = os.path.join(directory, file)
|
|
try:
|
|
module = _load_module(file_path)
|
|
if hasattr(module, "NODE_CLASS_MAPPINGS"):
|
|
NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS)
|
|
if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS"):
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
|
|
except Exception as e:
|
|
print(f"[LD-PromptNodes] Error loading {file}: {e}")
|
|
|
|
|
|
_load_directory(NODE_DIR, skip_names={"download", os.path.basename(__file__)[:-3]})
|
|
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|