Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 09:14:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(dicts): add curated dicts, pruning, and manifest generation

Browse Source 
Add 4 curated vocabulary dicts (art, photography, quality, negative)
that ship bundled for out-of-the-box autocomplete. Add manifest.json
listing all dicts available on HuggingFace.

CLI tools:
- cli/prune_dicts.py: per-category pruning with configurable thresholds
- cli/gen_manifest.py: generate manifest from on-disk dict files
- cli/fetch_dicts.py: unified 14-category scheme, idol source, auto
  manifest regeneration after fetch

Fetch and prune auto-update manifest.json when one already exists.
This commit is contained in:
CalamitousFelicitousness
2026-03-24 23:41:06 +00:00
parent f3e0b4a86b
commit 36b2ed6bb7
8 changed files with 837 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cli/fetch_dicts.py
+200 -99
View File
@@ -6,8 +6,12 @@ Usage:
python cli/fetch_dicts.py e621 [--output PATH] [--min-count N]
python cli/fetch_dicts.py rule34 --key USER_ID:API_KEY [--output PATH] [--min-count N]
python cli/fetch_dicts.py sankaku [--output PATH] [--min-count N]
python cli/fetch_dicts.py idol [--output PATH] [--min-count N]
python cli/fetch_dicts.py all --key USER_ID:API_KEY [--output-dir DIR] [--min-count N]
Fetches are checkpointed 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.
@@ -22,59 +26,115 @@ from datetime import date
import requests
DANBOORU_CATEGORIES = {
# 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": "#a800aa"},
"3": {"name": "copyright", "color": "#dd00dd"},
"4": {"name": "character", "color": "#00ab2c"},
"5": {"name": "meta", "color": "#ee8800"},
}
E621_CATEGORIES = {
"0": {"name": "general", "color": "#0075f8"},
"1": {"name": "artist", "color": "#a800aa"},
"3": {"name": "copyright", "color": "#dd00dd"},
"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": "invalid", "color": "#ff3d3d"},
"7": {"name": "meta", "color": "#ee8800"},
"8": {"name": "lore", "color": "#228b22"},
"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"},
}
# Rule34 uses the same category IDs as Danbooru (Gelbooru-compatible)
RULE34_CATEGORIES = {
"0": {"name": "general", "color": "#0075f8"},
"1": {"name": "artist", "color": "#a800aa"},
"3": {"name": "copyright", "color": "#dd00dd"},
"4": {"name": "character", "color": "#00ab2c"},
"5": {"name": "meta", "color": "#ee8800"},
}
# 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}
SANKAKU_CATEGORIES = {
"0": {"name": "general", "color": "#0075f8"},
"1": {"name": "artist", "color": "#a800aa"},
"2": {"name": "studio", "color": "#a800aa"},
"3": {"name": "copyright", "color": "#dd00dd"},
"4": {"name": "character", "color": "#00ab2c"},
"5": {"name": "species", "color": "#ed5d1f"},
"8": {"name": "medium", "color": "#ee8800"},
"9": {"name": "meta", "color": "#ee8800"},
# 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)"
CHECKPOINT_INTERVAL = 50 # save progress every N pages
MAX_RETRIES = 3
RETRY_BACKOFF = 5 # seconds, multiplied by attempt number
def fetch_danbooru(min_count: int = 10, **_kwargs) -> list:
# ── Checkpoint helpers ──
def save_checkpoint(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_checkpoint(path: str) -> tuple[int, list] | tuple[None, list]:
"""Load fetch 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 checkpoint: page {page}, {len(tags)} tags", file=sys.stderr)
return page, tags
except (json.JSONDecodeError, KeyError) as e:
print(f" Warning: corrupt checkpoint, starting fresh ({e})", file=sys.stderr)
return None, []
def clear_checkpoint(path: str):
"""Remove checkpoint 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, checkpoint_path: str = "", **_kwargs) -> list:
"""Fetch tags from Danbooru API, paginated."""
tags = []
page = 1
start_page, tags = load_checkpoint(checkpoint_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 = session.get(url, timeout=30)
resp.raise_for_status()
resp = fetch_with_retry(session, url)
except requests.RequestException as e:
print(f" Error on page {page}: {e}", file=sys.stderr)
break
@@ -90,23 +150,24 @@ def fetch_danbooru(min_count: int = 10, **_kwargs) -> list:
print(f" Page {page}: {len(data)} tags (total: {len(tags)})", file=sys.stderr)
if below_threshold:
break
if page % CHECKPOINT_INTERVAL == 0:
save_checkpoint(checkpoint_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, **_kwargs) -> list:
def fetch_e621(min_count: int = 10, checkpoint_path: str = "", **_kwargs) -> list:
"""Fetch tags from e621 API, paginated."""
tags = []
page = 1
start_page, tags = load_checkpoint(checkpoint_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 = session.get(url, timeout=30)
resp.raise_for_status()
resp = fetch_with_retry(session, url)
except requests.RequestException as e:
print(f" Error on page {page}: {e}", file=sys.stderr)
break
@@ -122,13 +183,15 @@ def fetch_e621(min_count: int = 10, **_kwargs) -> list:
print(f" Page {page}: {len(data)} tags (total: {len(tags)})", file=sys.stderr)
if below_threshold:
break
if page % CHECKPOINT_INTERVAL == 0:
save_checkpoint(checkpoint_path, page, tags)
page += 1
time.sleep(1.0) # e621 rate limit is stricter
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) -> list:
def fetch_gelbooru(base_url: str, min_count: int = 10, api_key: str | None = None, rate_limit: float = 0.5, checkpoint_path: str = "") -> list:
"""Fetch tags from a Gelbooru-compatible API (rule34, gelbooru, etc.).
Unlike Danbooru/e621, the Gelbooru tag endpoint doesn't support ordering
@@ -137,8 +200,8 @@ def fetch_gelbooru(base_url: str, min_count: int = 10, api_key: str | None = Non
"""
import xml.etree.ElementTree as ET
tags = []
page = 0
start_page, tags = load_checkpoint(checkpoint_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
@@ -157,18 +220,22 @@ def fetch_gelbooru(base_url: str, min_count: int = 10, api_key: str | None = Non
while True:
params["pid"] = str(page)
try:
resp = session.get(base_url, params=params, timeout=30)
resp.raise_for_status()
resp = fetch_with_retry(session, base_url, params=params)
except requests.RequestException as e:
print(f" Error on page {page}: {e}", file=sys.stderr)
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'):
# Non-XML response (auth error, empty, etc.)
if page == 0:
print(f" Error: {text[:200]}", file=sys.stderr)
break
root = ET.fromstring(text)
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"))
@@ -177,102 +244,132 @@ def fetch_gelbooru(base_url: str, min_count: int = 10, api_key: str | None = Non
print(f" Page {page}: {len(elements)} tags (total: {len(tags)})", file=sys.stderr)
if len(elements) < page_size:
break
if page % CHECKPOINT_INTERVAL == 0:
save_checkpoint(checkpoint_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, **_kwargs) -> list:
def fetch_rule34(min_count: int = 10, api_key: str | None = None, checkpoint_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)
return fetch_gelbooru("https://api.rule34.xxx/index.php", min_count=min_count, api_key=api_key, checkpoint_path=checkpoint_path)
def fetch_sankaku(min_count: int = 10, **_kwargs) -> list:
def fetch_sankaku(min_count: int = 10, checkpoint_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 underscores.
Tag names come as English with spaces - converted to lowercase.
"""
tags = []
page = 1
page_size = 200 # API max
start_page, tags = load_checkpoint(checkpoint_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 = session.get(
"https://sankakuapi.com/tags",
params={"limit": page_size, "page": page, "order": "count"},
timeout=30,
)
resp.raise_for_status()
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}: {e}", file=sys.stderr)
print(f" Error on page {page} after {MAX_RETRIES} retries: {e}", file=sys.stderr)
break
data = resp.json()
if not data:
break
below_threshold = True
for tag in data:
count = tag.get("post_count", 0)
if count >= min_count:
below_threshold = False
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])
else:
below_threshold = True
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 % CHECKPOINT_INTERVAL == 0:
save_checkpoint(checkpoint_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, checkpoint_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_checkpoint(checkpoint_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 % CHECKPOINT_INTERVAL == 0:
save_checkpoint(checkpoint_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,
"categories": DANBOORU_CATEGORIES,
},
"e621": {
"fetch": fetch_e621,
"categories": E621_CATEGORIES,
},
"rule34": {
"fetch": fetch_rule34,
"categories": RULE34_CATEGORIES,
},
"sankaku": {
"fetch": fetch_sankaku,
"categories": SANKAKU_CATEGORIES,
},
"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, categories: dict, output_path: str, separator: str = "_"):
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)
"""
if separator == "_":
tags = [[t[0].replace(" ", "_"), t[1], t[2]] for t in tags]
else:
tags = [[t[0].replace("_", " "), t[1], t[2]] for t in tags]
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": categories,
"tags": tags,
"categories": UNIFIED_CATEGORIES,
"tags": normalized,
}
os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True)
tmp_path = output_path + ".tmp"
@@ -289,12 +386,16 @@ def fetch_source(name: str, output: str, min_count: int, api_key: str | None = N
print(f"Unknown source: {name}. Available: {', '.join(SOURCES.keys())}", file=sys.stderr)
sys.exit(1)
source = SOURCES[name]
checkpoint_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)
tags = source["fetch"](min_count=min_count, api_key=api_key, checkpoint_path=checkpoint_path)
if not tags:
print(f" No tags fetched for {name}", file=sys.stderr)
return
write_dict(name, tags, source["categories"], output, separator=separator)
write_dict(name, tags, source["type_map"], output, separator=separator)
clear_checkpoint(checkpoint_path)
from gen_manifest import update_manifest
update_manifest(os.path.dirname(output) or ".")
def main():