148 lines
5.1 KiB
Python
148 lines
5.1 KiB
Python
import torch
|
|
from safetensors.torch import load_file
|
|
import logging
|
|
import gc
|
|
import os
|
|
from .utils import get_llm_adapters, get_llm_adapter_path
|
|
from .llm_to_sdxl_adapter import LLMToSDXLAdapter
|
|
|
|
logger = logging.getLogger("LLM-SDXL-Adapter")
|
|
|
|
class LLMAdapterLoaderCustom:
|
|
"""
|
|
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_name": (get_llm_adapters(), {
|
|
"default": get_llm_adapters()[0] if get_llm_adapters() else None
|
|
}),
|
|
},
|
|
"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_name, 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
|
|
|
|
adapter_path = get_llm_adapter_path(adapter_name)
|
|
|
|
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)}")
|
|
|
|
|
|
|
|
# Node mapping for ComfyUI registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LLMAdapterLoaderCustom": LLMAdapterLoaderCustom,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LLMAdapterLoaderCustom": "LLM Adapter Loader Custom",
|
|
} |