From 3fdb5dc3b864698022d19c19db5d4b4cd5bf73e4 Mon Sep 17 00:00:00 2001 From: NeuroSenko Date: Wed, 9 Jul 2025 05:20:56 +0300 Subject: [PATCH] init --- .gitignore | 1 + README.md | 165 +++++++++++++++++++++++++ __init__.py | 169 ++++++++++++++++++++++++++ apply_llm_to_sdxl_adapter.py | 61 ++++++++++ llm_adapter_loader.py | 151 +++++++++++++++++++++++ llm_model_loader.py | 99 +++++++++++++++ llm_text_encoder.py | 107 ++++++++++++++++ llm_to_sdxl_adapter.py | 228 +++++++++++++++++++++++++++++++++++ utils.py | 47 ++++++++ 9 files changed, 1028 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 __init__.py create mode 100644 apply_llm_to_sdxl_adapter.py create mode 100644 llm_adapter_loader.py create mode 100644 llm_model_loader.py create mode 100644 llm_text_encoder.py create mode 100644 llm_to_sdxl_adapter.py create mode 100644 utils.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ba0430d --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__/ \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..1d7a5c0 --- /dev/null +++ b/README.md @@ -0,0 +1,165 @@ +# ComfyUI LLM SDXL Adapter + +![Version](https://img.shields.io/badge/version-1.0.0-blue.svg) +![License](https://img.shields.io/badge/license-MIT-green.svg) +![Python](https://img.shields.io/badge/python-3.8+-blue.svg) +![ComfyUI](https://img.shields.io/badge/ComfyUI-compatible-orange.svg) + +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. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..dc92012 --- /dev/null +++ b/__init__.py @@ -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") \ No newline at end of file diff --git a/apply_llm_to_sdxl_adapter.py b/apply_llm_to_sdxl_adapter.py new file mode 100644 index 0000000..56182d5 --- /dev/null +++ b/apply_llm_to_sdxl_adapter.py @@ -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" +} \ No newline at end of file diff --git a/llm_adapter_loader.py b/llm_adapter_loader.py new file mode 100644 index 0000000..f63888d --- /dev/null +++ b/llm_adapter_loader.py @@ -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", +} \ No newline at end of file diff --git a/llm_model_loader.py b/llm_model_loader.py new file mode 100644 index 0000000..a146a2e --- /dev/null +++ b/llm_model_loader.py @@ -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", +} \ No newline at end of file diff --git a/llm_text_encoder.py b/llm_text_encoder.py new file mode 100644 index 0000000..3d77df1 --- /dev/null +++ b/llm_text_encoder.py @@ -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" +} \ No newline at end of file diff --git a/llm_to_sdxl_adapter.py b/llm_to_sdxl_adapter.py new file mode 100644 index 0000000..9dbb15e --- /dev/null +++ b/llm_to_sdxl_adapter.py @@ -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 \ No newline at end of file diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..4ea8714 --- /dev/null +++ b/utils.py @@ -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