""" DJZ Prompt Extractor V2 - ComfyUI Custom Node Extracts text prompts from PNG image metadata generated by ComfyUI. V2 Features: - Returns all strings sorted by length (longest first) - Offset parameter to select which string to return (1 = longest, 2 = second longest, etc.) - Improved system_prompt filtering with prefix matching - Support for widgets_values segmented prompts Author: Drift Johnson Repository: https://github.com/MushroomFleet/ComfyUI_PromptExtractor_nodes """ import json import os from PIL import Image import numpy as np import torch import folder_paths class DJZ_PromptExtractorV2: """ A ComfyUI node that extracts text prompts from PNG metadata with offset selection. Returns strings sorted by length (longest first). Use the offset parameter to select which string to return: - offset=1: longest string (default) - offset=2: second longest string - offset=N: Nth longest string This approach solves issues with mislabeled or duplicated metadata by allowing users to iterate through all available strings. Uses Load Image style file selector for full metadata access. """ def __init__(self): pass @classmethod def INPUT_TYPES(cls): input_dir = folder_paths.get_input_directory() files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] # Filter for image files (primarily PNG but allow others) image_extensions = ('.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp') files = sorted([f for f in files if f.lower().endswith(image_extensions)]) return { "required": { "image": (files, {"image_upload": True}), }, "optional": { "offset": ("INT", { "default": 1, "min": 1, "max": 100, "step": 1, "display": "number", "tooltip": "Which string to return: 1=longest, 2=second longest, etc." }), "ignore_system_prompt": ("BOOLEAN", { "default": True, "tooltip": "Filter out system_prompt content from results" }), } } RETURN_TYPES = ("STRING", "STRING", "INT", "STRING", "IMAGE",) RETURN_NAMES = ("selected_string", "all_strings_json", "total_count", "debug_info", "image",) FUNCTION = "extract_prompt" CATEGORY = "DJZ-Nodes" OUTPUT_NODE = False @classmethod def IS_CHANGED(cls, image, offset=1, ignore_system_prompt=True): image_path = folder_paths.get_annotated_filepath(image) return os.path.getmtime(image_path) @classmethod def VALIDATE_INPUTS(cls, image, offset=1, ignore_system_prompt=True): if not folder_paths.exists_annotated_filepath(image): return f"Invalid image file: {image}" return True @staticmethod def find_all_strings(obj, strings=None): """Recursively find all strings in a JSON object/array.""" if strings is None: strings = [] if isinstance(obj, str): strings.append(obj) elif isinstance(obj, list): for item in obj: DJZ_PromptExtractorV2.find_all_strings(item, strings) elif isinstance(obj, dict): for value in obj.values(): DJZ_PromptExtractorV2.find_all_strings(value, strings) return strings @staticmethod def collect_system_prompts(obj, prefixes, full_strings, current_key=""): """Recursively collect system_prompt values from nested objects.""" if isinstance(obj, str): # If we're inside a key that looks like system_prompt, collect it if 'system_prompt' in current_key.lower() or 'systemprompt' in current_key.lower(): if len(obj) >= 20: prefixes.append(obj[:20]) full_strings.add(obj) elif isinstance(obj, list): for item in obj: DJZ_PromptExtractorV2.collect_system_prompts(item, prefixes, full_strings, current_key) elif isinstance(obj, dict): for key, value in obj.items(): DJZ_PromptExtractorV2.collect_system_prompts(value, prefixes, full_strings, key) @staticmethod def find_widgets_values(obj, results=None): """Extract and join widgets_values arrays (segmented prompts).""" if results is None: results = [] if isinstance(obj, list): for item in obj: DJZ_PromptExtractorV2.find_widgets_values(item, results) elif isinstance(obj, dict): # Check if this object has widgets_values if 'widgets_values' in obj and isinstance(obj['widgets_values'], list): # Filter to only string values and join them string_values = [ v.strip() for v in obj['widgets_values'] if isinstance(v, str) and v.strip() and v.strip() != ',' ] if len(string_values) > 1: # Join the values - this is likely a segmented prompt joined = ' '.join(string_values) joined = ' '.join(joined.split()) # Normalize whitespace if joined: results.append(joined) # Recurse into object values for value in obj.values(): DJZ_PromptExtractorV2.find_widgets_values(value, results) return results def extract_all_sorted_strings_from_info(self, img_info, ignore_system_prompt=True): """ Extract all unique strings from PNG metadata info dict, sorted by length (longest first). Args: img_info: The img.info dict from PIL Image Returns: tuple: (sorted_strings_list, metadata_dict, debug_info) """ # Get metadata from PNG info metadata = {} system_prompt_prefixes = [] system_prompt_full_strings = set() all_strings_set = set() # First pass: collect system_prompt content if filtering if ignore_system_prompt: for key, value in img_info.items(): key_lower = key.lower() # Direct system_prompt key if key_lower == 'system_prompt' or key_lower == 'systemprompt': if isinstance(value, str): try: parsed = json.loads(value) for s in self.find_all_strings(parsed): if len(s) >= 20: system_prompt_prefixes.append(s[:20]) system_prompt_full_strings.add(s) except json.JSONDecodeError: if len(value) >= 20: system_prompt_prefixes.append(value[:20]) system_prompt_full_strings.add(value) else: # Scan inside other chunks for embedded system_prompt keys if isinstance(value, str): try: parsed = json.loads(value) self.collect_system_prompts(parsed, system_prompt_prefixes, system_prompt_full_strings) except json.JSONDecodeError: pass def matches_system_prompt(s): """Check if a string matches any collected system_prompt.""" if not ignore_system_prompt: return False if s in system_prompt_full_strings: return True if len(s) >= 20 and system_prompt_prefixes: s_prefix = s[:20] if any(prefix == s_prefix for prefix in system_prompt_prefixes): return True return False # Second pass: collect all strings for key, value in img_info.items(): key_lower = key.lower() # Skip system_prompt key entirely if filtering if ignore_system_prompt and (key_lower == 'system_prompt' or key_lower == 'systemprompt'): if isinstance(value, str): try: metadata[key] = json.loads(value) except json.JSONDecodeError: metadata[key] = value continue if isinstance(value, str): try: parsed = json.loads(value) metadata[key] = parsed # Find widgets_values joined strings for s in self.find_widgets_values(parsed): if s.strip() and not matches_system_prompt(s): all_strings_set.add(s) # Find regular strings for s in self.find_all_strings(parsed): if s.strip() and not matches_system_prompt(s): all_strings_set.add(s) except json.JSONDecodeError: # Not JSON, add raw string metadata[key] = value if value.strip() and not matches_system_prompt(value): all_strings_set.add(value) elif isinstance(value, bytes): try: decoded = value.decode('utf-8') parsed = json.loads(decoded) metadata[key] = parsed for s in self.find_widgets_values(parsed): if s.strip() and not matches_system_prompt(s): all_strings_set.add(s) for s in self.find_all_strings(parsed): if s.strip() and not matches_system_prompt(s): all_strings_set.add(s) except: pass # Convert to list and sort by length (longest first) sorted_strings = sorted( [s for s in all_strings_set if s], key=len, reverse=True ) debug_info = f"Found {len(sorted_strings)} unique strings" if ignore_system_prompt and system_prompt_full_strings: debug_info += f", filtered {len(system_prompt_full_strings)} system_prompt values" return sorted_strings, metadata, debug_info def extract_prompt(self, image, offset=1, ignore_system_prompt=True): """ Main extraction function. Args: image: Filename of the image in the input directory offset: Which string to return (1=longest, 2=second longest, etc.) ignore_system_prompt: Whether to filter out system_prompt content Returns: tuple: (selected_string, all_strings_json, total_count, debug_info, image_tensor) """ # Get full path from the selected filename image_path = folder_paths.get_annotated_filepath(image) # Open image and get metadata try: img = Image.open(image_path) img_info = img.info except Exception as e: # Return empty tensor on error empty_tensor = torch.zeros((1, 64, 64, 3), dtype=torch.float32) return ( f"[Error opening image: {str(e)}]", "[]", 0, f"Error: {str(e)}", empty_tensor ) # Load the image as tensor for output img_rgb = img.convert("RGB") image_tensor = torch.from_numpy(np.array(img_rgb).astype(np.float32) / 255.0).unsqueeze(0) # Extract strings from metadata sorted_strings, metadata, debug_info = self.extract_all_sorted_strings_from_info( img_info, ignore_system_prompt ) if not sorted_strings: return ( "[No strings found]", "[]", 0, debug_info or "No text content found in metadata", image_tensor ) # Adjust offset to 0-based index index = offset - 1 # Clamp index to valid range if index < 0: index = 0 elif index >= len(sorted_strings): index = len(sorted_strings) - 1 debug_info += f" | Offset {offset} exceeds count, using last string (index {index + 1})" selected_string = sorted_strings[index] # Create JSON array of all strings with their lengths strings_with_info = [ {"index": i + 1, "length": len(s), "preview": s[:100] + "..." if len(s) > 100 else s} for i, s in enumerate(sorted_strings) ] all_strings_json = json.dumps(strings_with_info, indent=2) return ( selected_string, all_strings_json, len(sorted_strings), debug_info, image_tensor ) class DJZ_PromptExtractorV2_Batch: """ Batch version of V2 extractor that processes multiple images from a directory. """ def __init__(self): pass @classmethod def INPUT_TYPES(cls): return { "required": { "directory_path": ("STRING", { "default": "", "multiline": False, "placeholder": "Path to directory containing PNG files" }), }, "optional": { "offset": ("INT", { "default": 1, "min": 1, "max": 100, "step": 1, "display": "number", "tooltip": "Which string to extract from each image: 1=longest, 2=second longest, etc." }), "max_images": ("INT", { "default": 10, "min": 1, "max": 100, "step": 1 }), "ignore_system_prompt": ("BOOLEAN", { "default": True, }), } } RETURN_TYPES = ("STRING", "STRING", "INT",) RETURN_NAMES = ("prompts_list", "filenames_list", "processed_count",) FUNCTION = "extract_batch" CATEGORY = "DJZ-Nodes" OUTPUT_NODE = False def extract_from_path(self, filepath, offset=1, ignore_system_prompt=True): """Extract prompt from a file path (for batch processing).""" try: img = Image.open(filepath) img_info = img.info except Exception as e: return f"[Error: {str(e)}]" extractor = DJZ_PromptExtractorV2() sorted_strings, _, _ = extractor.extract_all_sorted_strings_from_info( img_info, ignore_system_prompt ) if not sorted_strings: return "[No strings found]" # Adjust offset to 0-based index index = offset - 1 if index < 0: index = 0 elif index >= len(sorted_strings): index = len(sorted_strings) - 1 return sorted_strings[index] def extract_batch(self, directory_path, offset=1, max_images=10, ignore_system_prompt=True): """Extract prompts from all PNG files in a directory.""" if not os.path.isdir(directory_path): return (f"[Error: Directory not found: {directory_path}]", "", 0) # Find all PNG files png_files = [f for f in os.listdir(directory_path) if f.lower().endswith('.png')] png_files = png_files[:max_images] if not png_files: return ("[No PNG files found in directory]", "", 0) prompts = [] filenames = [] for filename in png_files: filepath = os.path.join(directory_path, filename) selected_string = self.extract_from_path(filepath, offset, ignore_system_prompt) if selected_string and not selected_string.startswith("["): prompts.append(selected_string) filenames.append(filename) # Join with separators prompts_output = "\n---\n".join(prompts) if prompts else "[No prompts extracted]" filenames_output = "\n".join(filenames) if filenames else "" return (prompts_output, filenames_output, len(prompts)) class DJZ_PromptExtractorV2_StringNavigator: """ Helper node that takes the all_strings_json output and allows navigation through the strings with offset selection. """ def __init__(self): pass @classmethod def INPUT_TYPES(cls): return { "required": { "all_strings_json": ("STRING", { "multiline": True, "placeholder": "Connect to all_strings_json output from PromptExtractorV2" }), "offset": ("INT", { "default": 1, "min": 1, "max": 100, "step": 1, "display": "number", }), }, } RETURN_TYPES = ("STRING", "INT", "STRING",) RETURN_NAMES = ("selected_string", "string_length", "preview",) FUNCTION = "navigate" CATEGORY = "DJZ-Nodes" OUTPUT_NODE = False def navigate(self, all_strings_json, offset=1): """Navigate through the strings list using offset.""" try: strings_info = json.loads(all_strings_json) except json.JSONDecodeError: return ("[Invalid JSON input]", 0, "") if not strings_info: return ("[No strings available]", 0, "") # Adjust offset to 0-based index index = offset - 1 if index < 0: index = 0 elif index >= len(strings_info): index = len(strings_info) - 1 item = strings_info[index] # The preview contains the actual string (possibly truncated) # We need to return the full string, which is in preview # For full string access, user should use the main extractor preview = item.get("preview", "") length = item.get("length", 0) return (preview, length, f"String {offset} of {len(strings_info)}: {length} chars") # Node registration for ComfyUI NODE_CLASS_MAPPINGS = { "DJZ_PromptExtractorV2": DJZ_PromptExtractorV2, "DJZ_PromptExtractorV2_Batch": DJZ_PromptExtractorV2_Batch, "DJZ_PromptExtractorV2_StringNavigator": DJZ_PromptExtractorV2_StringNavigator, } NODE_DISPLAY_NAME_MAPPINGS = { "DJZ_PromptExtractorV2": "DJZ Prompt Extractor V2", "DJZ_PromptExtractorV2_Batch": "DJZ Prompt Extractor V2 (Batch)", "DJZ_PromptExtractorV2_StringNavigator": "DJZ String Navigator V2", }