58 lines
2.4 KiB
Python
58 lines
2.4 KiB
Python
# ComfyUI/custom_nodes/ComfyUI_LocalLLMNodes/__init__.py
|
|
# --- Conditional Import Logic ---
|
|
LOCAL_LLM_NODES_AVAILABLE = False
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
# --- Check for core dependencies ---
|
|
# We need transformers and torch for the LLM nodes to work.
|
|
try:
|
|
import transformers
|
|
import torch
|
|
# Optionally check for bitsandbytes here if you plan to use quantization immediately
|
|
# import bitsandbytes as bnb
|
|
CORE_DEPS_AVAILABLE = True
|
|
except ImportError as e:
|
|
print(f"[ComfyUI_LocalLLMNodes] Core dependencies (transformers, torch) not found or not importable: {e}")
|
|
CORE_DEPS_AVAILABLE = False
|
|
|
|
if CORE_DEPS_AVAILABLE:
|
|
try:
|
|
# --- Import Node Classes ---
|
|
from .local_llm_connector import SetLocalLLMServiceConnector
|
|
from .local_prompt_generator import (
|
|
LocalKontextPromptGenerator,
|
|
AddUserLocalKontextPreset,
|
|
RemoveUserLocalKontextPreset
|
|
)
|
|
|
|
# --- Define Mappings ---
|
|
# Using the class names directly as keys is standard.
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SetLocalLLMServiceConnector": SetLocalLLMServiceConnector,
|
|
"LocalKontextPromptGenerator": LocalKontextPromptGenerator,
|
|
"AddUserLocalKontextPreset": AddUserLocalKontextPreset,
|
|
"RemoveUserLocalKontextPreset": RemoveUserLocalKontextPreset,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SetLocalLLMServiceConnector": "Set Local LLM Service Connector 🐑",
|
|
"LocalKontextPromptGenerator": "Local Kontext Prompt Generator 🐑",
|
|
"AddUserLocalKontextPreset": "Add User Local Kontext Preset 🐑",
|
|
"RemoveUserLocalKontextPreset": "Remove User Local Kontext Preset 🐑",
|
|
}
|
|
|
|
LOCAL_LLM_NODES_AVAILABLE = True
|
|
print("[ComfyUI_LocalLLMNodes] All nodes are available.")
|
|
|
|
except Exception as e:
|
|
print(f"[ComfyUI_LocalLLMNodes] Error importing node classes: {e}")
|
|
# If import fails, NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS remain empty,
|
|
# and ComfyUI won't register any nodes from this package.
|
|
else:
|
|
print("[ComfyUI_LocalLLMNodes] Nodes will NOT be available due to missing core dependencies.")
|
|
|
|
# --- Define what ComfyUI sees ---
|
|
# ComfyUI looks for these specific dictionaries in __init__.py
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|