# /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() def resolve_inline_choices(text: str, rng) -> str: """ Resolves standard dynamic prompt inline choices like {red|blue|green}. Supports weighted choices like {0.2::red|0.8::blue} and nested ones like {a|{b|c}}. """ if not text: return text pattern = re.compile(r'\{([^{}]+)\}') def repl(match): choices_str = match.group(1) choices = choices_str.split('|') weighted_choices =[] for c in choices: c = c.strip() if '::' in c: parts = c.split('::', 1) try: weight = float(parts[0]) choice = parts[1] except ValueError: weight = 1.0 choice = c else: weight = 1.0 choice = c weighted_choices.append((weight, choice)) total_weight = sum(w for w, c in weighted_choices) if total_weight <= 0: picked = rng.choice([c for w, c in weighted_choices]) if weighted_choices else "" else: r = rng.uniform(0, total_weight) upto = 0 picked = weighted_choices[-1][1] if weighted_choices else "" for w, c in weighted_choices: if upto + w >= r: picked = c break upto += w return picked prev_text = None # We loop it to resolve nested brackets from the innermost layer recursively while text != prev_text: prev_text = text text = pattern.sub(repl, text) return text # --- 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