mirror of
https://github.com/vladmandic/automatic
synced 2026-09-18 16:54:33 +02:00
feat(dicts): add curated dicts, pruning, and manifest generation
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
+200
-99
@@ -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():
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Generate manifest.json for the HuggingFace dict repository.
|
||||
|
||||
Reads dict JSON files and writes a manifest with accurate tag counts,
|
||||
file sizes, and versions. Only listed files are included - curated dicts
|
||||
that ship bundled with the repo should not be passed as arguments.
|
||||
|
||||
Usage:
|
||||
python cli/gen_manifest.py data/dicts/danbooru.json data/dicts/e621.json ...
|
||||
python cli/gen_manifest.py data/dicts/danbooru.json -o data/dicts/manifest.json
|
||||
python cli/gen_manifest.py data/dicts/*.json --exclude art photography quality negative
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Human-readable descriptions keyed by dict name.
|
||||
# Add entries here when new sources are added to fetch_dicts.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 dict 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(dicts_dir: str) -> bool:
|
||||
"""Regenerate manifest.json in dicts_dir if one already exists.
|
||||
|
||||
Only updates entries for dicts already listed in the manifest - does not
|
||||
add new dicts. Returns True if the manifest was updated, False if no
|
||||
manifest exists to update.
|
||||
"""
|
||||
manifest_path = os.path.join(dicts_dir, "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("dicts", [])}
|
||||
entries = []
|
||||
for name in sorted(existing_names):
|
||||
filepath = os.path.join(dicts_dir, f"{name}.json")
|
||||
if not os.path.isfile(filepath):
|
||||
continue
|
||||
try:
|
||||
entries.append(build_entry(filepath))
|
||||
except (json.JSONDecodeError, KeyError):
|
||||
continue
|
||||
manifest["dicts"] = entries
|
||||
with open(manifest_path, "w", encoding="utf-8") as f:
|
||||
json.dump(manifest, f, indent=2, ensure_ascii=False)
|
||||
f.write("\n")
|
||||
print(f" Manifest updated: {manifest_path} ({len(entries)} dicts)", file=sys.stderr)
|
||||
return True
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Generate manifest.json for HF dict repo")
|
||||
parser.add_argument("files", nargs="+", help="Dict 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="Dict 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 = {"dicts": 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, indent=2, ensure_ascii=False)
|
||||
f.write("\n")
|
||||
|
||||
print(f"\nManifest written: {args.output} ({len(entries)} dicts)", file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,137 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Prune tag dictionaries with per-category minimum post counts.
|
||||
|
||||
Usage:
|
||||
python cli/prune_dicts.py data/dicts/danbooru.json
|
||||
python cli/prune_dicts.py data/dicts/*.json --general 500 --artist 20
|
||||
python cli/prune_dicts.py data/dicts/sankaku.json -o data/dicts/sankaku-pruned.json
|
||||
python cli/prune_dicts.py data/dicts/*.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 fetch_dicts.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}
|
||||
|
||||
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)
|
||||
from gen_manifest import update_manifest
|
||||
update_manifest(os.path.dirname(output_path) or ".")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user