From 8a2bc6f44da967c2962fe4031e91807568a827ff Mon Sep 17 00:00:00 2001 From: Legende <68666984+xLegende@users.noreply.github.com> Date: Tue, 10 Jun 2025 19:29:04 +0200 Subject: [PATCH] Add files via upload refactor --- __init__.py | 54 ++-- categorized_prompt_analyzer.py | 311 +++++-------------- categorized_prompt_formatter.py | 414 +++++-------------------- categorized_random_prompt_formatter.py | 345 +++++---------------- prompt_formatter_utils.py | 156 ++++++++++ 5 files changed, 416 insertions(+), 864 deletions(-) create mode 100644 prompt_formatter_utils.py diff --git a/__init__.py b/__init__.py index 8b7710b..8220db7 100644 --- a/__init__.py +++ b/__init__.py @@ -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") \ No newline at end of file diff --git a/categorized_prompt_analyzer.py b/categorized_prompt_analyzer.py index f4b0a8a..d371ea3 100644 --- a/categorized_prompt_analyzer.py +++ b/categorized_prompt_analyzer.py @@ -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) \ No newline at end of file + return (summary, total_count, details_str, unmatched_str) \ No newline at end of file diff --git a/categorized_prompt_formatter.py b/categorized_prompt_formatter.py index 310c47b..e8f4b75 100644 --- a/categorized_prompt_formatter.py +++ b/categorized_prompt_formatter.py @@ -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)) \ No newline at end of file diff --git a/categorized_random_prompt_formatter.py b/categorized_random_prompt_formatter.py index 4669caf..268df0d 100644 --- a/categorized_random_prompt_formatter.py +++ b/categorized_random_prompt_formatter.py @@ -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) \ No newline at end of file + \ No newline at end of file diff --git a/prompt_formatter_utils.py b/prompt_formatter_utils.py new file mode 100644 index 0000000..ecd5d17 --- /dev/null +++ b/prompt_formatter_utils.py @@ -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 \ No newline at end of file