diff --git a/cli/tags-fetch.py b/cli/tags-fetch.py new file mode 100644 index 000000000..6bef29e7f --- /dev/null +++ b/cli/tags-fetch.py @@ -0,0 +1,429 @@ +#!/usr/bin/env python3 +"""Fetch and convert booru tag databases to SD.Next autocomplete format. + +Usage: + python cli/tags-fetch.py danbooru [--output PATH] [--min-count N] + python cli/tags-fetch.py e621 [--output PATH] [--min-count N] + python cli/tags-fetch.py rule34 --key USER_ID:API_KEY [--output PATH] [--min-count N] + python cli/tags-fetch.py sankaku [--output PATH] [--min-count N] + python cli/tags-fetch.py idol [--output PATH] [--min-count N] + python cli/tags-fetch.py all --key USER_ID:API_KEY [--output-dir DIR] [--min-count N] + +Progress is saved every 50 pages to a .partial file so interrupted +runs can be resumed. Transient HTTP errors are retried with backoff. + +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 + +import requests + +# Unified category scheme - all sources map their native type IDs to these. +# Every dict file uses these same IDs and colors. +UNIFIED_CATEGORIES = { + "0": {"name": "general", "color": "#0075f8"}, + "1": {"name": "artist", "color": "#cc0000"}, + "2": {"name": "studio", "color": "#ff4500"}, + "3": {"name": "copyright", "color": "#9900ff"}, + "4": {"name": "character", "color": "#00ab2c"}, + "5": {"name": "species", "color": "#ed5d1f"}, + "6": {"name": "genre", "color": "#8a66ff"}, + "7": {"name": "medium", "color": "#00cccc"}, + "8": {"name": "meta", "color": "#6b7280"}, + "9": {"name": "lore", "color": "#228b22"}, + "10": {"name": "lens", "color": "#e67e22"}, + "11": {"name": "lighting", "color": "#f1c40f"}, + "12": {"name": "composition", "color": "#1abc9c"}, + "13": {"name": "color", "color": "#e84393"}, +} + +# Source → unified type maps. Each maps the source's native category IDs +# to the unified IDs above. Unmapped IDs default to 0 (general). +DANBOORU_TYPE_MAP = {0: 0, 1: 1, 3: 3, 4: 4, 5: 8} +E621_TYPE_MAP = {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 0, 7: 8, 8: 9} +RULE34_TYPE_MAP = {0: 0, 1: 1, 3: 3, 4: 4, 5: 8} +SANKAKU_TYPE_MAP = {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 6, 8: 7, 9: 8} + +# Idol Complex uses Sankaku's granular tag type system (21 subcategories). +# API type → unified category: +# 0=studio, 1=artist, 2=studio, 3=franchise, 4=character, +# 5=photoset, 6=genre, 8=medium, 9=meta, 10=fashion, +# 11=anatomy, 12=pose, 13=activity, 14=role, 15=flora, +# 17=fauna/entity, 18=object/setting, 19=substance, 20=general, +# 21=language, 22=automatic +IDOL_TYPE_MAP = { + 0: 2, 1: 1, 2: 2, 3: 3, 4: 4, 5: 8, 6: 6, 8: 7, 9: 8, + 10: 0, 11: 0, 12: 0, 13: 0, 14: 0, 15: 0, 17: 0, 18: 0, + 19: 0, 20: 0, 21: 8, 22: 8, +} + +USER_AGENT = "SDNext-DictFetcher/1.0 (tag autocomplete)" +SAVE_INTERVAL = 50 # save progress every N pages +MAX_RETRIES = 3 +RETRY_BACKOFF = 5 # seconds, multiplied by attempt number + + +# -- Partial save/resume -- + +def save_partial(path: str, page: int, tags: list): + """Save fetch progress to a .partial file.""" + if not path: + return + tmp = path + ".tmp" + with open(tmp, "w", encoding="utf-8") as f: + json.dump({"page": page, "tags": tags}, f, separators=(",", ":")) + os.replace(tmp, path) + + +def load_partial(path: str) -> tuple[int, list] | tuple[None, list]: + """Load saved progress from a .partial file. Returns (page, tags) or (None, []).""" + if not path or not os.path.isfile(path): + return None, [] + try: + with open(path, encoding="utf-8") as f: + data = json.load(f) + page = data["page"] + tags = data["tags"] + print(f" Resuming from partial: page {page}, {len(tags)} tags", file=sys.stderr) + return page, tags + except (json.JSONDecodeError, KeyError) as e: + print(f" Warning: corrupt .partial file, starting fresh ({e})", file=sys.stderr) + return None, [] + + +def clear_partial(path: str): + """Remove .partial file after successful completion.""" + if path and os.path.isfile(path): + os.remove(path) + + +# -- HTTP retry helper -- + +def fetch_with_retry(session: requests.Session, url: str, params: dict | None = None, timeout: int = 30) -> requests.Response: + """GET with retry and exponential backoff for transient errors.""" + for attempt in range(MAX_RETRIES): + try: + resp = session.get(url, params=params, timeout=timeout) + resp.raise_for_status() + return resp + except requests.RequestException as e: + if attempt < MAX_RETRIES - 1: + wait = RETRY_BACKOFF * (attempt + 1) + print(f" Retry {attempt + 1}/{MAX_RETRIES} in {wait}s: {e}", file=sys.stderr) + time.sleep(wait) + else: + raise + + +# -- Fetchers -- + +def fetch_danbooru(min_count: int = 10, partial_path: str = "", **_kwargs) -> list: + """Fetch tags from Danbooru API, paginated.""" + start_page, tags = load_partial(partial_path) + page = (start_page + 1) if start_page is not None else 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 = fetch_with_retry(session, url) + 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 + if page % SAVE_INTERVAL == 0: + save_partial(partial_path, page, tags) + page += 1 + time.sleep(0.5) + tags.sort(key=lambda t: t[2], reverse=True) + return tags + + +def fetch_e621(min_count: int = 10, partial_path: str = "", **_kwargs) -> list: + """Fetch tags from e621 API, paginated.""" + start_page, tags = load_partial(partial_path) + page = (start_page + 1) if start_page is not None else 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 = fetch_with_retry(session, url) + 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 + if page % SAVE_INTERVAL == 0: + save_partial(partial_path, page, tags) + page += 1 + time.sleep(1.0) + tags.sort(key=lambda t: t[2], reverse=True) + return tags + + +def fetch_gelbooru(base_url: str, min_count: int = 10, api_key: str | None = None, rate_limit: float = 0.5, partial_path: str = "") -> list: + """Fetch tags from a Gelbooru-compatible API (rule34, gelbooru, etc.). + + Unlike Danbooru/e621, the Gelbooru tag endpoint doesn't support ordering + by count, so we paginate through all tags and filter client-side. + The tag endpoint returns XML (json=1 is not supported for tags). + """ + import xml.etree.ElementTree as ET + + start_page, tags = load_partial(partial_path) + page = (start_page + 1) if start_page is not None else 0 + page_size = 1000 + session = requests.Session() + session.headers["User-Agent"] = USER_AGENT + params: dict[str, str] = { + "page": "dapi", "s": "tag", "q": "index", + "limit": str(page_size), + } + if api_key: + if ":" not in api_key: + print(" Error: --key must be USER_ID:API_KEY format", file=sys.stderr) + return [] + uid, key = api_key.split(":", 1) + params["user_id"] = uid + params["api_key"] = key + + while True: + params["pid"] = str(page) + try: + resp = fetch_with_retry(session, base_url, params=params) + except requests.RequestException as e: + print(f" Error on page {page} after {MAX_RETRIES} retries: {e}", file=sys.stderr) + break + text = resp.text.strip() + if not text or text.startswith('"') or not text.startswith('= min_count: + tags.append([el.get("name", ""), int(el.get("type", "0")), count]) + print(f" Page {page}: {len(elements)} tags (total: {len(tags)})", file=sys.stderr) + if len(elements) < page_size: + break + if page % SAVE_INTERVAL == 0: + save_partial(partial_path, page, tags) + page += 1 + time.sleep(rate_limit) + tags.sort(key=lambda t: t[2], reverse=True) + return tags + + +def fetch_rule34(min_count: int = 10, api_key: str | None = None, partial_path: str = "", **_kwargs) -> list: + """Fetch tags from rule34.xxx.""" + if not api_key: + print(" Warning: no --key provided, rule34 may rate-limit aggressively", file=sys.stderr) + return fetch_gelbooru("https://api.rule34.xxx/index.php", min_count=min_count, api_key=api_key, partial_path=partial_path) + + +def fetch_sankaku(min_count: int = 10, partial_path: str = "", **_kwargs) -> list: + """Fetch tags from Sankaku Complex (chan.sankakucomplex.com). + + Uses the public JSON API at sankakuapi.com which supports order=count, + so we can stop early when counts drop below min_count. + Tag names come as English with spaces - converted to lowercase. + """ + start_page, tags = load_partial(partial_path) + page = (start_page + 1) if start_page is not None else 1 + page_size = 200 + session = requests.Session() + session.headers["User-Agent"] = USER_AGENT + while True: + try: + resp = fetch_with_retry(session, "https://sankakuapi.com/tags", + params={"limit": page_size, "page": page, "order": "count"}) + except requests.RequestException as e: + print(f" Error on page {page} after {MAX_RETRIES} retries: {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 + name = tag.get("name_en") or tag.get("name_ja", "") + if not name: + continue + name = name.strip().lower() + tags.append([name, tag.get("type", 0), count]) + below_threshold = all(tag.get("post_count", 0) < min_count for tag in data) + print(f" Page {page}: {len(data)} tags (total: {len(tags)})", file=sys.stderr) + if below_threshold: + break + if page % SAVE_INTERVAL == 0: + save_partial(partial_path, page, tags) + page += 1 + time.sleep(0.5) + tags.sort(key=lambda t: t[2], reverse=True) + return tags + + +def fetch_idol(min_count: int = 10, partial_path: str = "", **_kwargs) -> list: + """Fetch tags from Idol Complex (idol.sankakucomplex.com). + + Uses the legacy JSON API at iapi.sankakucomplex.com. Supports order=count. + Max 50 tags per page. Type IDs are non-standard - remapped in write_dict. + """ + start_page, tags = load_partial(partial_path) + page = (start_page + 1) if start_page is not None else 1 + page_size = 50 + session = requests.Session() + session.headers["User-Agent"] = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/138.0.0.0 Safari/537.36" + while True: + try: + resp = fetch_with_retry(session, "https://iapi.sankakucomplex.com/tags.json", + params={"limit": page_size, "page": page, "order": "count"}) + except requests.RequestException as e: + print(f" Error on page {page} after {MAX_RETRIES} retries: {e}", file=sys.stderr) + break + data = resp.json() + if not data: + break + for tag in data: + count = tag.get("count", 0) + if count < min_count: + continue + name = tag.get("name", "") + if not name: + continue + tags.append([name, tag.get("type", 0), count]) + below_threshold = all(tag.get("count", 0) < min_count for tag in data) + print(f" Page {page}: {len(data)} tags (total: {len(tags)})", file=sys.stderr) + if below_threshold: + break + if page % SAVE_INTERVAL == 0: + save_partial(partial_path, page, tags) + page += 1 + time.sleep(1.0) + tags.sort(key=lambda t: t[2], reverse=True) + return tags + + +SOURCES = { + "danbooru": {"fetch": fetch_danbooru, "type_map": DANBOORU_TYPE_MAP}, + "e621": {"fetch": fetch_e621, "type_map": E621_TYPE_MAP}, + "rule34": {"fetch": fetch_rule34, "type_map": RULE34_TYPE_MAP}, + "sankaku": {"fetch": fetch_sankaku, "type_map": SANKAKU_TYPE_MAP}, + "idol": {"fetch": fetch_idol, "type_map": IDOL_TYPE_MAP}, +} + + +def write_dict(name: str, tags: list, type_map: dict, output_path: str, separator: str = "_"): + """Write dict JSON file atomically. + + type_map remaps source-native category IDs to unified IDs. + separator controls the word separator in tag names: + "_" (default) → "high_resolution" (booru convention, anime/illustration models) + " " → "high resolution" (natural language, SDXL/Flux-style models) + """ + normalized = [] + for t in tags: + tag_name = t[0].replace(" ", "_") if separator == "_" else t[0].replace("_", " ") + category = type_map.get(t[1], 0) + normalized.append([tag_name, category, t[2]]) + data = { + "name": name, + "version": date.today().isoformat(), + "categories": UNIFIED_CATEGORIES, + "tags": normalized, + } + 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, api_key: str | None = None, separator: str = "_"): + """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] + partial_path = output + ".partial" + print(f"Fetching {name} (min_count={min_count})...", file=sys.stderr) + tags = source["fetch"](min_count=min_count, api_key=api_key, partial_path=partial_path) + if not tags: + print(f" No tags fetched for {name}", file=sys.stderr) + return + write_dict(name, tags, source["type_map"], output, separator=separator) + clear_partial(partial_path) + from importlib.util import spec_from_file_location, module_from_spec + spec = spec_from_file_location("tags_manifest", os.path.join(os.path.dirname(__file__), "tags-manifest.py")) + mod = module_from_spec(spec) + spec.loader.exec_module(mod) + update_manifest = mod.update_manifest + update_manifest(os.path.dirname(output) or ".") + + +def main(): + parser = argparse.ArgumentParser(description="Fetch booru tag databases for SD.Next 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)") + parser.add_argument("--key", "-k", help="API key as USER_ID:API_KEY (required for rule34)") + parser.add_argument("--spaces", action="store_true", help="Use spaces instead of underscores in tag names (for natural language models like SDXL/Flux)") + args = parser.parse_args() + separator = " " if args.spaces else "_" + + 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, api_key=args.key, separator=separator) + else: + output = args.output or f"{args.source}.json" + fetch_source(args.source, output, args.min_count, api_key=args.key, separator=separator) + + print("Done.", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/cli/tags-manifest.py b/cli/tags-manifest.py new file mode 100644 index 000000000..2958e55d5 --- /dev/null +++ b/cli/tags-manifest.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""Generate manifest.json for the HuggingFace autocomplete repository. + +Reads tag JSON files and writes a manifest with accurate tag counts, +file sizes, and versions. + +Usage: + python cli/tags-manifest.py models/autocomplete/danbooru.json models/autocomplete/e621.json ... + python cli/tags-manifest.py models/autocomplete/*.json -o models/autocomplete/manifest.json +""" + +import argparse +import json +import os +import sys + +# Human-readable descriptions keyed by file name. +# Add entries here when new sources are added to tags-fetch.py. +DESCRIPTIONS = { + "art": "Art movements, styles, and techniques", + "danbooru": "Danbooru image board tags - anime/illustration focused", + "e621": "e621 tags - furry/animal art focused", + "idol": "Idol Complex tags - Japanese idol photography", + "negative": "Common negative prompt terms", + "photography": "Photography terms - lens, lighting, composition, color", + "quality": "Quality and aesthetic meta tags", + "rule34": "Rule34.xxx tags - multi-fandom", + "sankaku": "Sankaku Complex tags - anime/illustration with granular categories", +} + + +def build_entry(filepath: str) -> dict: + """Build a manifest entry from a tag JSON file.""" + with open(filepath, encoding="utf-8") as f: + data = json.load(f) + name = data.get("name") or os.path.splitext(os.path.basename(filepath))[0] + size_mb = round(os.path.getsize(filepath) / (1024 * 1024), 1) + return { + "name": name, + "description": DESCRIPTIONS.get(name, ""), + "version": data.get("version", ""), + "tag_count": len(data.get("tags", [])), + "size_mb": size_mb, + } + + +def update_manifest(directory: str) -> bool: + """Regenerate manifest.json in directory if one already exists. + + Only updates entries already listed in the manifest - does not + add new entries. Returns True if the manifest was updated, False + if no manifest exists to update. + """ + manifest_path = os.path.join(directory, "manifest.json") + if not os.path.isfile(manifest_path): + return False + with open(manifest_path, encoding="utf-8") as f: + manifest = json.load(f) + existing_names = {e["name"] for e in manifest.get("entries", [])} + entries = [] + for name in sorted(existing_names): + filepath = os.path.join(directory, f"{name}.json") + if not os.path.isfile(filepath): + continue + try: + entries.append(build_entry(filepath)) + except (json.JSONDecodeError, KeyError): + continue + manifest["entries"] = entries + with open(manifest_path, "w", encoding="utf-8") as f: + json.dump(manifest, f, ensure_ascii=False, separators=(",", ":")) + f.write("\n") + print(f" Manifest updated: {manifest_path} ({len(entries)} entries)", file=sys.stderr) + return True + + +def main(): + parser = argparse.ArgumentParser(description="Generate manifest.json for HF autocomplete repo") + parser.add_argument("files", nargs="+", help="Tag JSON files to include") + parser.add_argument("-o", "--output", default="manifest.json", help="Output path (default: manifest.json)") + parser.add_argument("--exclude", nargs="*", default=[], help="Names to exclude (e.g. art quality)") + args = parser.parse_args() + + exclude = set(args.exclude) + entries = [] + for filepath in args.files: + basename = os.path.splitext(os.path.basename(filepath))[0] + if basename in exclude or basename == "manifest": + continue + if not os.path.isfile(filepath): + print(f" Skipping {filepath}: not found", file=sys.stderr) + continue + try: + entry = build_entry(filepath) + entries.append(entry) + print(f" {entry['name']:>12}: {entry['tag_count']:>10,} tags, {entry['size_mb']:>6.1f} MB, v={entry['version']}", file=sys.stderr) + except (json.JSONDecodeError, KeyError) as e: + print(f" Skipping {filepath}: {e}", file=sys.stderr) + + entries.sort(key=lambda e: e["name"]) + manifest = {"entries": entries} + + os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True) + with open(args.output, "w", encoding="utf-8") as f: + json.dump(manifest, f, ensure_ascii=False, separators=(",", ":")) + f.write("\n") + + print(f"\nManifest written: {args.output} ({len(entries)} entries)", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/cli/tags-prune.py b/cli/tags-prune.py new file mode 100644 index 000000000..68c130eaf --- /dev/null +++ b/cli/tags-prune.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +"""Prune tag autocomplete files with per-category minimum post counts. + +Usage: + python cli/tags-prune.py models/autocomplete/danbooru.json + python cli/tags-prune.py models/autocomplete/*.json --general 500 --artist 20 + python cli/tags-prune.py models/autocomplete/sankaku.json -o sankaku-pruned.json + python cli/tags-prune.py models/autocomplete/*.json --dry-run + +Category-aware pruning: artists and characters are kept at lower thresholds +(proper nouns that need autocomplete), while general tags are pruned more +aggressively (common vocabulary you'd type naturally). +""" + +import argparse +import json +import os +import sys + +# Category ID -> name (must match UNIFIED_CATEGORIES in tags-fetch.py) +CATEGORY_NAMES = { + 0: "general", + 1: "artist", + 2: "studio", + 3: "copyright", + 4: "character", + 5: "species", + 6: "genre", + 7: "medium", + 8: "meta", + 9: "lore", +} + +# Default minimum post counts per category. +# Low for proper nouns (hard to guess), high for common vocabulary (easy to type). +DEFAULTS = { + "general": 200, + "artist": 10, + "studio": 10, + "copyright": 20, + "character": 10, + "species": 50, + "genre": 100, + "medium": 100, + "meta": 1000, + "lore": 50, +} + + +def prune(tags: list, thresholds: dict[str, int]) -> tuple[list, dict[str, tuple[int, int]]]: + """Prune tags by per-category thresholds. + + Returns (pruned_tags, stats) where stats maps category name + to (before_count, after_count). + """ + stats: dict[str, tuple[int, int]] = {} + kept = [] + for tag in tags: + cat_id = tag[1] + cat_name = CATEGORY_NAMES.get(cat_id, "general") + threshold = thresholds.get(cat_name, thresholds.get("general", 10)) + before, after = stats.get(cat_name, (0, 0)) + before += 1 + if tag[2] >= threshold: + kept.append(tag) + after += 1 + stats[cat_name] = (before, after) + return kept, stats + + +def main(): + parser = argparse.ArgumentParser( + description="Prune tag dictionaries with per-category minimum post counts", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog="Defaults: " + ", ".join(f"{k}={v}" for k, v in DEFAULTS.items()), + ) + parser.add_argument("files", nargs="+", help="Dict JSON files to prune") + parser.add_argument("-o", "--output", help="Output file (single file mode only)") + parser.add_argument("--dry-run", action="store_true", help="Show what would be pruned without writing") + parser.add_argument("--in-place", "-i", action="store_true", help="Overwrite input files") + + for name, default in DEFAULTS.items(): + parser.add_argument(f"--{name}", type=int, default=default, help=f"Min posts for {name} (default: {default})") + + args = parser.parse_args() + + if args.output and len(args.files) > 1: + print("Error: --output can only be used with a single input file", file=sys.stderr) + sys.exit(1) + + if not args.output and not args.in_place and not args.dry_run: + print("Error: specify --in-place, --output, or --dry-run", file=sys.stderr) + sys.exit(1) + + thresholds = {name: getattr(args, name) for name in DEFAULTS} + + from importlib.util import spec_from_file_location, module_from_spec + spec = spec_from_file_location("tags_manifest", os.path.join(os.path.dirname(__file__), "tags-manifest.py")) + mod = module_from_spec(spec) + spec.loader.exec_module(mod) + update_manifest = mod.update_manifest + + for filepath in args.files: + if not os.path.isfile(filepath): + print(f" Skipping {filepath}: not found", file=sys.stderr) + continue + + with open(filepath, encoding="utf-8") as f: + data = json.load(f) + + tags = data.get("tags", []) + pruned, stats = prune(tags, thresholds) + total_before = sum(s[0] for s in stats.values()) + total_after = sum(s[1] for s in stats.values()) + + name = data.get("name", os.path.basename(filepath)) + print(f"\n{name}: {total_before:,} → {total_after:,} tags ({total_before - total_after:,} removed)", file=sys.stderr) + for cat_name in sorted(stats, key=lambda c: list(CATEGORY_NAMES.values()).index(c) if c in CATEGORY_NAMES.values() else 99): + before, after = stats[cat_name] + threshold = thresholds.get(cat_name, 10) + removed = before - after + if removed > 0: + print(f" {cat_name:>12}: {before:>8,} → {after:>8,} (min {threshold:>5}, -{removed:,})", file=sys.stderr) + else: + print(f" {cat_name:>12}: {before:>8,} (min {threshold:>5}, all kept)", file=sys.stderr) + + if args.dry_run: + continue + + data["tags"] = pruned + output_path = args.output or filepath + 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} ({size_mb:.1f} MB)", file=sys.stderr) + update_manifest(os.path.dirname(output_path) or ".") + + +if __name__ == "__main__": + main() diff --git a/html/locale_en.json b/html/locale_en.json index c785cc937..8f2ac1e48 100644 --- a/html/locale_en.json +++ b/html/locale_en.json @@ -54,9 +54,12 @@ {"id":"","label":"_Guidance scale","localized":"","hint":"","ui":"txt2img"}, {"id":"","label":"_Guidance rescale","localized":"","hint":"","ui":"txt2img"}, {"id":"","label":"_Guidance start","localized":"","hint":"","ui":"txt2img"}, - {"id":"","label":"_Guidance stop","localized":"","hint":"","ui":"txt2img"} + {"id":"","label":"_Guidance stop","localized":"","hint":"","ui":"txt2img"}, + {"id":"tag_autocomplete_refresh","label":"⟲","localized":"","hint":"Fetch list of available tag dictionaries from the remote repository","ui":"script_autocomplete"}, + {"id":"tag_autocomplete_update","label":"⇩","localized":"","hint":"Re-download enabled dictionaries if a newer version is available","ui":"script_autocomplete"} ], "a": [ + {"id":"","label":"Active dictionaries","localized":"","hint":"Select which tag dictionaries are used for prompt autocompletion.
Dictionaries not yet downloaded locally will be fetched automatically when the autocomplete engine loads them.","ui":"script_autocomplete"}, {"id":"txt2img_advanced","label":"Advanced","localized":"","hint":"Advanced settings used to run image generation","ui":"txt2img"}, {"id":"txt2img_adapters","label":"Adapters","localized":"","hint":"Settings related to IP Adapters","ui":"txt2img"}, {"id":"component-981","label":"Apply to model","localized":"","hint":"","ui":"script_layerdiffuse"}, @@ -71,6 +74,7 @@ {"id":"","label":"Answer","localized":"","hint":"","ui":"caption"}, {"id":"","label":"Adjust start","localized":"","hint":"Starting step when sigma adjust occurs","ui":"txt2img"}, {"id":"","label":"Adjust end","localized":"","hint":"Ending step when sigma adjust occurs","ui":"txt2img"}, + {"id":"","label":"Autocomplete","localized":"","hint":"Enable or disable Tag Autocomplete. Choose which dictionaries are used for prompt autocompletion in Extras","ui":"control"}, {"id":"","label":"AutoGuidance dropout","localized":"","hint":"","ui":"txt2img"}, {"id":"","label":"AutoGuidance layers","localized":"","hint":"","ui":"txt2img"}, {"id":"","label":"AutoGuidance config","localized":"","hint":"","ui":"txt2img"}, @@ -315,6 +319,7 @@ {"id":"","label":"CivitAI discard downloads with hash mismatch","localized":"","hint":"","ui":"settings_huggingface"}, {"id":"","label":"Cache text encoder results","localized":"","hint":"","ui":"settings_legacy_options"}, {"id":"","label":"contain","localized":"","hint":"","ui":"settings_legacy_options"}, + {"id":"","label":"Comma separator","localized":"","hint":"Automatically insert a comma between tags when accepting an autocomplete suggestion.
Disable for natural-language prompts where commas are not used as delimiters.","ui":"script_autocomplete"}, {"id":"","label":"Ctrl+up/down word delimiters","localized":"","hint":"","ui":"settings_legacy_options"}, {"id":"","label":"Ctrl+up/down precision when editing (attention:1.1)","localized":"","hint":"","ui":"settings_legacy_options"}, {"id":"","label":"Ctrl+up/down precision when editing ","localized":"","hint":"","ui":"settings_legacy_options"}, @@ -941,6 +946,7 @@ {"id":"","label":"max-autotune-no-cudagraphs","localized":"","hint":"","ui":"settings_compile"}, {"id":"","label":"Maximum image size (MP)","localized":"","hint":"","ui":"settings_saving-images"}, {"id":"","label":"Max words","localized":"","hint":"","ui":"settings_saving-paths"}, + {"id":"","label":"Min characters","localized":"","hint":"Number of characters that must be typed before autocomplete suggestions appear.
Lower values show suggestions sooner but may feel noisy; higher values wait for a more specific prefix.","ui":"script_autocomplete"}, {"id":"","label":"Modern","localized":"","hint":"","ui":"settings_ui"}, {"id":"","label":"Mount URL subpath","localized":"","hint":"","ui":"settings_ui"}, {"id":"","label":"Mobile scale","localized":"","hint":"","ui":"settings_ui"}, @@ -1229,6 +1235,7 @@ {"id":"","label":"RAS enabled","localized":"","hint":"","ui":"settings_advanced"}, {"id":"","label":"reduce-overhead","localized":"","hint":"","ui":"settings_compile"}, {"id":"","label":"repeated","localized":"","hint":"","ui":"settings_compile"}, + {"id":"","label":"Replace underscores","localized":"","hint":"Display underscores in tag names as spaces in the autocomplete suggestion list.
For example, long_hair appears as long hair.","ui":"script_autocomplete"}, {"id":"","label":"Root model folder","localized":"","hint":"","ui":"settings_system-paths"}, {"id":"","label":"Resize background color","localized":"","hint":"","ui":"settings_saving-images"}, {"id":"","label":"Restore from metadata: skip params","localized":"","hint":"","ui":"settings_image-metadata"}, @@ -1436,6 +1443,7 @@ {"id":"","label":"T2I Adapter","localized":"","hint":"","ui":"control"}, {"id":"","label":"Tagger","localized":"","hint":"Tag images using anime-focused classification models like WaifuDiffusion or DeepBooru.","ui":"caption"}, {"id":"btn_wd_tag","label":"Tag","localized":"","hint":"","ui":"caption"}, + {"id":"","label":"Tag Autocomplete","localized":"","hint":"Suggests matching tags from booru and other dictionaries as you type in prompt fields.
Use the refresh button to fetch the list of available dictionaries, then select which ones to enable.","ui":"script_autocomplete"}, {"id":"","label":"Text Encoder","localized":"","hint":"Settings related to text encoder and prompt encoding processing during generate"}, {"id":"","label":"Text","localized":"","hint":"Create image from text"}, {"id":"","label":"TorchAO","localized":"","hint":"","ui":"settings_quantization"}, diff --git a/javascript/autocomplete.js b/javascript/autocomplete.js new file mode 100644 index 000000000..9d4eab9c0 --- /dev/null +++ b/javascript/autocomplete.js @@ -0,0 +1,564 @@ +/* + * Tag autocomplete for SD.Next prompt textareas. + * + * Ported from Enso's CodeMirror-based autocomplete (autocomplete.ts). + * Uses binary search on sorted tag arrays for O(log n) prefix lookup, + * with substring fallback for 4+ char queries. + */ + +// -- Category colors (unified 14-category scheme) -- + +const CATEGORY_COLORS = { + 0: '#0075f8', // general + 1: '#cc0000', // artist + 2: '#ff4500', // studio + 3: '#9900ff', // copyright + 4: '#00ab2c', // character + 5: '#ed5d1f', // species + 6: '#8a66ff', // genre + 7: '#00cccc', // medium + 8: '#6b7280', // meta + 9: '#228b22', // lore + 10: '#e67e22', // lens + 11: '#f1c40f', // lighting + 12: '#1abc9c', // composition + 13: '#e84393', // color +}; + +const CATEGORY_NAMES = { + 0: 'general', + 1: 'artist', + 2: 'studio', + 3: 'copyright', + 4: 'character', + 5: 'species', + 6: 'genre', + 7: 'medium', + 8: 'meta', + 9: 'lore', + 10: 'lens', + 11: 'lighting', + 12: 'composition', + 13: 'color', +}; + +let active = false; + +// -- Utilities (ported from Enso) -- + +/** Binary search for the first tag where tag.name >= prefix. */ +function lowerBound(tags, prefix) { + let lo = 0; + let hi = tags.length; + while (lo < hi) { + const mid = (lo + hi) >>> 1; + if (tags[mid].name < prefix) lo = mid + 1; + else hi = mid; + } + return lo; +} + +/** Format post count as abbreviated string. */ +function formatCount(count) { + if (count >= 1_000_000) return `${(count / 1_000_000).toFixed(1)}M`; + if (count >= 1_000) return `${Math.round(count / 1_000)}k`; + return String(count); +} + +/** + * Estimate viewport Y of the bottom of the caret line using a persistent + * offscreen mirror div. Styles and width are re-read from the textarea on + * every call so resized textareas are handled correctly. + */ +let caretMirror = null; +let caretMarker = null; +const MIRROR_PROPS = ['fontFamily', 'fontSize', 'fontWeight', 'fontStyle', + 'lineHeight', 'letterSpacing', 'wordSpacing', 'textTransform', + 'padding', 'border', 'boxSizing']; + +function caretViewportY(textarea) { + if (!caretMirror) { + caretMirror = document.createElement('div'); + caretMirror.className = 'autocomplete-mirror'; + caretMirror.style.whiteSpace = 'pre-wrap'; + caretMirror.style.wordWrap = 'break-word'; + caretMirror.style.position = 'absolute'; + caretMirror.style.left = '-9999px'; + caretMirror.style.overflow = 'hidden'; + caretMarker = document.createElement('span'); + caretMarker.textContent = '\u200b'; + document.body.appendChild(caretMirror); + } + const cs = getComputedStyle(textarea); + for (const p of MIRROR_PROPS) caretMirror.style[p] = cs[p]; + caretMirror.style.width = `${textarea.offsetWidth}px`; + caretMirror.textContent = textarea.value.substring(0, textarea.selectionStart); + caretMirror.appendChild(caretMarker); + const offset = caretMarker.offsetTop + caretMarker.offsetHeight; + return textarea.getBoundingClientRect().top + offset - textarea.scrollTop; +} + +// -- TagIndex -- + +class TagIndex { + constructor(data) { + this.categories = data.categories || {}; + // Build sorted array of {name, category, count} from raw [name, catId, count] tuples + this.tags = data.tags.map(([name, category, count]) => ({ + name: name.toLowerCase(), + display: name, + category, + count, + })); + this.tags.sort((a, b) => a.name.localeCompare(b.name)); + } + + /** Prefix search with binary search. Returns matches sorted by count descending. */ + search(prefix, limit = 20) { + const query = prefix.toLowerCase().replace(/ /g, '_'); + if (!query) return []; + const start = lowerBound(this.tags, query); + const matches = []; + for (let i = start; i < this.tags.length && matches.length < limit * 5; i++) { + if (!this.tags[i].name.startsWith(query)) break; + matches.push(this.tags[i]); + } + // Substring fallback for 4+ chars if prefix found nothing + if (matches.length === 0 && query.length >= 4) { + for (let i = 0; i < this.tags.length && matches.length < limit * 5; i++) { + if (this.tags[i].name.includes(query)) matches.push(this.tags[i]); + } + } + matches.sort((a, b) => b.count - a.count); + return matches.slice(0, limit); + } +} + +// -- Engine -- + +const engine = { + indices: new Map(), // name -> TagIndex + categoryColors: { ...CATEGORY_COLORS }, + categoryNames: { ...CATEGORY_NAMES }, + + async loadEnabled() { + const enabled = window.opts?.autocomplete_enabled || []; + active = window.opts?.autocomplete_active || false; + if (!active) { + this.indices.clear(); + return; + } + const toLoad = enabled.filter((n) => !this.indices.has(n)); + const toRemove = [...this.indices.keys()].filter((n) => !enabled.includes(n)); + toRemove.forEach((n) => this.indices.delete(n)); + await Promise.all(toLoad.map(async (name) => { + try { + const resp = await fetch(`${window.api}/autocomplete/${name}`, { credentials: 'include' }); + if (!resp.ok) throw new Error(`${resp.status}`); + const data = await resp.json(); + this.indices.set(name, new TagIndex(data)); + // Extract category colors from first loaded file + if (data.categories) { + Object.entries(data.categories).forEach(([id, cat]) => { + if (cat.color) this.categoryColors[id] = cat.color; + if (cat.name) this.categoryNames[id] = cat.name; + }); + } + log('autoComplete', { loaded: name, tags: data.tags?.length || 0 }); + } catch (e) { + log('autoComplete', { failed: name, error: e }); + } + })); + }, + + searchAll(prefix, limit = 20) { + if (this.indices.size === 0) return []; + const all = []; + this.indices.forEach((index) => { + all.push(...index.search(prefix, limit)); + }); + // Deduplicate by name, keeping highest count + const seen = new Map(); + all.forEach((tag) => { + const existing = seen.get(tag.name); + if (!existing || tag.count > existing.count) seen.set(tag.name, tag); + }); + const results = [...seen.values()]; + results.sort((a, b) => b.count - a.count); + return results.slice(0, limit); + }, +}; + +// -- Textarea integration -- + +/** Extract the current word being typed at the cursor position. */ +function getCurrentWord(textarea) { + const { value, selectionStart } = textarea; + if (selectionStart !== textarea.selectionEnd) return null; // has selection + // Scan backward from cursor to find word start + let start = selectionStart; + while (start > 0) { + const ch = value[start - 1]; + if (ch === ',' || ch === '\n') break; + start--; + } + // Skip leading whitespace + while (start < selectionStart && value[start] === ' ') start++; + const word = value.slice(start, selectionStart); + if (!word) return null; + // Skip if inside angle brackets (LoRA/embedding syntax) + const before = value.slice(0, selectionStart); + const lastOpen = before.lastIndexOf('<'); + const lastClose = before.lastIndexOf('>'); + if (lastOpen > lastClose) return null; + // Skip if inside wildcard syntax + const wcBefore = before.slice(start); + if (wcBefore.startsWith('__') && !wcBefore.endsWith('__')) return null; + return { word, start, end: selectionStart }; +} + +/** Insert a tag at the current word position, replacing the typed prefix. */ +function insertTag(textarea, tagName) { + const info = getCurrentWord(textarea); + if (!info) return; + const { value } = textarea; + const before = value.slice(0, info.start); + const after = value.slice(info.end); + // Build insertion: tag + separator + const useComma = window.opts?.autocomplete_append_comma ?? true; + const sep = useComma ? ',' : ''; + const needsSepBefore = before.length > 0 && before.trimEnd().length > 0 && !before.trimEnd().endsWith(','); + const prefix = needsSepBefore ? `${sep} ` : ''; + let suffix = `${sep} `; + if (after.length > 0 && after.trimStart().startsWith(',')) suffix = ' '; + const insertion = `${prefix}${tagName}${suffix}`; + textarea.value = before.trimEnd() + (before.trimEnd().length > 0 ? ' ' : '') + insertion + after.trimStart(); + // Position cursor after the inserted tag + separator + const cursorPos = before.trimEnd().length + (before.trimEnd().length > 0 ? 1 : 0) + insertion.length; + textarea.selectionStart = cursorPos; + textarea.selectionEnd = cursorPos; + // Sync with Gradio + if (typeof updateInput === 'function') updateInput(textarea); +} + +// -- Dropdown -- + +const dropdown = { + el: null, + listEl: null, + selectedIndex: -1, + results: [], + textarea: null, + query: '', + visible: false, + + init() { + this.el = document.createElement('div'); + this.el.className = 'autocompleteResults'; + this.el.style.display = 'none'; + this.listEl = document.createElement('ul'); + this.listEl.className = 'autocompleteResultsList'; + this.el.appendChild(this.listEl); + document.body.appendChild(this.el); + this.el.addEventListener('mousedown', (e) => e.preventDefault()); // prevent blur on click + this.el.addEventListener('click', (e) => { + const li = e.target.closest('li'); + if (!li) return; + const idx = [...this.listEl.children].indexOf(li); + if (idx >= 0 && idx < this.results.length) { + this.selectedIndex = idx; + this.accept(); + } + }); + this.resizeObserver = new ResizeObserver(() => { + if (this.visible) this.position(); + }); + }, + + show(results, textarea, query) { + if (results.length === 0) { this.hide(); return; } + if (this.textarea !== textarea) { + if (this.textarea) this.resizeObserver.unobserve(this.textarea); + this.resizeObserver.observe(textarea); + } + this.results = results; + this.textarea = textarea; + this.query = query || ''; + this.selectedIndex = -1; + this.render(); + this.position(); + this.el.style.display = ''; + this.visible = true; + }, + + hide() { + if (this.textarea) this.resizeObserver.unobserve(this.textarea); + this.textarea = null; + this.el.style.display = 'none'; + this.visible = false; + this.results = []; + this.selectedIndex = -1; + }, + + render() { + const replaceUnderscores = window.opts?.autocomplete_replace_underscores ?? true; + const queryNorm = this.query.toLowerCase().replace(/ /g, '_'); + this.listEl.replaceChildren(); + this.results.forEach((tag, i) => { + const li = document.createElement('li'); + if (i === this.selectedIndex) li.classList.add('selected'); + const dot = document.createElement('span'); + dot.className = 'autocomplete-category'; + dot.style.color = engine.categoryColors[tag.category] || '#888'; + dot.textContent = '\u25CF'; + dot.title = engine.categoryNames[tag.category] || ''; + const name = document.createElement('span'); + name.className = 'autocomplete-tag'; + const tagText = replaceUnderscores ? tag.display.replace(/_/g, ' ') : tag.display; + const matchPos = tag.name.indexOf(queryNorm); + if (matchPos >= 0 && queryNorm.length > 0) { + const mark = document.createElement('mark'); + mark.textContent = tagText.slice(matchPos, matchPos + queryNorm.length); + name.append( + document.createTextNode(tagText.slice(0, matchPos)), + mark, + document.createTextNode(tagText.slice(matchPos + queryNorm.length)), + ); + } else { + name.textContent = tagText; + } + const count = document.createElement('span'); + count.className = 'autocomplete-count'; + count.textContent = tag.count > 0 ? formatCount(tag.count) : ''; + li.append(dot, name, count); + li.addEventListener('mouseenter', () => { + this.selectedIndex = i; + this.updateSelection(); + }); + this.listEl.appendChild(li); + }); + }, + + position() { + if (!this.textarea) return; + const rect = this.textarea.getBoundingClientRect(); + // Position near the caret line instead of the textarea bottom + const cursorBottom = caretViewportY(this.textarea); + const anchorY = Math.max(rect.top, Math.min(cursorBottom, rect.bottom)); + const spaceBelow = window.innerHeight - anchorY; + const dropHeight = Math.min(this.el.scrollHeight, 300); + if (spaceBelow >= dropHeight || spaceBelow >= anchorY - rect.top) { + this.el.style.top = `${anchorY + 2}px`; + } else { + this.el.style.top = `${anchorY - dropHeight - 2}px`; + } + this.el.style.left = `${rect.left}px`; + this.el.style.width = `${rect.width}px`; + }, + + updateSelection() { + [...this.listEl.children].forEach((li, i) => { + li.classList.toggle('selected', i === this.selectedIndex); + }); + const selected = this.listEl.children[this.selectedIndex]; + if (selected) selected.scrollIntoView({ block: 'nearest' }); + }, + + navigate(dir) { + if (this.results.length === 0) return; + if (this.selectedIndex === -1) { + this.selectedIndex = dir > 0 ? 0 : this.results.length - 1; + } else { + this.selectedIndex = (this.selectedIndex + dir + this.results.length) % this.results.length; + } + this.updateSelection(); + }, + + accept() { + if (this.selectedIndex < 0 || this.selectedIndex >= this.results.length) { + // Tab with no selection: select first + if (this.results.length > 0) { + this.selectedIndex = 0; + this.updateSelection(); + } + return; + } + const tag = this.results[this.selectedIndex]; + if (this.textarea) insertTag(this.textarea, tag.display); + this.hide(); + }, +}; + +// -- Event handlers -- + +let debounceTimer = null; + +function onInput(textarea) { + if (!active) return; + const minChars = window.opts?.autocomplete_min_chars ?? 3; + const info = getCurrentWord(textarea); + if (!info || info.word.length < minChars) { + dropdown.hide(); + return; + } + clearTimeout(debounceTimer); + debounceTimer = setTimeout(() => { + const results = engine.searchAll(info.word); + dropdown.show(results, textarea, info.word); + }, 150); +} + +function onKeyDown(e) { + if (!dropdown.visible) return; + switch (e.key) { + case 'ArrowDown': + e.preventDefault(); + e.stopPropagation(); + dropdown.navigate(1); + break; + case 'ArrowUp': + e.preventDefault(); + e.stopPropagation(); + dropdown.navigate(-1); + break; + case 'Enter': + if (dropdown.selectedIndex >= 0) { + e.preventDefault(); + e.stopPropagation(); + dropdown.accept(); + } + break; + case 'Tab': + e.preventDefault(); + e.stopPropagation(); + dropdown.accept(); + break; + case 'Escape': + e.preventDefault(); + e.stopPropagation(); + dropdown.hide(); + break; + default: + break; + } +} + +/** Attach autocomplete to a single textarea. */ +function attachAutocomplete(textarea) { + textarea.addEventListener('input', () => onInput(textarea)); + textarea.addEventListener('keydown', onKeyDown); + textarea.addEventListener('focusout', () => { + setTimeout(() => dropdown.hide(), 200); + }); +} + +// -- Prompt textarea IDs -- + +const PROMPT_IDS = [ + 'txt2img_prompt', 'txt2img_neg_prompt', + 'img2img_prompt', 'img2img_neg_prompt', + 'control_prompt', 'control_neg_prompt', + 'video_prompt', 'video_neg_prompt', +]; + +// -- Active button -- + +function patchActiveButton() { + const buttons = [...gradioApp().querySelectorAll('.autocomplete-active')]; + active = window.opts?.autocomplete_active || false; + buttons.forEach((btn) => { + btn.classList.toggle('autocomplete-active', active); + btn.classList.toggle('autocomplete-inactive', !active); + btn.parentElement.onclick = () => { + active = !active; + window.opts.autocomplete_active = !active; + btn.classList.toggle('autocomplete-active', active); + btn.classList.toggle('autocomplete-inactive', !active); + }; + }); +} + +// -- Config bridge -- + +/** Monkey-patch script config bridge textboxes to push autocomplete config changes to window.opts immediately. */ +function patchConfigBridge() { + const elements = gradioApp().querySelectorAll('[id$="_tag_autocomplete_config_json"]'); + for (const el of elements) { + const textarea = el.querySelector('textarea'); + if (!textarea || textarea.acBridgePatched) continue; + textarea.acBridgePatched = true; + const proto = Object.getOwnPropertyDescriptor(HTMLTextAreaElement.prototype, 'value'); + Object.defineProperty(textarea, 'value', { + set(newValue) { + const oldValue = proto.get.call(textarea); + proto.set.call(textarea, newValue); + if (oldValue !== newValue && newValue) { + try { + const cfg = JSON.parse(newValue); + for (const [key, val] of Object.entries(cfg)) window.opts[key] = val; + executeCallbacks(optionsChangedCallbacks); + } catch { /* ignore parse errors */ } + } + }, + get() { return proto.get.call(textarea); }, + }); + } +} + +// -- Initialization -- + +async function initAutocomplete() { + const enabled = window.opts?.autocomplete_enabled || []; + active = window.opts?.autocomplete_active || false; + log('autoComplete', { active, enabled }); + // Inject styles (CSS files in javascript/ are not auto-loaded) + const style = document.createElement('style'); + style.textContent = [ + '.autocompleteResults { position: fixed; z-index: 9999; max-height: 300px; overflow-y: auto;', + ' background: var(--sd-main-background-color, var(--background-fill-primary, #1f2937));', + ' border: 1px solid var(--sd-input-border-color, var(--border-color-primary, #374151));', + ' border-radius: var(--sd-border-radius, 6px); box-shadow: 0 4px 16px rgba(0,0,0,0.4);', + ' font-size: 13px; scrollbar-width: thin; }', + '.autocompleteResultsList { list-style: none; margin: 0; padding: 4px 0; }', + '.autocompleteResultsList > li { display: flex; align-items: center; padding: 6px 12px; cursor: pointer;', + ' gap: 8px; line-height: 1.4; transition: background 0.1s ease; border-bottom: 1px solid rgba(255,255,255,0.03); }', + '.autocompleteResultsList > li:last-child { border-bottom: none; }', + '.autocompleteResultsList > li:hover { background: var(--sd-panel-background-color, var(--input-background-fill-focus, #374151)); }', + '.autocompleteResultsList > li.selected { background: var(--sd-main-accent-color, var(--button-primary-background-fill, #4b5563)); }', + '.autocomplete-category { font-size: 10px; flex-shrink: 0; width: 10px; text-align: center; cursor: help; }', + '.autocomplete-tag { flex: 1; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }', + '.autocomplete-tag mark { background: transparent; color: inherit; font-weight: 700; }', + '.autocomplete-count { font-size: 0.75em; opacity: 0.45; flex-shrink: 0; font-variant-numeric: tabular-nums;', + ' background: rgba(255,255,255,0.06); padding: 1px 6px; border-radius: 8px; min-width: 28px; text-align: right; }', + ].join('\n'); + document.head.appendChild(style); + dropdown.init(); + await engine.loadEnabled(); + // Attach to all prompt textareas; even if no dictionaries loaded yet, they may be enabled later via script UI + let attached = 0; + PROMPT_IDS.forEach((id) => { + const textarea = gradioApp().querySelector(`#${id} > label > textarea`); + if (textarea) { + attachAutocomplete(textarea); + attached++; + } + }); + log('autoComplete', { attached, dicts: engine.indices.size }); + // Reload when settings change + onOptionsChanged(async () => { + const newActive = window.opts?.autocomplete_active || false; + const newEnabled = window.opts?.autocomplete_enabled || []; + const currentKeys = [...engine.indices.keys()].sort().join(','); + const newKeys = [...newEnabled].sort().join(','); + if ((currentKeys !== newKeys) || (active !== newActive)) { + log('autoComplete', { reload: newEnabled }); + await engine.loadEnabled(); + active = newActive; + patchActiveButton(); + } + }); + // Watch for config updates from the script UI bridge + patchConfigBridge(); + patchActiveButton(); + onAfterUiUpdate(() => patchConfigBridge()); +} diff --git a/javascript/setHints.js b/javascript/setHints.js index 8ae428ae8..b35fd0eb3 100644 --- a/javascript/setHints.js +++ b/javascript/setHints.js @@ -324,8 +324,11 @@ async function setHints() { for (const el of elements) { // localize elements text let found; - if (el.dataset.original) found = localeData.data.find((l) => l.label.toLowerCase().trim() === el.dataset.original.toLowerCase().trim()); - else found = localeData.data.find((l) => l.label.toLowerCase().trim() === el.textContent.toLowerCase().trim()); + if (el.id) found = localeData.data.find((l) => l.id && (l.id === el.id || el.id.endsWith(l.id))); // prefer id match for disambiguation + if (!found) { + if (el.dataset.original) found = localeData.data.find((l) => l.label.toLowerCase().trim() === el.dataset.original.toLowerCase().trim()); + else found = localeData.data.find((l) => l.label.toLowerCase().trim() === el.textContent.toLowerCase().trim()); + } if (found?.localized?.length > 0) { if (!el.dataset.original) el.dataset.original = el.textContent; replaceTextContent(el, found.localized); @@ -359,9 +362,12 @@ async function applyHintToElement(el) { || (el.tagName === 'SPAN' && (el.parentElement?.tagName === 'LABEL' || el.parentElement?.classList.contains('label-wrap'))); if (!isValidElement) return; - let found; // find matching hint data - if (el.dataset.original) found = localeData.data.find((l) => l.label.toLowerCase().trim() === el.dataset.original.toLowerCase().trim()); - else found = localeData.data.find((l) => l.label.toLowerCase().trim() === el.textContent.toLowerCase().trim()); + let found; // find matching hint data - prefer id match for disambiguation + if (el.id) found = localeData.data.find((l) => l.id && (l.id === el.id || el.id.endsWith(l.id))); + if (!found) { + if (el.dataset.original) found = localeData.data.find((l) => l.label.toLowerCase().trim() === el.dataset.original.toLowerCase().trim()); + else found = localeData.data.find((l) => l.label.toLowerCase().trim() === el.textContent.toLowerCase().trim()); + } if (found?.localized?.length > 0) { // apply localization if found if (!el.dataset.original) el.dataset.original = el.textContent; diff --git a/javascript/startup.js b/javascript/startup.js index 5f6b3f926..a51a0272d 100644 --- a/javascript/startup.js +++ b/javascript/startup.js @@ -59,6 +59,7 @@ async function initStartup() { // optinally wait for modern ui if (window.waitForUiReady) await waitForUiReady(); + initAutocomplete(); monitorConnection(); removeSplash(); diff --git a/modules/api/api.py b/modules/api/api.py index 41ec90876..42d6c912d 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) + # autocomplete api + from modules.api import autocomplete as autocomplete_api + autocomplete_api.init(getattr(shared.opts, 'autocomplete_dir', '') or os.path.join(paths.models_path, 'autocomplete')) + autocomplete_api.register_api(self) + # gallery api from modules.api import gallery gallery.register_api(self.app) @@ -139,9 +145,9 @@ class Api: # hide trailing-slash duplicates from OpenAPI schema from fastapi.routing import APIRoute - paths = {r.path for r in self.app.routes if hasattr(r, 'path')} + route_paths = {r.path for r in self.app.routes if hasattr(r, 'path')} for route in self.app.routes: - if isinstance(route, APIRoute) and len(route.path) > 1 and route.path.endswith('/') and route.path[:-1] in paths: + if isinstance(route, APIRoute) and len(route.path) > 1 and route.path.endswith('/') and route.path[:-1] in route_paths: route.include_in_schema = False # upload api diff --git a/modules/api/autocomplete.py b/modules/api/autocomplete.py new file mode 100644 index 000000000..33f4e6749 --- /dev/null +++ b/modules/api/autocomplete.py @@ -0,0 +1,254 @@ +"""V1 tag autocomplete endpoints. + +Serves pre-built tag files (Danbooru, e621, natural language, artists) +from JSON files in the configured autocomplete directory. Remote files +are hosted on HuggingFace and downloaded on demand. +""" + +import asyncio +import json +import os + +from fastapi.exceptions import HTTPException + +from modules.api.models import ItemAutocomplete, ItemAutocompleteContent, ItemAutocompleteRemote +from modules.logger import log + + +autocomplete_dir: str = "" +cache: dict[str, dict] = {} +HF_REPO = "CalamitousFelicitousness/prompt-vocab" +HF_BASE = f"https://huggingface.co/datasets/{HF_REPO}/resolve/main" +MANIFEST_CACHE_SEC = 300 # re-fetch manifest every 5 minutes +manifest_cache: dict = {} # {"data": [...], "fetched_at": float} + + +def init(path: str) -> None: + """Set the autocomplete directory path. Called once during API registration.""" + global autocomplete_dir # pylint: disable=global-statement + autocomplete_dir = path + + +def get_cached(name: str) -> dict: + """Load a tag 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 name") + path = os.path.join(autocomplete_dir, f"{name}.json") + if not os.path.isfile(path): + cache.pop(name, None) + # Auto-download from HF if available in manifest + try: + manifest = fetch_manifest_sync() + if any(e.get('name') == name for e in manifest): + log.info(f'Autocomplete: name="{name}" auto-download') + download_sync(name) + else: + raise HTTPException(status_code=404, detail=f"Not found: {name}") + except HTTPException: + raise + except Exception as e: + raise HTTPException(status_code=404, detail=f"Not found: {name} ({e})") from e + 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_all_sync() -> list[ItemAutocomplete]: + """Scan autocomplete directory and return metadata for each tag file.""" + if not autocomplete_dir or not os.path.isdir(autocomplete_dir): + return [] + items = [] + for filename in sorted(os.listdir(autocomplete_dir)): + if not filename.endswith('.json') or filename.startswith('.') or filename == 'manifest.json': + continue + name = filename.rsplit('.', 1)[0] + try: + entry = get_cached(name) + meta = entry['meta'] + items.append(ItemAutocomplete( + 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_all() -> list[ItemAutocomplete]: + """List available tag autocomplete files.""" + return await asyncio.to_thread(list_all_sync) + + +async def get_content(name: str) -> ItemAutocompleteContent: + """Get full tag file 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 ItemAutocompleteContent( + name=content.get('name', name), + version=content.get('version', ''), + categories=content.get('categories', {}), + tags=content.get('tags', []), + ) + + +# -- Remote management -- + +def fetch_manifest_sync() -> list[dict]: + """Fetch manifest.json from HuggingFace, with caching.""" + import time + import requests + now = time.time() + if manifest_cache.get('data') and now - manifest_cache.get('fetched_at', 0) < MANIFEST_CACHE_SEC: + return manifest_cache['data'] + url = f"{HF_BASE}/manifest.json" + try: + resp = requests.get(url, timeout=15) + resp.raise_for_status() + data = resp.json() + entries = data.get('entries', data) if isinstance(data, dict) else data + manifest_cache['data'] = entries + manifest_cache['fetched_at'] = now + return entries + except Exception as e: + log.warning(f"Autocomplete: Failed to fetch manifest: {e}") + return manifest_cache.get('data', []) + + +def local_names() -> set[str]: + """Return set of locally available autocomplete file names.""" + if not autocomplete_dir or not os.path.isdir(autocomplete_dir): + return set() + return { + f.rsplit('.', 1)[0] + for f in os.listdir(autocomplete_dir) + if f.endswith('.json') and not f.startswith('.') and f != 'manifest.json' + } + + +def local_version(name: str) -> str: + """Return the version string of a local tag file, or empty string.""" + path = os.path.join(autocomplete_dir, f"{name}.json") + if not os.path.isfile(path): + return "" + try: + entry = cache.get(name) + if entry: + return entry['meta'].get('version', '') + with open(path, encoding='utf-8') as f: + data = json.load(f) + return data.get('version', '') + except Exception: + return "" + + +async def list_remote() -> list[ItemAutocompleteRemote]: + """List tag files available for download from HuggingFace.""" + entries = await asyncio.to_thread(fetch_manifest_sync) + local = await asyncio.to_thread(local_names) + results = [] + for e in entries: + name = e['name'] + is_local = name in local + remote_version = e.get('version', '') + update = False + if is_local and remote_version: + lv = await asyncio.to_thread(local_version, name) + update = bool(lv and lv != remote_version) + results.append(ItemAutocompleteRemote( + name=name, + description=e.get('description', ''), + version=remote_version, + tag_count=e.get('tag_count', 0), + size_mb=e.get('size_mb', 0), + downloaded=is_local, + update_available=update, + )) + return results + + +def download_sync(name: str) -> str: + """Download a tag file from HuggingFace to the local autocomplete directory.""" + import requests + if '/' in name or '\\' in name or '..' in name: + raise HTTPException(status_code=400, detail="Invalid name") + os.makedirs(autocomplete_dir, exist_ok=True) + url = f"{HF_BASE}/{name}.json" + try: + resp = requests.get(url, timeout=120, stream=True) + resp.raise_for_status() + except requests.RequestException as e: + raise HTTPException(status_code=502, detail=f"Failed to download {name}: {e}") from e + target = os.path.join(autocomplete_dir, f"{name}.json") + tmp = target + ".tmp" + size = 0 + with open(tmp, 'wb') as f: + for chunk in resp.iter_content(chunk_size=1024 * 256): + f.write(chunk) + size += len(chunk) + os.replace(tmp, target) + cache.pop(name, None) + log.info(f'Autocomplete: name="{name}" url={url} ({size / 1024 / 1024:.2f}MB) downloaded') + return target + + +async def download(name: str): + """Download a tag file from HuggingFace.""" + await asyncio.to_thread(download_sync, name) + entry = await asyncio.to_thread(get_cached, name) + meta = entry['meta'] + return ItemAutocomplete( + name=meta['name'], + version=meta['version'], + tag_count=meta['tag_count'], + categories=meta['categories'], + size=entry['size'], + ) + + +async def delete(name: str): + """Delete a locally downloaded tag file.""" + if '/' in name or '\\' in name or '..' in name: + raise HTTPException(status_code=400, detail="Invalid name") + path = os.path.join(autocomplete_dir, f"{name}.json") + if not os.path.isfile(path): + raise HTTPException(status_code=404, detail=f"Not found: {name}") + await asyncio.to_thread(os.remove, path) + cache.pop(name, None) + return {"status": "deleted", "name": name} + + +def register_api(api): + api.add_api_route("/sdapi/v1/autocomplete", list_all, methods=["GET"], response_model=list[ItemAutocomplete], tags=["Enumerators"]) + api.add_api_route("/sdapi/v1/autocomplete/remote", list_remote, methods=["GET"], response_model=list[ItemAutocompleteRemote], tags=["Enumerators"]) + api.add_api_route("/sdapi/v1/autocomplete/{name}", get_content, methods=["GET"], response_model=ItemAutocompleteContent, tags=["Enumerators"]) + api.add_api_route("/sdapi/v1/autocomplete/{name}/download", download, methods=["POST"], response_model=ItemAutocomplete, tags=["Enumerators"]) + api.add_api_route("/sdapi/v1/autocomplete/{name}", delete, methods=["DELETE"], tags=["Enumerators"]) diff --git a/modules/api/models.py b/modules/api/models.py index aa2bb197d..9e6e6b9ed 100644 --- a/modules/api/models.py +++ b/modules/api/models.py @@ -509,6 +509,28 @@ 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 ItemAutocomplete(BaseModel): + name: str = Field(title="Name", description="Autocomplete file identifier (filename without extension)") + version: str = Field(default="", title="Version", description="Version string") + tag_count: int = Field(default=0, title="Tag count", description="Number of tags") + 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 ItemAutocompleteContent(BaseModel): + name: str = Field(title="Name", description="Autocomplete file identifier") + version: str = Field(default="", title="Version", description="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") + +class ItemAutocompleteRemote(BaseModel): + name: str = Field(title="Name", description="Autocomplete file identifier") + description: str = Field(default="", title="Description", description="Human-readable description") + version: str = Field(default="", title="Version", description="Version string") + tag_count: int = Field(default=0, title="Tag count", description="Number of tags") + size_mb: float = Field(default=0, title="Size (MB)", description="Approximate file size in megabytes") + downloaded: bool = Field(default=False, title="Downloaded", description="Whether available locally") + update_available: bool = Field(default=False, title="Update available", description="Whether a newer version exists remotely") + # 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..06851a6f5 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('autocomplete_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 9feee1e9f..47a69f5b1 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -389,6 +389,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), + "autocomplete_dir": OptionInfo(os.path.join(paths.models_path, 'autocomplete'), "Folder with tag autocomplete files", 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), @@ -657,6 +658,14 @@ def create_settings(cmd_opts): "disabled_extensions": OptionInfo([], "Disable these extensions", gr.Textbox, {"visible": False}), "sd_checkpoint_hash": OptionInfo("", "SHA256 hash of the current checkpoint", gr.Textbox, {"visible": False}), "tooltips": OptionInfo("UI Tooltips", "UI tooltips", gr.Radio, {"choices": ["None", "Browser default", "UI tooltips"], "visible": False}), + + # Autocomplete settings (controlled via Tag Autocomplete script UI) + "autocomplete_active": OptionInfo(False, "Enable Autocomplete", gr.Checkbox, {"visible": False}), + "autocomplete_enabled": OptionInfo([], "Enabled tag autocomplete files", gr.Dropdown, {"multiselect": True, "choices": [], "visible": False}), + "autocomplete_min_chars": OptionInfo(3, "Min autocomplete chars", gr.Slider, {"minimum": 2, "maximum": 6, "step": 1, "visible": False}), + "autocomplete_replace_underscores": OptionInfo(True, "Replace underscores in autocomplete", gr.Checkbox, {"visible": False}), + "autocomplete_append_comma": OptionInfo(True, "Append comma after autocomplete", gr.Checkbox, {"visible": False}), + # Caption settings (controlled via Caption Tab UI) "caption_default_type": OptionInfo("VLM", "Default caption type", gr.Radio, {"choices": ["VLM", "OpenCLiP", "Tagger"], "visible": False}), "tagger_show_scores": OptionInfo(False, "Tagger: show confidence scores in results", gr.Checkbox, {"visible": False}), diff --git a/scripts/autocomplete.py b/scripts/autocomplete.py new file mode 100644 index 000000000..fea760a08 --- /dev/null +++ b/scripts/autocomplete.py @@ -0,0 +1,193 @@ +"""Always-on script providing tag autocomplete dictionary management UI.""" + +import json +import gradio as gr +from modules import shared, scripts_manager +from modules.api import autocomplete as ac_api +from modules.ui_components import ToolButton +import modules.ui_symbols as symbols +from modules.logger import log + + +def get_all_names(): + """Merge local file names with cached remote manifest names.""" + local = ac_api.local_names() + remote = set() + cached = ac_api.manifest_cache.get('data') + if cached: + remote = {e['name'] for e in cached if 'name' in e} + return sorted(local | remote) + + +def get_config_json(): + """Serialize autocomplete opts for the JS config bridge.""" + return json.dumps({ + "autocomplete_active": bool(shared.opts.data.get('autocomplete_active', False)), + "autocomplete_enabled": list(shared.opts.data.get('autocomplete_enabled', [])), + "autocomplete_min_chars": shared.opts.data.get('autocomplete_min_chars', 3), + "autocomplete_replace_underscores": shared.opts.data.get('autocomplete_replace_underscores', True), + "autocomplete_append_comma": shared.opts.data.get('autocomplete_append_comma', True), + }) + + +def on_active_change(value): + shared.opts.data['autocomplete_active'] = bool(value) + shared.opts.save(silent=True) + return get_config_json(), "" + + +def on_enabled_change(selected): + shared.opts.data['autocomplete_enabled'] = list(selected) + shared.opts.save(silent=True) + return get_config_json(), "" + + +def on_min_chars_change(value): + shared.opts.data['autocomplete_min_chars'] = int(value) + shared.opts.save(silent=True) + return get_config_json() + + +def on_replace_underscores_change(value): + shared.opts.data['autocomplete_replace_underscores'] = bool(value) + shared.opts.save(silent=True) + return get_config_json() + + +def on_append_comma_change(value): + shared.opts.data['autocomplete_append_comma'] = bool(value) + shared.opts.save(silent=True) + return get_config_json() + + +def format_status(local, remote_entries, fetch_ok): + """Build status HTML showing available dictionaries.""" + lines = [] + remote_names = set() + for e in remote_entries: + name = e.get('name', '') + remote_names.add(name) + dl_status = '' if name in local else symbols.save + desc = e.get('description', '') + # size = e.get('size_mb', 0) + tags = e.get('tag_count', 0) + lines.append(f"{name} | {desc} | {tags:,} tags {dl_status}") + for name in sorted(local - remote_names): + lines.append(f"{name}") + if not fetch_ok: + lines.insert(0, "Remote fetch failed; showing local files only") + elif not lines: + lines.append("No dictionaries found") + return "
".join(lines) + + +def on_refresh(): + """Fetch remote manifest and update dropdown choices.""" + try: + ac_api.manifest_cache.pop('fetched_at', None) # force re-fetch by expiring cache + ac_api.fetch_manifest_sync() + fetch_ok = bool(ac_api.manifest_cache.get('fetched_at')) + names = get_all_names() + current = list(shared.opts.data.get('autocomplete_enabled', [])) + local = ac_api.local_names() + remote_entries = ac_api.manifest_cache.get('data', []) + msg = format_status(local, remote_entries, fetch_ok) + return gr.update(choices=names, value=current), msg + except Exception as e: + log.warning(f"Autocomplete refresh: {e}") + return gr.update(), f"Refresh failed: {e}" + + +def on_update(selected): + """Re-download enabled dictionaries if remote version is newer.""" + if not selected: + return "No dictionaries enabled" + try: + entries = ac_api.fetch_manifest_sync() + except Exception as e: + return f"Failed to fetch manifest: {e}" + updated = [] + for name in selected: + remote_entry = next((e for e in entries if e.get('name') == name), None) + if not remote_entry: + continue + remote_ver = remote_entry.get('version', '') + local_ver = ac_api.local_version(name) + if not local_ver or (remote_ver and local_ver != remote_ver): + try: + ac_api.download_sync(name) + updated.append(name) + except Exception as e: + log.warning(f"Autocomplete update {name}: {e}") + if updated: + return f"Updated: {', '.join(updated)}" + return "All dictionaries are up to date" + + +class AutocompleteScript(scripts_manager.Script): + + def show(self, is_img2img): + return scripts_manager.AlwaysVisible + + def title(self): + return "Tag Autocomplete" + + def ui(self, is_img2img): + initial_names = get_all_names() + initial_enabled = list(shared.opts.data.get('autocomplete_enabled', [])) + + with gr.Accordion('Tag Autocomplete', open=False, elem_id='autocomplete_settings'): + with gr.Row(): + active_cb = gr.Checkbox( + label="Enable Autocomplete", + value=bool(shared.opts.data.get('autocomplete_active', False)), + elem_id=self.elem_id("active"), + ) + with gr.Row(): + enabled_dd = gr.Dropdown( + label="Active dictionaries", + multiselect=True, + choices=initial_names, + value=initial_enabled, + interactive=True, + elem_id=self.elem_id("enabled"), + ) + refresh_btn = ToolButton(value=symbols.refresh, elem_id=self.elem_id("refresh")) + update_btn = ToolButton(value=symbols.save, elem_id=self.elem_id("update")) + with gr.Row(): + replace_underscores = gr.Checkbox( + label="Replace underscores", + value=shared.opts.data.get('autocomplete_replace_underscores', True), + elem_id=self.elem_id("replace_underscores"), + ) + append_comma = gr.Checkbox( + label="Comma separator", + value=shared.opts.data.get('autocomplete_append_comma', True), + elem_id=self.elem_id("append_comma"), + ) + min_chars = gr.Slider( + label="Min characters", + minimum=2, maximum=6, step=1, + value=shared.opts.data.get('autocomplete_min_chars', 3), + elem_id=self.elem_id("min_chars"), + ) + with gr.Row(): + status = gr.HTML(value="", elem_id=self.elem_id("status")) + config_json = gr.Textbox( + value=get_config_json, + visible=False, + elem_id=self.elem_id("config_json"), + ) + + active_cb.change(fn=on_active_change, inputs=[active_cb], outputs=[config_json, status]) + enabled_dd.change(fn=on_enabled_change, inputs=[enabled_dd], outputs=[config_json, status]) + min_chars.change(fn=on_min_chars_change, inputs=[min_chars], outputs=[config_json]) + replace_underscores.change(fn=on_replace_underscores_change, inputs=[replace_underscores], outputs=[config_json]) + append_comma.change(fn=on_append_comma_change, inputs=[append_comma], outputs=[config_json]) + refresh_btn.click(fn=on_refresh, inputs=[], outputs=[enabled_dd, status]) + update_btn.click(fn=on_update, inputs=[enabled_dd], outputs=[status]) + + for comp in [enabled_dd, min_chars, replace_underscores, append_comma, config_json, status]: + comp.do_not_save_to_config = True + + return [active_cb, enabled_dd, min_chars, replace_underscores, append_comma, config_json]