restore debug nodes

__init__.py

@@ -5,6 +5,7 @@ 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):
     """
@@ -33,6 +34,17 @@ def tokenize_preserving_weights(clip, text):
     
     found_any = False
     
+    # Try standard tokenization first to get a baseline for all parts
+    # This ensures we don't lose parts that we can't manually tokenize
+    try:
+        # Some tokenizers might not support llama_template, so we try with and without
+        try:
+            out = clip.tokenize(text, llama_template="{}")
+        except TypeError:
+            out = clip.tokenize(text)
+    except Exception:
+        out = {}
+    
     for attr in potential_parts:
         if hasattr(tokenizer, attr):
             sub_tok = getattr(tokenizer, attr)
@@ -187,7 +199,14 @@ def tokenize_preserving_weights(clip, text):
                     import traceback
                     traceback.print_exc()
 
-    if not found_any:
+    # If manual tokenization was attempted and succeeded for at least one part,
+    # 'out' now contains a mix of standard tokenization (for skipped parts)
+    # and manual tokenization (for parts that needed weight preservation).
+    # If no manual tokenization happened, we rely entirely on the initial clip.tokenize result.
+    
+    if not found_any and not out:
+        # If both failed (e.g. clip.tokenize failed earlier and no manual parts found)
+        # Try one last time or return empty
         return clip.tokenize(text, llama_template="{}")
         
     return out
@@ -207,7 +226,6 @@ class SaveTokenEmbeddings:
                 "text": ("STRING", {"multiline": True, "dynamicPrompts": True}),
                 "slice_bos_eos": ("BOOLEAN", {"default": False}),
                 "filename_prefix": ("STRING", {"default": "token_embeds"}),
-                "split_by_model": ("BOOLEAN", {"default": False}),
             }
         }
 
@@ -216,7 +234,7 @@ class SaveTokenEmbeddings:
     OUTPUT_NODE = True
     CATEGORY = "EmbeddingToolkit"
 
-    def save_token_embeddings(self, clip, text, slice_bos_eos, filename_prefix, split_by_model=False):
+    def save_token_embeddings(self, clip, text, slice_bos_eos, filename_prefix):
         if clip is None:
             raise RuntimeError("CLIP input is None.")
         if clip.cond_stage_model is None:
@@ -330,49 +348,28 @@ class SaveTokenEmbeddings:
             print("SaveTokenEmbeddings: Error: No unweighted embeddings generated.")
             return {"ui": {"text": ["Error: No unweighted embeddings generated."]}}
 
-        subfolder_in_prefix, filename_base_orig = os.path.split(filename_prefix)
+        subfolder_in_prefix, filename_base = os.path.split(filename_prefix)
         primary_output_dir = self.output_dir_list[0]
         full_output_folder = os.path.join(primary_output_dir, subfolder_in_prefix)
         if subfolder_in_prefix:
             os.makedirs(full_output_folder, exist_ok=True)
 
-        dicts_to_save = []
-        if split_by_model:
-            for key, tensor in tensors_to_save.items():
-                dicts_to_save.append({key: tensor})
-        else:
-            dicts_to_save.append(tensors_to_save)
-
-        saved_paths = []
         metadata = {}
-
-        for current_tensors in dicts_to_save:
-            keys = sorted(list(current_tensors.keys()))
-            current_filename_base = filename_base_orig
-            if keys:
-                current_filename_base = f"{'_'.join(keys)}_{filename_base_orig}"
-
-            save_filename = f"{current_filename_base}.safetensors"
+        counter = 1
+        max_counter = 99999
+        while True:
+            save_filename = f"{filename_base}_{counter:05}.safetensors"
             save_path = os.path.join(full_output_folder, save_filename)
+            if not os.path.exists(save_path):
+                break
+            counter += 1
+            if counter > max_counter:
+                print(f"SaveTokenEmbeddings: Warning: Max file attempts for {filename_base}.")
+                return {"ui": {"text": [f"Error: Max files for {filename_base}"]}}
 
-            if os.path.exists(save_path):
-                counter = 1
-                max_counter = 99999
-                while True:
-                    save_filename = f"{current_filename_base}_{counter:05}.safetensors"
-                    save_path = os.path.join(full_output_folder, save_filename)
-                    if not os.path.exists(save_path):
-                        break
-                    counter += 1
-                    if counter > max_counter:
-                        print(f"SaveTokenEmbeddings: Warning: Max file attempts for {current_filename_base}.")
-                        return {"ui": {"text": [f"Error: Max files for {current_filename_base}"]}}
-
-            print(f"SaveTokenEmbeddings: Saving unweighted to: {save_path}")
-            comfy.utils.save_torch_file(current_tensors, save_path, metadata=metadata)
-            saved_paths.append(save_path)
-
-        return {"ui": {"text": [f"Saved unweighted to {', '.join(saved_paths)}"]}}
+        print(f"SaveTokenEmbeddings: Saving unweighted to: {save_path}")
+        comfy.utils.save_torch_file(tensors_to_save, save_path, metadata=metadata)
+        return {"ui": {"text": [f"Saved unweighted to {save_path}"]}}
 
 
 class SaveWeightedEmbeddings:
@@ -390,7 +387,6 @@ class SaveWeightedEmbeddings:
                 "text": ("STRING", {"multiline": True, "dynamicPrompts": True}),
                 "slice_bos_eos": ("BOOLEAN", {"default": False}),
                 "filename_prefix": ("STRING", {"default": "weighted_embed"}),
-                "split_by_model": ("BOOLEAN", {"default": False}),
             }
         }
 
@@ -399,7 +395,7 @@ class SaveWeightedEmbeddings:
     OUTPUT_NODE = True
     CATEGORY = "EmbeddingToolkit"
 
-    def save_weighted_embeddings(self, clip, text, slice_bos_eos, filename_prefix, split_by_model=False):
+    def save_weighted_embeddings(self, clip, text, slice_bos_eos, filename_prefix):
         if clip is None:
             raise RuntimeError("CLIP input is None.")
         if clip.cond_stage_model is None:
@@ -553,49 +549,28 @@ class SaveWeightedEmbeddings:
             print("SaveWeightedEmbeddings: Error: No weighted embeddings generated.")
             return {"ui": {"text": ["Error: No weighted embeddings generated."]}}
 
-        subfolder_in_prefix, filename_base_orig = os.path.split(filename_prefix)
+        subfolder_in_prefix, filename_base = os.path.split(filename_prefix)
         primary_output_dir = self.output_dir_list[0]
         full_output_folder = os.path.join(primary_output_dir, subfolder_in_prefix)
         if subfolder_in_prefix:
             os.makedirs(full_output_folder, exist_ok=True)
 
-        dicts_to_save = []
-        if split_by_model:
-            for key, tensor in tensors_to_save.items():
-                dicts_to_save.append({key: tensor})
-        else:
-            dicts_to_save.append(tensors_to_save)
-
-        saved_paths = []
         metadata = {}
-
-        for current_tensors in dicts_to_save:
-            keys = sorted(list(current_tensors.keys()))
-            current_filename_base = filename_base_orig
-            if keys:
-                current_filename_base = f"{'_'.join(keys)}_{filename_base_orig}"
-
-            save_filename = f"{current_filename_base}.safetensors"
+        counter = 1
+        max_counter = 99999
+        while True:
+            save_filename = f"{filename_base}_{counter:05}.safetensors"
             save_path = os.path.join(full_output_folder, save_filename)
+            if not os.path.exists(save_path):
+                break
+            counter += 1
+            if counter > max_counter:
+                print(f"SaveWeightedEmbeddings: Warning: Max file attempts for {filename_base}.")
+                return {"ui": {"text": [f"Error: Max files for {filename_base}"]}}
 
-            if os.path.exists(save_path):
-                counter = 1
-                max_counter = 99999
-                while True:
-                    save_filename = f"{current_filename_base}_{counter:05}.safetensors"
-                    save_path = os.path.join(full_output_folder, save_filename)
-                    if not os.path.exists(save_path):
-                        break
-                    counter += 1
-                    if counter > max_counter:
-                        print(f"SaveWeightedEmbeddings: Warning: Max file attempts for {current_filename_base}.")
-                        return {"ui": {"text": [f"Error: Max files for {current_filename_base}"]}}
-
-            print(f"SaveWeightedEmbeddings: Saving weighted to: {save_path}")
-            comfy.utils.save_torch_file(current_tensors, save_path, metadata=metadata)
-            saved_paths.append(save_path)
-
-        return {"ui": {"text": [f"Saved weighted to {', '.join(saved_paths)}"]}}
+        print(f"SaveWeightedEmbeddings: Saving weighted to: {save_path}")
+        comfy.utils.save_torch_file(tensors_to_save, save_path, metadata=metadata)
+        return {"ui": {"text": [f"Saved weighted to {save_path}"]}}
 
 class SaveA1111WeightedEmbeddings:
     def __init__(self):
@@ -612,7 +587,6 @@ class SaveA1111WeightedEmbeddings:
                 "text": ("STRING", {"multiline": True, "dynamicPrompts": True}),
                 "slice_bos_eos": ("BOOLEAN", {"default": False}),
                 "filename_prefix": ("STRING", {"default": "a1111_weighted_embed"}),
-                "split_by_model": ("BOOLEAN", {"default": False}),
             }
         }
 
@@ -621,7 +595,7 @@ class SaveA1111WeightedEmbeddings:
     OUTPUT_NODE = True
     CATEGORY = "EmbeddingToolkit"
 
-    def save_a1111_weighted_embeddings(self, clip, text, slice_bos_eos, filename_prefix, split_by_model=False):
+    def save_a1111_weighted_embeddings(self, clip, text, slice_bos_eos, filename_prefix):
         if clip is None:
             raise RuntimeError("CLIP input is None.")
         if clip.cond_stage_model is None:
@@ -752,49 +726,28 @@ class SaveA1111WeightedEmbeddings:
             print("SaveA1111WeightedEmbeddings: Error: No A1111-style weighted embeddings generated.")
             return {"ui": {"text": ["Error: No A1111-style weighted embeddings generated."]}}
 
-        subfolder_in_prefix, filename_base_orig = os.path.split(filename_prefix)
+        subfolder_in_prefix, filename_base = os.path.split(filename_prefix)
         primary_output_dir = self.output_dir_list[0]
         full_output_folder = os.path.join(primary_output_dir, subfolder_in_prefix)
         if subfolder_in_prefix:
             os.makedirs(full_output_folder, exist_ok=True)
 
-        dicts_to_save = []
-        if split_by_model:
-            for key, tensor in tensors_to_save.items():
-                dicts_to_save.append({key: tensor})
-        else:
-            dicts_to_save.append(tensors_to_save)
-
-        saved_paths = []
         metadata = {}
-
-        for current_tensors in dicts_to_save:
-            keys = sorted(list(current_tensors.keys()))
-            current_filename_base = filename_base_orig
-            if keys:
-                current_filename_base = f"{'_'.join(keys)}_{filename_base_orig}"
-
-            save_filename = f"{current_filename_base}.safetensors"
+        counter = 1
+        max_counter = 99999
+        while True:
+            save_filename = f"{filename_base}_{counter:05}.safetensors"
             save_path = os.path.join(full_output_folder, save_filename)
+            if not os.path.exists(save_path):
+                break
+            counter += 1
+            if counter > max_counter:
+                print(f"SaveA1111WeightedEmbeddings: Warning: Max file attempts for {filename_base}.")
+                return {"ui": {"text": [f"Error: Max files for {filename_base}"]}}
 
-            if os.path.exists(save_path):
-                counter = 1
-                max_counter = 99999
-                while True:
-                    save_filename = f"{current_filename_base}_{counter:05}.safetensors"
-                    save_path = os.path.join(full_output_folder, save_filename)
-                    if not os.path.exists(save_path):
-                        break
-                    counter += 1
-                    if counter > max_counter:
-                        print(f"SaveA1111WeightedEmbeddings: Warning: Max file attempts for {current_filename_base}.")
-                        return {"ui": {"text": [f"Error: Max files for {current_filename_base}"]}}
-
-            print(f"SaveA1111WeightedEmbeddings: Saving A1111-style weighted to: {save_path}")
-            comfy.utils.save_torch_file(current_tensors, save_path, metadata=metadata)
-            saved_paths.append(save_path)
-
-        return {"ui": {"text": [f"Saved A1111-style weighted to {', '.join(saved_paths)}"]}}
+        print(f"SaveA1111WeightedEmbeddings: Saving A1111-style weighted to: {save_path}")
+        comfy.utils.save_torch_file(tensors_to_save, save_path, metadata=metadata)
+        return {"ui": {"text": [f"Saved A1111-style weighted to {save_path}"]}}
 
 
 class SliceExistingEmbedding:
@@ -900,39 +853,28 @@ class SliceExistingEmbedding:
 
         subfolder_in_prefix, filename_base_from_prefix = os.path.split(output_filename_prefix)
 
-        keys = sorted([k for k in sliced_tensors.keys() if k != "__metadata__"])
-        if keys:
-            filename_base_from_prefix = f"{'_'.join(keys)}_{filename_base_from_prefix}"
-
         output_target_folder = os.path.join(self.primary_embeddings_dir, subfolder_in_prefix)
         os.makedirs(output_target_folder, exist_ok=True)
 
         _, input_ext = os.path.splitext(embedding_file)
         output_ext = input_ext if input_ext.lower() in [".safetensors", ".pt"] else ".safetensors"
 
+        counter = 1
+        max_counter = 99999
         final_save_path = ""
+        while True:
+            save_filename_numbered = f"{filename_base_from_prefix}_{counter:05}{output_ext}"
+            current_save_path_candidate = os.path.join(output_target_folder, save_filename_numbered)
 
-        save_filename_numbered = f"{filename_base_from_prefix}{output_ext}"
-        current_save_path_candidate = os.path.join(output_target_folder, save_filename_numbered)
+            if not os.path.exists(current_save_path_candidate):
+                final_save_path = current_save_path_candidate
+                break
 
-        if not os.path.exists(current_save_path_candidate):
-            final_save_path = current_save_path_candidate
-        else:
-            counter = 1
-            max_counter = 99999
-            while True:
-                save_filename_numbered = f"{filename_base_from_prefix}_{counter:05}{output_ext}"
-                current_save_path_candidate = os.path.join(output_target_folder, save_filename_numbered)
-
-                if not os.path.exists(current_save_path_candidate):
-                    final_save_path = current_save_path_candidate
-                    break
-
-                counter += 1
-                if counter > max_counter:
-                    msg = f"Max file attempts for {output_filename_prefix}. Please check your embeddings folder or prefix."
-                    print(f"SliceExistingEmbedding: {msg}")
-                    return {"ui": {"text": [f"Error: {msg}"]}}
+            counter += 1
+            if counter > max_counter:
+                msg = f"Max file attempts for {output_filename_prefix}. Please check your embeddings folder or prefix."
+                print(f"SliceExistingEmbedding: {msg}")
+                return {"ui": {"text": [f"Error: {msg}"]}}
 
         try:
             if file_metadata:
@@ -952,6 +894,8 @@ NODE_CLASS_MAPPINGS = {
     "SaveWeightedEmbeddings": SaveWeightedEmbeddings,
     "SaveA1111WeightedEmbeddings": SaveA1111WeightedEmbeddings,
     "SliceExistingEmbedding": SliceExistingEmbedding,
+    "PromptEmbeddingFixer": PromptEmbeddingFixer,
+    "InspectEmbeddingForClip": InspectEmbeddingForClip,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -959,6 +903,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

@@ -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",
+}