import torch import os import folder_paths import comfy.utils class SaveTokenEmbeddings: def __init__(self): self.output_dir_list = folder_paths.get_folder_paths("embeddings") if not self.output_dir_list: print("SaveTokenEmbeddings: Warning: 'embeddings' folder type not found. Falling back to main output directory.") self.output_dir_list = [folder_paths.get_output_directory()] @classmethod def INPUT_TYPES(cls): return { "required": { "clip": ("CLIP",), "text": ("STRING", {"multiline": True, "dynamicPrompts": True}), "filename_prefix": ("STRING", {"default": "token_embeds"}), } } RETURN_TYPES = () FUNCTION = "save_token_embeddings" OUTPUT_NODE = True CATEGORY = "EmbeddingToolkit" def save_token_embeddings(self, clip, text, filename_prefix): if clip is None: raise RuntimeError("CLIP input is None.") if clip.cond_stage_model is None: 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) 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'} for key_suffix, attr_name in potential_clip_parts.items(): if hasattr(actual_clip_model_wrapper, attr_name): model_instance = getattr(actual_clip_model_wrapper, attr_name) if model_instance is not None and hasattr(model_instance, "process_tokens") and hasattr(model_instance, "transformer"): sd_clip_instances[key_suffix] = model_instance else: if model_instance is None: print(f"SaveTokenEmbeddings: Warning: Attribute '{attr_name}' is None. Skipping '{key_suffix}'.") else: print(f"SaveTokenEmbeddings: Warning: Attribute '{attr_name}' (type {type(model_instance)}) is not a valid SDClipModel. Skipping '{key_suffix}'.") if not sd_clip_instances and hasattr(actual_clip_model_wrapper, "process_tokens") and hasattr(actual_clip_model_wrapper, "transformer"): print(f"SaveTokenEmbeddings: Warning: No standard attributes. Fallback: cond_stage_model as 'l'.") sd_clip_instances['l'] = actual_clip_model_wrapper if not sd_clip_instances: print(f"SaveTokenEmbeddings: Error: No valid SDClipModel instances found.") return {"ui": {"text": ["Error: No valid SDClipModel instances found."]}} tensors_to_save = {} for key_suffix, sd_clip_model_inst in sd_clip_instances.items(): current_token_segments_with_weights = None if isinstance(tokenized_text_with_weights, dict): current_token_segments_with_weights = tokenized_text_with_weights.get(key_suffix) elif key_suffix == 'l' and not isinstance(tokenized_text_with_weights, dict): current_token_segments_with_weights = tokenized_text_with_weights if not current_token_segments_with_weights: print(f"SaveTokenEmbeddings: Warning: No token data for '{key_suffix}'. Skipping.") continue if not hasattr(sd_clip_model_inst, 'transformer') or sd_clip_model_inst.transformer is None: print(f"SaveTokenEmbeddings: Error: SDClipModel for '{key_suffix}' has invalid 'transformer'. Skipping.") continue device = sd_clip_model_inst.transformer.get_input_embeddings().weight.device tokens_for_process_tokens_method = [] for segment_with_weights in current_token_segments_with_weights: tokens_for_process_tokens_method.append([item[0] for item in segment_with_weights]) if not tokens_for_process_tokens_method: print(f"SaveTokenEmbeddings: Warning: No token ID segments for '{key_suffix}'. Skipping.") continue embeds_out_full, attention_mask_full, num_tokens, embeds_info = sd_clip_model_inst.process_tokens( tokens_for_process_tokens_method, device ) print(f"Returned values: num_tokens={num_tokens}, embeds_info={embeds_info}") all_actual_embeddings_for_this_clip_part = [] num_segments = embeds_out_full.shape[0] for s in range(num_segments): current_segment_embeds = embeds_out_full[s] current_segment_mask = attention_mask_full[s] actual_embeddings_this_segment = current_segment_embeds[current_segment_mask == 1] if actual_embeddings_this_segment.shape[0] > 0: all_actual_embeddings_for_this_clip_part.append(actual_embeddings_this_segment) if not all_actual_embeddings_for_this_clip_part: print(f"SaveTokenEmbeddings: Warning: No actual tokens for '{key_suffix}'. Skipping.") continue final_embeddings_for_this_clip_part = torch.cat(all_actual_embeddings_for_this_clip_part, dim=0) num_tokens_in_final = final_embeddings_for_this_clip_part.shape[0] if key_suffix == 'l' or key_suffix == 'g': if num_tokens_in_final >= 2: final_embeddings_for_this_clip_part = final_embeddings_for_this_clip_part[1:-1] elif num_tokens_in_final > 0: final_embeddings_for_this_clip_part = final_embeddings_for_this_clip_part[0:0] elif key_suffix.startswith('t5') or key_suffix.startswith('umt5') or key_suffix.startswith('pile'): if num_tokens_in_final >= 1: final_embeddings_for_this_clip_part = final_embeddings_for_this_clip_part[:-1] if final_embeddings_for_this_clip_part.shape[0] == 0: 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'): save_key_in_file = key_suffix else: save_key_in_file = f"clip_{key_suffix}" tensors_to_save[save_key_in_file] = final_embeddings_for_this_clip_part.cpu() print(f"SaveTokenEmbeddings: Processed unweighted for '{save_key_in_file}': shape {final_embeddings_for_this_clip_part.shape}") if not tensors_to_save: 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) 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}"]}} 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: def __init__(self): self.output_dir_list = folder_paths.get_folder_paths("embeddings") if not self.output_dir_list: print("SaveWeightedEmbeddings: Warning: 'embeddings' folder type not found. Falling back to main output directory.") self.output_dir_list = [folder_paths.get_output_directory()] @classmethod def INPUT_TYPES(cls): return { "required": { "clip": ("CLIP",), "text": ("STRING", {"multiline": True, "dynamicPrompts": True}), "filename_prefix": ("STRING", {"default": "weighted_embed"}), } } RETURN_TYPES = () FUNCTION = "save_weighted_embeddings" OUTPUT_NODE = True CATEGORY = "EmbeddingToolkit" def save_weighted_embeddings(self, clip, text, filename_prefix): if clip is None: raise RuntimeError("CLIP input is None.") if clip.cond_stage_model is None: 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) 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'} for key_suffix, attr_name in potential_clip_parts.items(): if hasattr(actual_clip_model_wrapper, attr_name): model_instance = getattr(actual_clip_model_wrapper, attr_name) if model_instance is not None and hasattr(model_instance, "process_tokens") and hasattr(model_instance, "transformer"): sd_clip_instances[key_suffix] = model_instance else: if model_instance is None: print(f"SaveWeightedEmbeddings: Warning: Attribute '{attr_name}' is None. Skipping '{key_suffix}'.") else: print(f"SaveWeightedEmbeddings: Warning: Attribute '{attr_name}' (type {type(model_instance)}) is not a valid SDClipModel. Skipping '{key_suffix}'.") if not sd_clip_instances and hasattr(actual_clip_model_wrapper, "process_tokens") and hasattr(actual_clip_model_wrapper, "transformer"): print(f"SaveWeightedEmbeddings: Warning: No standard attributes. Fallback: cond_stage_model as 'l'.") sd_clip_instances['l'] = actual_clip_model_wrapper if not sd_clip_instances: print(f"SaveWeightedEmbeddings: Error: No valid SDClipModel instances found.") return {"ui": {"text": ["Error: No valid SDClipModel instances found."]}} tensors_to_save = {} for key_suffix, sd_clip_model_inst in sd_clip_instances.items(): current_prompt_segments_with_weights = None if isinstance(tokenized_text_with_weights, dict): current_prompt_segments_with_weights = tokenized_text_with_weights.get(key_suffix) elif key_suffix == 'l' and not isinstance(tokenized_text_with_weights, dict): current_prompt_segments_with_weights = tokenized_text_with_weights if not current_prompt_segments_with_weights: print(f"SaveWeightedEmbeddings: Warning: No token data for '{key_suffix}'. Skipping.") continue if not hasattr(sd_clip_model_inst, 'transformer') or sd_clip_model_inst.transformer is None: print(f"SaveWeightedEmbeddings: Error: SDClipModel for '{key_suffix}' has invalid 'transformer'. Skipping.") continue device = sd_clip_model_inst.transformer.get_input_embeddings().weight.device prompt_token_id_segments = [] max_segment_len = 0 for segment_w_weights in current_prompt_segments_with_weights: prompt_token_id_segments.append([item[0] for item in segment_w_weights]) max_segment_len = max(max_segment_len, len(segment_w_weights)) if not prompt_token_id_segments: print(f"SaveWeightedEmbeddings: Warning: No token ID segments for '{key_suffix}'. Skipping.") continue empty_token_sequence = [] if hasattr(sd_clip_model_inst, "gen_empty_tokens") and hasattr(sd_clip_model_inst, "special_tokens"): empty_token_sequence = sd_clip_model_inst.gen_empty_tokens(sd_clip_model_inst.special_tokens, max_segment_len) else: pad_token_id = None if hasattr(sd_clip_model_inst, "special_tokens") and "pad" in sd_clip_model_inst.special_tokens: pad_token_id = sd_clip_model_inst.special_tokens["pad"] if pad_token_id is not None: print(f"SaveWeightedEmbeddings: Warning: '{key_suffix}' SDClipModel instance missing 'gen_empty_tokens' or 'special_tokens'. Using simple pad token sequence for empty.") empty_token_sequence = [pad_token_id] * max_segment_len else: print(f"SaveWeightedEmbeddings: Error: Cannot determine pad token for '{key_suffix}'. Skipping weighting for this part.") continue all_token_sequences_for_process_tokens = prompt_token_id_segments + [empty_token_sequence] embeds_combined, attention_mask_combined, num_tokens, embeds_info = sd_clip_model_inst.process_tokens( all_token_sequences_for_process_tokens, device ) print(f"Returned values: num_tokens={num_tokens}, embeds_info={embeds_info}") num_prompt_segments = len(prompt_token_id_segments) prompt_embeds_full_segments = embeds_combined[:num_prompt_segments] embeds_empty_segment_full = embeds_combined[num_prompt_segments:] prompt_attention_mask_full_segments = attention_mask_combined[:num_prompt_segments] all_weighted_actual_embeddings_for_this_clip_part = [] for s_idx in range(num_prompt_segments): current_segment_raw_embeds = prompt_embeds_full_segments[s_idx] segment_data_from_tokenizer = current_prompt_segments_with_weights[s_idx] current_segment_mask_from_process_tokens = prompt_attention_mask_full_segments[s_idx] weighted_segment_embeds = torch.zeros_like(current_segment_raw_embeds) for token_idx in range(current_segment_raw_embeds.shape[0]): is_active_token = current_segment_mask_from_process_tokens[token_idx] == 1 original_embedding = current_segment_raw_embeds[token_idx] embed_empty_for_this_pos = embeds_empty_segment_full[0, token_idx] if is_active_token: weight = segment_data_from_tokenizer[token_idx][1] if token_idx < len(segment_data_from_tokenizer) else 1.0 if weight == 1.0: weighted_segment_embeds[token_idx] = original_embedding else: weighted_segment_embeds[token_idx] = (original_embedding - embed_empty_for_this_pos) * weight + embed_empty_for_this_pos else: weighted_segment_embeds[token_idx] = original_embedding actual_weighted_embeddings_this_segment = weighted_segment_embeds[current_segment_mask_from_process_tokens == 1] if actual_weighted_embeddings_this_segment.shape[0] > 0: all_weighted_actual_embeddings_for_this_clip_part.append(actual_weighted_embeddings_this_segment) if not all_weighted_actual_embeddings_for_this_clip_part: print(f"SaveWeightedEmbeddings: Warning: No actual weighted tokens for '{key_suffix}'. Skipping.") continue final_embeddings_for_this_clip_part = torch.cat(all_weighted_actual_embeddings_for_this_clip_part, dim=0) num_tokens_in_final = final_embeddings_for_this_clip_part.shape[0] if key_suffix == 'l' or key_suffix == 'g': if num_tokens_in_final >= 2: final_embeddings_for_this_clip_part = final_embeddings_for_this_clip_part[1:-1] elif num_tokens_in_final > 0: final_embeddings_for_this_clip_part = final_embeddings_for_this_clip_part[0:0] elif key_suffix.startswith('t5') or key_suffix.startswith('umt5') or key_suffix.startswith('pile'): if num_tokens_in_final >= 1: final_embeddings_for_this_clip_part = final_embeddings_for_this_clip_part[:-1] if final_embeddings_for_this_clip_part.shape[0] == 0: 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'): save_key_in_file = key_suffix else: save_key_in_file = f"clip_{key_suffix}" tensors_to_save[save_key_in_file] = final_embeddings_for_this_clip_part.cpu() print(f"SaveWeightedEmbeddings: Processed weighted for '{save_key_in_file}': shape {final_embeddings_for_this_clip_part.shape}") if not tensors_to_save: 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) 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}"]}} 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): self.output_dir_list = folder_paths.get_folder_paths("embeddings") if not self.output_dir_list: print("SaveA1111WeightedEmbeddings: Warning: 'embeddings' folder type not found. Falling back to main output directory.") self.output_dir_list = [folder_paths.get_output_directory()] @classmethod def INPUT_TYPES(cls): return { "required": { "clip": ("CLIP",), "text": ("STRING", {"multiline": True, "dynamicPrompts": True}), "filename_prefix": ("STRING", {"default": "a1111_weighted_embed"}), } } RETURN_TYPES = () FUNCTION = "save_a1111_weighted_embeddings" OUTPUT_NODE = True CATEGORY = "EmbeddingToolkit" def save_a1111_weighted_embeddings(self, clip, text, filename_prefix): if clip is None: raise RuntimeError("CLIP input is None.") if clip.cond_stage_model is None: 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) 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'} for key_suffix, attr_name in potential_clip_parts.items(): if hasattr(actual_clip_model_wrapper, attr_name): model_instance = getattr(actual_clip_model_wrapper, attr_name) if model_instance is not None and hasattr(model_instance, "process_tokens") and hasattr(model_instance, "transformer"): sd_clip_instances[key_suffix] = model_instance else: if model_instance is None: print(f"SaveA1111WeightedEmbeddings: Warning: Attribute '{attr_name}' is None. Skipping '{key_suffix}'.") else: print(f"SaveA1111WeightedEmbeddings: Warning: Attribute '{attr_name}' (type {type(model_instance)}) is not a valid SDClipModel. Skipping '{key_suffix}'.") if not sd_clip_instances and hasattr(actual_clip_model_wrapper, "process_tokens") and hasattr(actual_clip_model_wrapper, "transformer"): print(f"SaveA1111WeightedEmbeddings: Warning: No standard attributes. Fallback: cond_stage_model as 'l'.") sd_clip_instances['l'] = actual_clip_model_wrapper if not sd_clip_instances: print(f"SaveA1111WeightedEmbeddings: Error: No valid SDClipModel instances found.") return {"ui": {"text": ["Error: No valid SDClipModel instances found."]}} tensors_to_save = {} for key_suffix, sd_clip_model_inst in sd_clip_instances.items(): current_prompt_segments_with_weights = None if isinstance(tokenized_text_with_weights, dict): current_prompt_segments_with_weights = tokenized_text_with_weights.get(key_suffix) elif key_suffix == 'l' and not isinstance(tokenized_text_with_weights, dict): current_prompt_segments_with_weights = tokenized_text_with_weights if not current_prompt_segments_with_weights: print(f"SaveA1111WeightedEmbeddings: Warning: No token data for '{key_suffix}'. Skipping.") continue if not hasattr(sd_clip_model_inst, 'transformer') or sd_clip_model_inst.transformer is None: print(f"SaveA1111WeightedEmbeddings: Error: SDClipModel for '{key_suffix}' has invalid 'transformer'. Skipping.") continue device = sd_clip_model_inst.transformer.get_input_embeddings().weight.device prompt_token_id_segments = [] for segment_w_weights in current_prompt_segments_with_weights: prompt_token_id_segments.append([item[0] for item in segment_w_weights]) if not prompt_token_id_segments: print(f"SaveA1111WeightedEmbeddings: Warning: No token ID segments for '{key_suffix}'. Skipping.") continue embeds_out_full, attention_mask_full, num_tokens, embeds_info = sd_clip_model_inst.process_tokens( prompt_token_id_segments, device ) print(f"Returned values: num_tokens={num_tokens}, embeds_info={embeds_info}") all_scaled_actual_embeddings_for_this_clip_part = [] num_prompt_segments = embeds_out_full.shape[0] for s_idx in range(num_prompt_segments): current_segment_raw_embeds = embeds_out_full[s_idx] segment_data_from_tokenizer = current_prompt_segments_with_weights[s_idx] current_segment_mask_from_process_tokens = attention_mask_full[s_idx] scaled_segment_embeds = torch.zeros_like(current_segment_raw_embeds) for token_idx in range(current_segment_raw_embeds.shape[0]): is_active_token = current_segment_mask_from_process_tokens[token_idx] == 1 original_embedding = current_segment_raw_embeds[token_idx] if is_active_token: weight = segment_data_from_tokenizer[token_idx][1] if token_idx < len(segment_data_from_tokenizer) else 1.0 scaled_segment_embeds[token_idx] = original_embedding * weight else: scaled_segment_embeds[token_idx] = original_embedding actual_scaled_embeddings_this_segment = scaled_segment_embeds[current_segment_mask_from_process_tokens == 1] if actual_scaled_embeddings_this_segment.shape[0] > 0: all_scaled_actual_embeddings_for_this_clip_part.append(actual_scaled_embeddings_this_segment) if not all_scaled_actual_embeddings_for_this_clip_part: print(f"SaveA1111WeightedEmbeddings: Warning: No actual scaled tokens found for '{key_suffix}' after processing. Skipping.") continue final_embeddings_for_this_clip_part = torch.cat(all_scaled_actual_embeddings_for_this_clip_part, dim=0) num_tokens_in_final = final_embeddings_for_this_clip_part.shape[0] if key_suffix == 'l' or key_suffix == 'g': if num_tokens_in_final >= 2: final_embeddings_for_this_clip_part = final_embeddings_for_this_clip_part[1:-1] elif num_tokens_in_final > 0: final_embeddings_for_this_clip_part = final_embeddings_for_this_clip_part[0:0] elif key_suffix.startswith('t5') or key_suffix.startswith('umt5') or key_suffix.startswith('pile'): if num_tokens_in_final >= 1: final_embeddings_for_this_clip_part = final_embeddings_for_this_clip_part[:-1] if final_embeddings_for_this_clip_part.shape[0] == 0: 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'): save_key_in_file = key_suffix else: save_key_in_file = f"clip_{key_suffix}" tensors_to_save[save_key_in_file] = final_embeddings_for_this_clip_part.cpu() print(f"SaveA1111WeightedEmbeddings: Processed A1111-style weighted for '{save_key_in_file}': shape {final_embeddings_for_this_clip_part.shape}") if not tensors_to_save: 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) 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}"]}} 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: def __init__(self): self.embeddings_dir_paths = folder_paths.get_folder_paths("embeddings") if not self.embeddings_dir_paths: print("SliceExistingEmbedding: Warning: 'embeddings' folder type not found. Falling back to main output directory.") self.embeddings_dir_paths = [folder_paths.get_output_directory()] self.primary_embeddings_dir = self.embeddings_dir_paths[0] os.makedirs(self.primary_embeddings_dir, exist_ok=True) @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 Exception as e: print(f"SliceExistingEmbedding: Error listing embedding files: {e}") embedding_files = ["None"] if not embedding_files: embedding_files = ["None"] return { "required": { "embedding_file": (embedding_files,), "output_filename_prefix": ("STRING", {"default": "sliced_embedding"}), } } RETURN_TYPES = () FUNCTION = "slice_and_save" OUTPUT_NODE = True CATEGORY = "EmbeddingToolkit/Utils" def slice_and_save(self, embedding_file, output_filename_prefix): if embedding_file == "None" or not embedding_file: print("SliceExistingEmbedding: No embedding file selected.") return {"ui": {"text": ["No embedding file selected."]}} full_file_path = folder_paths.get_full_path("embeddings", embedding_file) if full_file_path is None or not os.path.exists(full_file_path): potential_path = os.path.join(self.primary_embeddings_dir, embedding_file) if os.path.exists(potential_path): full_file_path = potential_path else: print(f"SliceExistingEmbedding: Error: File not found: {embedding_file}") return {"ui": {"text": [f"Error: File not found: {embedding_file}"]}} try: loaded_data = comfy.utils.load_torch_file(full_file_path) metadata_key = "__metadata__" file_metadata = {} if isinstance(loaded_data, dict) and metadata_key in loaded_data: file_metadata = loaded_data.pop(metadata_key) except Exception as e: print(f"SliceExistingEmbedding: Error loading {embedding_file}: {e}") return {"ui": {"text": [f"Error loading {embedding_file}: {e}"]}} if not isinstance(loaded_data, dict): print(f"SliceExistingEmbedding: Error: Loaded data from {embedding_file} is not a dictionary (unexpected format).") return {"ui": {"text": [f"Error: Unexpected format in {embedding_file}."]}} sliced_tensors = {} any_sliced = False for key, tensor in loaded_data.items(): if not isinstance(tensor, torch.Tensor) or tensor.ndim == 0: sliced_tensors[key] = tensor continue original_shape = tensor.shape num_tokens = original_shape[0] if tensor.ndim > 0 else 0 processed_tensor = tensor if num_tokens == 0: print(f"SliceExistingEmbedding: Tensor for key '{key}' is already empty (shape {original_shape}). Keeping as is.") elif key.startswith("clip_l") or key.startswith("clip_g"): print(f"SliceExistingEmbedding: Applying CLIP-style slicing (BOS & EOS) to key '{key}' (shape {original_shape})") if num_tokens >= 2: processed_tensor = tensor[1:-1] any_sliced = True elif num_tokens > 0 : processed_tensor = tensor[0:0] any_sliced = True print(f" Resulting shape: {processed_tensor.shape}") elif key.startswith("t5") or key.startswith("umt5") or key.startswith("pile"): print(f"SliceExistingEmbedding: Applying T5-style slicing (EOS) to key '{key}' (shape {original_shape})") if num_tokens >= 1: processed_tensor = tensor[:-1] any_sliced = True print(f" Resulting shape: {processed_tensor.shape}") else: print(f"SliceExistingEmbedding: Info: Key '{key}' does not match known slicing patterns. Tensor will not be sliced.") sliced_tensors[key] = processed_tensor if not any_sliced: msg = f"No changes made to {embedding_file} (no applicable slicing rules or tensors already short/empty)." print(f"SliceExistingEmbedding: {msg}") return {"ui": {"text": [msg]}} subfolder_in_prefix, filename_base_from_prefix = os.path.split(output_filename_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 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: sliced_tensors[metadata_key] = file_metadata comfy.utils.save_torch_file(sliced_tensors, final_save_path) msg = f"Saved sliced embedding to {final_save_path}" print(f"SliceExistingEmbedding: {msg}") return {"ui": {"text": [msg]}} except Exception as e: msg = f"Error saving to {final_save_path}: {e}" print(f"SliceExistingEmbedding: {msg}") return {"ui": {"text": [f"Error saving to {final_save_path}: {e}"]}} NODE_CLASS_MAPPINGS = { "SaveTokenEmbeddings": SaveTokenEmbeddings, "SaveWeightedEmbeddings": SaveWeightedEmbeddings, "SaveA1111WeightedEmbeddings": SaveA1111WeightedEmbeddings, "SliceExistingEmbedding": SliceExistingEmbedding, } NODE_DISPLAY_NAME_MAPPINGS = { "SaveTokenEmbeddings": "Save Token Embeddings", "SaveWeightedEmbeddings": "Save Weighted Embeddings", "SaveA1111WeightedEmbeddings": "Save A1111-style Weighted Embeddings", "SliceExistingEmbedding": "Slice Existing Embedding File", } __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']