445 lines
16 KiB
Python
445 lines
16 KiB
Python
"""
|
|
SDXL Long Context Node for ComfyUI
|
|
|
|
Custom checkpoint loader for SDXL models with extended positional
|
|
embeddings (up to 248 tokens) created by convert_spectral.py.
|
|
|
|
The converter extends positional embeddings via spectral extrapolation (FFT).
|
|
It does NOT touch position_ids, causal masks, or any other weights.
|
|
|
|
This loader:
|
|
1. Detects extended pos_emb in checkpoint
|
|
2. Truncates to 77 for ComfyUI's standard loader compatibility
|
|
3. Loads via standard pipeline
|
|
4. Patches back full embeddings + updates position_ids and max_length
|
|
5. Verifies everything
|
|
|
|
Use standard CLIPTextEncode nodes for text encoding afterward.
|
|
|
|
Install: copy this folder to ComfyUI/custom_nodes/sdxl_long_context/
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
import folder_paths
|
|
import comfy.sd
|
|
|
|
from safetensors import safe_open
|
|
|
|
logger = logging.getLogger("SDXLLongContext")
|
|
|
|
|
|
# ================================================================
|
|
# Exact keys matching the converter
|
|
# ================================================================
|
|
|
|
POS_EMB_KEYS = [
|
|
# CLIP-L
|
|
"conditioner.embedders.0.transformer.text_model.embeddings.position_embedding.weight",
|
|
"cond_stage_model.transformer.text_model.embeddings.position_embedding.weight",
|
|
"text_encoder.text_model.embeddings.position_embedding.weight",
|
|
# CLIP-G
|
|
"conditioner.embedders.1.model.positional_embedding",
|
|
"conditioner.embedders.1.transformer.text_model.embeddings.position_embedding.weight",
|
|
"text_encoder_2.text_model.embeddings.position_embedding.weight",
|
|
]
|
|
|
|
CLIP_L_PREFIXES = [
|
|
"conditioner.embedders.0.",
|
|
"cond_stage_model.",
|
|
"text_encoder.",
|
|
]
|
|
|
|
CLIP_G_PREFIXES = [
|
|
"conditioner.embedders.1.",
|
|
"text_encoder_2.",
|
|
]
|
|
|
|
MAX_LENGTH = 248
|
|
|
|
|
|
# ================================================================
|
|
# Utilities
|
|
# ================================================================
|
|
|
|
def detect_extended_keys(state_dict):
|
|
"""Find POS_EMB_KEYS that are present and extended beyond 77."""
|
|
found = {}
|
|
for key in POS_EMB_KEYS:
|
|
if key not in state_dict:
|
|
continue
|
|
tensor = state_dict[key]
|
|
if tensor.dim() == 2 and tensor.shape[0] > 77:
|
|
found[key] = tensor
|
|
return found
|
|
|
|
|
|
def get_max_length(extended_keys):
|
|
"""Get target max_length from extended embeddings."""
|
|
max_len = 77
|
|
for tensor in extended_keys.values():
|
|
if tensor.shape[0] > max_len:
|
|
max_len = tensor.shape[0]
|
|
return max_len
|
|
|
|
|
|
def classify_key(key):
|
|
"""Determine if a key belongs to clip_l or clip_g."""
|
|
for prefix in CLIP_G_PREFIXES:
|
|
if key.startswith(prefix):
|
|
return "clip_g"
|
|
for prefix in CLIP_L_PREFIXES:
|
|
if key.startswith(prefix):
|
|
return "clip_l"
|
|
return None
|
|
|
|
|
|
def get_nested_attr(obj, path):
|
|
"""Get attribute by dot-separated path."""
|
|
for part in path.split("."):
|
|
obj = getattr(obj, part)
|
|
return obj
|
|
|
|
|
|
def set_nested_attr(obj, path, value):
|
|
"""Set attribute by dot-separated path."""
|
|
parts = path.split(".")
|
|
if len(parts) > 1:
|
|
parent = get_nested_attr(obj, ".".join(parts[:-1]))
|
|
else:
|
|
parent = obj
|
|
setattr(parent, parts[-1], value)
|
|
|
|
|
|
# ================================================================
|
|
# CLIP Patching
|
|
# ================================================================
|
|
|
|
def find_position_embedding(clip_submodel):
|
|
"""
|
|
Find position embedding in a CLIP submodel.
|
|
Returns (path, module_or_none, is_embedding).
|
|
|
|
Handles both:
|
|
- HF format: nn.Embedding at transformer.text_model.embeddings.position_embedding
|
|
- OpenCLIP format: nn.Parameter at positional_embedding
|
|
"""
|
|
# HF-style
|
|
hf_path = "transformer.text_model.embeddings.position_embedding"
|
|
try:
|
|
module = get_nested_attr(clip_submodel, hf_path)
|
|
if isinstance(module, nn.Embedding):
|
|
return hf_path, module, True
|
|
except AttributeError:
|
|
pass
|
|
|
|
# OpenCLIP-style
|
|
if hasattr(clip_submodel, "positional_embedding"):
|
|
param = clip_submodel.positional_embedding
|
|
if isinstance(param, (nn.Parameter, torch.Tensor)):
|
|
return "positional_embedding", None, False
|
|
|
|
# Fallback search
|
|
for name, module in clip_submodel.named_modules():
|
|
if isinstance(module, nn.Embedding) and "position" in name.lower():
|
|
return name, module, True
|
|
|
|
return None, None, False
|
|
|
|
|
|
def patch_position_embedding(clip_submodel, full_tensor, max_length, clip_name):
|
|
"""Patch a CLIP submodel with full positional embedding tensor."""
|
|
path, module, is_embedding = find_position_embedding(clip_submodel)
|
|
|
|
if path is None:
|
|
print(f" WARNING: No position_embedding found in {clip_name}")
|
|
return False
|
|
|
|
if is_embedding:
|
|
# nn.Embedding
|
|
old_num = module.num_embeddings
|
|
old_dim = module.embedding_dim
|
|
|
|
if old_dim != full_tensor.shape[1]:
|
|
print(f" WARNING: dim mismatch {clip_name}.{path}: "
|
|
f"model={old_dim}, checkpoint={full_tensor.shape[1]}")
|
|
return False
|
|
|
|
new_emb = nn.Embedding(
|
|
max_length, old_dim,
|
|
dtype=full_tensor.dtype,
|
|
device=full_tensor.device,
|
|
)
|
|
new_emb.weight = nn.Parameter(full_tensor.clone())
|
|
set_nested_attr(clip_submodel, path, new_emb)
|
|
|
|
print(f" Patched: {clip_name}.{path}")
|
|
print(f" nn.Embedding({old_num}, {old_dim}) -> "
|
|
f"nn.Embedding({max_length}, {old_dim})")
|
|
else:
|
|
# nn.Parameter (OpenCLIP)
|
|
old_shape = getattr(clip_submodel, path).shape
|
|
setattr(clip_submodel, path, nn.Parameter(full_tensor.clone()))
|
|
|
|
print(f" Patched: {clip_name}.{path}")
|
|
print(f" Parameter{tuple(old_shape)} -> Parameter{tuple(full_tensor.shape)}")
|
|
|
|
return True
|
|
|
|
|
|
def patch_position_ids(clip_submodel, max_length, clip_name):
|
|
"""Update position_ids buffers to match new max_length."""
|
|
patched = False
|
|
for name, buf in list(clip_submodel.named_buffers()):
|
|
if "position_ids" not in name:
|
|
continue
|
|
if buf.shape[-1] == max_length:
|
|
continue
|
|
if buf.shape[-1] != 77:
|
|
continue
|
|
|
|
new_ids = torch.arange(max_length, dtype=buf.dtype, device=buf.device)
|
|
if buf.dim() == 2:
|
|
new_ids = new_ids.unsqueeze(0)
|
|
|
|
try:
|
|
parts = name.split(".")
|
|
if len(parts) > 1:
|
|
parent = get_nested_attr(clip_submodel, ".".join(parts[:-1]))
|
|
else:
|
|
parent = clip_submodel
|
|
parent.register_buffer(parts[-1], new_ids)
|
|
print(f" Updated: {clip_name}.{name}: {tuple(buf.shape)} -> {tuple(new_ids.shape)}")
|
|
patched = True
|
|
except Exception as e:
|
|
logger.warning(f" Could not update {clip_name}.{name}: {e}")
|
|
|
|
return patched
|
|
|
|
|
|
def update_max_length_settings(clip, max_length):
|
|
"""Update max_length in CLIP model and tokenizer."""
|
|
cond_model = clip.cond_stage_model
|
|
|
|
for attr_name in ["clip_l", "clip_g"]:
|
|
if not hasattr(cond_model, attr_name):
|
|
continue
|
|
sub = getattr(cond_model, attr_name)
|
|
if hasattr(sub, "max_length"):
|
|
old = sub.max_length
|
|
sub.max_length = max_length
|
|
print(f" {attr_name}.max_length: {old} -> {max_length}")
|
|
|
|
if hasattr(clip, "tokenizer"):
|
|
tokenizer = clip.tokenizer
|
|
for attr_name in ["clip_l", "clip_g"]:
|
|
if not hasattr(tokenizer, attr_name):
|
|
continue
|
|
sub_tok = getattr(tokenizer, attr_name)
|
|
for attr in ["max_length", "max_tokens", "model_max_length"]:
|
|
if hasattr(sub_tok, attr):
|
|
old = getattr(sub_tok, attr)
|
|
setattr(sub_tok, attr, max_length)
|
|
print(f" tokenizer.{attr_name}.{attr}: {old} -> {max_length}")
|
|
|
|
if hasattr(tokenizer, "max_length"):
|
|
old = tokenizer.max_length
|
|
tokenizer.max_length = max_length
|
|
print(f" tokenizer.max_length: {old} -> {max_length}")
|
|
|
|
|
|
def verify_patching(clip, max_length):
|
|
"""Verify all position embeddings have expected size."""
|
|
cond_model = clip.cond_stage_model
|
|
print(f" --- Verification ---")
|
|
all_ok = True
|
|
|
|
for clip_name in ["clip_l", "clip_g"]:
|
|
if not hasattr(cond_model, clip_name):
|
|
continue
|
|
sub = getattr(cond_model, clip_name)
|
|
|
|
# nn.Embedding
|
|
for name, module in sub.named_modules():
|
|
if isinstance(module, nn.Embedding) and "position" in name.lower():
|
|
ok = module.num_embeddings == max_length
|
|
status = "PASS" if ok else "FAIL"
|
|
print(f" {status}: {clip_name}.{name} = "
|
|
f"({module.num_embeddings}, {module.embedding_dim})")
|
|
if not ok:
|
|
all_ok = False
|
|
|
|
# nn.Parameter (OpenCLIP)
|
|
if hasattr(sub, "positional_embedding"):
|
|
param = sub.positional_embedding
|
|
if isinstance(param, (nn.Parameter, torch.Tensor)):
|
|
ok = param.shape[0] == max_length
|
|
status = "PASS" if ok else "FAIL"
|
|
print(f" {status}: {clip_name}.positional_embedding = {tuple(param.shape)}")
|
|
if not ok:
|
|
all_ok = False
|
|
|
|
# max_length
|
|
if hasattr(sub, "max_length"):
|
|
ok = sub.max_length == max_length
|
|
status = "PASS" if ok else "FAIL"
|
|
print(f" {status}: {clip_name}.max_length = {sub.max_length}")
|
|
if not ok:
|
|
all_ok = False
|
|
|
|
# position_ids
|
|
for name, buf in sub.named_buffers():
|
|
if "position_ids" in name:
|
|
ok = buf.shape[-1] == max_length
|
|
status = "PASS" if ok else "FAIL"
|
|
print(f" {status}: {clip_name}.{name} = {tuple(buf.shape)}")
|
|
if not ok:
|
|
all_ok = False
|
|
|
|
print(f" {'All checks passed!' if all_ok else 'ERRORS DETECTED!'}")
|
|
return all_ok
|
|
|
|
|
|
# ================================================================
|
|
# Node
|
|
# ================================================================
|
|
|
|
class SDXLLongCheckpointLoader:
|
|
"""
|
|
Checkpoint loader for SDXL models with extended context (up to 248 tokens).
|
|
|
|
For checkpoints converted by convert_spectral.py (spectral extrapolation).
|
|
Also loads standard 77-token checkpoints normally.
|
|
|
|
Use standard CLIPTextEncode nodes for text encoding.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
|
RETURN_NAMES = ("model", "clip", "vae")
|
|
FUNCTION = "load_checkpoint"
|
|
CATEGORY = "loaders/sdxl_long_context"
|
|
|
|
def load_checkpoint(self, ckpt_name):
|
|
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
|
print(f"\n[SDXLLongContext] Loading: {ckpt_name}")
|
|
|
|
# Step 1: Read state_dict
|
|
print("[SDXLLongContext] Step 1: Reading state_dict...")
|
|
state_dict = {}
|
|
metadata = {}
|
|
with safe_open(ckpt_path, framework="pt", device="cpu") as f:
|
|
try:
|
|
metadata = f.metadata() or {}
|
|
except Exception:
|
|
metadata = {}
|
|
for key in f.keys():
|
|
state_dict[key] = f.get_tensor(key)
|
|
|
|
print(f" Tensors: {len(state_dict)}")
|
|
|
|
if metadata.get("long_context"):
|
|
print(f" Converter metadata:")
|
|
print(f" method: {metadata.get('long_context')}")
|
|
print(f" original_length: {metadata.get('original_length')}")
|
|
print(f" target_length: {metadata.get('target_length')}")
|
|
print(f" damping: {metadata.get('damping', 'N/A')}")
|
|
try:
|
|
mod_keys = json.loads(metadata.get("modified_keys", "[]"))
|
|
for mk in mod_keys:
|
|
print(f" modified: {mk}")
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
# Step 2: Detect extended embeddings
|
|
print("[SDXLLongContext] Step 2: Detecting extended pos_emb...")
|
|
extended_keys = detect_extended_keys(state_dict)
|
|
|
|
if not extended_keys:
|
|
print(" Standard checkpoint, using normal loader")
|
|
return comfy.sd.load_checkpoint_guess_config(
|
|
ckpt_path, output_vae=True, output_clip=True,
|
|
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
|
)[:3]
|
|
|
|
max_length = min(get_max_length(extended_keys), MAX_LENGTH)
|
|
print(f" Found {len(extended_keys)} extended key(s), max_length={max_length}")
|
|
for key, tensor in extended_keys.items():
|
|
print(f" [{classify_key(key)}] {key}: {tuple(tensor.shape)}")
|
|
|
|
# Step 3: Save full tensors, truncate for standard loader
|
|
print("[SDXLLongContext] Step 3: Truncating to 77 for standard loader...")
|
|
full_tensors = {}
|
|
for key, tensor in extended_keys.items():
|
|
clip_type = classify_key(key)
|
|
full_tensors[key] = (clip_type, tensor.clone())
|
|
state_dict[key] = tensor[:77].contiguous()
|
|
print(f" {key}: ({tensor.shape[0]}, {tensor.shape[1]}) -> (77, {tensor.shape[1]})")
|
|
|
|
# Step 4: Standard loading
|
|
print("[SDXLLongContext] Step 4: Loading via standard ComfyUI loader...")
|
|
out = comfy.sd.load_state_dict_guess_config(
|
|
state_dict, output_vae=True, output_clip=True,
|
|
output_clipvision=False,
|
|
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
|
output_model=True,
|
|
)
|
|
model, clip, vae = out[0], out[1], out[2]
|
|
print(" Standard loading complete")
|
|
|
|
# Step 5: Patch CLIP
|
|
print(f"[SDXLLongContext] Step 5: Patching CLIP to {max_length} tokens...")
|
|
cond_model = clip.cond_stage_model
|
|
|
|
clips_to_patch = {}
|
|
for key, (clip_type, full_tensor) in full_tensors.items():
|
|
if clip_type is None:
|
|
print(f" WARNING: Cannot classify {key}, skipping")
|
|
continue
|
|
if clip_type not in clips_to_patch:
|
|
clips_to_patch[clip_type] = full_tensor
|
|
|
|
for clip_name, full_tensor in clips_to_patch.items():
|
|
if not hasattr(cond_model, clip_name):
|
|
print(f" WARNING: cond_stage_model has no '{clip_name}'")
|
|
continue
|
|
|
|
sub = getattr(cond_model, clip_name)
|
|
patch_position_embedding(sub, full_tensor, max_length, clip_name)
|
|
patch_position_ids(sub, max_length, clip_name)
|
|
|
|
# Step 6: Update max_length
|
|
print(f"[SDXLLongContext] Step 6: Updating max_length settings...")
|
|
update_max_length_settings(clip, max_length)
|
|
|
|
# Step 7: Verify
|
|
print(f"[SDXLLongContext] Step 7: Verification")
|
|
verify_patching(clip, max_length)
|
|
|
|
print(f"[SDXLLongContext] Done! Context: {max_length} tokens\n")
|
|
return (model, clip, vae)
|
|
|
|
|
|
# ================================================================
|
|
# Registration
|
|
# ================================================================
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SDXLLongCheckpointLoader": SDXLLongCheckpointLoader,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SDXLLongCheckpointLoader": "Load Checkpoint (SDXL Long 248)",
|
|
}
|