Add files via upload
refactor
This commit is contained in:
Legende
+29
-25
@@ -1,34 +1,38 @@
|
||||
# /ComfyUI-Prompt-Formatter/__init__.py
|
||||
# /ComfyUI-Prompt-Formatter/__init__.py
|
||||
|
||||
# Import the node class
|
||||
from .categorized_prompt_formatter import CategorizedPromptFormatter, NODE_NAME as NODE_NAME_FORMATTER
|
||||
from .categorized_random_prompt_formatter import CategorizedRandomPromptFormatter, NODE_NAME_RANDOM
|
||||
from .categorized_prompt_analyzer import CategorizedPromptAnalyzer, NODE_NAME_ANALYZER
|
||||
from .wildcard_importer import WildcardImporter, NODE_NAME_IMPORTER
|
||||
# Ensure utils are loaded first to check dependencies
|
||||
from . import prompt_formatter_utils
|
||||
|
||||
# Mapping node class names to implementations
|
||||
# Import node classes
|
||||
from .categorized_prompt_formatter import CategorizedPromptFormatter
|
||||
from .categorized_random_prompt_formatter import CategorizedRandomPromptFormatter
|
||||
from .categorized_prompt_analyzer import CategorizedPromptAnalyzer
|
||||
from .wildcard_importer import WildcardImporter
|
||||
|
||||
# Standardize NODE_NAME attribute in WildcardImporter for consistency
|
||||
if not hasattr(WildcardImporter, 'NODE_NAME'):
|
||||
WildcardImporter.NODE_NAME = "Wildcard Importer"
|
||||
|
||||
# Mapping node implementations
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
NODE_NAME_FORMATTER: CategorizedPromptFormatter,
|
||||
NODE_NAME_RANDOM: CategorizedRandomPromptFormatter,
|
||||
NODE_NAME_ANALYZER: CategorizedPromptAnalyzer,
|
||||
NODE_NAME_IMPORTER: WildcardImporter,
|
||||
CategorizedPromptFormatter.NODE_NAME: CategorizedPromptFormatter,
|
||||
CategorizedRandomPromptFormatter.NODE_NAME: CategorizedRandomPromptFormatter,
|
||||
CategorizedPromptAnalyzer.NODE_NAME: CategorizedPromptAnalyzer,
|
||||
WildcardImporter.NODE_NAME: WildcardImporter,
|
||||
}
|
||||
|
||||
# Mapping node class names to display names for the UI
|
||||
# Mapping display names for the UI
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
NODE_NAME_FORMATTER: "📝 Categorized Prompt Formatter",
|
||||
NODE_NAME_RANDOM: "🎲 Categorized Random Prompt Formatter",
|
||||
NODE_NAME_ANALYZER: "📊 Categorized Prompt Analyzer",
|
||||
NODE_NAME_IMPORTER: "📂 Wildcard Importer",
|
||||
CategorizedPromptFormatter.NODE_NAME: "📝 Categorized Prompt Formatter",
|
||||
CategorizedRandomPromptFormatter.NODE_NAME: "🎲 Categorized Random Prompt Formatter",
|
||||
CategorizedPromptAnalyzer.NODE_NAME: "📊 Categorized Prompt Analyzer",
|
||||
WildcardImporter.NODE_NAME: "📂 Wildcard Importer",
|
||||
}
|
||||
# A dictionary that contains all nodes in this file
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
print("...")
|
||||
print("### Loading: ComfyUI-Prompt-Formatter ###")
|
||||
print("### Version: 1.4.1")
|
||||
print(f"### Node: {NODE_NAME_FORMATTER} -> {NODE_DISPLAY_NAME_MAPPINGS[NODE_NAME_FORMATTER]}")
|
||||
print(f"### Node: {NODE_NAME_RANDOM} -> {NODE_DISPLAY_NAME_MAPPINGS[NODE_NAME_RANDOM]}")
|
||||
print(f"### Node: {NODE_NAME_ANALYZER} -> {NODE_DISPLAY_NAME_MAPPINGS[NODE_NAME_ANALYZER]}")
|
||||
print(f"### Node: {NODE_NAME_IMPORTER} -> {NODE_DISPLAY_NAME_MAPPINGS[NODE_NAME_IMPORTER]}")
|
||||
print("...")
|
||||
# Print loading info to console
|
||||
print("\n### Loading: ComfyUI-Prompt-Formatter (Version: 1.4.2) ###")
|
||||
for name, display_name in NODE_DISPLAY_NAME_MAPPINGS.items():
|
||||
print(f" - {name} -> {display_name}")
|
||||
print("###-------------------------------------------###\n")
|
||||
+71
-240
@@ -2,130 +2,24 @@
|
||||
|
||||
import yaml
|
||||
import re
|
||||
import os
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
# --- Dependency Check ---
|
||||
try:
|
||||
import yaml
|
||||
except ImportError:
|
||||
print("PyYAML is not installed. Please run install.py or install it manually.", file=sys.stderr)
|
||||
|
||||
# --- Constants ---
|
||||
NODE_NAME_ANALYZER = "Categorized Prompt Analyzer" # Specific name
|
||||
INCLUDE_DIRECTIVE = "$include"
|
||||
TAGS_KEY = "tags"
|
||||
|
||||
|
||||
def get_node_directory_analyzer():
|
||||
return Path(__file__).parent
|
||||
|
||||
def find_yaml_file_analyzer(filename):
|
||||
if not filename: return None
|
||||
if Path(filename).is_absolute():
|
||||
if Path(filename).is_file(): return Path(filename)
|
||||
else: print(f"Warning [Analyzer]: Absolute path specified but not found: {filename}"); return None
|
||||
node_dir = get_node_directory_analyzer()
|
||||
relative_to_node = node_dir / filename
|
||||
if relative_to_node.is_file(): return relative_to_node
|
||||
try:
|
||||
current_dir = Path.cwd()
|
||||
input_dir = current_dir / 'input'
|
||||
if input_dir.exists():
|
||||
input_dir_path = input_dir / filename
|
||||
if input_dir_path.is_file(): return input_dir_path
|
||||
search_dir = node_dir
|
||||
for _ in range(5):
|
||||
input_dir = search_dir / 'input'
|
||||
if input_dir.exists():
|
||||
input_dir_path = input_dir / filename
|
||||
if input_dir_path.is_file(): return input_dir_path
|
||||
if (search_dir / 'ComfyUI').exists() or (search_dir / 'main.py').exists():
|
||||
input_dir = search_dir / 'input'
|
||||
if input_dir.exists():
|
||||
input_dir_path = input_dir / filename
|
||||
if input_dir_path.is_file(): return input_dir_path
|
||||
break
|
||||
parent = search_dir.parent
|
||||
if parent == search_dir: break
|
||||
search_dir = parent
|
||||
except Exception as e: print(f"Warning [Analyzer]: Error searching for ComfyUI input directory: {e}")
|
||||
print(f"Warning [Analyzer]: YAML file '{filename}' not found as absolute, relative to node, or in input dir.")
|
||||
return None
|
||||
|
||||
def resolve_category_tags_analyzer(category_name, yaml_data, resolved_cache, recursion_guard=None):
|
||||
category_name = str(category_name).strip()
|
||||
if recursion_guard is None: recursion_guard = set()
|
||||
if len(recursion_guard) > 20: return set() # Depth limit
|
||||
if category_name in recursion_guard: print(f"Warning [Analyzer]: Circular dependency: '{category_name}'."); return set()
|
||||
if category_name in resolved_cache: return resolved_cache[category_name]
|
||||
if category_name not in yaml_data: return set()
|
||||
recursion_guard.add(category_name)
|
||||
category_data = yaml_data[category_name]
|
||||
final_tags = set()
|
||||
if isinstance(category_data, list):
|
||||
for item in category_data:
|
||||
item_str = str(item).strip();
|
||||
if not item_str: continue
|
||||
is_full_include = item_str.startswith('$') and len(item_str) > 1 and ' ' not in item_str and item_str[1:].replace('_', '').isalnum()
|
||||
if is_full_include:
|
||||
included_category = item_str[1:]
|
||||
final_tags.update(resolve_category_tags_analyzer(included_category, yaml_data, resolved_cache, recursion_guard.copy()))
|
||||
else:
|
||||
match = re.search(r'\$(\w+)', item_str)
|
||||
if match:
|
||||
ref_category_name = match.group(1); placeholder = match.group(0)
|
||||
resolved_ref_tags = resolve_category_tags_analyzer(ref_category_name, yaml_data, resolved_cache, recursion_guard.copy())
|
||||
if not resolved_ref_tags: pass # Optionally add literal item_str here
|
||||
else:
|
||||
for resolved_tag in resolved_ref_tags:
|
||||
new_tag = item_str.replace(placeholder, str(resolved_tag).strip(), 1)
|
||||
final_tags.add(new_tag)
|
||||
else: final_tags.add(item_str)
|
||||
elif isinstance(category_data, dict):
|
||||
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:
|
||||
final_tags.update(resolve_category_tags_analyzer(included_category, yaml_data, resolved_cache, recursion_guard.copy()))
|
||||
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())
|
||||
resolved_cache[category_name] = final_tags
|
||||
return final_tags
|
||||
|
||||
def parse_tag_analyzer(tag, handle_weights=True):
|
||||
original_tag = str(tag).strip(); base_tag = original_tag
|
||||
if not handle_weights: return original_tag, base_tag
|
||||
patterns = [ r"^\s*\(\s*(.*?)\s*(?::\s*[\d.]+\s*)?\)\s*$", r"^\s*\[\s*(.*?)\s*\]\s*$", r"^\s*\{\s*(.*?)\s*\}\s*$" ]
|
||||
for pattern in patterns:
|
||||
match = re.match(pattern, original_tag)
|
||||
if match:
|
||||
base_tag = match.group(1).strip()
|
||||
nested_match = re.match(r"^\s*[\(\[\{](.*?)[\)\]\}]\s*$", base_tag)
|
||||
if nested_match: base_tag = nested_match.group(1).strip()
|
||||
return original_tag, base_tag
|
||||
return original_tag, base_tag.strip()
|
||||
|
||||
def clean_output_string_analyzer(text, delimiter=", "):
|
||||
if not text: return ""
|
||||
text = text.strip().strip(delimiter.strip()).strip()
|
||||
delimiter_pattern = r'\s*' + re.escape(delimiter.strip()) + r'\s*'
|
||||
text = re.sub(f'({delimiter_pattern})+', delimiter, text)
|
||||
if text.startswith(delimiter): text = text[len(delimiter):].lstrip()
|
||||
if text.endswith(delimiter.rstrip()): text = text[:-len(delimiter.rstrip())].rstrip()
|
||||
return text
|
||||
|
||||
# --- The Analyzer Node Class ---
|
||||
# Local utility imports
|
||||
from .prompt_formatter_utils import (
|
||||
find_yaml_file,
|
||||
parse_tag,
|
||||
clean_output_string,
|
||||
resolve_category_tags,
|
||||
INCLUDE_DIRECTIVE,
|
||||
TAGS_KEY
|
||||
)
|
||||
|
||||
class CategorizedPromptAnalyzer:
|
||||
"""
|
||||
A ComfyUI node to analyze tag occurrences in a prompt based on specific
|
||||
tags or categories defined in a YAML file.
|
||||
A node to analyze tag occurrences in a prompt against YAML-defined categories.
|
||||
"""
|
||||
NODE_NAME = "Categorized Prompt Analyzer"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
@@ -137,8 +31,7 @@ class CategorizedPromptAnalyzer:
|
||||
"optional": {
|
||||
"input_delimiter": ("STRING", {"default": ","}),
|
||||
"target_delimiter": ("STRING", {"default": ","}),
|
||||
"output_delimiter": ("STRING", {"default": ", "}), # For list outputs
|
||||
"strip_whitespace": ("BOOLEAN", {"default": True}),
|
||||
"output_delimiter": ("STRING", {"default": ", "}),
|
||||
"case_sensitive_matching": ("BOOLEAN", {"default": False}),
|
||||
"handle_weights": ("BOOLEAN", {"default": True}),
|
||||
"match_underscores_spaces": ("BOOLEAN", {"default": True}),
|
||||
@@ -152,138 +45,76 @@ class CategorizedPromptAnalyzer:
|
||||
FUNCTION = "analyze_prompt"
|
||||
CATEGORY = "text/analysis"
|
||||
|
||||
def analyze_prompt(self, prompt, analyze_targets, category_definition_file,
|
||||
input_delimiter=",", target_delimiter=",", output_delimiter=", ",
|
||||
strip_whitespace=True, case_sensitive_matching=False,
|
||||
handle_weights=True, match_underscores_spaces=True,
|
||||
generate_detailed_output=False, generate_unmatched_output=False):
|
||||
|
||||
def analyze_prompt(self, prompt, analyze_targets, category_definition_file, **kwargs):
|
||||
# --- 1. Load & Resolve YAML ---
|
||||
resolved_category_tags = {}
|
||||
yaml_path = find_yaml_file(category_definition_file, self.NODE_NAME)
|
||||
raw_yaml_data = None
|
||||
yaml_path = find_yaml_file_analyzer(category_definition_file)
|
||||
|
||||
if yaml_path and yaml_path.exists():
|
||||
|
||||
if yaml_path:
|
||||
try:
|
||||
with open(yaml_path, 'r', encoding='utf-8') as f: raw_yaml_data = yaml.safe_load(f)
|
||||
if not isinstance(raw_yaml_data, dict): raw_yaml_data = None
|
||||
except Exception as e: print(f"Error [Analyzer]: Loading YAML {yaml_path}: {e}")
|
||||
except Exception as e: print(f"Error [{self.NODE_NAME}]: Loading YAML {yaml_path}: {e}")
|
||||
|
||||
case_sensitive = kwargs.get("case_sensitive_matching", False)
|
||||
if raw_yaml_data:
|
||||
resolved_tags_cache = {}
|
||||
all_category_names = list(raw_yaml_data.keys())
|
||||
for category_name in all_category_names:
|
||||
if str(category_name).strip() != INCLUDE_DIRECTIVE and str(category_name).strip() != TAGS_KEY:
|
||||
# Resolve tags using the helper
|
||||
resolved_tags_set = resolve_category_tags_analyzer(category_name, raw_yaml_data, resolved_tags_cache)
|
||||
# Store resolved tags, applying case sensitivity NOW for matching later
|
||||
normalized_category_name = str(category_name).strip() if case_sensitive_matching else str(category_name).strip().lower()
|
||||
normalized_tags = set()
|
||||
for tag in resolved_tags_set:
|
||||
# Also normalize tags based on case sensitivity for matching
|
||||
normalized_tags.add(tag if case_sensitive_matching else tag.lower())
|
||||
resolved_category_tags[normalized_category_name] = normalized_tags
|
||||
# else: Warning already printed by find_yaml_file
|
||||
cache = {}
|
||||
for cat_name in list(raw_yaml_data.keys()):
|
||||
if str(cat_name).strip() not in [INCLUDE_DIRECTIVE, TAGS_KEY]:
|
||||
tags_set = resolve_category_tags(cat_name, raw_yaml_data, cache, self.NODE_NAME)
|
||||
key = str(cat_name).strip() if case_sensitive else str(cat_name).strip().lower()
|
||||
resolved_category_tags[key] = {t if case_sensitive else t.lower() for t in tags_set}
|
||||
|
||||
# --- 2. Parse & Process Targets ---
|
||||
targets_info = [] # List of tuples: (original_target, normalized_target, is_category, {normalized_tags_to_match})
|
||||
original_targets_list = [t.strip() for t in analyze_targets.split(target_delimiter) if t.strip()]
|
||||
# --- 2. Parse Targets ---
|
||||
targets_info = []
|
||||
original_targets = [t.strip() for t in analyze_targets.split(kwargs.get("target_delimiter", ",")) if t.strip()]
|
||||
match_spaces = kwargs.get("match_underscores_spaces", True)
|
||||
|
||||
for original_target in original_targets_list:
|
||||
normalized_target = original_target if case_sensitive_matching else original_target.lower()
|
||||
is_category = False
|
||||
normalized_tags_to_match = set()
|
||||
for target in original_targets:
|
||||
norm_target = target if case_sensitive else target.lower()
|
||||
tags_to_match = resolved_category_tags.get(norm_target, set())
|
||||
if not tags_to_match: # It's a literal tag, not a category
|
||||
tags_to_match.add(norm_target)
|
||||
if match_spaces:
|
||||
tags_to_match.add(norm_target.replace('_', ' '))
|
||||
tags_to_match.add(norm_target.replace(' ', '_'))
|
||||
targets_info.append((target, tags_to_match))
|
||||
|
||||
if normalized_target in resolved_category_tags:
|
||||
# It's a category target
|
||||
is_category = True
|
||||
normalized_tags_to_match = resolved_category_tags[normalized_target]
|
||||
else:
|
||||
# It's a literal target
|
||||
is_category = False
|
||||
# Generate variants for the literal target itself for matching
|
||||
literal_variants = {normalized_target}
|
||||
if match_underscores_spaces:
|
||||
literal_variants.add(normalized_target.replace('_', ' '))
|
||||
literal_variants.add(normalized_target.replace(' ', '_'))
|
||||
normalized_tags_to_match = literal_variants
|
||||
|
||||
targets_info.append((original_target, normalized_target, is_category, normalized_tags_to_match))
|
||||
|
||||
# --- 3. Parse Input Prompt ---
|
||||
input_tags_processed = [] # List of tuples: (original_tag, base_tag_normalized, {variants})
|
||||
all_input_tags_original = set()
|
||||
|
||||
if prompt:
|
||||
raw_tags = prompt.split(input_delimiter)
|
||||
for raw_tag in raw_tags:
|
||||
tag_original = raw_tag.strip() if strip_whitespace else raw_tag
|
||||
if not tag_original: continue
|
||||
all_input_tags_original.add(tag_original)
|
||||
|
||||
# Parse tag, get base form
|
||||
_, base_parsed = parse_tag_analyzer(tag_original, handle_weights)
|
||||
base_tag_normalized = base_parsed if case_sensitive_matching else base_parsed.lower()
|
||||
|
||||
# Generate variants for matching
|
||||
variants = {base_tag_normalized}
|
||||
if match_underscores_spaces:
|
||||
variants.add(base_tag_normalized.replace('_', ' '))
|
||||
variants.add(base_tag_normalized.replace(' ', '_'))
|
||||
|
||||
input_tags_processed.append((tag_original, base_tag_normalized, variants))
|
||||
|
||||
|
||||
# --- 4. Perform Counting ---
|
||||
counts = defaultdict(int) # original_target -> count
|
||||
details_data = defaultdict(list) # original_target -> list_of_matched_input_tags
|
||||
matched_input_originals_set = set()
|
||||
|
||||
for tag_original, _, input_variants in input_tags_processed:
|
||||
tag_matched_at_least_once = False
|
||||
for original_target, normalized_target, is_category, target_match_set in targets_info:
|
||||
# Check if any input variant matches any tag in the target's match set
|
||||
found_match = False
|
||||
for variant in input_variants:
|
||||
if variant in target_match_set:
|
||||
found_match = True
|
||||
break # Found a match for this input tag against this target
|
||||
|
||||
if found_match:
|
||||
counts[original_target] += 1
|
||||
if generate_detailed_output:
|
||||
details_data[original_target].append(tag_original)
|
||||
matched_input_originals_set.add(tag_original)
|
||||
tag_matched_at_least_once = True
|
||||
# Note: If an input tag matches multiple targets (e.g., literal and category),
|
||||
# it will correctly increment the count for *each* target it matches.
|
||||
|
||||
|
||||
# --- 5. Calculate Unmatched ---
|
||||
unmatched_tags_list = []
|
||||
if generate_unmatched_output:
|
||||
unmatched_tags_list = list(all_input_tags_original - matched_input_originals_set)
|
||||
|
||||
|
||||
# --- 6. Format Outputs ---
|
||||
summary_parts = []
|
||||
for original_target in original_targets_list: # Keep original order
|
||||
summary_parts.append(f"{original_target}: {counts[original_target]}")
|
||||
analysis_summary = output_delimiter.join(summary_parts)
|
||||
|
||||
total_matched_count = sum(counts.values())
|
||||
# --- 3. Parse and Count Prompt ---
|
||||
counts = defaultdict(int)
|
||||
details = defaultdict(list)
|
||||
matched_input_tags = set()
|
||||
all_input_tags = set()
|
||||
handle_weights = kwargs.get("handle_weights", True)
|
||||
|
||||
for raw_tag in prompt.split(kwargs.get("input_delimiter", ",")):
|
||||
original_tag = raw_tag.strip()
|
||||
if not original_tag: continue
|
||||
all_input_tags.add(original_tag)
|
||||
_, base = parse_tag(original_tag, handle_weights)
|
||||
norm_base = base if case_sensitive else base.lower()
|
||||
|
||||
variants = {norm_base}
|
||||
if match_spaces: variants.update({norm_base.replace('_', ' '), norm_base.replace(' ', '_')})
|
||||
|
||||
for t_original, t_match_set in targets_info:
|
||||
if not variants.isdisjoint(t_match_set):
|
||||
counts[t_original] += 1
|
||||
details[t_original].append(original_tag)
|
||||
matched_input_tags.add(original_tag)
|
||||
|
||||
# --- 4. Format Outputs ---
|
||||
output_delimiter = kwargs.get("output_delimiter", ", ")
|
||||
summary = output_delimiter.join([f"{t}: {counts[t]}" for t in original_targets])
|
||||
total_count = sum(counts.values())
|
||||
|
||||
details_str = ""
|
||||
if generate_detailed_output:
|
||||
detail_parts = []
|
||||
for original_target in original_targets_list:
|
||||
matched_list_str = output_delimiter.join(details_data[original_target])
|
||||
detail_parts.append(f"{original_target}: {counts[original_target]} [{matched_list_str}]")
|
||||
details_str = "\n".join(detail_parts) # Use newline for readability
|
||||
if kwargs.get("generate_detailed_output", False):
|
||||
details_str = "\n".join([f"{t}: {counts[t]} [{output_delimiter.join(details[t])}]" for t in original_targets])
|
||||
|
||||
unmatched_tags_str = ""
|
||||
if generate_unmatched_output:
|
||||
unmatched_tags_str = output_delimiter.join(unmatched_tags_list)
|
||||
unmatched_str = ""
|
||||
if kwargs.get("generate_unmatched_output", False):
|
||||
unmatched_str = output_delimiter.join(list(all_input_tags - matched_input_tags))
|
||||
|
||||
|
||||
return (analysis_summary, total_matched_count, details_str, unmatched_tags_str)
|
||||
return (summary, total_count, details_str, unmatched_str)
|
||||
+81
-333
@@ -2,192 +2,31 @@
|
||||
|
||||
import yaml
|
||||
import re
|
||||
import os
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
# --- Dependency Check ---
|
||||
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" # 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."""
|
||||
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
|
||||
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
|
||||
node_dir = get_node_directory()
|
||||
relative_to_node = node_dir / filename
|
||||
if relative_to_node.is_file(): return relative_to_node
|
||||
try:
|
||||
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
|
||||
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():
|
||||
if (search_dir / 'input').exists():
|
||||
input_dir_path = search_dir / 'input' / filename
|
||||
if input_dir_path.is_file(): return input_dir_path
|
||||
break
|
||||
parent = search_dir.parent
|
||||
if parent == search_dir: break
|
||||
search_dir = parent
|
||||
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. """
|
||||
original_tag = str(tag).strip() # Ensure string type
|
||||
base_tag = 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 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
|
||||
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 ""
|
||||
text = text.strip().strip(delimiter.strip()).strip()
|
||||
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): # Or _random suffix
|
||||
"""
|
||||
Recursively resolves tags for a category, handling $include directives,
|
||||
$category list includes, and inline $category expansion within strings.
|
||||
Returns a set of unique tags (strings) for the category.
|
||||
"""
|
||||
category_name = str(category_name).strip()
|
||||
if recursion_guard is None: recursion_guard = set()
|
||||
# Basic recursion depth limit as extra safety
|
||||
if len(recursion_guard) > 20:
|
||||
print(f"Warning: Recursion depth limit exceeded processing '{category_name}'. Returning empty set.")
|
||||
return set()
|
||||
if category_name in recursion_guard:
|
||||
print(f"Warning: Circular dependency detected involving category '{category_name}'.")
|
||||
return set()
|
||||
if category_name in resolved_cache: return resolved_cache[category_name]
|
||||
if category_name not in yaml_data:
|
||||
# Don't warn here, allows optional categories in references
|
||||
return set()
|
||||
|
||||
recursion_guard.add(category_name)
|
||||
category_data = yaml_data[category_name]
|
||||
final_tags = set()
|
||||
|
||||
if isinstance(category_data, list):
|
||||
for item in category_data:
|
||||
item_str = str(item).strip()
|
||||
if not item_str: continue
|
||||
|
||||
# Check for full category include first (starts with $, no spaces, alphanumeric name)
|
||||
is_full_include = item_str.startswith('$') and len(item_str) > 1 and ' ' not in item_str and item_str[1:].replace('_', '').isalnum()
|
||||
|
||||
if is_full_include:
|
||||
# --- Handle Full Category Include: $other_category ---
|
||||
included_category = item_str[1:]
|
||||
final_tags.update(resolve_category_tags(included_category, yaml_data, resolved_cache, recursion_guard.copy())) # Use correct function name
|
||||
else:
|
||||
# --- Handle Inline Expansion: "prefix $category suffix" ---
|
||||
# Use regex to find the first $word pattern
|
||||
# Limit to one expansion per string for simplicity
|
||||
match = re.search(r'\$(\w+)', item_str) # \w+ matches letters, numbers, underscore
|
||||
|
||||
if match:
|
||||
ref_category_name = match.group(1)
|
||||
placeholder = match.group(0) # The full '$word' matched
|
||||
|
||||
# Resolve the referenced category's tags
|
||||
resolved_ref_tags = resolve_category_tags(ref_category_name, yaml_data, resolved_cache, recursion_guard.copy())
|
||||
|
||||
if not resolved_ref_tags:
|
||||
print(f"Warning: Inline reference ${ref_category_name} in '{item_str}' resolved to an empty set or category not found. Skipping expansion for this item.")
|
||||
else:
|
||||
# Substitute each resolved tag into the original string pattern
|
||||
for resolved_tag in resolved_ref_tags:
|
||||
# Perform replacement. Use count=1 if placeholder could appear multiple times.
|
||||
new_tag = item_str.replace(placeholder, str(resolved_tag).strip(), 1)
|
||||
final_tags.add(new_tag)
|
||||
else:
|
||||
# No $category reference found, treat as a literal tag
|
||||
final_tags.add(item_str)
|
||||
|
||||
# --- Keep supporting the old dictionary format ---
|
||||
elif isinstance(category_data, dict):
|
||||
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:
|
||||
final_tags.update(resolve_category_tags(included_category, yaml_data, resolved_cache, recursion_guard.copy()))
|
||||
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())
|
||||
|
||||
else:
|
||||
print(f"Warning: Category '{category_name}' definition ignored (not list or dictionary).")
|
||||
|
||||
resolved_cache[category_name] = final_tags
|
||||
# We DON'T remove from recursion_guard here because the cache handles it
|
||||
return final_tags
|
||||
|
||||
|
||||
# --- The Node Class ---
|
||||
# Local utility imports
|
||||
from .prompt_formatter_utils import (
|
||||
find_yaml_file,
|
||||
parse_tag,
|
||||
clean_output_string,
|
||||
resolve_category_tags,
|
||||
INCLUDE_DIRECTIVE,
|
||||
TAGS_KEY
|
||||
)
|
||||
|
||||
class CategorizedPromptFormatter:
|
||||
"""
|
||||
A ComfyUI node to categorize tags from an input prompt based on a YAML file
|
||||
(supporting $include directives) and format them using a template.
|
||||
A ComfyUI node to categorize and format tags from a prompt using a YAML file.
|
||||
"""
|
||||
NODE_NAME = "Categorized Prompt Formatter"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
||||
"category_definition_file": ("STRING", {"default": "prompt_categories.yaml"}),
|
||||
"output_template": ("STRING", {"multiline": True, "default": "<|quality|>, <|character_num|>, <|person_details|>, <|clothing|>, <|setting|>"}),
|
||||
"output_template": ("STRING", {"multiline": True, "default": "<|quality|>, <|character|>, <|clothing|>, <|setting|>"}),
|
||||
},
|
||||
"optional": {
|
||||
"input_delimiter": ("STRING", {"default": ","}),
|
||||
@@ -206,194 +45,103 @@ class CategorizedPromptFormatter:
|
||||
FUNCTION = "format_prompt"
|
||||
CATEGORY = "text/filtering"
|
||||
|
||||
def format_prompt(self, prompt, category_definition_file, output_template,
|
||||
input_delimiter=",", output_delimiter=", ", strip_whitespace=True,
|
||||
case_sensitive_matching=False, handle_weights=True,
|
||||
match_underscores_spaces=True, disable_duplicates=False,
|
||||
unmatched_tag_handling="discard"):
|
||||
|
||||
def format_prompt(self, prompt, category_definition_file, output_template, **kwargs):
|
||||
# --- 1. Load & Resolve Category Definitions ---
|
||||
tag_to_categories_map = defaultdict(list)
|
||||
yaml_path = find_yaml_file(category_definition_file)
|
||||
yaml_path = find_yaml_file(category_definition_file, self.NODE_NAME)
|
||||
raw_yaml_data = None
|
||||
|
||||
if yaml_path and yaml_path.exists():
|
||||
if yaml_path:
|
||||
try:
|
||||
with open(yaml_path, 'r', encoding='utf-8') as f:
|
||||
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}")
|
||||
except yaml.YAMLError as e:
|
||||
print(f"Error parsing YAML file {yaml_path}: {e}")
|
||||
print(f"Warning [{self.NODE_NAME}]: YAML file '{yaml_path}' is not a dictionary.")
|
||||
raw_yaml_data = None
|
||||
except Exception as e:
|
||||
print(f"Error loading YAML file {yaml_path}: {e}")
|
||||
print(f"Error [{self.NODE_NAME}]: Loading YAML file {yaml_path}: {e}")
|
||||
|
||||
if raw_yaml_data:
|
||||
resolved_tags_cache = {}
|
||||
all_category_names = list(raw_yaml_data.keys())
|
||||
for cat_name in list(raw_yaml_data.keys()):
|
||||
if str(cat_name).strip() not in [INCLUDE_DIRECTIVE, TAGS_KEY]:
|
||||
resolve_category_tags(cat_name, raw_yaml_data, resolved_tags_cache, self.NODE_NAME)
|
||||
|
||||
# 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 or empty/invalid. Proceeding without categorization.")
|
||||
# tag_to_categories_map remains empty
|
||||
case_sensitive = kwargs.get("case_sensitive_matching", False)
|
||||
for cat_name, tags_set in resolved_tags_cache.items():
|
||||
for tag in tags_set:
|
||||
key = tag if case_sensitive else tag.lower()
|
||||
if key: tag_to_categories_map[key].append(cat_name)
|
||||
|
||||
for key in tag_to_categories_map:
|
||||
tag_to_categories_map[key] = list(set(tag_to_categories_map[key]))
|
||||
|
||||
# --- 2. Parse Input Prompt & Categorize ---
|
||||
categorized_tags = defaultdict(list)
|
||||
all_input_tags_original = set()
|
||||
processed_input_tags = [] # Store (original_tag, base_tag for matching)
|
||||
all_input_tags = set()
|
||||
matched_tags = set()
|
||||
|
||||
if prompt:
|
||||
raw_tags = prompt.split(input_delimiter)
|
||||
for raw_tag in raw_tags:
|
||||
tag_original = raw_tag.strip() if strip_whitespace else raw_tag
|
||||
if not tag_original: continue
|
||||
handle_weights = kwargs.get("handle_weights", True)
|
||||
match_spaces = kwargs.get("match_underscores_spaces", True)
|
||||
case_sensitive = kwargs.get("case_sensitive_matching", False)
|
||||
|
||||
all_input_tags_original.add(tag_original)
|
||||
original_parsed, base_parsed = parse_tag(tag_original, handle_weights)
|
||||
for raw_tag in prompt.split(kwargs.get("input_delimiter", ",")):
|
||||
tag_original = raw_tag.strip() if kwargs.get("strip_whitespace", True) else raw_tag
|
||||
if not tag_original: continue
|
||||
all_input_tags.add(tag_original)
|
||||
|
||||
# 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))
|
||||
_, base_parsed = parse_tag(tag_original, handle_weights)
|
||||
lookup_key = base_parsed if case_sensitive else base_parsed.lower()
|
||||
|
||||
variants = {lookup_key}
|
||||
if match_spaces:
|
||||
variants.add(lookup_key.replace(' ', '_'))
|
||||
variants.add(lookup_key.replace('_', ' '))
|
||||
|
||||
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 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 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) # Mark the ORIGINAL tag as matched
|
||||
for variant in variants:
|
||||
if variant in tag_to_categories_map:
|
||||
for category in tag_to_categories_map[variant]:
|
||||
categorized_tags[category].append(tag_original)
|
||||
matched_tags.add(tag_original)
|
||||
break
|
||||
|
||||
# --- 3. Process Template ---
|
||||
placeholder_regex = r"<\|([^:]+?)(?::(-?\d+))?\|>"
|
||||
|
||||
formatted_output_string = ""
|
||||
already_added_tags = set()
|
||||
used_in_template_tags = set() # Tracks tags referenced by template (before dedup)
|
||||
|
||||
output_delimiter = kwargs.get("output_delimiter", ", ")
|
||||
disable_dupes = kwargs.get("disable_duplicates", False)
|
||||
added_tags = set()
|
||||
used_tags = set()
|
||||
result_parts = []
|
||||
last_end = 0
|
||||
# Use the updated regex with finditer
|
||||
for match in re.finditer(placeholder_regex, output_template):
|
||||
category_name = match.group(1).strip()
|
||||
limit_str = match.group(2) # Captures signed number string (e.g., "5", "-3") or None
|
||||
start, end = match.span()
|
||||
|
||||
# Append the literal text segment before the placeholder
|
||||
for match in re.finditer(r"<\|([^:]+?)(?::(-?\d+))?\|>", output_template):
|
||||
cat_name, limit_str, start, end = match.group(1).strip(), match.group(2), *match.span()
|
||||
result_parts.append(output_template[last_end:start])
|
||||
|
||||
limit = int(limit_str) if limit_str else None
|
||||
tags_for_cat = categorized_tags.get(cat_name, [])
|
||||
|
||||
tags_to_process = tags_for_cat
|
||||
if limit is not None:
|
||||
tags_to_process = tags_for_cat[:limit] if limit > 0 else (tags_for_cat[limit:] if limit < 0 else [])
|
||||
|
||||
limit = None
|
||||
if limit_str:
|
||||
try:
|
||||
limit = int(limit_str)
|
||||
except ValueError:
|
||||
print(f"Warning: Invalid limit format '{limit_str}' for category '{category_name}'. Ignoring limit.")
|
||||
limit = None # Revert to no limit on error
|
||||
|
||||
# Get the list of tags associated with this category name from the input prompt
|
||||
tags_for_category = categorized_tags.get(category_name, [])
|
||||
tags_to_process = [] # Initialize empty
|
||||
|
||||
# --- Apply Slicing Logic Based on Limit Sign ---
|
||||
if limit is None:
|
||||
# No limit specified, take all
|
||||
tags_to_process = tags_for_category
|
||||
elif limit == 0:
|
||||
# Limit is zero, take none
|
||||
tags_to_process = []
|
||||
elif limit > 0:
|
||||
# Positive limit: take the first N tags
|
||||
tags_to_process = tags_for_category[:limit]
|
||||
else: # limit < 0
|
||||
# Negative limit: take the last N (abs(limit)) tags
|
||||
tags_to_process = tags_for_category[limit:]
|
||||
|
||||
# --- Process the selected tags ---
|
||||
if tags_to_process:
|
||||
# Mark only the tags we are potentially using (after limit/slicing) as 'used'
|
||||
used_in_template_tags.update(tags_to_process)
|
||||
|
||||
tags_to_join = []
|
||||
if disable_duplicates:
|
||||
# Filter the selected list against already added tags
|
||||
for tag in tags_to_process:
|
||||
if tag not in already_added_tags:
|
||||
tags_to_join.append(tag)
|
||||
already_added_tags.add(tag) # Mark as added
|
||||
else:
|
||||
# Use all tags from the (potentially sliced) list
|
||||
tags_to_join = tags_to_process
|
||||
# already_added_tags.update(tags_to_join) # Optional if disable_duplicates is False
|
||||
|
||||
if tags_to_join:
|
||||
joined_tags = output_delimiter.join(tags_to_join)
|
||||
result_parts.append(joined_tags)
|
||||
|
||||
used_tags.update(tags_to_process)
|
||||
tags_to_join = [t for t in tags_to_process if not (disable_dupes and t in added_tags and not added_tags.add(t))]
|
||||
if tags_to_join: result_parts.append(output_delimiter.join(tags_to_join))
|
||||
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)
|
||||
formatted_prompt = "".join(result_parts)
|
||||
|
||||
# --- 4. Handle Unmatched Tags ---
|
||||
# This section remains the same
|
||||
rejected_tags_list = []
|
||||
if unmatched_tag_handling != "discard":
|
||||
rejected_tags_list = list(all_input_tags_original - used_in_template_tags)
|
||||
unmatched_handling = kwargs.get("unmatched_tag_handling", "discard")
|
||||
rejected_prompt = ""
|
||||
if unmatched_handling != "discard":
|
||||
rejected_tags = list(all_input_tags - used_tags)
|
||||
if rejected_tags:
|
||||
rejected_prompt = output_delimiter.join(rejected_tags)
|
||||
if unmatched_handling == "append_end":
|
||||
if formatted_prompt: formatted_prompt += output_delimiter
|
||||
formatted_prompt += rejected_prompt
|
||||
rejected_prompt = ""
|
||||
|
||||
rejected_output_string = ""
|
||||
if rejected_tags_list:
|
||||
rejected_output_string = output_delimiter.join(rejected_tags_list)
|
||||
if unmatched_tag_handling == "append_end":
|
||||
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:
|
||||
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
|
||||
|
||||
|
||||
# --- 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 (clean_output_string(formatted_prompt, output_delimiter), clean_output_string(rejected_prompt, output_delimiter))
|
||||
@@ -1,281 +1,94 @@
|
||||
# /ComfyUI-Prompt-Formatter/categorized_random_prompt_formatter.py
|
||||
#/ComfyUI-Prompt-Formatter/categorized_random_prompt_formatter.py
|
||||
|
||||
import yaml
|
||||
import re
|
||||
import os
|
||||
import sys
|
||||
import random
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
Local utility imports
|
||||
|
||||
# --- Dependency Check ---
|
||||
try:
|
||||
import yaml
|
||||
except ImportError:
|
||||
print("PyYAML is not installed. Please run install.py or install it manually.", file=sys.stderr)
|
||||
|
||||
# --- Constants ---
|
||||
NODE_NAME_RANDOM = "Categorized Random Prompt Formatter" # Specific name for this node
|
||||
INCLUDE_DIRECTIVE = "$include"
|
||||
TAGS_KEY = "tags"
|
||||
|
||||
# --- Helper Functions (Copied/Adapted from previous node) ---
|
||||
|
||||
def get_node_directory_random():
|
||||
"""Gets the directory path of the current node."""
|
||||
return Path(__file__).parent
|
||||
|
||||
def find_yaml_file_random(filename):
|
||||
"""Finds the YAML file (adapted naming)."""
|
||||
if not filename: return None
|
||||
if Path(filename).is_absolute():
|
||||
if Path(filename).is_file(): return Path(filename)
|
||||
else: print(f"Warning [RandomNode]: Absolute path specified but not found: {filename}"); return None
|
||||
node_dir = get_node_directory_random()
|
||||
relative_to_node = node_dir / filename
|
||||
if relative_to_node.is_file(): return relative_to_node
|
||||
try:
|
||||
current_dir = Path.cwd()
|
||||
input_dir = current_dir / 'input'
|
||||
if input_dir.exists():
|
||||
input_dir_path = input_dir / filename
|
||||
if input_dir_path.is_file(): return input_dir_path
|
||||
search_dir = node_dir
|
||||
for _ in range(5):
|
||||
input_dir = search_dir / 'input'
|
||||
if input_dir.exists():
|
||||
input_dir_path = input_dir / filename
|
||||
if input_dir_path.is_file(): return input_dir_path
|
||||
if (search_dir / 'ComfyUI').exists() or (search_dir / 'main.py').exists():
|
||||
input_dir = search_dir / 'input'
|
||||
if input_dir.exists():
|
||||
input_dir_path = input_dir / filename
|
||||
if input_dir_path.is_file(): return input_dir_path
|
||||
break
|
||||
parent = search_dir.parent
|
||||
if parent == search_dir: break
|
||||
search_dir = parent
|
||||
except Exception as e: print(f"Warning [RandomNode]: Error searching for ComfyUI input directory: {e}")
|
||||
print(f"Warning [RandomNode]: YAML file '{filename}' not found as absolute, relative to node, or in input dir.")
|
||||
return None
|
||||
|
||||
def resolve_category_tags_random(category_name, yaml_data, resolved_cache, recursion_guard=None): # Or _random suffix
|
||||
"""
|
||||
Recursively resolves tags for a category, handling $include directives,
|
||||
$category list includes, and inline $category expansion within strings.
|
||||
Returns a set of unique tags (strings) for the category.
|
||||
"""
|
||||
category_name = str(category_name).strip()
|
||||
if recursion_guard is None: recursion_guard = set()
|
||||
# Basic recursion depth limit as extra safety
|
||||
if len(recursion_guard) > 20:
|
||||
print(f"Warning: Recursion depth limit exceeded processing '{category_name}'. Returning empty set.")
|
||||
return set()
|
||||
if category_name in recursion_guard:
|
||||
print(f"Warning: Circular dependency detected involving category '{category_name}'.")
|
||||
return set()
|
||||
if category_name in resolved_cache: return resolved_cache[category_name]
|
||||
if category_name not in yaml_data:
|
||||
# Don't warn here, allows optional categories in references
|
||||
return set()
|
||||
|
||||
recursion_guard.add(category_name)
|
||||
category_data = yaml_data[category_name]
|
||||
final_tags = set()
|
||||
|
||||
if isinstance(category_data, list):
|
||||
for item in category_data:
|
||||
item_str = str(item).strip()
|
||||
if not item_str: continue
|
||||
|
||||
# Check for full category include first (starts with $, no spaces, alphanumeric name)
|
||||
is_full_include = item_str.startswith('$') and len(item_str) > 1 and ' ' not in item_str and item_str[1:].replace('_', '').isalnum()
|
||||
|
||||
if is_full_include:
|
||||
# --- Handle Full Category Include: $other_category ---
|
||||
included_category = item_str[1:]
|
||||
final_tags.update(resolve_category_tags_random(included_category, yaml_data, resolved_cache, recursion_guard.copy())) # Use correct function name
|
||||
else:
|
||||
# --- Handle Inline Expansion: "prefix $category suffix" ---
|
||||
# Use regex to find the first $word pattern
|
||||
# Limit to one expansion per string for simplicity
|
||||
match = re.search(r'\$(\w+)', item_str) # \w+ matches letters, numbers, underscore
|
||||
|
||||
if match:
|
||||
ref_category_name = match.group(1)
|
||||
placeholder = match.group(0) # The full '$word' matched
|
||||
|
||||
# Resolve the referenced category's tags
|
||||
resolved_ref_tags = resolve_category_tags_random(ref_category_name, yaml_data, resolved_cache, recursion_guard.copy())
|
||||
|
||||
if not resolved_ref_tags:
|
||||
print(f"Warning: Inline reference ${ref_category_name} in '{item_str}' resolved to an empty set or category not found. Skipping expansion for this item.")
|
||||
else:
|
||||
# Substitute each resolved tag into the original string pattern
|
||||
for resolved_tag in resolved_ref_tags:
|
||||
# Perform replacement. Use count=1 if placeholder could appear multiple times.
|
||||
new_tag = item_str.replace(placeholder, str(resolved_tag).strip(), 1)
|
||||
final_tags.add(new_tag)
|
||||
else:
|
||||
# No $category reference found, treat as a literal tag
|
||||
final_tags.add(item_str)
|
||||
|
||||
# --- Keep supporting the old dictionary format ---
|
||||
elif isinstance(category_data, dict):
|
||||
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:
|
||||
final_tags.update(resolve_category_tags_random(included_category, yaml_data, resolved_cache, recursion_guard.copy()))
|
||||
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())
|
||||
|
||||
else:
|
||||
print(f"Warning: Category '{category_name}' definition ignored (not list or dictionary).")
|
||||
|
||||
resolved_cache[category_name] = final_tags
|
||||
# We DON'T remove from recursion_guard here because the cache handles it
|
||||
return final_tags
|
||||
|
||||
def clean_output_string_random(text, delimiter=", "):
|
||||
"""Cleans up the final output string (adapted naming)."""
|
||||
if not text: return ""
|
||||
text = text.strip().strip(delimiter.strip()).strip()
|
||||
delimiter_pattern = r'\s*' + re.escape(delimiter.strip()) + r'\s*'
|
||||
text = re.sub(f'({delimiter_pattern})+', delimiter, text)
|
||||
if text.startswith(delimiter): text = text[len(delimiter):].lstrip()
|
||||
if text.endswith(delimiter.rstrip()): text = text[:-len(delimiter.rstrip())].rstrip()
|
||||
return text
|
||||
|
||||
# --- The Random Node Class ---
|
||||
from .prompt_formatter_utils import (
|
||||
find_yaml_file,
|
||||
clean_output_string,
|
||||
resolve_category_tags,
|
||||
INCLUDE_DIRECTIVE,
|
||||
TAGS_KEY
|
||||
)
|
||||
|
||||
class CategorizedRandomPromptFormatter:
|
||||
"""
|
||||
A ComfyUI node to generate random prompts by selecting tags from categories
|
||||
defined in a YAML file, based on a template and a seed.
|
||||
"""
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"category_definition_file": ("STRING", {"default": "prompt_categories.yaml"}),
|
||||
"output_template": ("STRING", {
|
||||
"multiline": True,
|
||||
"default": "<|quality:1|>, <|character_num:1|>, <|person_details:3|>, <|clothing:1|>, <|setting:1|>, <|style:1|>"
|
||||
}),
|
||||
"seed": ("INT", {"default": -1, "min": -1, "max": 0xffffffffffffffff}),
|
||||
},
|
||||
"optional": {
|
||||
"output_delimiter": ("STRING", {"default": ", "}),
|
||||
}
|
||||
}
|
||||
"""
|
||||
A ComfyUI node to generate random prompts from categories defined in a YAML file.
|
||||
"""
|
||||
NODE_NAME = "Categorized Random Prompt Formatter"
|
||||
|
||||
RETURN_TYPES = ("STRING", "INT")
|
||||
RETURN_NAMES = ("random_prompt", "used_seed")
|
||||
FUNCTION = "generate_prompt"
|
||||
CATEGORY = "text/generation"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"category_definition_file": ("STRING", {"default": "prompt_categories.yaml"}),
|
||||
"output_template": ("STRING", {
|
||||
"multiline": True,
|
||||
"default": "<|quality:1|>, <|character:1|>, <|details:3|>, <|setting:1|>"
|
||||
}),
|
||||
"seed": ("INT", {"default": -1, "min": -1, "max": 0xffffffffffffffff}),
|
||||
},
|
||||
"optional": { "output_delimiter": ("STRING", {"default": ", "}) }
|
||||
}
|
||||
|
||||
def generate_prompt(self, category_definition_file, output_template, seed,
|
||||
output_delimiter=", ",
|
||||
):
|
||||
RETURN_TYPES = ("STRING", "INT")
|
||||
RETURN_NAMES = ("random_prompt", "used_seed")
|
||||
FUNCTION = "generate_prompt"
|
||||
CATEGORY = "text/generation"
|
||||
|
||||
# --- 1. Handle Seed ---
|
||||
if seed == -1:
|
||||
used_seed = random.randint(0, 0xffffffffffffffff)
|
||||
else:
|
||||
used_seed = seed
|
||||
# Initialize a dedicated random number generator for this execution
|
||||
rng = random.Random(used_seed)
|
||||
def generate_prompt(self, category_definition_file, output_template, seed, output_delimiter=", "):
|
||||
# --- 1. Handle Seed ---
|
||||
used_seed = random.randint(0, 0xffffffffffffffff) if seed == -1 else seed
|
||||
rng = random.Random(used_seed)
|
||||
|
||||
# --- 2. Load & Resolve Category Definitions ---
|
||||
resolved_categories = {} # Store category_name -> set_of_tags
|
||||
yaml_path = find_yaml_file_random(category_definition_file)
|
||||
raw_yaml_data = None
|
||||
# --- 2. Load & Resolve Categories ---
|
||||
resolved_categories = {}
|
||||
yaml_path = find_yaml_file(category_definition_file, self.NODE_NAME)
|
||||
raw_yaml_data = None
|
||||
|
||||
if yaml_path and yaml_path.exists():
|
||||
try:
|
||||
with open(yaml_path, 'r', encoding='utf-8') as f:
|
||||
raw_yaml_data = yaml.safe_load(f)
|
||||
if not isinstance(raw_yaml_data, dict):
|
||||
print(f"Warning [RandomNode]: YAML file '{yaml_path}' is not a dictionary. Cannot process.")
|
||||
raw_yaml_data = None
|
||||
except Exception as e:
|
||||
print(f"Error [RandomNode]: Loading YAML file {yaml_path}: {e}")
|
||||
if yaml_path:
|
||||
try:
|
||||
with open(yaml_path, 'r', encoding='utf-8') as f:
|
||||
raw_yaml_data = yaml.safe_load(f)
|
||||
if not isinstance(raw_yaml_data, dict):
|
||||
print(f"Warning [{self.NODE_NAME}]: YAML file '{yaml_path}' is not a dictionary.")
|
||||
raw_yaml_data = None
|
||||
except Exception as e:
|
||||
print(f"Error [{self.NODE_NAME}]: Loading YAML file {yaml_path}: {e}")
|
||||
|
||||
if raw_yaml_data:
|
||||
resolved_tags_cache = {}
|
||||
all_category_names = list(raw_yaml_data.keys())
|
||||
for category_name in all_category_names:
|
||||
if str(category_name).strip() != INCLUDE_DIRECTIVE and str(category_name).strip() != TAGS_KEY:
|
||||
# Resolve tags for each category using the helper
|
||||
resolved_categories[str(category_name).strip()] = resolve_category_tags_random(
|
||||
category_name, raw_yaml_data, resolved_tags_cache
|
||||
)
|
||||
else:
|
||||
print(f"Warning [RandomNode]: Category file '{category_definition_file}' not found or empty. Cannot generate prompt.")
|
||||
return ("", used_seed) # Return empty if no categories loaded
|
||||
if raw_yaml_data:
|
||||
cache = {}
|
||||
for cat_name in list(raw_yaml_data.keys()):
|
||||
if str(cat_name).strip() not in [INCLUDE_DIRECTIVE, TAGS_KEY]:
|
||||
resolved_categories[str(cat_name).strip()] = resolve_category_tags(cat_name, raw_yaml_data, cache, self.NODE_NAME)
|
||||
else:
|
||||
return ("", used_seed)
|
||||
|
||||
# --- 3. Process Template and Generate Random Tags ---
|
||||
# Regex to capture category name (group 1) and optional count (group 2)
|
||||
# Only positive counts are meaningful here. Default count is 1.
|
||||
placeholder_regex = r"<\|([^:]+?)(?::(\d+))?\|>"
|
||||
# --- 3. Process Template & Generate Prompt ---
|
||||
result_parts = []
|
||||
last_end = 0
|
||||
for match in re.finditer(r"<\|([^:]+?)(?::(\d+))?\|>", output_template):
|
||||
cat_name, count_str, start, end = match.group(1).strip(), match.group(2), *match.span()
|
||||
result_parts.append(output_template[last_end:start])
|
||||
|
||||
num_to_pick = int(count_str) if count_str and int(count_str) >= 0 else 1
|
||||
if num_to_pick > 0:
|
||||
available_tags = list(resolved_categories.get(cat_name, []))
|
||||
if available_tags:
|
||||
sample_count = min(num_to_pick, len(available_tags))
|
||||
tags_to_join = rng.sample(available_tags, sample_count)
|
||||
result_parts.append(output_delimiter.join(tags_to_join))
|
||||
else:
|
||||
print(f"Warning [{self.NODE_NAME}]: Category '{cat_name}' not found or empty.")
|
||||
|
||||
last_end = end
|
||||
|
||||
result_parts.append(output_template[last_end:])
|
||||
final_prompt = clean_output_string("".join(result_parts), output_delimiter)
|
||||
|
||||
return (final_prompt, used_seed)
|
||||
|
||||
generated_prompt_string = ""
|
||||
result_parts = []
|
||||
last_end = 0
|
||||
|
||||
for match in re.finditer(placeholder_regex, output_template):
|
||||
category_name = match.group(1).strip()
|
||||
count_str = match.group(2) # Optional count N
|
||||
start, end = match.span()
|
||||
|
||||
# Append literal text before placeholder
|
||||
result_parts.append(output_template[last_end:start])
|
||||
|
||||
# Determine the number of tags to pick (default 1)
|
||||
num_to_pick = 1 # Default
|
||||
if count_str:
|
||||
try:
|
||||
count = int(count_str)
|
||||
if count >= 0: # Allow 0 to explicitly pick none
|
||||
num_to_pick = count
|
||||
else:
|
||||
print(f"Warning [RandomNode]: Negative count '{count_str}' for category '{category_name}' invalid. Defaulting to 1.")
|
||||
num_to_pick = 1 # Default to 1 if negative
|
||||
except ValueError:
|
||||
print(f"Warning [RandomNode]: Invalid count format '{count_str}' for category '{category_name}'. Defaulting to 1.")
|
||||
num_to_pick = 1
|
||||
|
||||
tags_to_join = []
|
||||
if num_to_pick > 0:
|
||||
# Get the set of available tags for this category
|
||||
available_tags_set = resolved_categories.get(category_name)
|
||||
|
||||
if available_tags_set and len(available_tags_set) > 0:
|
||||
# Ensure we don't try to sample more tags than available
|
||||
actual_num_to_sample = min(num_to_pick, len(available_tags_set))
|
||||
|
||||
# Convert set to list for sampling, then sample using the seeded RNG
|
||||
tags_to_join = rng.sample(list(available_tags_set), actual_num_to_sample)
|
||||
elif not available_tags_set:
|
||||
print(f"Warning [RandomNode]: Category '{category_name}' requested in template but not found in YAML.")
|
||||
|
||||
|
||||
# Join the randomly selected tags
|
||||
if tags_to_join:
|
||||
joined_tags = output_delimiter.join(tags_to_join)
|
||||
result_parts.append(joined_tags)
|
||||
# If num_to_pick was 0 or category empty/not found, nothing is appended for this placeholder
|
||||
|
||||
last_end = end
|
||||
|
||||
# Append any remaining literal text
|
||||
result_parts.append(output_template[last_end:])
|
||||
generated_prompt_string = "".join(result_parts)
|
||||
|
||||
# --- 4. Cleanup and Return ---
|
||||
final_prompt = clean_output_string_random(generated_prompt_string, output_delimiter)
|
||||
|
||||
return (final_prompt, used_seed)
|
||||
|
||||
@@ -0,0 +1,156 @@
|
||||
# /ComfyUI-Prompt-Formatter/prompt_formatter_utils.py
|
||||
|
||||
import yaml
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# --- Dependency Check ---
|
||||
try:
|
||||
import yaml
|
||||
except ImportError:
|
||||
# This will be printed on the console when ComfyUI loads the nodes
|
||||
print("PyYAML is not installed. Please run install.py or install it manually.", file=sys.stderr)
|
||||
# Re-raise the error to prevent nodes from being registered if the dependency is missing
|
||||
raise ImportError("PyYAML is required for ComfyUI-Prompt-Formatter nodes.")
|
||||
|
||||
# --- Constants ---
|
||||
INCLUDE_DIRECTIVE = "$include"
|
||||
TAGS_KEY = "tags"
|
||||
|
||||
# --- Path Utilities ---
|
||||
|
||||
def get_node_directory():
|
||||
"""Gets the directory path of the custom node pack."""
|
||||
return Path(__file__).parent
|
||||
|
||||
def find_yaml_file(filename: str, node_name: str = "Prompt Formatter"):
|
||||
"""
|
||||
Finds a YAML file by checking for an absolute path, a path relative to the node,
|
||||
and a path within various possible 'input' directories.
|
||||
"""
|
||||
if not filename:
|
||||
return None
|
||||
|
||||
if Path(filename).is_absolute():
|
||||
if Path(filename).is_file(): return Path(filename)
|
||||
else:
|
||||
print(f"Warning [{node_name}]: 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
|
||||
|
||||
try:
|
||||
search_dir = Path.cwd()
|
||||
# Look in standard ComfyUI input dir first
|
||||
input_dir_path = search_dir / 'input' / filename
|
||||
if input_dir_path.is_file(): return input_dir_path
|
||||
|
||||
# Fallback search moving up from the node directory to find ComfyUI root
|
||||
search_dir = node_dir
|
||||
for _ in range(5):
|
||||
is_comfy_root = (search_dir / 'ComfyUI').exists() or (search_dir / 'main.py').exists()
|
||||
if is_comfy_root:
|
||||
input_dir_path = search_dir / 'input' / filename
|
||||
if input_dir_path.is_file(): return input_dir_path
|
||||
break
|
||||
|
||||
parent = search_dir.parent
|
||||
if parent == search_dir: break
|
||||
search_dir = parent
|
||||
except Exception as e:
|
||||
print(f"Warning [{node_name}]: Error searching for ComfyUI input directory: {e}")
|
||||
|
||||
print(f"Warning [{node_name}]: YAML file '{filename}' not found as absolute, relative to node, or in standard input directories.")
|
||||
return None
|
||||
|
||||
# --- String & Tag Processing ---
|
||||
|
||||
def clean_output_string(text: str, delimiter: str = ", "):
|
||||
"""Cleans up the final output string by normalizing delimiters and trimming whitespace."""
|
||||
if not text: return ""
|
||||
text = text.strip().strip(delimiter.strip()).strip()
|
||||
delimiter_pattern = r'\s*' + re.escape(delimiter.strip()) + r'\s*'
|
||||
text = re.sub(f'({delimiter_pattern})+', delimiter, text)
|
||||
return text.strip(delimiter.strip()).strip()
|
||||
|
||||
def parse_tag(tag: str, handle_weights: bool = True):
|
||||
"""
|
||||
Parses a tag to extract the base tag (without weights/emphasis).
|
||||
"""
|
||||
original_tag = str(tag).strip()
|
||||
if not handle_weights: return original_tag, original_tag
|
||||
|
||||
base_tag = original_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()
|
||||
nested_match = re.match(r"^\s*[\(\[\{](.*?)[\)\]\}]\s*$", base_tag)
|
||||
if nested_match: base_tag = nested_match.group(1).strip()
|
||||
return original_tag, base_tag
|
||||
|
||||
return original_tag, base_tag.strip()
|
||||
|
||||
# --- YAML Category Resolution ---
|
||||
|
||||
def resolve_category_tags(category_name: str, yaml_data: dict, resolved_cache: dict, node_name: str = "Prompt Formatter", recursion_guard: set = None):
|
||||
"""
|
||||
Recursively resolves tags for a category, handling includes and inline expansion.
|
||||
"""
|
||||
category_name = str(category_name).strip()
|
||||
if recursion_guard is None: recursion_guard = set()
|
||||
|
||||
if len(recursion_guard) > 20:
|
||||
print(f"Warning [{node_name}]: Recursion depth limit exceeded for '{category_name}'.")
|
||||
return set()
|
||||
if category_name in recursion_guard:
|
||||
print(f"Warning [{node_name}]: Circular dependency for category '{category_name}'.")
|
||||
return set()
|
||||
if category_name in resolved_cache: return resolved_cache[category_name]
|
||||
if category_name not in yaml_data: return set()
|
||||
|
||||
recursion_guard.add(category_name)
|
||||
category_data = yaml_data[category_name]
|
||||
final_tags = set()
|
||||
|
||||
if isinstance(category_data, list):
|
||||
for item in category_data:
|
||||
item_str = str(item).strip()
|
||||
if not item_str: continue
|
||||
|
||||
is_full_include = item_str.startswith('$') and len(item_str) > 1 and ' ' not in item_str and item_str[1:].replace('_', '').isalnum()
|
||||
if is_full_include:
|
||||
final_tags.update(resolve_category_tags(item_str[1:], yaml_data, resolved_cache, node_name, recursion_guard.copy()))
|
||||
else:
|
||||
match = re.search(r'\$(\w+)', item_str)
|
||||
if match:
|
||||
ref_cat_name = match.group(1)
|
||||
resolved_ref_tags = resolve_category_tags(ref_cat_name, yaml_data, resolved_cache, node_name, recursion_guard.copy())
|
||||
if not resolved_ref_tags:
|
||||
print(f"Warning [{node_name}]: Inline reference ${ref_cat_name} in '{item_str}' is empty or not found.")
|
||||
for r_tag in resolved_ref_tags:
|
||||
final_tags.add(item_str.replace(match.group(0), str(r_tag).strip(), 1))
|
||||
else:
|
||||
final_tags.add(item_str)
|
||||
elif isinstance(category_data, dict):
|
||||
if INCLUDE_DIRECTIVE in category_data:
|
||||
includes = category_data[INCLUDE_DIRECTIVE]
|
||||
include_list = [str(inc).strip() for inc in includes] if isinstance(includes, list) else [str(includes).strip()]
|
||||
for included_category in include_list:
|
||||
final_tags.update(resolve_category_tags(included_category, yaml_data, resolved_cache, node_name, recursion_guard.copy()))
|
||||
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())
|
||||
else:
|
||||
print(f"Warning [{node_name}]: Category '{category_name}' type '{type(category_data).__name__}' is not a list or dict.")
|
||||
|
||||
resolved_cache[category_name] = final_tags
|
||||
return final_tags
|
||||
Reference in New Issue
Block a user