diff --git a/README.md b/README.md index 498cc8f..a976d0b 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,11 @@ Provides the "manual" for available generation models. * **Best for**: Learning the prompting style, recommended settings, and available resolution presets for a specific model. * **Workflow**: Call without arguments to see the catalog; call with `model_name` for the detailed guide. +### 🏷️ `search_tags` +Searches the Danbooru tag database for recognized tags and aliases. +* **Best for**: Finding the correct booru-style tags, checking tag popularity, and resolving aliases (e.g., 'lesbian' → 'yuri'). +* **Workflow**: Provide a query string to get a list of the most popular matching tags. + ### 🌐 `browse_wikipedia` Allows the model to browse Wikipedia using its API. * **Best for**: Quickly retrieving summaries, structural maps (ToC), or specific section content from Wikipedia without dumping the entire page. diff --git a/config.py b/config.py index 8cb5f3f..6ff030d 100644 --- a/config.py +++ b/config.py @@ -16,6 +16,7 @@ USER_AGENT = "Mozilla/5.0 (X11; Linux x86_64; rv:151.0) Gecko/20100101 Firefox/1 SD_URL = "http://127.0.0.1:7860" MODEL_PRESETS_PATH = str(ROOT_DIR / "model_presets.toml") RES_PRESETS_PATH = str(ROOT_DIR / "resolution_presets.toml") +TAG_DATABASE_PATH = "/home/matt/stable-diffusion-webui/extensions/a1111-sd-webui-tagcomplete/tags/danbooru.csv" # --- Model Specific Token Tuning (Tuned for Gemma 4) --- # Patch size is typically (clip.vision.patch_size * n_merge) diff --git a/tools/__init__.py b/tools/__init__.py index fca7060..12a02ed 100644 --- a/tools/__init__.py +++ b/tools/__init__.py @@ -25,6 +25,7 @@ from .get_text_context import handle as get_text_context_handler from .wikipedia import handle as wikipedia_handler from .get_model_info import handle as get_model_info_handler from .generate_image import handle as generate_image_handler +from .search_tags import handle as search_tags_handler # Central registry of all available tools TOOL_REGISTRY = [ @@ -154,4 +155,16 @@ TOOL_REGISTRY = [ }, handler=generate_image_handler ), + Tool( + name="search_tags", + description="Searches the Danbooru tag database for tags matching a query. Returns the most popular tags including alias matches. Useful for finding the correct booru-style tags for anime models.", + schema={ + "type": "object", + "properties": { + "query": {"type": "string", "description": "The tag or partial tag to search for (e.g., 'white hair' or 'lesbian')."} + }, + "required": ["query"], + }, + handler=search_tags_handler + ), ] diff --git a/tools/search_tags.py b/tools/search_tags.py new file mode 100644 index 0000000..5c0b70a --- /dev/null +++ b/tools/search_tags.py @@ -0,0 +1,68 @@ +import csv +import os +from typing import List, Dict, Any +from tools.utils import ToolError, format_count, get_type_suffix + +# Tag database path is now managed in config.py +async def handle(args: Dict[str, Any]) -> Dict[str, Any]: + """ + Searches the Danbooru tag database for tags matching a query. + Returns the most popular tags including alias matches. + """ + query = args.get("query", "").lower() + if not query: + raise ToolError("The 'query' argument is required.") + + tag_file_path = config.TAG_DATABASE_PATH + if not os.path.exists(tag_file_path): + raise ToolError(f"Tag database not found at {tag_file_path}. Please ensure the tag-autocomplete extension is installed.") + + matches = [] + try: + with open(tag_file_path, mode='r', encoding='utf-8') as f: + reader = csv.reader(f) + for row in reader: + if not row or len(row) < 3: + continue + + name = row[0].lower() + tag_type = row[1] + count = int(row[2]) if row[2].isdigit() else 0 + aliases = row[3].lower().split(',') if len(row) > 3 else [] + + # Check for match in name or aliases + is_direct = query in name + is_alias = any(query in alias.strip() for alias in aliases) + + if is_direct or is_alias: + matches.append({ + "name": row[0], + "type": tag_type, + "count": count, + "matched_via": name if is_direct else "alias", + "alias_match": "" if is_direct else next((a.strip() for a in aliases if query in a.strip()), query) + }) + + except Exception as e: + raise ToolError(f"Error reading tag database: {str(e)}") + + # Sort by count descending + matches.sort(key=lambda x: x["count"], reverse=True) + + # Format top 20 results + results = [] + for m in matches[:20]: + count_fmt = format_count(str(m["count"])) + type_sfx = get_type_suffix(m["type"]) + + if m["matched_via"] == "alias": + line = f"{m['alias_match']} → {m['name']} ({count_fmt}){type_sfx}" + else: + line = f"{m['name']} ({count_fmt}){type_sfx}" + + results.append(line) + + if not results: + return {"text": f"No tags found matching '{query}'."} + + return {"text": "Top matches:\n" + "\n".join(results)} diff --git a/tools/utils.py b/tools/utils.py index ad35198..411f427 100644 --- a/tools/utils.py +++ b/tools/utils.py @@ -117,3 +117,26 @@ def parse_indices(indices_str: str) -> List[int]: except ValueError: raise ToolError(f"Invalid index format: {part}") return indices + +def format_count(count_str: str) -> str: + """Formats large numbers into a human-readable string (e.g., 1.2M, 218k).""" + try: + count = int(count_str) + if count >= 1_000_000: + return f"{count / 1_000_000:.1f}M".replace(".0", "") + if count >= 1_000: + return f"{count / 1_000:.0f}k" + return str(count) + except (ValueError, TypeError): + return "0" + +def get_type_suffix(type_val: str) -> str: + """Returns the human-readable suffix for the tag type.""" + mapping = { + "0": "", # General + "1": " [Artist]", + "3": " [Copyright]", + "4": " [Character]", + "5": " [Meta]" + } + return mapping.get(str(type_val), f" [Type {type_val}]")