init
This commit is contained in:
NeuroSenko
@@ -0,0 +1 @@
|
||||
__pycache__/
|
||||
@@ -0,0 +1,165 @@
|
||||
# ComfyUI LLM SDXL Adapter
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
A comprehensive set of ComfyUI nodes for using Large Language Models (LLM) as text encoders for SDXL image generation through a trained adapter.
|
||||
|
||||
## 🚀 Key Features
|
||||
|
||||
- **Multi-LLM Support**: Gemma and other compatible models (extensible architecture)
|
||||
- **Flexible Architecture**: Modular nodes for customizable workflows
|
||||
- **Transformer Adapter**: Convert LLM embeddings to SDXL format
|
||||
- **Chat Templates**: Support for system prompts and conversation formatting
|
||||
- **Memory Optimization**: Model loading management and caching
|
||||
- **Configurable Parameters**: Flexible adapter architecture configuration
|
||||
|
||||
## 📦 Installation
|
||||
|
||||
### Requirements
|
||||
- Python 3.8+
|
||||
- ComfyUI
|
||||
- CUDA (recommended)
|
||||
- Latest transformers library (tested on 4.53.1)
|
||||
|
||||
### Install Dependencies
|
||||
```bash
|
||||
pip install transformers>=4.53.1 safetensors einops torch
|
||||
```
|
||||
|
||||
### Install Nodes
|
||||
1. Clone the repository to `ComfyUI/custom_nodes/`:
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes/
|
||||
git clone https://github.com/NeuroSenko/comfyui_llm_sdxl_adapter.git
|
||||
```
|
||||
|
||||
2. Restart ComfyUI
|
||||
|
||||
## 🗂️ Project Structure
|
||||
|
||||
```
|
||||
comfyui_llm_sdxl_adapter/
|
||||
├── __init__.py # Entry point and node registration
|
||||
├── llm_model_loader.py # LLM model loader
|
||||
├── llm_text_encoder.py # LLM text encoder
|
||||
├── lll_adapter_loader.py # Adapter loader
|
||||
├── llm_to_sdxl_adapter.py # Adapter architecture
|
||||
├── apply_llm_to_sdxl_adapter.py # Adapter application
|
||||
├── utils.py # Utility functions
|
||||
└── README.md # Documentation
|
||||
```
|
||||
|
||||
## 🔧 ComfyUI Nodes
|
||||
|
||||
### LLM Model Loader
|
||||
Loads language model and tokenizer.
|
||||
|
||||
**Inputs:**
|
||||
- `model_path`: Path to LLM model
|
||||
- `device`: Device (auto/cuda:0/cuda:1/cpu)
|
||||
- `force_reload`: Force reload model
|
||||
|
||||
**Outputs:**
|
||||
- `model`: LLM model instance
|
||||
- `tokenizer`: Tokenizer
|
||||
- `info`: Model information
|
||||
|
||||
### LLM Text Encoder
|
||||
Encodes text using loaded LLM.
|
||||
|
||||
**Inputs:**
|
||||
- `model`: LLM model
|
||||
- `tokenizer`: Tokenizer
|
||||
- `text`: Text to encode
|
||||
- `system_prompt`: System prompt (optional)
|
||||
- `skip_first`: Number of tokens to skip
|
||||
|
||||
**Outputs:**
|
||||
- `hidden_states`: LLM hidden states
|
||||
- `info`: Encoding information
|
||||
|
||||
### LLM Adapter Loader
|
||||
Loads trained LLM -> SDXL adapter.
|
||||
|
||||
**Inputs:**
|
||||
- `adapter_path`: Path to adapter file (.safetensors)
|
||||
- `llm_dim`: LLM dimension (default 1152)
|
||||
- `sdxl_seq_dim`: SDXL sequence dimension (2048)
|
||||
- `sdxl_pooled_dim`: SDXL pooled output dimension (1280)
|
||||
- `target_seq_len`: Target sequence length (308)
|
||||
- `n_wide_blocks`: Number of wide blocks (2)
|
||||
- `n_narrow_blocks`: Number of narrow blocks (3)
|
||||
|
||||
**Outputs:**
|
||||
- `adapter`: Adapter instance
|
||||
- `info`: Adapter information
|
||||
|
||||
### Apply LLM To SDXL Adapter
|
||||
Applies adapter to LLM hidden states.
|
||||
|
||||
**Inputs:**
|
||||
- `llm_hidden_states`: LLM hidden states
|
||||
- `adapter`: Loaded adapter
|
||||
|
||||
**Outputs:**
|
||||
- `conditioning`: SDXL conditioning
|
||||
- `info`: Application information
|
||||
|
||||
## 🎯 Basic Workflow
|
||||
|
||||
```
|
||||
LLM Model Loader → LLM Text Encoder → LLM Adapter Loader → Apply LLM To SDXL Adapter → KSampler
|
||||
```
|
||||
|
||||
### Usage Example:
|
||||
|
||||
1. **Load Model**: Use `LLM Model Loader` to load an LLM (e.g., Gemma-3-1b-it)
|
||||
2. **Encode Text**: Connect model to `LLM Text Encoder` and input your prompt
|
||||
3. **Load Adapter**: Use `LLM Adapter Loader` to load trained adapter
|
||||
4. **Apply Adapter**: Connect hidden states and adapter through `Apply LLM To SDXL Adapter`
|
||||
5. **Generate**: Connect the resulting conditioning to standard SDXL pipeline
|
||||
|
||||
## 📁 File Structure
|
||||
|
||||
### LLM Models
|
||||
Place LLM models in the folder:
|
||||
```
|
||||
ComfyUI/models/LLM/
|
||||
└── gemma-3-1b-it/
|
||||
```
|
||||
|
||||
### Adapters
|
||||
Place trained adapters in the folder:
|
||||
```
|
||||
ComfyUI/models/llm_adapters/
|
||||
└── gemma_to_sdxl_adapter.safetensors
|
||||
```
|
||||
|
||||
## ⚙️ Adapter Architecture
|
||||
|
||||
The `LLMToSDXLAdapter` uses a transformer architecture with two stages:
|
||||
|
||||
1. **Wide Processing**: Process full sequence (512 tokens)
|
||||
2. **Compression**: Compress to target length (308 tokens) via cross-attention
|
||||
3. **Narrow Processing**: Final processing of compressed sequence
|
||||
4. **Pooling**: Create pooled output for SDXL
|
||||
|
||||
### Key Features:
|
||||
- Positional embeddings for input and output sequences
|
||||
- Learnable compression queries for controlled compression
|
||||
- Attention masking for proper padding handling
|
||||
- Dropout and LayerNorm for training stability
|
||||
|
||||
## 🔍 Debugging
|
||||
|
||||
To enable detailed logging:
|
||||
|
||||
```python
|
||||
import logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
```
|
||||
|
||||
Information about loaded nodes is displayed during ComfyUI initialization.
|
||||
+169
@@ -0,0 +1,169 @@
|
||||
"""
|
||||
ComfyUI LLM SDXL Adapter v1.0.0
|
||||
|
||||
Transform any Language Model into a powerful text encoder for SDXL image generation.
|
||||
This plugin provides ability to use trained neural adapter that bridges LLM embeddings to SDXL's
|
||||
conditioning format, enabling richer and more nuanced prompt understanding.
|
||||
|
||||
🧠 Core Components:
|
||||
- LLMModelLoader: Load and manage Language Models (Gemma, extensible to others)
|
||||
- LLMTextEncoder: Process text prompts through LLM with chat templates
|
||||
- AdapterLoader: Load trained LLM→SDXL transformation adapters
|
||||
- ApplyLLMToSDXLAdapter: Convert LLM embeddings to SDXL conditioning
|
||||
|
||||
🔧 Key Features:
|
||||
- Supports Gemma models with extensible architecture
|
||||
- Advanced transformer-based adapter with compression stages
|
||||
- Memory-efficient model loading and caching
|
||||
- Configurable chat templates and system prompts
|
||||
- Full integration with ComfyUI's SDXL pipeline
|
||||
|
||||
⚡ Quick Start:
|
||||
1. Install: pip install transformers>=4.53.1 safetensors einops
|
||||
2. Place models in ComfyUI/models/LLM/
|
||||
3. Place adapters in ComfyUI/models/llm_adapters/
|
||||
4. Workflow: LLMModelLoader → LLMTextEncoder → LLMAdapterLoader → ApplyLLMToSDXLAdapter → KSampler
|
||||
|
||||
Author: NeuroSenko | License: MIT | Tested with transformers 4.53.1
|
||||
"""
|
||||
|
||||
import logging
|
||||
import sys
|
||||
import os
|
||||
from typing import Dict, Tuple
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 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.error("Please install: pip install transformers safetensors einops")
|
||||
raise
|
||||
|
||||
# Import all node modules from separate files
|
||||
try:
|
||||
from .llm_model_loader import NODE_CLASS_MAPPINGS as LLM_MODEL_LOADER_MAPPINGS
|
||||
from .llm_model_loader import NODE_DISPLAY_NAME_MAPPINGS as LLM_MODEL_LOADER_DISPLAY_MAPPINGS
|
||||
|
||||
from .llm_text_encoder import NODE_CLASS_MAPPINGS as LLM_ENCODER_MAPPINGS
|
||||
from .llm_text_encoder import NODE_DISPLAY_NAME_MAPPINGS as LLM_ENCODER_DISPLAY_MAPPINGS
|
||||
|
||||
from .llm_adapter_loader import NODE_CLASS_MAPPINGS as ADAPTER_LOADER_MAPPINGS
|
||||
from .llm_adapter_loader import NODE_DISPLAY_NAME_MAPPINGS as ADAPTER_LOADER_DISPLAY_MAPPINGS
|
||||
|
||||
from .apply_llm_to_sdxl_adapter import NODE_CLASS_MAPPINGS as ADAPTER_NODE_MAPPINGS
|
||||
from .apply_llm_to_sdxl_adapter import NODE_DISPLAY_NAME_MAPPINGS as ADAPTER_NODE_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_MODEL_LOADER_MAPPINGS,
|
||||
LLM_ENCODER_MAPPINGS,
|
||||
ADAPTER_LOADER_MAPPINGS,
|
||||
ADAPTER_NODE_MAPPINGS,
|
||||
]
|
||||
|
||||
all_display_mappings = [
|
||||
LLM_MODEL_LOADER_DISPLAY_MAPPINGS,
|
||||
LLM_ENCODER_DISPLAY_MAPPINGS,
|
||||
ADAPTER_LOADER_DISPLAY_MAPPINGS,
|
||||
ADAPTER_NODE_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__ = "NeuroSenko"
|
||||
__description__ = "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 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_HIDDEN_STATES": "LLM model hidden states",
|
||||
"LLM_ADAPTER": "Adapter model instance",
|
||||
"VECTOR_CONDITIONING": "SDXL vector conditioning",
|
||||
|
||||
}
|
||||
|
||||
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"
|
||||
]
|
||||
|
||||
# Print initialization message
|
||||
print(f"\n{'='*60}")
|
||||
print(f" ComfyUI LLM SDXL Adapter v{__version__} - Loaded Successfully")
|
||||
print(f"{'='*60}")
|
||||
print(f" Available Nodes: {len(NODE_CLASS_MAPPINGS)}")
|
||||
print(f" Main Workflow: LLMModelLoader -> LLMTextEncoder -> LLMAdapterLoader -> ApplyLLMToSDXLAdapter -> KSampler")
|
||||
print(f" Supports: Gemma, Llama, Mistral, and other compatible LLMs")
|
||||
print(f" Quick Start: Use modular nodes for flexible workflows")
|
||||
print(f"{'='*60}\n")
|
||||
@@ -0,0 +1,61 @@
|
||||
import torch
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ApplyLLMToSDXLAdapter:
|
||||
"""
|
||||
ComfyUI node that applies loaded LLM to SDXL adapter transformation
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"llm_hidden_states": ("LLM_HIDDEN_STATES",),
|
||||
"llm_adapter": ("LLM_ADAPTER",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING", "STRING")
|
||||
RETURN_NAMES = ("conditioning", "info")
|
||||
FUNCTION = "apply_adapter"
|
||||
CATEGORY = "llm_sdxl"
|
||||
|
||||
def apply_adapter(self, llm_hidden_states, llm_adapter):
|
||||
"""Apply the LLM to SDXL adapter transformation"""
|
||||
try:
|
||||
# Apply adapter
|
||||
with torch.no_grad():
|
||||
conditioning, pooled_output = llm_adapter(llm_hidden_states)
|
||||
|
||||
# Format conditioning for ComfyUI
|
||||
# ComfyUI expects conditioning as a list of [cond_tensor, metadata_dict] tuples
|
||||
comfy_conditioning = [[conditioning, {"pooled_output": pooled_output}]]
|
||||
|
||||
# Prepare info
|
||||
info = f"Conditioning shape: {conditioning.shape}"
|
||||
|
||||
logger.info(f"Applied LLM to SDXL adapter: {info}")
|
||||
|
||||
return (comfy_conditioning, info)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to apply adapter: {str(e)}")
|
||||
raise Exception(f"Adapter application failed: {str(e)}")
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, llm_hidden_states, llm_adapter):
|
||||
"""ComfyUI change detection"""
|
||||
return hash(str(llm_hidden_states.shape) + str(id(llm_adapter)))
|
||||
|
||||
|
||||
# Node mapping for ComfyUI registration
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ApplyLLMToSDXLAdapter": ApplyLLMToSDXLAdapter
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ApplyLLMToSDXLAdapter": "Apply LLM To SDXL Adapter"
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
import torch
|
||||
from safetensors.torch import load_file
|
||||
import logging
|
||||
import gc
|
||||
import os
|
||||
from .utils import get_llm_adapters
|
||||
from .llm_to_sdxl_adapter import LLMToSDXLAdapter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LLMAdapterLoader:
|
||||
"""
|
||||
ComfyUI node that loads LLM to SDXL adapter
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.adapter = None
|
||||
self.current_adapter_path = None
|
||||
self.device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"adapter_path": (get_llm_adapters(), {
|
||||
"default": get_llm_adapters()[0] if get_llm_adapters() else "checkpoint_epoch_10.safetensors"
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"llm_dim": ("INT", {
|
||||
"default": 1152,
|
||||
"min": 512,
|
||||
"max": 4096
|
||||
}),
|
||||
"sdxl_seq_dim": ("INT", {
|
||||
"default": 2048,
|
||||
"min": 1024,
|
||||
"max": 4096
|
||||
}),
|
||||
"sdxl_pooled_dim": ("INT", {
|
||||
"default": 1280,
|
||||
"min": 512,
|
||||
"max": 2048
|
||||
}),
|
||||
"target_seq_len": ("INT", {
|
||||
"default": 308,
|
||||
"min": 64,
|
||||
"max": 1024
|
||||
}),
|
||||
"n_wide_blocks": ("INT", {
|
||||
"default": 2,
|
||||
"min": 1,
|
||||
"max": 8
|
||||
}),
|
||||
"n_narrow_blocks": ("INT", {
|
||||
"default": 3,
|
||||
"min": 1,
|
||||
"max": 8
|
||||
}),
|
||||
"num_heads": ("INT", {
|
||||
"default": 16,
|
||||
"min": 4,
|
||||
"max": 32
|
||||
}),
|
||||
"dropout": ("FLOAT", {
|
||||
"default": 0.1,
|
||||
"min": 0.0,
|
||||
"max": 0.5,
|
||||
"step": 0.01
|
||||
}),
|
||||
"device": (["auto", "cuda:0", "cuda:1", "cpu"], {
|
||||
"default": "auto"
|
||||
}),
|
||||
"force_reload": ("BOOLEAN", {
|
||||
"default": False
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LLM_ADAPTER", "STRING")
|
||||
RETURN_NAMES = ("llm_adapter", "info")
|
||||
FUNCTION = "load_adapter"
|
||||
CATEGORY = "llm_sdxl"
|
||||
|
||||
def load_adapter(self, adapter_path, llm_dim=1152, sdxl_seq_dim=2048, sdxl_pooled_dim=1280,
|
||||
target_seq_len=308, n_wide_blocks=2, n_narrow_blocks=3, num_heads=16, dropout=0.1, device="auto", force_reload=False):
|
||||
"""Load and initialize the LLM to SDXL adapter"""
|
||||
if device == "auto":
|
||||
device = self.device
|
||||
|
||||
try:
|
||||
# Check if we need to reload
|
||||
if force_reload or self.adapter is None or self.current_adapter_path != adapter_path:
|
||||
# Clear previous adapter
|
||||
if self.adapter is not None:
|
||||
del self.adapter
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
logger.info(f"Loading LLM to SDXL adapter from {adapter_path}")
|
||||
|
||||
# Initialize adapter with specified parameters
|
||||
self.adapter = LLMToSDXLAdapter(
|
||||
llm_dim=llm_dim,
|
||||
sdxl_seq_dim=sdxl_seq_dim,
|
||||
sdxl_pooled_dim=sdxl_pooled_dim,
|
||||
target_seq_len=target_seq_len,
|
||||
n_wide_blocks=n_wide_blocks,
|
||||
n_narrow_blocks=n_narrow_blocks,
|
||||
num_heads=num_heads,
|
||||
dropout=dropout
|
||||
)
|
||||
|
||||
# Load checkpoint if file exists
|
||||
if os.path.exists(adapter_path):
|
||||
checkpoint = load_file(adapter_path)
|
||||
self.adapter.load_state_dict(checkpoint)
|
||||
logger.info(f"Loaded adapter weights from {adapter_path}")
|
||||
else:
|
||||
logger.warning(f"Adapter file not found: {adapter_path}, using initialized weights")
|
||||
|
||||
# Move to device
|
||||
self.adapter.to(device)
|
||||
self.adapter.eval()
|
||||
|
||||
self.current_adapter_path = adapter_path
|
||||
logger.info("LLM to SDXL adapter loaded successfully")
|
||||
|
||||
info = f"Adapter: {adapter_path}\nDevice: {device}\nLLM dim: {llm_dim}\nSDXL seq dim: {sdxl_seq_dim}\nTarget seq len: {target_seq_len}"
|
||||
|
||||
return (self.adapter, info)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load adapter: {str(e)}")
|
||||
raise Exception(f"Adapter loading failed: {str(e)}")
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, adapter_path, **kwargs):
|
||||
"""ComfyUI change detection"""
|
||||
return hash(adapter_path + str(kwargs))
|
||||
|
||||
|
||||
# Node mapping for ComfyUI registration
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"LLMAdapterLoader": LLMAdapterLoader,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"LLMAdapterLoader": "LLM Adapter Loader",
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
import torch
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
import gc
|
||||
import logging
|
||||
from .utils import get_llm_checkpoints
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LLMModelLoader:
|
||||
"""
|
||||
ComfyUI node that loads Language Model and tokenizer
|
||||
Supports various LLM architectures (Gemma, Llama, Mistral, etc.)
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.current_model_path = None
|
||||
self.device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model_path": (get_llm_checkpoints(), {
|
||||
"default": get_llm_checkpoints()[0] if get_llm_checkpoints() else "./gemma-3-1b-it"
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"device": (["auto", "cuda:0", "cuda:1", "cpu"], {
|
||||
"default": "auto"
|
||||
}),
|
||||
"force_reload": ("BOOLEAN", {
|
||||
"default": False
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LLM_MODEL", "LLM_TOKENIZER", "STRING")
|
||||
RETURN_NAMES = ("model", "tokenizer", "info")
|
||||
FUNCTION = "load_model"
|
||||
CATEGORY = "llm_sdxl"
|
||||
|
||||
def load_model(self, model_path, device="auto", force_reload=False):
|
||||
"""Load Language Model and tokenizer"""
|
||||
if device == "auto":
|
||||
device = self.device
|
||||
|
||||
try:
|
||||
# Check if we need to reload
|
||||
if force_reload or self.model is None or self.current_model_path != model_path:
|
||||
# Clear previous model
|
||||
if self.model is not None:
|
||||
del self.model
|
||||
del self.tokenizer
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
logger.info(f"Loading Language Model from {model_path}")
|
||||
|
||||
self.model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map=device,
|
||||
output_hidden_states=True,
|
||||
trust_remote_code=True
|
||||
)
|
||||
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=True
|
||||
)
|
||||
|
||||
self.current_model_path = model_path
|
||||
logger.info("Language Model loaded successfully")
|
||||
|
||||
info = f"Model: {model_path}\nDevice: {device}\nLoaded: {self.model is not None}"
|
||||
|
||||
return (self.model, self.tokenizer, info)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load Language Model: {str(e)}")
|
||||
raise Exception(f"Model loading failed: {str(e)}")
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, model_path, device="auto", force_reload=False):
|
||||
"""ComfyUI change detection"""
|
||||
return hash(model_path + device + str(force_reload))
|
||||
|
||||
|
||||
# Node mapping for ComfyUI registration
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"LLMModelLoader": LLMModelLoader,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"LLMModelLoader": "LLM Model Loader",
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
import torch
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LLMTextEncoder:
|
||||
"""
|
||||
ComfyUI node that encodes text using a loaded Language Model
|
||||
Supports various LLM architectures with chat templates
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("LLM_MODEL",),
|
||||
"tokenizer": ("LLM_TOKENIZER",),
|
||||
"text": ("STRING", {
|
||||
"multiline": True,
|
||||
"default": "masterpiece, best quality, 1girl, anime style"
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"system_prompt": ("STRING", {
|
||||
"multiline": True,
|
||||
"default": "You are expert in understanding of user prompts for image generations. Create an image according to the prompt from user."
|
||||
}),
|
||||
"skip_first": ("INT", {
|
||||
"default": 27,
|
||||
"min": 0,
|
||||
"max": 100,
|
||||
"step": 1
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LLM_HIDDEN_STATES", "STRING")
|
||||
RETURN_NAMES = ("hidden_states", "info")
|
||||
FUNCTION = "encode_text"
|
||||
CATEGORY = "llm_sdxl"
|
||||
|
||||
def encode_text(self, model, tokenizer, text, system_prompt="You are expert in understanding of user prompts for image generations. Create an image according to the prompt from user.", skip_first=27):
|
||||
"""
|
||||
Encode text using Language Model and return hidden states
|
||||
"""
|
||||
try:
|
||||
# Get model device
|
||||
device = next(model.parameters()).device
|
||||
|
||||
# Prepare chat template
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": system_prompt}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": text}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
# Apply chat template
|
||||
inputs = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
return_dict=True,
|
||||
return_tensors="pt",
|
||||
add_generation_prompt=True,
|
||||
).to(device)
|
||||
|
||||
# Generate hidden states
|
||||
with torch.no_grad():
|
||||
outputs = model(**inputs)
|
||||
|
||||
# Extract hidden states, skipping first tokens
|
||||
hidden_states = outputs['hidden_states'][-1][:, skip_first:, :].to(torch.float)
|
||||
|
||||
# Prepare info
|
||||
info = f"Text: {text[:50]}...\nTokens after skip: {hidden_states.shape[1]}\nShape: {hidden_states.shape}"
|
||||
|
||||
logger.info(f"Encoded text with shape: {hidden_states.shape}")
|
||||
|
||||
return (hidden_states, info)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to encode text: {str(e)}")
|
||||
raise Exception(f"Text encoding failed: {str(e)}")
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, model, tokenizer, text, system_prompt="You are expert in understanding of user prompts for image generations. Create an image according to the prompt from user.", skip_first=27):
|
||||
"""ComfyUI change detection"""
|
||||
return hash(text + system_prompt + str(skip_first))
|
||||
|
||||
|
||||
# Node mapping for ComfyUI registration
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"LLMTextEncoder": LLMTextEncoder
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"LLMTextEncoder": "LLM Text Encoder"
|
||||
}
|
||||
@@ -0,0 +1,228 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def pad_to_length(tensor, target_length, dim=1, value=0):
|
||||
"""Universal function for padding tensors"""
|
||||
current_length = tensor.size(dim)
|
||||
|
||||
if current_length >= target_length:
|
||||
return tensor.narrow(dim, 0, target_length)
|
||||
|
||||
pad_size = list(tensor.shape)
|
||||
pad_size[dim] = target_length - current_length
|
||||
|
||||
padding = torch.full(
|
||||
pad_size,
|
||||
value,
|
||||
device=tensor.device,
|
||||
dtype=tensor.dtype
|
||||
)
|
||||
|
||||
return torch.cat([tensor, padding], dim=dim)
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, dim, num_heads=16, mlp_ratio=4.0, dropout=0.0):
|
||||
super().__init__()
|
||||
|
||||
self.norm1 = nn.LayerNorm(dim)
|
||||
self.attn = nn.MultiheadAttention(
|
||||
dim, num_heads, batch_first=True, dropout=dropout
|
||||
)
|
||||
|
||||
self.norm2 = nn.LayerNorm(dim)
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(dim, int(dim * mlp_ratio)),
|
||||
nn.GELU(),
|
||||
nn.Linear(int(dim * mlp_ratio), dim)
|
||||
)
|
||||
|
||||
def forward(self, x, mask=None):
|
||||
# Self-attention
|
||||
normed = self.norm1(x)
|
||||
|
||||
# Use key_padding_mask instead of attn_mask
|
||||
if mask is not None:
|
||||
# key_padding_mask: True means "ignore this token"
|
||||
# Our mask: 1 = real token, 0 = padding
|
||||
# So we invert
|
||||
key_padding_mask = ~mask.bool()
|
||||
else:
|
||||
key_padding_mask = None
|
||||
|
||||
attn_out, _ = self.attn(
|
||||
normed, normed, normed,
|
||||
key_padding_mask=key_padding_mask,
|
||||
need_weights=False
|
||||
)
|
||||
x = x + attn_out
|
||||
|
||||
# MLP
|
||||
x = x + self.mlp(self.norm2(x))
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class LLMToSDXLAdapter(nn.Module):
|
||||
"""
|
||||
Universal adapter for converting any LLM embeddings to SDXL format
|
||||
Supports various LLM architectures (Gemma, Llama, Mistral, etc.)
|
||||
"""
|
||||
def __init__(self,
|
||||
llm_dim=1152, # Changed from gemma_dim to llm_dim
|
||||
sdxl_seq_dim=2048,
|
||||
sdxl_pooled_dim=1280,
|
||||
max_input_len=512,
|
||||
target_seq_len=308,
|
||||
n_wide_blocks=3, # Blocks BEFORE compression
|
||||
n_narrow_blocks=3, # Blocks AFTER compression
|
||||
num_heads=16,
|
||||
dropout=0):
|
||||
super().__init__()
|
||||
|
||||
self.max_input_len = max_input_len
|
||||
self.target_seq_len = target_seq_len
|
||||
self.num_heads = num_heads
|
||||
|
||||
# Projections
|
||||
self.seq_projection = nn.Linear(llm_dim, sdxl_seq_dim)
|
||||
|
||||
# Positional embeddings for full sequence
|
||||
self.input_position_embeddings = nn.Parameter(
|
||||
torch.randn(1, max_input_len, sdxl_seq_dim)
|
||||
)
|
||||
# Positional embeddings for compressed sequence
|
||||
self.output_position_embeddings = nn.Parameter(
|
||||
torch.randn(1, target_seq_len, sdxl_seq_dim)
|
||||
)
|
||||
|
||||
# Wide blocks - processing full sequence (512 tokens)
|
||||
self.wide_attention_blocks = nn.ModuleList([
|
||||
TransformerBlock(sdxl_seq_dim, num_heads=num_heads, dropout=dropout)
|
||||
for _ in range(n_wide_blocks)
|
||||
])
|
||||
|
||||
# Compression: Cross-attention with learnable queries
|
||||
self.compression_queries = nn.Parameter(
|
||||
torch.randn(1, target_seq_len, sdxl_seq_dim)
|
||||
)
|
||||
self.compression_attention = nn.MultiheadAttention(
|
||||
embed_dim=sdxl_seq_dim,
|
||||
num_heads=num_heads,
|
||||
batch_first=True,
|
||||
dropout=dropout
|
||||
)
|
||||
# Norm layer after compression for stability
|
||||
self.compression_norm = nn.LayerNorm(sdxl_seq_dim)
|
||||
# Optional gate mechanism for weighting information
|
||||
self.compression_gate = nn.Sequential(
|
||||
nn.Linear(sdxl_seq_dim * 2, sdxl_seq_dim),
|
||||
nn.Sigmoid()
|
||||
)
|
||||
|
||||
# Narrow blocks - processing compressed sequence (308 tokens)
|
||||
self.narrow_attention_blocks = nn.ModuleList([
|
||||
TransformerBlock(sdxl_seq_dim, num_heads=num_heads, dropout=dropout)
|
||||
for _ in range(n_narrow_blocks)
|
||||
])
|
||||
|
||||
# Pooling head - now works with processed sequence
|
||||
self.pooling_attention = nn.MultiheadAttention(
|
||||
embed_dim=sdxl_seq_dim,
|
||||
num_heads=num_heads,
|
||||
batch_first=True,
|
||||
dropout=dropout
|
||||
)
|
||||
|
||||
# Learnable [CLS]-like token for pooling
|
||||
self.pooling_token = nn.Parameter(torch.randn(1, 1, sdxl_seq_dim))
|
||||
|
||||
# Final projection for pooled embeddings
|
||||
self.pooled_projection = nn.Sequential(
|
||||
nn.Linear(sdxl_seq_dim, sdxl_seq_dim),
|
||||
nn.LayerNorm(sdxl_seq_dim),
|
||||
nn.GELU(),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(sdxl_seq_dim, sdxl_pooled_dim)
|
||||
)
|
||||
|
||||
def forward(self, llm_hidden_states, attention_mask=None):
|
||||
batch_size, seq_len, _ = llm_hidden_states.shape
|
||||
|
||||
# Project to target dimension
|
||||
hidden_states = self.seq_projection(llm_hidden_states)
|
||||
|
||||
# Padding/truncation to max_input_len
|
||||
if seq_len > self.max_input_len:
|
||||
hidden_states = hidden_states[:, :self.max_input_len, :]
|
||||
if attention_mask is not None:
|
||||
attention_mask = attention_mask[:, :self.max_input_len]
|
||||
else:
|
||||
if seq_len < self.max_input_len:
|
||||
hidden_states = pad_to_length(hidden_states, self.max_input_len, dim=1)
|
||||
if attention_mask is not None:
|
||||
attention_mask = pad_to_length(attention_mask, self.max_input_len, dim=1, value=0)
|
||||
else:
|
||||
attention_mask = torch.ones(batch_size, self.max_input_len, device=hidden_states.device)
|
||||
attention_mask[:, seq_len:] = 0
|
||||
|
||||
# Add positional embeddings
|
||||
hidden_states = hidden_states + self.input_position_embeddings
|
||||
|
||||
# ===== STAGE 1: Wide Processing (full sequence) =====
|
||||
for block in self.wide_attention_blocks:
|
||||
hidden_states = block(hidden_states, attention_mask)
|
||||
|
||||
# ===== STAGE 2: Compression (512 -> 308) =====
|
||||
# Prepare queries for compression
|
||||
queries = self.compression_queries.expand(batch_size, -1, -1)
|
||||
|
||||
# Cross-attention for compression
|
||||
if attention_mask is not None:
|
||||
key_padding_mask = ~attention_mask.bool()
|
||||
else:
|
||||
key_padding_mask = None
|
||||
|
||||
compressed_sequence, compression_weights = self.compression_attention(
|
||||
queries,
|
||||
hidden_states,
|
||||
hidden_states,
|
||||
key_padding_mask=key_padding_mask,
|
||||
need_weights=True,
|
||||
average_attn_weights=False # Get weights for each head
|
||||
)
|
||||
|
||||
# Optional: Gate mechanism for mixing with queries
|
||||
gate_input = torch.cat([queries, compressed_sequence], dim=-1)
|
||||
gate_weights = self.compression_gate(gate_input)
|
||||
compressed_sequence = gate_weights * compressed_sequence + (1 - gate_weights) * queries
|
||||
|
||||
# Apply normalization
|
||||
compressed_sequence = self.compression_norm(compressed_sequence)
|
||||
|
||||
# Add output positional embeddings
|
||||
compressed_sequence = compressed_sequence + self.output_position_embeddings
|
||||
|
||||
# ===== STAGE 3: Narrow Processing (compressed sequence) =====
|
||||
for block in self.narrow_attention_blocks:
|
||||
compressed_sequence = block(compressed_sequence)
|
||||
|
||||
# ===== STAGE 4: Pooling for Vector Embeddings =====
|
||||
# Pool the compressed sequence for vector embeddings
|
||||
pooling_tokens = self.pooling_token.expand(batch_size, -1, -1)
|
||||
pooled_output, pooling_weights = self.pooling_attention(
|
||||
pooling_tokens,
|
||||
compressed_sequence,
|
||||
compressed_sequence,
|
||||
need_weights=False
|
||||
)
|
||||
pooled_output = pooled_output.squeeze(1) # Remove sequence dimension
|
||||
|
||||
# Final projection for pooled embeddings
|
||||
pooled_output = self.pooled_projection(pooled_output)
|
||||
|
||||
return compressed_sequence, pooled_output
|
||||
@@ -0,0 +1,47 @@
|
||||
import os
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_llm_checkpoints():
|
||||
"""
|
||||
Scan models/LLM directory for available checkpoints
|
||||
Returns list of checkpoint directories/files
|
||||
"""
|
||||
llm_path = os.path.join("models", "LLM")
|
||||
checkpoints = []
|
||||
|
||||
if os.path.exists(llm_path):
|
||||
for item in os.listdir(llm_path):
|
||||
item_path = os.path.join(llm_path, item)
|
||||
if os.path.isdir(item_path):
|
||||
# Check if it's a valid model directory (contains config.json or similar)
|
||||
if any(f in os.listdir(item_path) for f in ['config.json', 'model.safetensors', 'pytorch_model.bin']):
|
||||
checkpoints.append(item_path)
|
||||
elif item.endswith(('.safetensors', '.bin', '.pt')):
|
||||
checkpoints.append(item_path)
|
||||
|
||||
if not checkpoints:
|
||||
checkpoints = ["./gemma-3-1b-it"] # Default fallback
|
||||
|
||||
return checkpoints
|
||||
|
||||
|
||||
def get_llm_adapters():
|
||||
"""
|
||||
Scan models/llm_adapters directory for available adapters
|
||||
Returns list of adapter files
|
||||
"""
|
||||
adapters_path = os.path.join("models", "llm_adapters")
|
||||
adapters = []
|
||||
|
||||
if os.path.exists(adapters_path):
|
||||
for item in os.listdir(adapters_path):
|
||||
if item.endswith('.safetensors'):
|
||||
adapters.append(os.path.join(adapters_path, item))
|
||||
|
||||
if not adapters:
|
||||
adapters = ["checkpoint_epoch_10.safetensors"] # Default fallback
|
||||
|
||||
return adapters
|
||||
Reference in New Issue
Block a user