nodes/xLegende-ComfyUI-Prompt-Formatter
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
Legende
2025-04-13 08:13:30 +02:00
committed by GitHub
parent 8ba5d7c8f7
commit 1bae5e2eb4
2 changed files with 172 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
categorized_prompt_formatter.py
+163 -144
View File
@@ -12,135 +12,153 @@ try:
import yaml
except ImportError:
print("PyYAML is not installed. Please run install.py or install it manually.", file=sys.stderr)
# Consider adding instructions or a more robust check here
# For now, we assume it's installed if this point is reached without error
# --- Constants ---
NODE_NAME = "Categorized Prompt Formatter"
NODE_REPO = "ComfyUI-Prompt-Formatter"
NODE_REPO = "ComfyUI-Prompt-Formatter" # For finding relative paths
INCLUDE_DIRECTIVE = "$include" # Define the directive key
TAGS_KEY = "tags" # Optional key for direct tags within an include dict
# --- Helper Functions ---
def get_node_directory():
"""Gets the directory path of the current node."""
# This assumes the script is directly inside the node's folder in custom_nodes
return Path(__file__).parent
def find_yaml_file(filename):
"""Finds the YAML file, checking relative to node, inputs, and absolute."""
if not filename:
return None
# 1. Check if it's an absolute path
if Path(filename).is_absolute():
if Path(filename).is_file():
return Path(filename)
else:
print(f"Warning: Absolute path specified but not found: {filename}")
return None
# 2. Check relative to the custom node's directory
if Path(filename).is_file(): return Path(filename)
else: print(f"Warning: Absolute path specified but not found: {filename}"); return None
node_dir = get_node_directory()
relative_to_node = node_dir / filename
if relative_to_node.is_file():
return relative_to_node
# 3. Check relative to ComfyUI's input directory (common practice)
# Need to find the base ComfyUI directory first.
# This is a bit heuristic, go up until we find 'ComfyUI'.
# This might fail in unusual setups.
if relative_to_node.is_file(): return relative_to_node
try:
current_dir = Path.cwd() # Usually ComfyUI base dir when run normally
current_dir = Path.cwd()
if (current_dir / 'input').exists():
input_dir_path = current_dir / 'input' / filename
if input_dir_path.is_file():
return input_dir_path
# Fallback search upwards from node dir if cwd isn't reliable
if input_dir_path.is_file(): return input_dir_path
search_dir = node_dir
for _ in range(5):
if (search_dir / 'input').exists():
input_dir_path = search_dir / 'input' / filename
if input_dir_path.is_file():
return input_dir_path
if (search_dir / 'ComfyUI').exists() or (search_dir / 'main.py').exists(): # Found potential base
if input_dir_path.is_file(): return input_dir_path
if (search_dir / 'ComfyUI').exists() or (search_dir / 'main.py').exists():
if (search_dir / 'input').exists():
input_dir_path = search_dir / 'input' / filename
if input_dir_path.is_file():
return input_dir_path
break
if input_dir_path.is_file(): return input_dir_path
break
parent = search_dir.parent
if parent == search_dir: break
if parent == search_dir: break
search_dir = parent
except Exception as e:
print(f"Warning: Error searching for ComfyUI input directory: {e}")
except Exception as e: print(f"Warning: Error searching for ComfyUI input directory: {e}")
print(f"Warning: YAML file '{filename}' not found as absolute, relative to node, or in input dir.")
return None
def parse_tag(tag, handle_weights=True):
"""
Parses a tag to extract the base tag and preserve the original form.
Handles common weighting and emphasis syntax like (tag:1.2), [tag], {tag}.
Returns (original_tag, base_tag).
"""
original_tag = tag.strip()
""" Parses a tag to extract the base tag and preserve the original form. """
original_tag = str(tag).strip() # Ensure string type
base_tag = original_tag
if handle_weights:
# More robust regex to handle variations and potential nesting (simple cases)
# Pattern for (tag:weight) or (tag)
match = re.match(r"^\s*\(\s*(.*?)\s*(?::\s*[\d.]+\s*)?\)\s*$", original_tag)
if not handle_weights: return original_tag, base_tag
patterns = [
r"^\s*\(\s*(.*?)\s*(?::\s*[\d.]+\s*)?\)\s*$", # (tag:weight) or (tag)
r"^\s*\[\s*(.*?)\s*\]\s*$", # [tag]
r"^\s*\{\s*(.*?)\s*\}\s*$" # {tag}
]
for pattern in patterns:
match = re.match(pattern, original_tag)
if match:
base_tag = match.group(1).strip()
# Handle nested cases like ((tag)) - recursively call? Or just strip outer? Let's strip outer.
nested_match = re.match(r"^\s*\(\s*(.*?)\s*\)\s*$", base_tag)
# Handle simple nesting like ((tag))
nested_match = re.match(r"^\s*[\(\[\{](.*?)[\)\]\}]\s*$", base_tag)
if nested_match:
base_tag = nested_match.group(1).strip()
return original_tag, base_tag
# Pattern for [tag]
match = re.match(r"^\s*\[\s*(.*?)\s*\]\s*$", original_tag)
if match:
base_tag = match.group(1).strip()
# Handle nested [[tag]]
nested_match = re.match(r"^\s*\[\s*(.*?)\s*\]\s*$", base_tag)
if nested_match:
base_tag = nested_match.group(1).strip()
return original_tag, base_tag
# Pattern for {tag}
match = re.match(r"^\s*\{\s*(.*?)\s*\}\s*$", original_tag)
if match:
base_tag = match.group(1).strip()
# Handle nested {{tag}}
nested_match = re.match(r"^\s*\{\s*(.*?)\s*\}\s*$", base_tag)
if nested_match:
base_tag = nested_match.group(1).strip()
return original_tag, base_tag
# If no specific weight/emphasis pattern matched, return the stripped original
return original_tag, base_tag.strip()
return original_tag, base_tag.strip() # Return stripped original if no pattern matched
def clean_output_string(text, delimiter=", "):
"""Cleans up the final output string."""
if not text:
return ""
# Remove leading/trailing whitespace and delimiters
if not text: return ""
text = text.strip().strip(delimiter.strip()).strip()
# Replace multiple delimiters (and surrounding whitespace) with a single one
delimiter_pattern = r'\s*' + re.escape(delimiter.strip()) + r'\s*'
text = re.sub(f'({delimiter_pattern})+', delimiter, text)
# Remove delimiter if it's the very start or end after replacements
if text.startswith(delimiter): text = text[len(delimiter):].lstrip()
if text.endswith(delimiter.rstrip()): text = text[:-len(delimiter.rstrip())].rstrip()
return text
# --- YAML Include Resolution Helper ---
def resolve_category_tags(category_name, yaml_data, resolved_cache, recursion_guard=None):
"""
Recursively resolves tags for a category, handling $include directives.
Returns a set of unique tags (strings) for the category.
"""
category_name = str(category_name).strip() # Ensure consistent key type
if recursion_guard is None:
recursion_guard = set()
if category_name in recursion_guard:
print(f"Warning: Circular dependency detected involving category '{category_name}'. Skipping recursive include.")
return set()
if category_name in resolved_cache:
return resolved_cache[category_name]
if category_name not in yaml_data:
print(f"Warning: Included category '{category_name}' not found in YAML data.")
return set()
recursion_guard.add(category_name) # Mark as visiting
category_data = yaml_data[category_name]
final_tags = set()
if isinstance(category_data, list):
# Simple list of tags
final_tags.update(str(tag).strip() for tag in category_data if str(tag).strip())
elif isinstance(category_data, dict):
# Dictionary format, check for $include and direct tags
# Check for $include directive first
if INCLUDE_DIRECTIVE in category_data:
includes = category_data[INCLUDE_DIRECTIVE]
include_list = []
if isinstance(includes, list):
include_list = [str(inc).strip() for inc in includes if str(inc).strip()]
elif isinstance(includes, str):
include_list = [includes.strip()] if includes.strip() else []
for included_category in include_list:
# Pass a copy of the guard to prevent siblings affecting each other's cycle detection
final_tags.update(resolve_category_tags(included_category, yaml_data, resolved_cache, recursion_guard.copy()))
# Check for direct tags (e.g., under 'tags:' key or others if structure allows)
if TAGS_KEY in category_data and isinstance(category_data[TAGS_KEY], list):
final_tags.update(str(tag).strip() for tag in category_data[TAGS_KEY] if str(tag).strip())
# Optional: Handle case where dict has no $include or 'tags', maybe treat other keys?
# else:
# print(f"Warning: Dict category '{category_name}' lacks '{INCLUDE_DIRECTIVE}' or '{TAGS_KEY}'. Tags ignored.")
# pass
else:
# Handle cases where a category definition is neither list nor dict (e.g., just a string)
print(f"Warning: Category '{category_name}' definition is not a list or dictionary. Ignoring.")
resolved_cache[category_name] = final_tags
# No need to remove from recursion_guard here as we pass copies down
return final_tags
# --- The Node Class ---
class CategorizedPromptFormatter:
"""
A ComfyUI node to categorize tags from an input prompt based on a YAML file
and format them into a new prompt string using a template.
(supporting $include directives) and format them using a template.
"""
@classmethod
def INPUT_TYPES(cls):
@@ -165,46 +183,64 @@ class CategorizedPromptFormatter:
RETURN_TYPES = ("STRING", "STRING")
RETURN_NAMES = ("formatted_prompt", "rejected_prompt")
FUNCTION = "format_prompt"
CATEGORY = "text/filtering"
CATEGORY = "text/filtering"
def format_prompt(self, prompt, category_definition_file, output_template,
input_delimiter=",", output_delimiter=", ", strip_whitespace=True,
case_sensitive_matching=False, match_underscores_spaces=True, handle_weights=True,
disable_duplicates=False, unmatched_tag_handling="discard"):
case_sensitive_matching=False, handle_weights=True,
match_underscores_spaces=True, disable_duplicates=False,
unmatched_tag_handling="discard"):
# --- 1. Load Category Definitions ---
# --- 1. Load & Resolve Category Definitions ---
tag_to_categories_map = defaultdict(list)
yaml_path = find_yaml_file(category_definition_file)
raw_yaml_data = None
if yaml_path and yaml_path.exists():
try:
with open(yaml_path, 'r', encoding='utf-8') as f:
yaml_data = yaml.safe_load(f)
if yaml_data:
for category, tags in yaml_data.items():
if isinstance(tags, list):
for tag in tags:
processed_tag = str(tag).strip()
if not case_sensitive_matching:
processed_tag = processed_tag.lower()
if processed_tag:
tag_to_categories_map[processed_tag].append(str(category).strip())
raw_yaml_data = yaml.safe_load(f)
if not isinstance(raw_yaml_data, dict):
print(f"Warning: YAML file '{yaml_path}' does not contain a dictionary at the top level. Cannot process categories.")
raw_yaml_data = None # Treat as empty
except FileNotFoundError:
print(f"Error: YAML file not found at resolved path: {yaml_path}")
return ("", prompt)
except yaml.YAMLError as e:
print(f"Error parsing YAML file {yaml_path}: {e}")
return ("", prompt)
except Exception as e:
print(f"Error loading YAML file {yaml_path}: {e}")
return ("", prompt)
if raw_yaml_data:
resolved_tags_cache = {}
all_category_names = list(raw_yaml_data.keys())
# Resolve all categories first using the helper function
for category_name in all_category_names:
# Don't try to resolve the directive key itself if used at top level
if str(category_name).strip() != INCLUDE_DIRECTIVE and str(category_name).strip() != TAGS_KEY:
resolve_category_tags(category_name, raw_yaml_data, resolved_tags_cache)
# Build the final tag -> [categories] map from the resolved cache
for category_name, resolved_tags_set in resolved_tags_cache.items():
for tag in resolved_tags_set: # Tags are already stripped strings here
processed_tag_key = tag if case_sensitive_matching else tag.lower()
if processed_tag_key: # Ensure not empty after potential lowercasing
# A tag belongs to the category under which it was resolved
# (This includes parent categories that included it)
tag_to_categories_map[processed_tag_key].append(category_name)
# Deduplicate the list of categories for each tag
for tag_key in tag_to_categories_map:
tag_to_categories_map[tag_key] = list(set(tag_to_categories_map[tag_key]))
else:
print(f"Warning: Category definition file '{category_definition_file}' not found. Proceeding without categorization.")
print(f"Warning: Category definition file '{category_definition_file}' not found or empty/invalid. Proceeding without categorization.")
# tag_to_categories_map remains empty
# --- 2. Parse Input Prompt & Categorize ---
categorized_tags = defaultdict(list)
all_input_tags_original = set()
processed_input_tags = []
processed_input_tags = [] # Store (original_tag, base_tag for matching)
if prompt:
raw_tags = prompt.split(input_delimiter)
@@ -215,114 +251,97 @@ class CategorizedPromptFormatter:
all_input_tags_original.add(tag_original)
original_parsed, base_parsed = parse_tag(tag_original, handle_weights)
# Prepare the tag used for dictionary lookups (case normalization)
# Prepare the key used for dictionary lookups
lookup_key_base = base_parsed if case_sensitive_matching else base_parsed.lower()
processed_input_tags.append((original_parsed, lookup_key_base))
processed_input_tags.append((original_parsed, lookup_key_base))
matched_original_tags = set()
for original_tag, lookup_key_base in processed_input_tags:
found_categories = None
variants_to_check = {lookup_key_base} # Start with the base lookup key
if match_underscores_spaces:
# Generate variants for lookup
variants_to_check = {
lookup_key_base,
lookup_key_base.replace('_', ' '),
lookup_key_base.replace(' ', '_')
}
# Check each variant against the map
for variant in variants_to_check:
categories = tag_to_categories_map.get(variant)
if categories:
found_categories = categories
break #
else:
# Direct lookup if matching is disabled
found_categories = tag_to_categories_map.get(lookup_key_base)
# Generate variants only if needed
underscore_variant = lookup_key_base.replace(' ', '_')
space_variant = lookup_key_base.replace('_', ' ')
if underscore_variant != lookup_key_base: variants_to_check.add(underscore_variant)
if space_variant != lookup_key_base: variants_to_check.add(space_variant)
# Check each variant against the map
for variant in variants_to_check:
categories = tag_to_categories_map.get(variant)
if categories:
found_categories = categories
break # Found a match, stop checking variants for this tag
# If categories were found
if found_categories:
for category in found_categories:
# Append the ORIGINAL tag (with weights etc.) to the list for each category it maps to
categorized_tags[category].append(original_tag)
matched_original_tags.add(original_tag)
matched_original_tags.add(original_tag) # Mark the ORIGINAL tag as matched
# --- 3. Process Template ---
formatted_output_string = "" # Will be built piece by piece
placeholders = re.findall(r"<\|(.*?)\|>|", output_template)
placeholders = [p for p in placeholders if p]
# Track tags added to the output string to prevent duplicates if requested
formatted_output_string = ""
already_added_tags = set()
# Track tags considered "used" by the template for rejection logic
used_in_template_tags = set()
used_in_template_tags = set() # Tracks tags referenced by template (before dedup)
result_parts = []
last_end = 0
for match in re.finditer(r"<\|(.*?)\|>|", output_template):
placeholder_content = match.group(1)
start, end = match.span()
# Append the literal text segment before the placeholder
result_parts.append(output_template[last_end:start])
if placeholder_content:
category_name = placeholder_content.strip()
# Use the categorized_tags dict built earlier
tags_for_category = categorized_tags.get(category_name, [])
if tags_for_category:
# Mark all tags associated with this placeholder as 'used' by the template
# This happens *before* deduplication, as the template *did* reference them.
used_in_template_tags.update(tags_for_category)
tags_to_join = []
if disable_duplicates:
# Filter out tags already added
for tag in tags_for_category:
if tag not in already_added_tags:
tags_to_join.append(tag)
already_added_tags.add(tag) # Mark as added
else:
# Include all tags if duplicates are allowed
tags_to_join = tags_for_category
# already_added_tags.update(tags_to_join) # Optional
if tags_to_join:
joined_tags = output_delimiter.join(tags_to_join)
result_parts.append(joined_tags)
last_end = end
# Append any remaining literal text after the last placeholder
result_parts.append(output_template[last_end:])
formatted_output_string = "".join(result_parts)
# --- 4. Handle Unmatched Tags ---
# This section remains the same
rejected_tags_list = []
if unmatched_tag_handling != "discard":
# Find tags that were in the input but NOT marked as 'used' by the template placeholders
rejected_tags_list = list(all_input_tags_original - used_in_template_tags)
rejected_output_string = ""
if rejected_tags_list:
rejected_output_string = output_delimiter.join(rejected_tags_list)
if unmatched_tag_handling == "append_end":
# Ensure there's a delimiter if the main string isn't empty and doesn't end with one
if formatted_output_string and not formatted_output_string.endswith(output_delimiter):
if formatted_output_string and not formatted_output_string.rstrip().endswith(output_delimiter.strip()):
# Add delimiter only if needed and formatted string isn't empty
formatted_output_string += output_delimiter
elif not formatted_output_string: # If main string is empty, don't start with delimiter
pass
else: # Main string ends with delimiter, add space if needed
if not formatted_output_string.endswith(" "):
formatted_output_string += " "
elif not formatted_output_string:
pass # Avoid leading delimiter if formatted is empty
# Append rejected tags
formatted_output_string += rejected_output_string
rejected_output_string = "" # Clear rejected as it was appended
formatted_output_string += output_delimiter.join(rejected_tags_list)
elif unmatched_tag_handling == "output_separately":
rejected_output_string = output_delimiter.join(rejected_tags_list)
# --- 5. Cleanup and Return ---
final_formatted = clean_output_string(formatted_output_string, output_delimiter)
final_rejected = clean_output_string(rejected_output_string, output_delimiter)
return (final_formatted, final_rejected)
return (final_formatted, final_rejected)
prompt_categories.yaml
+9 -4
View File
@@ -280,10 +280,9 @@ person_details:
- looking at viewer
- serious expression
- open mouth
# --- Clothing ---
clothing:
# Tops
# --- Tops ---
tops:
- shirt
- t-shirt
- blouse
@@ -297,6 +296,12 @@ clothing:
- blazer
- vest
- cardigan
# --- Clothing ---
clothing:
$include:
- tops
tags:
# Bottoms
- pants
- jeans