diff --git a/__init__.py b/__init__.py index 2d50000..17ed4a4 100644 --- a/__init__.py +++ b/__init__.py @@ -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}" diff --git a/pyproject.toml b/pyproject.toml index e996c6e..7bd8c8a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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]