mirror of
https://github.com/vladmandic/automatic
synced 2026-09-18 16:54:33 +02:00
Merge pull request #4707 from vladmandic/feat/tag-dictionaries
Feat/tag dictionaries
This commit is contained in:
Vladimir Mandic
GitHub
@@ -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('<?xml'):
|
||||
if page == 0:
|
||||
print(f" Error: {text[:200]}", file=sys.stderr)
|
||||
break
|
||||
try:
|
||||
root = ET.fromstring(text)
|
||||
except ET.ParseError as e:
|
||||
print(f" Skipping page {page}: malformed XML ({e})", file=sys.stderr)
|
||||
page += 1
|
||||
time.sleep(rate_limit)
|
||||
continue
|
||||
elements = root.findall('tag')
|
||||
for el in elements:
|
||||
count = int(el.get("count", "0"))
|
||||
if count >= 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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
+9
-1
@@ -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.<br>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.<br>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 <extra networks:0.9>","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.<br>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.<br>For example, <i>long_hair</i> appears as <i>long hair</i>.","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.<br>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"},
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
+11
-5
@@ -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;
|
||||
|
||||
@@ -59,6 +59,7 @@ async function initStartup() {
|
||||
|
||||
// optinally wait for modern ui
|
||||
if (window.waitForUiReady) await waitForUiReady();
|
||||
initAutocomplete();
|
||||
monitorConnection();
|
||||
removeSplash();
|
||||
|
||||
|
||||
+9
-3
@@ -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
|
||||
|
||||
@@ -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"])
|
||||
@@ -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]:
|
||||
|
||||
@@ -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', '')
|
||||
|
||||
@@ -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}),
|
||||
|
||||
@@ -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"<b>{name}</b> | {desc} | {tags:,} tags {dl_status}")
|
||||
for name in sorted(local - remote_names):
|
||||
lines.append(f"<b>{name}</b>")
|
||||
if not fetch_ok:
|
||||
lines.insert(0, "<i>Remote fetch failed; showing local files only</i>")
|
||||
elif not lines:
|
||||
lines.append("No dictionaries found")
|
||||
return "<br>".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]
|
||||
Reference in New Issue
Block a user