diff --git a/README.md b/README.md new file mode 100644 index 0000000..6aeaa7a --- /dev/null +++ b/README.md @@ -0,0 +1,80 @@ +# ComfyUI - Categorized Prompt Formatter Node + +A custom node for ComfyUI that allows you to filter and restructure text prompts based on tag categories defined in a YAML file. + +This node takes a standard comma-separated prompt, categorizes each tag according to your definitions, and then rebuilds the prompt based on a template you provide. It's useful for enforcing prompt structure, extracting specific concepts, or reordering elements dynamically. + +![Example Node UI (Conceptual - Actual UI may vary slightly)](placeholder_image.png) + +## Features + +* **YAML-based Categorization:** Define which tags belong to which categories in a simple `.yaml` file. +* **Multi-Category Tags:** A single tag (e.g., "blue eyes") can belong to multiple categories (e.g., `eyes`, `person_details`). +* **Template-Driven Output:** Structure your output prompt using placeholders like `<|category_name|>`. +* **Flexible Delimiters:** Configure input and output delimiters. +* **Weight/Emphasis Handling:** Attempts to preserve common syntax like `(tag:1.2)` and `[tag]` while using the base tag for categorization. +* **Case Sensitivity Option:** Choose whether matching should be case-sensitive or not. +* **Unmatched Tag Handling:** Decide whether to discard tags not used in the template, append them to the end, or output them separately. + +## Installation + +1. Clone this repository into your `ComfyUI/custom_nodes/` directory: + ```bash + cd ComfyUI/custom_nodes/ + git clone https://github.com/xLegende/ComfyUI-Prompt-Formatter.git + ``` + +2. Install dependencies: + * Navigate to the node's directory: `cd ComfyUI-Prompt-Formatter` + * Run the install script using your ComfyUI Python environment: + ```bash + # Example using a virtual environment: + # path/to/your/ComfyUI/python_env/bin/python install.py + # Or if your system python is the one used by ComfyUI: + python install.py + ``` + * This will install the `PyYAML` library if it's not already present. +3. Restart ComfyUI. + +The node "📝 Categorized Prompt Formatter" should now appear under the "text/filtering" category when you right-click or double-click on the ComfyUI canvas. + +## Usage + +1. **Prepare your YAML file:** + * Create a `.yaml` file (e.g., `my_categories.yaml`) or modify the included `prompt_categories.yaml`. + * The format is: + ```yaml + category_name_1: + - tag1 + - tag2 + - tag_shared + category_name_2: + - tag3 + - tag_shared # This tag now belongs to both categories + - "(tag with weight:1.1)" # Define tags exactly as they might appear + # ... etc + ``` + * Place this file somewhere ComfyUI can find it (e.g., in the `ComfyUI/input/` directory, or provide an absolute path, or place it in the `custom_nodes/ComfyUI-Prompt-Formatter/` folder). +2. **Add the Node:** Add the "📝 Categorized Prompt Formatter" node to your workflow. +3. **Connect Input:** Connect your source prompt string to the `prompt` input. +4. **Configure Widgets:** + * `category_definition_file`: Enter the name (e.g., `my_categories.yaml`) or full path to your YAML file. The node will search relative to itself, in `ComfyUI/input/`, and absolute paths. + * `output_template`: Define your desired output structure using `<|category_name|>` placeholders. Example: `<|quality|>, <|character_num|>, <|eyes|>, <|clothing|>, <|setting|>` + * Adjust delimiters, case sensitivity, weight handling, and unmatched tag handling as needed. +5. **Connect Output:** Connect the `formatted_prompt` output to the next node in your workflow (e.g., a CLIP Text Encode node). Use the optional `rejected_prompt` output for debugging or alternative workflows. + +## Example + +* **Input Prompt:** `"masterpiece, 1girl, (red eyes:1.1), long hair, wearing a blue dress, outdoors, detailed background"` +* **YAML (`prompt_categories.yaml` from this repo):** (Contains categories like `quality`, `character_num`, `eyes`, `hair_style`, `person_details`, `clothing`, `setting`) +* **Output Template:** `<|quality|>, <|character_num|>, <|eyes|>, <|clothing|>, <|setting|>` +* **Unmatched Handling:** `output_separately` + +* **Result (`formatted_prompt`):** `"masterpiece, 1girl, (red eyes:1.1), wearing a blue dress, outdoors"` +* **Result (`rejected_prompt`):** `"long hair, detailed background"` (Because `<|hair_style|>` wasn't in the template, and "detailed background" might not be defined or its category wasn't used) + +## Notes + +* The YAML file path finding tries to be robust, checking relative to the node, the ComfyUI `input` directory, and absolute paths. If it can't find the file, it will print a warning and proceed without categorization. +* Weight handling uses regular expressions to parse common patterns. Very complex nested structures might not parse correctly. The base tag extracted is used for lookup. +* If a tag belongs to multiple categories, and the template requests those multiple categories, the tag **will appear multiple times** in the output string by default. \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..34beb1d --- /dev/null +++ b/__init__.py @@ -0,0 +1,25 @@ +# /ComfyUI-Prompt-Formatter/__init__.py + +# Import the node class +from .categorized_prompt_formatter import CategorizedPromptFormatter, NODE_NAME + +# Mapping node class names to implementations +NODE_CLASS_MAPPINGS = { + NODE_NAME: CategorizedPromptFormatter +} + +# Mapping node class names to display names for the UI +NODE_DISPLAY_NAME_MAPPINGS = { + NODE_NAME: "📝 Categorized Prompt Formatter" # Add an emoji for fun! +} + +# A dictionary that contains all nodes in this file +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] + +print("...") +print("...") +print("### Loading: ComfyUI-Prompt-Formatter ###") +print("### Version: 1.0") # Optional version print +print(f"### Node: {NODE_NAME} -> {NODE_DISPLAY_NAME_MAPPINGS[NODE_NAME]}") +print("...") +print("...") \ No newline at end of file diff --git a/categorized_prompt_formatter.py b/categorized_prompt_formatter.py new file mode 100644 index 0000000..bb92e70 --- /dev/null +++ b/categorized_prompt_formatter.py @@ -0,0 +1,298 @@ +# /ComfyUI-Prompt-Formatter/categorized_prompt_formatter.py + +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) + # Optionally, try to run install.py automatically if possible, + # but this can be complex depending on environment. Best to rely on user action. + # raise ImportError("PyYAML is required but not installed. Run install.py.") from None + +# --- Constants --- +NODE_NAME = "Categorized Prompt Formatter" +NODE_REPO = "ComfyUI-Prompt-Formatter" # For finding relative paths + +# --- 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 # Explicit absolute path not found + + # 2. Check relative to the custom node's directory + 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. + try: + current_dir = Path.cwd() # Usually ComfyUI base dir when run normally + 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 + search_dir = node_dir + for _ in range(5): # Limit search depth + 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 (search_dir / 'input').exists(): + input_dir_path = search_dir / 'input' / filename + if input_dir_path.is_file(): + return input_dir_path + break # Stop searching up + parent = search_dir.parent + if parent == search_dir: break # Reached root + 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. + Handles common weighting and emphasis syntax like (tag:1.2), [tag], {tag}. + Returns (original_tag, base_tag). + """ + original_tag = tag.strip() + 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 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) + 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() + + +def clean_output_string(text, delimiter=", "): + """Cleans up the final output string.""" + if not text: + return "" + # Remove leading/trailing whitespace and delimiters + 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) + return text + +# --- 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. + """ + @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|>"}), + }, + "optional": { + "input_delimiter": ("STRING", {"default": ","}), + "output_delimiter": ("STRING", {"default": ", "}), + "strip_whitespace": ("BOOLEAN", {"default": True}), + "case_sensitive_matching": ("BOOLEAN", {"default": False}), + "handle_weights": ("BOOLEAN", {"default": True}), + "unmatched_tag_handling": (["discard", "append_end", "output_separately"], {"default": "discard"}), + } + } + + RETURN_TYPES = ("STRING", "STRING") + RETURN_NAMES = ("formatted_prompt", "rejected_prompt") + FUNCTION = "format_prompt" + CATEGORY = "text/filtering" # Or your preferred category + + def format_prompt(self, prompt, category_definition_file, output_template, + input_delimiter=",", output_delimiter=", ", strip_whitespace=True, + case_sensitive_matching=False, handle_weights=True, + unmatched_tag_handling="discard"): + + # --- 1. Load Category Definitions --- + tag_to_categories_map = defaultdict(list) + yaml_path = find_yaml_file(category_definition_file) + + 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() # Ensure string and strip + if not case_sensitive_matching: + processed_tag = processed_tag.lower() + if processed_tag: # Avoid empty tags + tag_to_categories_map[processed_tag].append(str(category).strip()) + except FileNotFoundError: + print(f"Error: YAML file not found at resolved path: {yaml_path}") + # Return empty prompts or raise an error? Let's return empty. + return ("", prompt) # Return original prompt as rejected if file fails + 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) + else: + print(f"Warning: Category definition file '{category_definition_file}' not found. Proceeding without categorization.") + # If no categories, all tags will be unmatched + + # --- 2. Parse Input Prompt & Categorize --- + categorized_tags = defaultdict(list) + all_input_tags_original = set() + processed_input_tags = [] # Store (original_tag, base_tag) + + 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 # Skip empty tags resulting from split + + all_input_tags_original.add(tag_original) + original_parsed, base_parsed = parse_tag(tag_original, handle_weights) + + lookup_tag = base_parsed if not case_sensitive_matching else base_parsed.lower() + processed_input_tags.append((original_parsed, lookup_tag)) + + + matched_original_tags = set() + for original_tag, lookup_tag in processed_input_tags: + categories = tag_to_categories_map.get(lookup_tag) + if categories: + for category in categories: + categorized_tags[category].append(original_tag) + matched_original_tags.add(original_tag) # Mark as matched if found in *any* category + + # --- 3. Process Template --- + formatted_output_string = output_template + placeholders = re.findall(r"<\|(.*?)\|>|", output_template) # Find placeholders like <|name|> + placeholders = [p for p in placeholders if p] # Filter out empty strings from regex artifact + + used_in_template_tags = set() + + # Use a temporary string builder approach to avoid issues with replacing overlapping parts + result_parts = [] + last_end = 0 + for match in re.finditer(r"<\|(.*?)\|>|", output_template): + placeholder_content = match.group(1) + start, end = match.span() + + # Append the text segment before the placeholder + result_parts.append(output_template[last_end:start]) + + if placeholder_content: + category_name = placeholder_content.strip() + tags_for_category = categorized_tags.get(category_name, []) + if tags_for_category: + joined_tags = output_delimiter.join(tags_for_category) + result_parts.append(joined_tags) + used_in_template_tags.update(tags_for_category) # Keep track of used tags + # If category not found or has no tags, effectively replaces with empty string + else: + # Append the matched placeholder itself if it wasn't valid (shouldn't happen with this regex but safer) + result_parts.append(match.group(0)) + + + last_end = end + + # Append any remaining text after the last placeholder + result_parts.append(output_template[last_end:]) + + formatted_output_string = "".join(result_parts) + + + # --- 4. Handle Unmatched Tags --- + rejected_tags_list = [] + if unmatched_tag_handling != "discard": + # Find tags that were in the input but *not* used in the template substitution + rejected_tags_list = list(all_input_tags_original - used_in_template_tags) + + rejected_output_string = "" + if 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): + 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 += " " + + + 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 diff --git a/install.py b/install.py new file mode 100644 index 0000000..9b081d7 --- /dev/null +++ b/install.py @@ -0,0 +1,36 @@ +# /ComfyUI-Prompt-Formatter/install.py + +import subprocess +import sys +import os + +def install_dependencies(): + print("...") + print("...") + print("Installing dependencies for ComfyUI-Prompt-Formatter") + print("...") + print("...") + try: + print("Installing PyYAML...") + # Use pip directly from the Python executable path + python_exe = sys.executable + subprocess.check_call([python_exe, '-m', 'pip', 'install', 'PyYAML']) + print("PyYAML installed successfully.") + except subprocess.CalledProcessError as e: + print(f"Error installing PyYAML: {e}") + print("Please try installing PyYAML manually:") + print(f"'{sys.executable}' -m pip install PyYAML") + except Exception as e: + print(f"An unexpected error occurred during installation: {e}") + +# Check if run as main script +if __name__ == "__main__": + install_dependencies() + + # Check if running in Colab or Paperspace to potentially avoid restart prompt + if 'COLAB_GPU' in os.environ or 'PAPERSPACE_NOTEBOOK_ID' in os.environ: + print("\nDependency installation complete. Please proceed.") + else: + # Basic prompt for non-automated environments + print("\nDependency installation finished.") + print("If ComfyUI is running, you may need to restart it for the changes to take effect.") \ No newline at end of file diff --git a/prompt_categories.yaml b/prompt_categories.yaml new file mode 100644 index 0000000..bf82ead --- /dev/null +++ b/prompt_categories.yaml @@ -0,0 +1,112 @@ +# /ComfyUI-Prompt-Formatter/prompt_categories.yaml +# Example category definitions for the Categorized Prompt Formatter node. + +# --- Quality Tags --- +quality: + - masterpiece + - best quality + - high quality + - absurdres + - highres + - ultra detailed + +# --- Style Tags --- +style: + - photorealistic + - photograph + - anime style + - manga style + - illustration + - concept art + - watercolor + - oil painting + +# --- Character Basics --- +character_num: + - 1girl + - 1boy + - 2girls + - solo + - multiple boys + +# --- Appearance Details --- +eyes: + - blue eyes + - green eyes + - brown eyes + - red eyes + - heterochromia + - glowing eyes + - closed eyes + +hair_color: + - blonde hair + - brown hair + - black hair + - red hair + - pink hair + - silver hair + +hair_style: + - long hair + - short hair + - medium hair + - ponytail + - twintails + - bob cut + - messy hair + +# --- General Person Details +person_details: + - blue eyes + - red eyes + - blonde hair + - long hair + - smiling + - looking at viewer + - serious expression + - open mouth + +# --- Clothing --- +clothing: + - dress + - red dress + - blue shirt + - t-shirt + - jeans + - skirt + - school uniform + - jacket + - hoodie + - hat + - glasses + +# --- Actions --- +actions: + - smiling + - looking at viewer + - running + - sitting + - standing + - lying down + - holding object + +# --- Setting/Background --- +setting: + - outdoors + - indoors + - nature + - forest + - city + - street + - room + - white background + - simple background + - detailed background + +# --- NSFW Example (Illustrative - Use Responsibly) --- +# nsfw_body: +# - large breasts +# - small breasts +# - nude +# - explicit \ No newline at end of file