From 897856baf073125836e67728aa1d5ff6fa501304 Mon Sep 17 00:00:00 2001 From: toyxyz Date: Tue, 29 Sep 2026 15:30:15 +0900 Subject: [PATCH] update booru tag prompter node --- .gitignore | 2 + README.md | 2 +- nodes/booru_tag_prompter.py | 19 +++++++- nodes/booru_wiki.py | 81 ++++++++++++++++++++++++++++++++- nodes/data/README.md | 7 +-- web/booru_tag_prompter.css | 14 ++++++ web/util/booru_wiki.js | 89 ++++++++++++++++++++++++++++++++++--- 7 files changed, 201 insertions(+), 13 deletions(-) diff --git a/.gitignore b/.gitignore index 85411e8..9fa1e3a 100644 --- a/.gitignore +++ b/.gitignore @@ -17,6 +17,8 @@ tests/ /.wildcard_settings.json /favorites/ /.favorite_settings.json +/nodes/data/danbooru_wiki.sqlite3 +/nodes/data/.danbooru_wiki_*.download # Local agent instructions /AGENTS.md diff --git a/README.md b/README.md index 885c953..e69acbd 100644 --- a/README.md +++ b/README.md @@ -284,7 +284,7 @@ Enter tags in `tags` and connect the `tags` output to a text encoder. Autocomple | --- | --- | | **W** (Wildcards) | Lists files in `wildcards/`. Click a file to insert `__name__`; each execution replaces it with one random non-empty line from that file. **Folder** opens the folder and **↻** refreshes the list. | | **F** (Favorites) | **Save** stores the entire current prompt. Click an entry to insert it at the cursor; right-click to edit or delete it. **↻** refreshes the list. | -| **Wiki** | Opens a movable offline browser. Browse categories, search, follow wiki links, and use **Insert** on a tag entry to add it at the cursor. Drag the title bar to move the window or its bottom-right corner to resize it. | +| **Wiki** | On first use, asks to download the offline database from Hugging Face (about 272 MiB). Later uses work offline. Browse categories, search, follow wiki links, and use **Insert** on a tag entry to add it at the cursor. Drag the title bar to move the window or its bottom-right corner to resize it. | | `camera` | Optionally connect `booru tag camera`; its selected camera guidance is appended after your text. | The node does not use Qwen or invent additional tags. Wildcard file contents and manually entered prompt text are preserved as written. diff --git a/nodes/booru_tag_prompter.py b/nodes/booru_tag_prompter.py index 7a9cba0..6f8d0f9 100644 --- a/nodes/booru_tag_prompter.py +++ b/nodes/booru_tag_prompter.py @@ -11,7 +11,7 @@ from .booru_wildcards import TOKEN, expand_wildcards, wildcard_files, wildcard_p from .booru_favorites import (FavoriteConflictError, active_favorite_root, browse_favorite_folder, change_favorite, list_favorites, save_favorite, set_favorite_path) -from .booru_wiki import search_wiki, wiki_categories, wiki_page +from .booru_wiki import download_wiki_db, search_wiki, wiki_categories, wiki_download_status, wiki_page LOG = logging.getLogger(__name__) @@ -115,6 +115,23 @@ def register_routes(): origin = request.headers.get("Origin") return request.remote in ("127.0.0.1", "::1") and (not origin or urlsplit(origin).netloc == request.host) + @PromptServer.instance.routes.get("/toyxyz/booru-tags/wiki/status") + async def wiki_status_route(request): + if not local_request(request): + return web.json_response({"error": "Wiki status requires local access."}, status=403) + return web.json_response(await asyncio.to_thread(wiki_download_status)) + + @PromptServer.instance.routes.post("/toyxyz/booru-tags/wiki/download") + async def wiki_download_route(request): + if not local_request(request): + return web.json_response({"error": "Wiki download requires local access."}, status=403) + try: + await asyncio.to_thread(download_wiki_db) + except OSError as exc: + LOG.warning("Wiki download failed: %s", exc) + return web.json_response({"error": str(exc)}, status=503) + return web.json_response({"ready": True}) + @PromptServer.instance.routes.get("/toyxyz/booru-tags/wiki/categories") async def wiki_categories_route(request): if not local_request(request): diff --git a/nodes/booru_wiki.py b/nodes/booru_wiki.py index 7e522c8..6dc691a 100644 --- a/nodes/booru_wiki.py +++ b/nodes/booru_wiki.py @@ -5,19 +5,96 @@ from __future__ import annotations import sqlite3 import json import re +import hashlib +import os +import tempfile +import threading from contextlib import closing from pathlib import Path from urllib.parse import quote +from urllib.request import Request, urlopen DEFAULT_DB = Path(__file__).resolve().parent / "data" / "danbooru_wiki.sqlite3" +WIKI_DOWNLOAD_URL = "https://huggingface.co/toyxyz/backup_models/resolve/main/danbooru_wiki.sqlite3" +WIKI_SHA256 = "c17b31e0f6468d2a0d5094452085925643d8442a2433f798f74f26923758bac6" +WIKI_SIZE = 284_766_208 +MAX_DOWNLOAD_SIZE = 400 * 1024 * 1024 +_download_lock = threading.Lock() +_progress_lock = threading.Lock() +_download_progress = {"downloading": False, "bytes": 0} PAGE_SIZE = 30 WIKI_LINK = re.compile(r"\[\[([^\]]+)\]\]") +def wiki_db_ready() -> bool: + """Reject missing files and Git LFS pointers without reading the entire DB.""" + try: + if DEFAULT_DB.stat().st_size < 100: + return False + with DEFAULT_DB.open("rb") as source: + return source.read(16) == b"SQLite format 3\x00" + except OSError: + return False + + +def wiki_download_status() -> dict: + ready = wiki_db_ready() + with _progress_lock: + downloading = _download_progress["downloading"] + received = _download_progress["bytes"] + return {"ready": ready, "downloading": downloading, + "percent": 100 if ready else min(99, received * 100 // WIKI_SIZE) if downloading else 0} + + +def download_wiki_db() -> Path: + """Install the pinned snapshot atomically after verifying its SHA-256.""" + with _download_lock: + if wiki_db_ready(): + return DEFAULT_DB + DEFAULT_DB.parent.mkdir(parents=True, exist_ok=True) + temporary = None + with _progress_lock: + _download_progress.update(downloading=True, bytes=0) + try: + with tempfile.NamedTemporaryFile(dir=DEFAULT_DB.parent, prefix=".danbooru_wiki_", + suffix=".download", delete=False) as target: + temporary = Path(target.name) + digest = hashlib.sha256() + size = 0 + request = Request(WIKI_DOWNLOAD_URL, headers={"User-Agent": "toyxyz-booru-wiki/1.0"}) + with urlopen(request, timeout=60) as source: + if source.status != 200: + raise OSError(f"Wiki download failed (HTTP {source.status}).") + while chunk := source.read(1024 * 1024): + size += len(chunk) + if size > MAX_DOWNLOAD_SIZE: + raise OSError("Wiki download exceeded the expected size.") + digest.update(chunk) + target.write(chunk) + with _progress_lock: + _download_progress["bytes"] = size + target.flush() + os.fsync(target.fileno()) + if size != WIKI_SIZE or digest.hexdigest() != WIKI_SHA256: + raise OSError("Downloaded wiki database does not match the expected snapshot.") + with closing(sqlite3.connect(f"file:{quote(temporary.resolve().as_posix(), safe='/:')}?mode=ro", + uri=True)) as con: + con.execute("SELECT name FROM sqlite_master LIMIT 1").fetchone() + os.replace(temporary, DEFAULT_DB) + return DEFAULT_DB + except (OSError, sqlite3.Error, ValueError) as exc: + raise OSError(str(exc) or "Wiki download failed. Please retry.") from exc + finally: + with _progress_lock: + _download_progress.update(downloading=False, bytes=0) + if temporary is not None: + temporary.unlink(missing_ok=True) + + def active_wiki_db() -> Path: - if not DEFAULT_DB.is_file(): - raise OSError("Bundled wiki database is missing from nodes/data.") + if not wiki_db_ready(): + raise OSError("Wiki database is unavailable. Open Wiki to download it.") return DEFAULT_DB diff --git a/nodes/data/README.md b/nodes/data/README.md index f79f509..f3611e5 100644 --- a/nodes/data/README.md +++ b/nodes/data/README.md @@ -10,9 +10,10 @@ This pinned snapshot contains 125,312 rows. It validates generated tag names offline; it does not establish that every tag is understood by a given Anima checkpoint. Refresh the snapshot and revision together when updating it. -`danbooru_wiki.sqlite3` is the bundled offline Danbooru wiki snapshot used by -the `booru tag prompter` Wiki panel. It is opened read-only and is not modified -by node execution. SHA-256: +`danbooru_wiki.sqlite3` is downloaded on first Wiki use from +`https://huggingface.co/toyxyz/backup_models/resolve/main/danbooru_wiki.sqlite3` +into this directory. It is not tracked in Git. Subsequent Wiki use is offline; +the file is opened read-only and is not modified by node execution. SHA-256: `c17b31e0f6468d2a0d5094452085925643d8442a2433f798f74f26923758bac6`. The snapshot is separate from the tag autocomplete CSV and does not update automatically. diff --git a/web/booru_tag_prompter.css b/web/booru_tag_prompter.css index 9d4f791..50dfda1 100644 --- a/web/booru_tag_prompter.css +++ b/web/booru_tag_prompter.css @@ -73,6 +73,20 @@ .toyxyz-wildcard-tabs button.toyxyz-tab-active { background:#315b84; border-color:#71b8ff; color:#fff; box-shadow:inset 0 0 0 1px #71b8ff; } .toyxyz-wiki-panel { position:fixed; z-index:10001; display:flex; flex-direction:column; box-sizing:border-box; height:min(760px,calc(100vh - 16px)); overflow:hidden; border:1px solid #647991; border-radius:9px; background:#222a34; color:#e9f0f7; box-shadow:0 16px 38px #000b; font:13px/1.5 sans-serif; } .toyxyz-wiki-panel[hidden] { display:none; } +.toyxyz-wiki-download-backdrop { position:fixed; z-index:10002; inset:0; display:flex; align-items:center; justify-content:center; padding:16px; background:#0009; } +.toyxyz-wiki-download-backdrop[hidden] { display:none; } +.toyxyz-wiki-download-box { box-sizing:border-box; width:min(440px,100%); padding:20px; border:1px solid #647991; border-radius:9px; background:#222a34; color:#e9f0f7; box-shadow:0 16px 38px #000b; font:13px/1.5 sans-serif; } +.toyxyz-wiki-download-box h2 { margin:0 0 10px; font-size:18px; } +.toyxyz-wiki-download-box p { margin:0 0 12px; } +.toyxyz-wiki-download-progress { display:flex; align-items:center; gap:10px; margin:0 0 12px; } +.toyxyz-wiki-download-progress[hidden] { display:none; } +.toyxyz-wiki-download-progress progress { flex:1; height:14px; accent-color:#75b8f5; } +.toyxyz-wiki-download-progress output { min-width:36px; text-align:right; font-variant-numeric:tabular-nums; } +.toyxyz-wiki-download-error { color:#ffbd9d; overflow-wrap:anywhere; } +.toyxyz-wiki-download-error:empty { display:none; } +.toyxyz-wiki-download-actions { display:flex; justify-content:flex-end; gap:8px; } +.toyxyz-wiki-download-actions button { padding:7px 12px; border:1px solid #637b91; border-radius:5px; background:#35485a; color:#eee; cursor:pointer; } +.toyxyz-wiki-download-actions button:disabled { opacity:.5; cursor:default; } .toyxyz-wiki-panel button { border:1px solid #536a80; border-radius:5px; padding:5px 8px; background:#304155; color:#edf4fa; font:inherit; cursor:pointer; } .toyxyz-wiki-panel button:hover { background:#415c77; } .toyxyz-wiki-panel button:disabled { opacity:.45; cursor:default; } diff --git a/web/util/booru_wiki.js b/web/util/booru_wiki.js index 2b3e47e..c767f39 100644 --- a/web/util/booru_wiki.js +++ b/web/util/booru_wiki.js @@ -116,13 +116,33 @@ export function installWiki(nodeRef,widget,input,api) { resizeGrip.tabIndex=0; layout.append(sidebar,browse,detail);panel.append(header,layout,resizeGrip);document.body.append(panel); + const downloadDialog=node('div','toyxyz-wiki-download-backdrop');downloadDialog.hidden=true; + const downloadBox=node('section','toyxyz-wiki-download-box'); + downloadBox.setAttribute('role','dialog');downloadBox.setAttribute('aria-modal','true'); + downloadBox.setAttribute('aria-label','Download Tag Wiki'); + downloadBox.append(node('h2','','Download Tag Wiki')); + const downloadMessage=node('p','','The offline Tag Wiki database is not installed. Download approximately 272 MiB from Hugging Face? The file will be saved in this custom node folder.'); + const progressRow=node('div','toyxyz-wiki-download-progress');progressRow.hidden=true; + const progressBar=node('progress');progressBar.max=100;progressBar.value=0; + progressBar.setAttribute('aria-label','Wiki download progress'); + const progressLabel=node('output','','0%'); + progressRow.append(progressBar,progressLabel); + const downloadError=node('p','toyxyz-wiki-download-error');downloadError.setAttribute('role','alert'); + const downloadActions=node('div','toyxyz-wiki-download-actions'); + const dismissDownload=node('button','','Cancel'),startDownload=node('button','','Download'); + dismissDownload.type=startDownload.type='button'; + downloadActions.append(dismissDownload,startDownload); + downloadBox.append(downloadMessage,progressRow,downloadError,downloadActions); + downloadDialog.append(downloadBox);document.body.append(downloadDialog); + const state={open:false,disposed:false,category:'',query:'',page:1, searchToken:0,detailToken:0,history:[],frame:0,lastRect:'',timer:0, - manualPosition:false,dragging:null,resizing:null}; - const readJSON=async url=>{ - const response=await api.fetchApi(url); + manualPosition:false,dragging:null,resizing:null,opening:false,downloading:false, + progressTimer:0,progressPending:false}; + const readJSON=async(url,options)=>{ + const response=await api.fetchApi(url,options); const data=await response.json().catch(()=>({})); - if(!response.ok)throw new Error(response.status===404&&url.includes('/categories') + if(!response.ok)throw new Error(response.status===404&&url.includes('/wiki/') ?'Restart ComfyUI to activate the Tag Wiki Browser.' :data.error||`Wiki request failed (HTTP ${response.status}).`); return data; @@ -353,6 +373,63 @@ export function installWiki(nodeRef,widget,input,api) { ++state.searchToken;++state.detailToken; } }; + const showDownload=()=>{ + downloadDialog.hidden=false; + downloadMessage.textContent='The offline Tag Wiki database is not installed. Download approximately 272 MiB from Hugging Face? The file will be saved in this custom node folder.'; + downloadError.textContent=''; + dismissDownload.disabled=false;startDownload.disabled=false; + startDownload.textContent='Download'; + progressBar.value=0;progressLabel.textContent='0%';progressRow.hidden=true; + }; + dismissDownload.addEventListener('click',()=>{if(!state.downloading)downloadDialog.hidden=true;}); + startDownload.addEventListener('click',async()=>{ + if(state.downloading||state.disposed)return; + state.downloading=true;dismissDownload.disabled=true;startDownload.disabled=true; + downloadError.textContent=''; + downloadMessage.textContent='Downloading and verifying the Tag Wiki database. Please wait; the wiki will open when ready.'; + progressBar.value=0;progressLabel.textContent='0%';progressRow.hidden=false; + const refreshProgress=async()=>{ + if(state.progressPending||!state.downloading||state.disposed)return; + state.progressPending=true; + try{ + const result=await readJSON('/toyxyz/booru-tags/wiki/status'); + if(!state.downloading||state.disposed)return; + const percent=Math.max(0,Math.min(99,Number(result.percent)||0)); + progressBar.value=percent;progressLabel.textContent=`${percent}%`; + }catch{ /* The download request reports a final error if status is unavailable. */ } + finally{state.progressPending=false;} + }; + state.progressTimer=setInterval(refreshProgress,250); + try{ + await readJSON('/toyxyz/booru-tags/wiki/download', {method:'POST'}); + if(state.disposed)return; + progressBar.value=100;progressLabel.textContent='100%'; + downloadDialog.hidden=true;setOpen(true); + }catch(error){ + if(!state.disposed){ + downloadMessage.textContent='The download did not complete. Your prompt remains available.'; + downloadError.textContent=error.message; + dismissDownload.disabled=false;startDownload.disabled=false; + startDownload.textContent='Retry'; + } + }finally{ + state.downloading=false; + clearInterval(state.progressTimer);state.progressTimer=0; + } + }); + const openWiki=async()=>{ + if(state.open){setOpen(false);return;} + if(state.opening||state.downloading)return; + state.opening=true;toggle.disabled=true; + try{ + const result=await readJSON('/toyxyz/booru-tags/wiki/status'); + if(state.disposed)return; + if(result.ready)setOpen(true); + else showDownload(); + }catch(error){ + if(!state.disposed){showDownload();downloadError.textContent=error.message;} + }finally{state.opening=false;toggle.disabled=false;} + }; search.addEventListener('input',()=>{ clearTimeout(state.timer);state.timer=setTimeout(()=>{ state.query=search.value.trim();state.page=1;loadResults(); @@ -367,7 +444,7 @@ export function installWiki(nodeRef,widget,input,api) { }); previous.addEventListener('click',()=>{if(state.page>1){--state.page;loadResults();}}); next.addEventListener('click',()=>{++state.page;loadResults();}); - toggle.addEventListener('click',()=>setOpen(!state.open));close.addEventListener('click',()=>setOpen(false)); + toggle.addEventListener('click',openWiki);close.addEventListener('click',()=>setOpen(false)); const removed=nodeRef.onRemoved; - nodeRef.onRemoved=function(){state.disposed=true;setOpen(false);panel.remove();return removed?.apply(this,arguments);}; + nodeRef.onRemoved=function(){state.disposed=true;clearInterval(state.progressTimer);setOpen(false);panel.remove();downloadDialog.remove();return removed?.apply(this,arguments);}; }