Merge pull request #4 from silveroxides/fix/better-filenames

Fix/better filenames

__init__.py

@@ -5,7 +5,6 @@ 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):
     """
@@ -208,6 +207,7 @@ 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 +216,7 @@ class SaveTokenEmbeddings:
     OUTPUT_NODE = True
     CATEGORY = "EmbeddingToolkit"
 
-    def save_token_embeddings(self, clip, text, slice_bos_eos, filename_prefix):
+    def save_token_embeddings(self, clip, text, slice_bos_eos, filename_prefix, split_by_model=False):
         if clip is None:
             raise RuntimeError("CLIP input is None.")
         if clip.cond_stage_model is None:
@@ -330,28 +330,49 @@ class SaveTokenEmbeddings:
             print("SaveTokenEmbeddings: Error: No unweighted embeddings generated.")
             return {"ui": {"text": ["Error: No unweighted embeddings generated."]}}
 
-        subfolder_in_prefix, filename_base = os.path.split(filename_prefix)
+        subfolder_in_prefix, filename_base_orig = 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)
 
-        metadata = {}
-        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}"]}}
+        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)
 
-        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}"]}}
+        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"
+            save_path = os.path.join(full_output_folder, save_filename)
+
+            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)}"]}}
 
 
 class SaveWeightedEmbeddings:
@@ -369,6 +390,7 @@ 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}),
             }
         }
 
@@ -377,7 +399,7 @@ class SaveWeightedEmbeddings:
     OUTPUT_NODE = True
     CATEGORY = "EmbeddingToolkit"
 
-    def save_weighted_embeddings(self, clip, text, slice_bos_eos, filename_prefix):
+    def save_weighted_embeddings(self, clip, text, slice_bos_eos, filename_prefix, split_by_model=False):
         if clip is None:
             raise RuntimeError("CLIP input is None.")
         if clip.cond_stage_model is None:
@@ -531,28 +553,49 @@ class SaveWeightedEmbeddings:
             print("SaveWeightedEmbeddings: Error: No weighted embeddings generated.")
             return {"ui": {"text": ["Error: No weighted embeddings generated."]}}
 
-        subfolder_in_prefix, filename_base = os.path.split(filename_prefix)
+        subfolder_in_prefix, filename_base_orig = 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)
 
-        metadata = {}
-        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}"]}}
+        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)
 
-        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}"]}}
+        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"
+            save_path = os.path.join(full_output_folder, save_filename)
+
+            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)}"]}}
 
 class SaveA1111WeightedEmbeddings:
     def __init__(self):
@@ -569,6 +612,7 @@ 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}),
             }
         }
 
@@ -577,7 +621,7 @@ class SaveA1111WeightedEmbeddings:
     OUTPUT_NODE = True
     CATEGORY = "EmbeddingToolkit"
 
-    def save_a1111_weighted_embeddings(self, clip, text, slice_bos_eos, filename_prefix):
+    def save_a1111_weighted_embeddings(self, clip, text, slice_bos_eos, filename_prefix, split_by_model=False):
         if clip is None:
             raise RuntimeError("CLIP input is None.")
         if clip.cond_stage_model is None:
@@ -708,28 +752,49 @@ 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 = os.path.split(filename_prefix)
+        subfolder_in_prefix, filename_base_orig = 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)
 
-        metadata = {}
-        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}"]}}
+        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)
 
-        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}"]}}
+        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"
+            save_path = os.path.join(full_output_folder, save_filename)
+
+            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)}"]}}
 
 
 class SliceExistingEmbedding:
@@ -835,28 +900,39 @@ 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)
 
-            if not os.path.exists(current_save_path_candidate):
-                final_save_path = current_save_path_candidate
-                break
+        save_filename_numbered = f"{filename_base_from_prefix}{output_ext}"
+        current_save_path_candidate = os.path.join(output_target_folder, save_filename_numbered)
 
-            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}"]}}
+        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}"]}}
 
         try:
             if file_metadata:
@@ -876,8 +952,6 @@ NODE_CLASS_MAPPINGS = {
     "SaveWeightedEmbeddings": SaveWeightedEmbeddings,
     "SaveA1111WeightedEmbeddings": SaveA1111WeightedEmbeddings,
     "SliceExistingEmbedding": SliceExistingEmbedding,
-    "PromptEmbeddingFixer": PromptEmbeddingFixer,
-    "InspectEmbeddingForClip": InspectEmbeddingForClip,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -885,8 +959,6 @@ 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

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

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]