519 lines
19 KiB
Python
519 lines
19 KiB
Python
"""
|
|
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",
|
|
}
|