Merge feature/llm-embeddings
This commit is contained in:
silveroxides
+224
-9
@@ -2,6 +2,195 @@ 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
|
||||
|
||||
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 +222,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 +317,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 +383,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 +518,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 +583,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 +695,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}"
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "ComfyUI_EmbeddingToolkit"
|
||||
description = "Toolkit for creating embeddings for various models in ComfyUI."
|
||||
version = "1.0.3"
|
||||
version = "1.0.4"
|
||||
license = {file = "LICENSE"}
|
||||
|
||||
[project.urls]
|
||||
|
||||
Reference in New Issue
Block a user