From 9343768d78a6c1a44e8349ce5c5a32a36fe3d8a2 Mon Sep 17 00:00:00 2001
From: CalamitousFelicitousness <iowasovereign@gmail.com>
Date: Mon, 23 Mar 2026 02:02:14 +0000
Subject: [PATCH] feat(api): add tag dictionary API and settings for prompt
 autocomplete

---
 cli/fetch_dicts.py        | 173 ++++++++++++++++++++++++++++++++++++++
 modules/api/api.py        |   8 +-
 modules/api/dicts.py      | 108 ++++++++++++++++++++++++
 modules/api/models.py     |  13 +++
 modules/paths.py          |   1 +
 modules/ui_definitions.py |  17 ++++
 6 files changed, 319 insertions(+), 1 deletion(-)
 create mode 100644 cli/fetch_dicts.py
 create mode 100644 modules/api/dicts.py

diff --git a/cli/fetch_dicts.py b/cli/fetch_dicts.py
new file mode 100644
index 000000000..4fb191667
--- /dev/null
+++ b/cli/fetch_dicts.py
@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+"""Fetch and convert booru tag databases to SD.Next dict format.
+
+Usage:
+    python cli/fetch_dicts.py danbooru [--output PATH] [--min-count N]
+    python cli/fetch_dicts.py e621 [--output PATH] [--min-count N]
+    python cli/fetch_dicts.py all [--output-dir DIR] [--min-count N]
+
+Output format:
+    JSON with { name, version, categories, tags: [[name, category_id, post_count], ...] }
+    Tags sorted by post_count descending.
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+from datetime import date, timezone
+
+import requests
+
+DANBOORU_CATEGORIES = {
+    "0": {"name": "general", "color": "#0075f8"},
+    "1": {"name": "artist", "color": "#a800aa"},
+    "3": {"name": "copyright", "color": "#dd00dd"},
+    "4": {"name": "character", "color": "#00ab2c"},
+    "5": {"name": "meta", "color": "#ee8800"},
+}
+
+E621_CATEGORIES = {
+    "0": {"name": "general", "color": "#0075f8"},
+    "1": {"name": "artist", "color": "#a800aa"},
+    "3": {"name": "copyright", "color": "#dd00dd"},
+    "4": {"name": "character", "color": "#00ab2c"},
+    "5": {"name": "species", "color": "#ed5d1f"},
+    "6": {"name": "invalid", "color": "#ff3d3d"},
+    "7": {"name": "meta", "color": "#ee8800"},
+    "8": {"name": "lore", "color": "#228b22"},
+}
+
+USER_AGENT = "SDNext-DictFetcher/1.0 (tag autocomplete)"
+
+
+def fetch_danbooru(min_count: int = 10) -> list:
+    """Fetch tags from Danbooru API, paginated."""
+    tags = []
+    page = 1
+    session = requests.Session()
+    session.headers["User-Agent"] = USER_AGENT
+    while True:
+        url = f"https://danbooru.donmai.us/tags.json?limit=1000&page={page}&search[order]=count"
+        try:
+            resp = session.get(url, timeout=30)
+            resp.raise_for_status()
+        except requests.RequestException as e:
+            print(f"  Error on page {page}: {e}", file=sys.stderr)
+            break
+        data = resp.json()
+        if not data:
+            break
+        for tag in data:
+            count = tag.get("post_count", 0)
+            if count < min_count:
+                continue
+            tags.append([tag["name"], tag["category"], count])
+        below_threshold = all(t.get("post_count", 0) < min_count for t in data)
+        print(f"  Page {page}: {len(data)} tags (total: {len(tags)})", file=sys.stderr)
+        if below_threshold:
+            break
+        page += 1
+        time.sleep(0.5)
+    tags.sort(key=lambda t: t[2], reverse=True)
+    return tags
+
+
+def fetch_e621(min_count: int = 10) -> list:
+    """Fetch tags from e621 API, paginated."""
+    tags = []
+    page = 1
+    session = requests.Session()
+    session.headers["User-Agent"] = USER_AGENT
+    while True:
+        url = f"https://e621.net/tags.json?limit=320&page={page}&search[order]=count"
+        try:
+            resp = session.get(url, timeout=30)
+            resp.raise_for_status()
+        except requests.RequestException as e:
+            print(f"  Error on page {page}: {e}", file=sys.stderr)
+            break
+        data = resp.json()
+        if not data:
+            break
+        for tag in data:
+            count = tag.get("post_count", 0)
+            if count < min_count:
+                continue
+            tags.append([tag["name"], tag["category"], count])
+        below_threshold = all(t.get("post_count", 0) < min_count for t in data)
+        print(f"  Page {page}: {len(data)} tags (total: {len(tags)})", file=sys.stderr)
+        if below_threshold:
+            break
+        page += 1
+        time.sleep(1.0)  # e621 rate limit is stricter
+    tags.sort(key=lambda t: t[2], reverse=True)
+    return tags
+
+
+SOURCES = {
+    "danbooru": {
+        "fetch": fetch_danbooru,
+        "categories": DANBOORU_CATEGORIES,
+    },
+    "e621": {
+        "fetch": fetch_e621,
+        "categories": E621_CATEGORIES,
+    },
+}
+
+
+def write_dict(name: str, tags: list, categories: dict, output_path: str):
+    """Write dict JSON file atomically."""
+    data = {
+        "name": name,
+        "version": date.today().isoformat(),
+        "categories": categories,
+        "tags": tags,
+    }
+    os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True)
+    tmp_path = output_path + ".tmp"
+    with open(tmp_path, "w", encoding="utf-8") as f:
+        json.dump(data, f, ensure_ascii=False, separators=(",", ":"))
+    os.replace(tmp_path, output_path)
+    size_mb = os.path.getsize(output_path) / (1024 * 1024)
+    print(f"  Written: {output_path} ({len(tags)} tags, {size_mb:.1f} MB)", file=sys.stderr)
+
+
+def fetch_source(name: str, output: str, min_count: int):
+    """Fetch and write a single source."""
+    if name not in SOURCES:
+        print(f"Unknown source: {name}. Available: {', '.join(SOURCES.keys())}", file=sys.stderr)
+        sys.exit(1)
+    source = SOURCES[name]
+    print(f"Fetching {name} (min_count={min_count})...", file=sys.stderr)
+    tags = source["fetch"](min_count=min_count)
+    if not tags:
+        print(f"  No tags fetched for {name}", file=sys.stderr)
+        return
+    write_dict(name, tags, source["categories"], output)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Fetch booru tag databases for SD.Next dict autocomplete")
+    parser.add_argument("source", choices=list(SOURCES.keys()) + ["all"], help="Tag source to fetch")
+    parser.add_argument("--output", "-o", help="Output file path (for single source)")
+    parser.add_argument("--output-dir", "-d", help="Output directory (for 'all')")
+    parser.add_argument("--min-count", "-m", type=int, default=10, help="Minimum post count to include (default: 10)")
+    args = parser.parse_args()
+
+    if args.source == "all":
+        output_dir = args.output_dir or "."
+        for name in SOURCES:
+            output = os.path.join(output_dir, f"{name}.json")
+            fetch_source(name, output, args.min_count)
+    else:
+        output = args.output or f"{args.source}.json"
+        fetch_source(args.source, output, args.min_count)
+
+    print("Done.", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/modules/api/api.py b/modules/api/api.py
index f460fea51..4505abbf7 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -1,9 +1,10 @@
+import os
 from threading import Lock
 from secrets import compare_digest
 from fastapi import FastAPI, APIRouter, Depends, Request
 from fastapi.security import HTTPBasic, HTTPBasicCredentials
 from fastapi.exceptions import HTTPException
-from modules import errors, shared
+from modules import errors, shared, paths
 from modules.logger import log
 from modules.api import models, endpoints, script, helpers, server, generate, process, control, docs, gpu
 
@@ -117,6 +118,11 @@ class Api:
         from modules.api import loras
         loras.register_api(self.app)
 
+        # dicts api
+        from modules.api import dicts as dicts_api
+        dicts_api.init(getattr(shared.opts, 'dicts_dir', '') or os.path.join(paths.models_path, 'dicts'))
+        dicts_api.register_api(self.app)
+
         # gallery api
         from modules.api import gallery
         gallery.register_api(self.app)
diff --git a/modules/api/dicts.py b/modules/api/dicts.py
new file mode 100644
index 000000000..1d6400717
--- /dev/null
+++ b/modules/api/dicts.py
@@ -0,0 +1,108 @@
+"""V1 dictionary / tag autocomplete endpoints.
+
+Serves pre-built tag dictionaries (Danbooru, e621, natural language, artists)
+from JSON files in the configured dicts directory.
+"""
+
+import asyncio
+import json
+import os
+from datetime import datetime
+
+from fastapi.exceptions import HTTPException
+
+from modules.api.models import ItemDict, ItemDictContent
+
+dicts_dir: str = ""
+cache: dict[str, dict] = {}
+
+
+def init(path: str) -> None:
+    """Set the dicts directory path. Called once during API registration."""
+    global dicts_dir  # noqa: PLW0603
+    dicts_dir = path
+
+
+def get_cached(name: str) -> dict:
+    """Load a dict file, returning cached version if file hasn't changed."""
+    if '/' in name or '\\' in name or '..' in name:
+        raise HTTPException(status_code=400, detail="Invalid dict name")
+    path = os.path.join(dicts_dir, f"{name}.json")
+    if not os.path.isfile(path):
+        cache.pop(name, None)
+        raise HTTPException(status_code=404, detail=f"Dict not found: {name}")
+    stat = os.stat(path)
+    entry = cache.get(name)
+    if entry and entry['mtime'] == stat.st_mtime:
+        return entry
+    with open(path, encoding='utf-8') as f:
+        data = json.load(f)
+    entry = {
+        'mtime': stat.st_mtime,
+        'size': stat.st_size,
+        'meta': {
+            'name': data.get('name', name),
+            'version': data.get('version', ''),
+            'tag_count': len(data.get('tags', [])),
+            'categories': {
+                str(k): v.get('name', str(k)) if isinstance(v, dict) else str(v)
+                for k, v in data.get('categories', {}).items()
+            },
+        },
+        'content': data,
+    }
+    cache[name] = entry
+    return entry
+
+
+def list_dicts_sync() -> list[ItemDict]:
+    """Scan dicts directory and return metadata for each dict file."""
+    if not dicts_dir or not os.path.isdir(dicts_dir):
+        return []
+    items = []
+    for filename in sorted(os.listdir(dicts_dir)):
+        if not filename.endswith('.json') or filename.startswith('.'):
+            continue
+        name = filename.rsplit('.', 1)[0]
+        try:
+            entry = get_cached(name)
+            meta = entry['meta']
+            items.append(ItemDict(
+                name=meta['name'],
+                version=meta['version'],
+                tag_count=meta['tag_count'],
+                categories=meta['categories'],
+                size=entry['size'],
+            ))
+        except Exception:
+            pass
+    return items
+
+
+async def list_dicts() -> list[ItemDict]:
+    """List available tag dictionaries."""
+    return await asyncio.to_thread(list_dicts_sync)
+
+
+async def get_dict(name: str) -> ItemDictContent:
+    """Get full dict content by name."""
+    def _load():
+        return get_cached(name)
+    try:
+        entry = await asyncio.to_thread(_load)
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e)) from e
+    content = entry['content']
+    return ItemDictContent(
+        name=content.get('name', name),
+        version=content.get('version', ''),
+        categories=content.get('categories', {}),
+        tags=content.get('tags', []),
+    )
+
+
+def register_api(app):
+    app.add_api_route("/sdapi/v1/dicts", list_dicts, methods=["GET"], response_model=list[ItemDict], tags=["Enumerators"])
+    app.add_api_route("/sdapi/v1/dicts/{name}", get_dict, methods=["GET"], response_model=ItemDictContent, tags=["Enumerators"])
diff --git a/modules/api/models.py b/modules/api/models.py
index 47ccecd6b..d3d030dec 100644
--- a/modules/api/models.py
+++ b/modules/api/models.py
@@ -509,6 +509,19 @@ class ItemLoadedModel(BaseModel):
     dtype: Optional[str] = Field(default=None, title="Dtype", description="Effective data type (e.g., float16, nf4)")
     extra: Optional[dict] = Field(default=None, title="Extra metadata", description="Additional metadata (role, class, quantization method, etc.)")
 
