nodes/silveroxides-ComfyUI_EmbeddingToolkit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

initial llm stuff and addressing llama templates exclusion

Browse Source
This commit is contained in:
silveroxides
2026-02-02 03:24:53 +01:00
parent db00c2ab81
commit d3dbd3165d
2 changed files with 381 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+229 -9
View File
@@ -2,6 +2,196 @@ import torch
import os
import folder_paths
import comfy.utils
import logging
import re
from comfy.sd1_clip import token_weights, escape_important, unescape_important
from .nodes_debug import PromptEmbeddingFixer, InspectEmbeddingForClip
def tokenize_preserving_weights(clip, text):
"""
Custom tokenizer helper that bypasses 'disable_weights=True' enforced by some
tokenizers (e.g., Flux/Mistral/Qwen) by manually implementing the tokenization loop
and weight parsing, ensuring weights are always respected.
"""
tokenizer = clip.tokenizer
out = {}
potential_parts = [
'clip_l', 'clip_g', 'clip_h',
't5xxl', 't5xl', 't5base', 'pile_t5xl', 'umt5xxl',
'qwen25_7b', 'qwen25_3b', 'qwen3_4b', 'qwen3_8b', 'qwen3_2b', 'qwen3_06b',
'mistral3_24b',
'gemma2_2b', 'gemma3_4b', 'gemma3_12b',
'jina_clip_2', 'gemma', 'jina', 'byt5_small',
'llama'
]
# Dynamic discovery for SD1Tokenizer subclasses (like Flux2, Klein)
# which store the sub-tokenizer attribute name in 'self.clip'
if hasattr(tokenizer, "clip") and isinstance(tokenizer.clip, str):
if tokenizer.clip not in potential_parts:
potential_parts.append(tokenizer.clip)
found_any = False
for attr in potential_parts:
if hasattr(tokenizer, attr):
sub_tok = getattr(tokenizer, attr)
# Check for essential attributes to perform manual tokenization
if hasattr(sub_tok, "tokenizer") and hasattr(sub_tok, "tokens_start"):
try:
# Manual parsing to ensure weights are respected
# We skip the 'disable_weights' check entirely by implementing the logic ourselves
# 1. Parse weights
text_processed = escape_important(text)
parsed_weights = token_weights(text_processed, 1.0)
# 2. Tokenize segments
tokens = []
embedding_identifier = getattr(sub_tok, "embedding_identifier", "embedding:")
embedding_directory = getattr(sub_tok, "embedding_directory", None)
for weighted_segment, weight in parsed_weights:
to_tokenize = unescape_important(weighted_segment)
# Handle embeddings (copied from SDTokenizer logic)
split = re.split(' {0}|\n{0}'.format(embedding_identifier), to_tokenize)
to_tokenize_list = [split[0]]
for i in range(1, len(split)):
to_tokenize_list.append("{}{}".format(embedding_identifier, split[i]))
to_tokenize_list = [x for x in to_tokenize_list if x != ""]
for word in to_tokenize_list:
# Embedding check
if word.startswith(embedding_identifier) and embedding_directory is not None:
if hasattr(sub_tok, "_try_get_embedding"):
embedding_name = word[len(embedding_identifier):].strip('\n')
embed, leftover = sub_tok._try_get_embedding(embedding_name)
if embed is None:
logging.warning(f"EmbeddingToolkit: warning, embedding:{embedding_name} does not exist, ignoring")
else:
if len(embed.shape) == 1:
tokens.append([(embed, weight)])
else:
tokens.append([(embed[x], weight) for x in range(embed.shape[0])])
if leftover != "":
word = leftover
else:
continue
# Tokenize
end = 999999999999
if getattr(sub_tok, "tokenizer_adds_end_token", False):
end = -1
start = getattr(sub_tok, "tokens_start", 0)
# HF Tokenizer call
hf_tokens = sub_tok.tokenizer(word)["input_ids"]
# Slice
sliced = hf_tokens[start:end] if end != 999999999999 else hf_tokens[start:]
tokens.append([(t, weight) for t in sliced])
# 3. Batching (replicating SDTokenizer logic)
batched_tokens = []
batch = []
start_token = getattr(sub_tok, "start_token", None)
end_token = getattr(sub_tok, "end_token", None)
pad_token = getattr(sub_tok, "pad_token", 0)
max_length = getattr(sub_tok, "max_length", 77)
max_word_length = getattr(sub_tok, "max_word_length", 8)
pad_to_max_length = getattr(sub_tok, "pad_to_max_length", True)
min_length = getattr(sub_tok, "min_length", None)
min_padding = getattr(sub_tok, "min_padding", None)
pad_left = getattr(sub_tok, "pad_left", False) # Mistral uses left padding!
if start_token is not None:
batch.append((start_token, 1.0, 0))
for i, t_group in enumerate(tokens):
is_large = len(t_group) >= max_word_length
has_end_token = 1 if end_token is not None else 0
while len(t_group) > 0:
if len(t_group) + len(batch) > max_length - has_end_token:
remaining_length = max_length - len(batch) - has_end_token
if is_large:
batch.extend([(t,w,i+1) for t,w in t_group[:remaining_length]])
if end_token is not None:
batch.append((end_token, 1.0, 0))
t_group = t_group[remaining_length:]
else:
if end_token is not None:
batch.append((end_token, 1.0, 0))
# Padding
if pad_to_max_length:
pad_amt = max_length - len(batch)
if pad_left:
for _ in range(pad_amt): batch.insert(0, (pad_token, 1.0, 0))
else:
batch.extend([(pad_token, 1.0, 0)] * pad_amt)
batched_tokens.append(batch)
batch = []
if start_token is not None:
batch.append((start_token, 1.0, 0))
else:
batch.extend([(t,w,i+1) for t,w in t_group])
t_group = []
# Last batch
if end_token is not None:
batch.append((end_token, 1.0, 0))
# Min padding
if min_padding is not None:
pad_amt = min_padding
if pad_left:
for _ in range(pad_amt): batch.insert(0, (pad_token, 1.0, 0))
else:
batch.extend([(pad_token, 1.0, 0)] * pad_amt)
# Pad to max
if pad_to_max_length and len(batch) < max_length:
pad_amt = max_length - len(batch)
if pad_left:
for _ in range(pad_amt): batch.insert(0, (pad_token, 1.0, 0))
else:
batch.extend([(pad_token, 1.0, 0)] * pad_amt)
# Min length
if min_length is not None and len(batch) < min_length:
pad_amt = min_length - len(batch)
if pad_left:
for _ in range(pad_amt): batch.insert(0, (pad_token, 1.0, 0))
else:
batch.extend([(pad_token, 1.0, 0)] * pad_amt)
batched_tokens.append(batch)
# Map key
key = attr
if attr.startswith("clip_") and len(attr) == 6:
key = attr[5:]
out[key] = batched_tokens
found_any = True
except Exception as e:
print(f"EmbeddingToolkit: Error manually tokenizing for '{attr}': {e}")
import traceback
traceback.print_exc()
if not found_any:
return clip.tokenize(text, llama_template="{}")
return out
class SaveTokenEmbeddings:
def __init__(self):
@@ -33,11 +223,19 @@ class SaveTokenEmbeddings:
print("SaveTokenEmbeddings: Error: clip.cond_stage_model is None.")
return {"ui": {"text": ["Error: CLIP's cond_stage_model is None."]}}
tokenized_text_with_weights = clip.tokenize(text)
tokenized_text_with_weights = clip.tokenize(text, llama_template="{}")
actual_clip_model_wrapper = clip.cond_stage_model
sd_clip_instances = {}
potential_clip_parts = {'l': 'clip_l', 'g': 'clip_g', 'pile_t5xl': 'pile_t5xl', 't5xl': 't5xl', 't5xxl': 't5xxl', 'umt5xxl': 'umt5xxl', 't5base': 't5base', 'qwen25_7b': 'qwen25_7b', 'qwen3_4b': 'qwen3_4b', 'qwen3_8b': 'qwen3_8b', 'mistral3_24b': 'mistral3_24b', 'gemma2_2b': 'gemma2_2b', 'gemma3_4b': 'gemma3_4b'}
potential_clip_parts = {
'l': 'clip_l', 'g': 'clip_g', 'h': 'clip_h',
'pile_t5xl': 'pile_t5xl', 't5xl': 't5xl', 't5xxl': 't5xxl', 'umt5xxl': 'umt5xxl', 't5base': 't5base',
'qwen25_7b': 'qwen25_7b', 'qwen25_3b': 'qwen25_3b', 'qwen3_4b': 'qwen3_4b', 'qwen3_8b': 'qwen3_8b', 'qwen3_2b': 'qwen3_2b', 'qwen3_06b': 'qwen3_06b',
'mistral3_24b': 'mistral3_24b',
'gemma2_2b': 'gemma2_2b', 'gemma3_4b': 'gemma3_4b', 'gemma3_12b': 'gemma3_12b',
'jina_clip_2': 'jina_clip_2', 'gemma': 'gemma', 'jina': 'jina', 'byt5': 'byt5_small',
'llama': 'llama',
}
for key_suffix, attr_name in potential_clip_parts.items():
if hasattr(actual_clip_model_wrapper, attr_name):
@@ -120,7 +318,7 @@ class SaveTokenEmbeddings:
print(f"SaveTokenEmbeddings: Warning: No embeddings left for '{key_suffix}' after slicing. Skipping.")
continue
if key_suffix.startswith('t5') or key_suffix.startswith('umt5') or key_suffix.startswith('pile') or key_suffix.startswith('qwen') or key_suffix.startswith('mistral') or key_suffix.startswith('gemma'):
if key_suffix.startswith('t5') or key_suffix.startswith('umt5') or key_suffix.startswith('pile') or key_suffix.startswith('qwen') or key_suffix.startswith('mistral') or key_suffix.startswith('gemma') or key_suffix.startswith('jina') or key_suffix.startswith('byt5') or key_suffix.startswith('llama'):
save_key_in_file = key_suffix
else:
save_key_in_file = f"clip_{key_suffix}"
@@ -186,11 +384,20 @@ class SaveWeightedEmbeddings:
print("SaveWeightedEmbeddings: Error: clip.cond_stage_model is None.")
return {"ui": {"text": ["Error: CLIP's cond_stage_model is None."]}}
tokenized_text_with_weights = clip.tokenize(text)
# Use custom tokenizer to ensure we get weights even for Flux/Mistral/Qwen
tokenized_text_with_weights = tokenize_preserving_weights(clip, text)
actual_clip_model_wrapper = clip.cond_stage_model
sd_clip_instances = {}
potential_clip_parts = {'l': 'clip_l', 'g': 'clip_g', 'pile_t5xl': 'pile_t5xl', 't5xl': 't5xl', 't5xxl': 't5xxl', 'umt5xxl': 'umt5xxl', 't5base': 't5base', 'qwen25_7b': 'qwen25_7b', 'qwen3_4b': 'qwen3_4b', 'qwen3_8b': 'qwen3_8b', 'mistral3_24b': 'mistral3_24b', 'gemma2_2b': 'gemma2_2b', 'gemma3_4b': 'gemma3_4b'}
potential_clip_parts = {
'l': 'clip_l', 'g': 'clip_g', 'h': 'clip_h',
'pile_t5xl': 'pile_t5xl', 't5xl': 't5xl', 't5xxl': 't5xxl', 'umt5xxl': 'umt5xxl', 't5base': 't5base',
'qwen25_7b': 'qwen25_7b', 'qwen25_3b': 'qwen25_3b', 'qwen3_4b': 'qwen3_4b', 'qwen3_8b': 'qwen3_8b', 'qwen3_2b': 'qwen3_2b', 'qwen3_06b': 'qwen3_06b',
'mistral3_24b': 'mistral3_24b',
'gemma2_2b': 'gemma2_2b', 'gemma3_4b': 'gemma3_4b', 'gemma3_12b': 'gemma3_12b',
'jina_clip_2': 'jina_clip_2', 'gemma': 'gemma', 'jina': 'jina', 'byt5': 'byt5_small',
'llama': 'llama',
}
for key_suffix, attr_name in potential_clip_parts.items():
if hasattr(actual_clip_model_wrapper, attr_name):
@@ -312,7 +519,7 @@ class SaveWeightedEmbeddings:
print(f"SaveWeightedEmbeddings: Warning: No embeddings left for '{key_suffix}' after slicing. Skipping.")
continue
if key_suffix.startswith('t5') or key_suffix.startswith('umt5') or key_suffix.startswith('pile') or key_suffix.startswith('qwen') or key_suffix.startswith('mistral') or key_suffix.startswith('gemma'):
if key_suffix.startswith('t5') or key_suffix.startswith('umt5') or key_suffix.startswith('pile') or key_suffix.startswith('qwen') or key_suffix.startswith('mistral') or key_suffix.startswith('gemma') or key_suffix.startswith('jina') or key_suffix.startswith('byt5') or key_suffix.startswith('llama'):
save_key_in_file = key_suffix
else:
save_key_in_file = f"clip_{key_suffix}"
@@ -377,11 +584,20 @@ class SaveA1111WeightedEmbeddings:
print("SaveA1111WeightedEmbeddings: Error: clip.cond_stage_model is None.")
return {"ui": {"text": ["Error: CLIP's cond_stage_model is None."]}}
tokenized_text_with_weights = clip.tokenize(text)
# Use custom tokenizer to ensure we get weights even for Flux/Mistral/Qwen
tokenized_text_with_weights = tokenize_preserving_weights(clip, text)
actual_clip_model_wrapper = clip.cond_stage_model
sd_clip_instances = {}
potential_clip_parts = {'l': 'clip_l', 'g': 'clip_g', 'pile_t5xl': 'pile_t5xl', 't5xl': 't5xl', 't5xxl': 't5xxl', 'umt5xxl': 'umt5xxl', 't5base': 't5base', 'qwen25_7b': 'qwen25_7b', 'qwen3_4b': 'qwen3_4b', 'qwen3_8b': 'qwen3_8b', 'mistral3_24b': 'mistral3_24b', 'gemma2_2b': 'gemma2_2b', 'gemma3_4b': 'gemma3_4b'}
potential_clip_parts = {
'l': 'clip_l', 'g': 'clip_g', 'h': 'clip_h',
'pile_t5xl': 'pile_t5xl', 't5xl': 't5xl', 't5xxl': 't5xxl', 'umt5xxl': 'umt5xxl', 't5base': 't5base',
'qwen25_7b': 'qwen25_7b', 'qwen25_3b': 'qwen25_3b', 'qwen3_4b': 'qwen3_4b', 'qwen3_8b': 'qwen3_8b', 'qwen3_2b': 'qwen3_2b', 'qwen3_06b': 'qwen3_06b',
'mistral3_24b': 'mistral3_24b',
'gemma2_2b': 'gemma2_2b', 'gemma3_4b': 'gemma3_4b', 'gemma3_12b': 'gemma3_12b',
'jina_clip_2': 'jina_clip_2', 'gemma': 'gemma', 'jina': 'jina', 'byt5': 'byt5_small',
'llama': 'llama',
}
for key_suffix, attr_name in potential_clip_parts.items():
if hasattr(actual_clip_model_wrapper, attr_name):
@@ -480,7 +696,7 @@ class SaveA1111WeightedEmbeddings:
print(f"SaveA1111WeightedEmbeddings: Warning: No embeddings left for '{key_suffix}' after slicing. Skipping.")
continue
if key_suffix.startswith('t5') or key_suffix.startswith('umt5') or key_suffix.startswith('pile') or key_suffix.startswith('qwen') or key_suffix.startswith('mistral') or key_suffix.startswith('gemma'):
if key_suffix.startswith('t5') or key_suffix.startswith('umt5') or key_suffix.startswith('pile') or key_suffix.startswith('qwen') or key_suffix.startswith('mistral') or key_suffix.startswith('gemma') or key_suffix.startswith('jina') or key_suffix.startswith('byt5') or key_suffix.startswith('llama'):
save_key_in_file = key_suffix
else:
save_key_in_file = f"clip_{key_suffix}"
@@ -660,6 +876,8 @@ NODE_CLASS_MAPPINGS = {
"SaveWeightedEmbeddings": SaveWeightedEmbeddings,
"SaveA1111WeightedEmbeddings": SaveA1111WeightedEmbeddings,
"SliceExistingEmbedding": SliceExistingEmbedding,
"PromptEmbeddingFixer": PromptEmbeddingFixer,
"InspectEmbeddingForClip": InspectEmbeddingForClip,
}
NODE_DISPLAY_NAME_MAPPINGS = {
@@ -667,6 +885,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"SaveWeightedEmbeddings": "Save Weighted Embeddings",
"SaveA1111WeightedEmbeddings": "Save A1111-style Weighted Embeddings",
"SliceExistingEmbedding": "Slice Existing Embedding File",
"PromptEmbeddingFixer": "Prompt Embedding Fixer",
"InspectEmbeddingForClip": "Inspect Embedding For Clip",
}
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
nodes_debug.py
+152
View File
@@ -0,0 +1,152 @@
import torch
import os
import folder_paths
import comfy.utils
import re
class PromptEmbeddingFixer:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"text": ("STRING", {"multiline": True, "dynamicPrompts": True}),
}
}
RETURN_TYPES = ("STRING",)
FUNCTION = "fix_prompt"
CATEGORY = "EmbeddingToolkit/Utils"
def fix_prompt(self, text):
# Add space before embedding if missing
text = re.sub(r'(?<!\s)(embedding:[^\s]+)', r' \1', text)
# Add space after embedding if missing
text = re.sub(r'(embedding:[^\s]+)(?!\s)', r'\1 ', text)
return (text,)
class InspectEmbeddingForClip:
@classmethod
def INPUT_TYPES(cls):
try:
embedding_files = [f for f in folder_paths.get_filename_list("embeddings") if f.lower().endswith((".safetensors", ".pt"))]
except:
embedding_files = ["None"]
return {
"required": {
"clip": ("CLIP",),
"embedding_file": (embedding_files,),
}
}
RETURN_TYPES = ("STRING",)
FUNCTION = "inspect_embedding"
CATEGORY = "EmbeddingToolkit/Debug"
OUTPUT_NODE = True
def inspect_embedding(self, clip, embedding_file):
output_text = []
output_text.append(f"Inspecting File: {embedding_file}")
# Load embedding
full_path = folder_paths.get_full_path("embeddings", embedding_file)
if not full_path:
output_text.append("Error: Embedding file not found.")
return {"ui": {"text": ["\n".join(output_text)]}}
try:
embed_data = comfy.utils.load_torch_file(full_path)
output_text.append(f"Keys found in file: {list(embed_data.keys())}")
except Exception as e:
output_text.append(f"Error loading file: {e}")
return {"ui": {"text": ["\n".join(output_text)]}}
# Inspect CLIP Wrapper
cond_model = clip.cond_stage_model
output_text.append(f"\nCLIP Model Wrapper: {type(cond_model).__name__}")
# Inspect Tokenizer Wrapper
tokenizer = clip.tokenizer
output_text.append(f"Tokenizer Wrapper: {type(tokenizer).__name__}")
if hasattr(tokenizer, "llama_template"):
output_text.append(f" Template: {repr(tokenizer.llama_template)}")
# Check sub-models and sub-tokenizers
potential_parts = [
'clip_l', 'clip_g', 'clip_h',
't5xxl', 't5xl', 't5base',
'qwen25_7b', 'qwen25_3b', 'qwen3_4b', 'qwen3_8b', 'qwen3_2b', 'qwen3_06b',
'mistral3_24b', 'gemma2_2b', 'gemma3_4b', 'gemma3_12b',
'jina_clip_2', 'gemma', 'jina', 'byt5_small', 'llama'
]
# Add dynamic parts from tokenizer
if hasattr(tokenizer, "clip") and isinstance(tokenizer.clip, str):
if tokenizer.clip not in potential_parts:
potential_parts.append(tokenizer.clip)
for part in potential_parts:
# Check Model
if hasattr(cond_model, part):
sub_model = getattr(cond_model, part)
output_text.append(f"\n--- Component: {part} ---")
# Check Tokenizer for this component
if hasattr(tokenizer, part):
sub_tok = getattr(tokenizer, part)
output_text.append(f" Tokenizer: Found")
output_text.append(f" Embedding Dir: {getattr(sub_tok, 'embedding_directory', 'N/A')}")
output_text.append(f" Embedding ID: {repr(getattr(sub_tok, 'embedding_identifier', 'N/A'))}")
else:
output_text.append(f" Tokenizer: NOT FOUND on wrapper")
# Check Model Embedding Size
if hasattr(sub_model, "transformer") and sub_model.transformer is not None:
try:
weight = sub_model.transformer.get_input_embeddings().weight
expected_size = weight.shape[-1]
output_text.append(f" Model Embedding Size: {expected_size}")
# Check compatibility with file
match_found = False
# 1. Check exact key match
if part in embed_data:
emb = embed_data[part]
output_text.append(f" [MATCH] Key '{part}' found in file.")
output_text.append(f" File Shape: {emb.shape}")
if emb.shape[-1] != expected_size:
output_text.append(f" [FAIL] Dimension mismatch! {emb.shape[-1]} vs {expected_size}")
else:
output_text.append(f" [OK] Dimensions match.")
match_found = True
# 2. Check for single tensor fallback
if not match_found and len(embed_data) == 1:
val = list(embed_data.values())[0]
if hasattr(val, "shape") and val.ndim >= 2:
output_text.append(f" [CHECK] Checking single tensor fallback...")
if val.shape[-1] == expected_size:
output_text.append(f" [OK] Fallback match! Shape: {val.shape}")
match_found = True
else:
output_text.append(f" [FAIL] Fallback mismatch. {val.shape[-1]} vs {expected_size}")
if not match_found:
output_text.append(f" [INFO] No compatible embedding found for this component.")
except Exception as e:
output_text.append(f" Error inspecting model component: {e}")
else:
output_text.append(" (No transformer found)")
return {"ui": {"text": ["\n".join(output_text)]}}
NODE_CLASS_MAPPINGS = {
"PromptEmbeddingFixer": PromptEmbeddingFixer,
"InspectEmbeddingForClip": InspectEmbeddingForClip,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"PromptEmbeddingFixer": "Prompt Embedding Fixer",
"InspectEmbeddingForClip": "Inspect Embedding For Clip",
}