125 lines
4.0 KiB
Python
125 lines
4.0 KiB
Python
import logging
|
|
import sys
|
|
import os
|
|
from typing import Dict, Tuple
|
|
# Copied and modified from NeuroSenko/ComfyUI_LLM_SDXL_Adapter
|
|
# Setup logging
|
|
logger = logging.getLogger("LLM-SDXL-Adapter-Additions")
|
|
logger.setLevel(logging.WARN)
|
|
|
|
# Add custom formatter with module prefix
|
|
if not logger.handlers:
|
|
handler = logging.StreamHandler()
|
|
formatter = logging.Formatter('[LLM-SDXL-Adapter-Additions] %(levelname)s: %(message)s')
|
|
handler.setFormatter(formatter)
|
|
logger.addHandler(handler)
|
|
logger.propagate = False # Prevent duplicate logs from parent loggers
|
|
|
|
# Add current directory to path for imports
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
if current_dir not in sys.path:
|
|
sys.path.append(current_dir)
|
|
|
|
# Check dependencies
|
|
try:
|
|
import torch
|
|
import transformers
|
|
import safetensors
|
|
import einops
|
|
logger.info("All required dependencies found")
|
|
except ImportError as e:
|
|
logger.error(f"Missing dependency: {e}")
|
|
logger
|
|
# Import all node modules from separate files
|
|
try:
|
|
from .llm_adapter_loader_explicit import NODE_CLASS_MAPPINGS as LLM_ADAPTER_LOADER_UNFUSED_MAPPINGS
|
|
from .llm_adapter_loader_explicit import NODE_DISPLAY_NAME_MAPPINGS as LLM_ADAPTER_LOADER_UNFUSED_DISPLAY_MAPPINGS
|
|
|
|
from .t5gemma_text_encode_v2 import NODE_CLASS_MAPPINGS as T5_GEMMA_TEXT_ENCODE_ADDITIONAL_MAPPINGS
|
|
from .t5gemma_text_encode_v2 import NODE_DISPLAY_NAME_MAPPINGS as T5_GEMMA_TEXT_ENCODE_ADDITIONAL_DISPLAY_MAPPINGS
|
|
|
|
|
|
logger.info("Successfully imported all node modules from separate files")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to import node modules: {e}")
|
|
raise
|
|
|
|
# Combine all node mappings
|
|
NODE_CLASS_MAPPINGS: Dict[str, type] = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {}
|
|
|
|
# Add all mappings from separate files
|
|
all_class_mappings = [
|
|
LLM_ADAPTER_LOADER_UNFUSED_MAPPINGS,
|
|
T5_GEMMA_TEXT_ENCODE_ADDITIONAL_MAPPINGS,
|
|
]
|
|
|
|
all_display_mappings = [
|
|
LLM_ADAPTER_LOADER_UNFUSED_DISPLAY_MAPPINGS,
|
|
T5_GEMMA_TEXT_ENCODE_ADDITIONAL_DISPLAY_MAPPINGS,
|
|
]
|
|
|
|
for mapping in all_class_mappings:
|
|
NODE_CLASS_MAPPINGS.update(mapping)
|
|
|
|
for mapping in all_display_mappings:
|
|
NODE_DISPLAY_NAME_MAPPINGS.update(mapping)
|
|
|
|
# Version and metadata
|
|
__version__ = "1.0.0"
|
|
__author__ = "Remix"
|
|
__description__ = "Additional ComfyUI nodes for LLM to SDXL adapter workflow"
|
|
|
|
# Export information for ComfyUI
|
|
WEB_DIRECTORY = "./web" # For any web UI components (if needed)
|
|
|
|
# Log successful initialization
|
|
logger.info(f"LLM SDXL Adapter Additions v{__version__} initialized successfully")
|
|
logger.info(f"Registered {len(NODE_CLASS_MAPPINGS)} nodes from separate files:")
|
|
for node_name in sorted(NODE_CLASS_MAPPINGS.keys()):
|
|
display_name = NODE_DISPLAY_NAME_MAPPINGS.get(node_name, node_name)
|
|
logger.info(f" - {node_name} ({display_name})")
|
|
|
|
# Custom type definitions for ComfyUI
|
|
CUSTOM_TYPES = {
|
|
"LLM_MODEL": "Language Model instance",
|
|
"LLM_TOKENIZER": "Language Model tokenizer instance",
|
|
"LLM_ADAPTER": "Adapter model instance",
|
|
}
|
|
|
|
def get_node_info() -> Dict[str, any]:
|
|
"""
|
|
Return information about available nodes for debugging/documentation
|
|
"""
|
|
return {
|
|
"version": __version__,
|
|
"author": __author__,
|
|
"description": __description__,
|
|
"nodes": {
|
|
name: {
|
|
"display_name": NODE_DISPLAY_NAME_MAPPINGS.get(name, name),
|
|
"class": cls.__name__,
|
|
"category": getattr(cls, "CATEGORY", "unknown") if hasattr(cls, "CATEGORY") else "unknown",
|
|
"function": getattr(cls, "FUNCTION", "unknown") if hasattr(cls, "FUNCTION") else "unknown"
|
|
}
|
|
for name, cls in NODE_CLASS_MAPPINGS.items()
|
|
},
|
|
"custom_types": CUSTOM_TYPES
|
|
}
|
|
|
|
# Optional: Setup hook for ComfyUI initialization
|
|
def setup_js():
|
|
"""
|
|
Setup any JavaScript/web components if needed
|
|
"""
|
|
pass
|
|
|
|
# Export what ComfyUI expects
|
|
__all__ = [
|
|
"NODE_CLASS_MAPPINGS",
|
|
"NODE_DISPLAY_NAME_MAPPINGS",
|
|
"WEB_DIRECTORY",
|
|
"get_node_info"
|
|
]
|
|
# Above from NeuroSenko/ComfyUI_LLM_SDXL_Adapter |