+class ItemDict(BaseModel):
+    name: str = Field(title="Name", description="Dictionary identifier (filename without extension)")
+    version: str = Field(default="", title="Version", description="Dictionary format version string")
+    tag_count: int = Field(default=0, title="Tag count", description="Number of tags in this dictionary")
+    categories: dict = Field(default_factory=dict, title="Categories", description="Category ID to display name mapping")
+    size: int = Field(default=0, title="Size", description="File size in bytes")
+
+class ItemDictContent(BaseModel):
+    name: str = Field(title="Name", description="Dictionary identifier")
+    version: str = Field(default="", title="Version", description="Dictionary format version string")
+    categories: dict = Field(default_factory=dict, title="Categories", description="Category definitions with name and color")
+    tags: list = Field(default_factory=list, title="Tags", description="Tag entries as [name, category_id, post_count] tuples")
+
 # helper function
 
 def create_model_from_signature(func: Callable, model_name: str, base_model: type[BaseModel] = BaseModel, additional_fields: list | None = None, exclude_fields: list[str] | None = None) -> type[BaseModel]:
diff --git a/modules/paths.py b/modules/paths.py
index a2f460c3b..cdb30ff30 100644
--- a/modules/paths.py
+++ b/modules/paths.py
@@ -129,6 +129,7 @@ def create_paths(opts):
     create_path(fix_path('styles_dir'))
     create_path(fix_path('yolo_dir'))
     create_path(fix_path('wildcards_dir'))
