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