+    create_path(fix_path('dicts_dir'))
 
     # Create resolved output paths (base + specific)
     base_samples = opts.data.get('outdir_samples', '')
diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py
index d43b45e3a..ea8c6e68c 100644
--- a/modules/ui_definitions.py
+++ b/modules/ui_definitions.py
@@ -51,6 +51,19 @@ def get_openvino_device_list():
         return []
 
 
+def list_dict_names():
+    """Return list of available tag dictionary names (JSON filenames without extension)."""
+    from modules import shared, paths as paths_module
+    dicts_dir = getattr(shared.opts, 'dicts_dir', None) or os.path.join(paths_module.models_path, 'dicts')
+    if not os.path.isdir(dicts_dir):
+        return []
+    return sorted(
+        os.path.splitext(f)[0]
+        for f in os.listdir(dicts_dir)
+        if f.endswith('.json') and not f.startswith('.')
+    )
+
+
 def create_settings(cmd_opts):
 
     # Calculate default modes
@@ -387,6 +400,7 @@ def create_settings(cmd_opts):
         "lora_dir": OptionInfo(os.path.join(paths.models_path, 'Lora'), "Folder with LoRA network(s)", folder=True),
         "styles_dir": OptionInfo(os.path.join(paths.models_path, 'styles'), "File or Folder with user-defined styles", folder=True),
         "wildcards_dir": OptionInfo(os.path.join(paths.models_path, 'wildcards'), "Folder with user-defined wildcards", folder=True),
+        "dicts_dir": OptionInfo(os.path.join(paths.models_path, 'dicts'), "Folder with tag dictionaries", folder=True),
         "embeddings_dir": OptionInfo(os.path.join(paths.models_path, 'embeddings'), "Folder with textual inversion embeddings", folder=True),
         "control_dir": OptionInfo(os.path.join(paths.models_path, 'control'), "Folder with Control models", folder=True),
         "yolo_dir": OptionInfo(os.path.join(paths.models_path, 'yolo'), "Folder with Yolo models", folder=True),
@@ -637,6 +651,9 @@ def create_settings(cmd_opts):
 
         "extra_networks_wildcard_sep": OptionInfo("<h2>Wildcards</h2>", "", gr.HTML),
         "wildcards_enabled": OptionInfo(True, "Enable file wildcards support"),
+
+        "extra_networks_dicts_sep": OptionInfo("<h2>Tag Dictionaries</h2>", "", gr.HTML),
+        "dicts_enabled": OptionInfo([], "Enabled tag dictionaries for prompt autocomplete", gr.CheckboxGroup, lambda: {"choices": list_dict_names()}),
     }))
 
     # --- Extensions ---