From ff0466957e370e8cfafac75756a82d0c8fee06d6 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 09:49:31 +0000 Subject: [PATCH] feat(gallery): search only generation prompts; scan off the event loop; cache results - New prompt_search.py (no ComfyUI imports): reads only the PNG 'prompt' chunk (tEXt/zTXt/iTXt), skips the large 'workflow' chunk, and extracts only text that feeds text-encoder nodes (following links to string/primitive/concat nodes). Model names, sampler names, filenames and other widgets are no longer matched, and filenames are no longer searched. - Per-file results are cached by (mtime, size); repeat searches only touch new/changed files. - gallery/list runs in a worker thread so large searches don't block ComfyUI's server; the scan stops when the client disconnects. - Frontend aborts superseded search requests and ignores stale results. - UI strings/locales and README updated; unit tests in tests/. --- README.md | 2 +- __init__.py | 210 +++++++++++++++------------------- js/locales/ar/main.json | 4 +- js/locales/en/main.json | 4 +- js/locales/es/main.json | 4 +- js/locales/fa/main.json | 4 +- js/locales/fr/main.json | 4 +- js/locales/he/main.json | 4 +- js/locales/it/main.json | 4 +- js/locales/ja/main.json | 4 +- js/locales/ko/main.json | 4 +- js/locales/pt-br/main.json | 4 +- js/locales/ru/main.json | 4 +- js/locales/tr/main.json | 4 +- js/locales/zh-tw/main.json | 4 +- js/locales/zh/main.json | 4 +- js/ui/galleryView.js | 32 ++++-- prompt_search.py | 218 ++++++++++++++++++++++++++++++++++++ tests/test_prompt_search.py | 180 +++++++++++++++++++++++++++++ 19 files changed, 541 insertions(+), 157 deletions(-) create mode 100644 prompt_search.py create mode 100644 tests/test_prompt_search.py diff --git a/README.md b/README.md index fa59f8e..77316ab 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,7 @@ A compact and powerful layout overhaul for ComfyUI that makes generation results ComfySidebar compare -* **[NEW] Output Gallery sidebar:** Browse files and subfolders directly from your ComfyUI `output` directory without leaving the interface. Features dedicated cards for images, videos, audio, and 3D assets (`.glb`/`.gltf`), breadcrumb navigation, prompt metadata search across subfolders, and persistent sorting. +* **[NEW] Output Gallery sidebar:** Browse files and subfolders directly from your ComfyUI `output` directory without leaving the interface. Features dedicated cards for images, videos, audio, and 3D assets (`.glb`/`.gltf`), breadcrumb navigation, prompt text search across subfolders (matches the positive/negative prompts stored in PNG outputs, not filenames), and persistent sorting. * **Hotkeys on any keyboard layout:** Shortcuts such as `Q`, `Ctrl+Q` and `Alt+R` keep working while a non-Latin layout (for example Russian) is active. Non-Latin keys are treated as the Latin key in the same physical position, and typing in text fields is not affected. Can be turned off with **Fix non-ASCII keybinds** in the settings. diff --git a/__init__.py b/__init__.py index 21e19f0..5051666 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,14 @@ +import asyncio import os import platform import subprocess -import json +import threading from aiohttp import web from server import PromptServer import folder_paths +from . import prompt_search + NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} WEB_DIRECTORY = "./js" @@ -143,41 +146,81 @@ def open_file_in_os(target_path): raise RuntimeError("Headless Linux server detected (no desktop GUI available).") -def check_png_prompt_match(filepath, search_term): - """Checks PNG prompt metadata strictly for text inputs matching search.""" - try: - with open(filepath, "rb") as f: - if f.read(8) != b"\x89PNG\r\n\x1a\n": - return False - while True: - len_b = f.read(4) - if len(len_b) < 4: - break - length = int.from_bytes(len_b, "big") - ctype = f.read(4) - if ctype in (b"tEXt", b"iTXt"): - data = f.read(length) - if data.startswith(b"prompt\x00") or b'"inputs"' in data: - text = data.decode("utf-8", errors="ignore") - idx = text.find("{") - if idx != -1: - try: - prompt_dict = json.loads(text[idx:]) - for node in prompt_dict.values(): - if isinstance(node, dict) and "inputs" in node: - for val in node["inputs"].values(): - if isinstance(val, str) and search_term in val.lower(): - return True - except Exception: - pass - elif ctype in (b"IDAT", b"IEND"): - break - else: - f.seek(length, 1) - f.seek(4, 1) - return False - except Exception: - return False +_GALLERY_EXTS = { + ".png", ".jpg", ".jpeg", ".webp", ".gif", ".avif", + ".mp4", ".webm", ".mov", ".mkv", + ".mp3", ".wav", ".ogg", ".flac", ".m4a", ".aac", + ".glb", ".gltf", ".obj", ".fbx", ".stl", + ".json" +} + + +def _relative_subfolder(directory, output_dir): + rel_sub = os.path.relpath(directory, output_dir).replace("\\", "/") + return "" if rel_sub == "." else rel_sub + + +def _list_gallery(subfolder, search, cancel): + """Blocking directory listing / prompt search. Runs in a worker thread. + + Returns (http_status, json_payload). + """ + output_dir = folder_paths.get_output_directory() + target_dir = os.path.realpath(os.path.join(output_dir, subfolder)) if subfolder else os.path.realpath(output_dir) + + if not is_path_safe(output_dir, target_dir): + return 403, {"error": "Unsafe path"} + + if not os.path.exists(target_dir) or not os.path.isdir(target_dir): + return 200, {"subfolder": subfolder, "folders": [], "files": []} + + folders = [] + files = [] + + if search: + # Recursive prompt-text search (PNG `prompt` metadata only; filenames are not matched) + for fpath, stat in prompt_search.iter_prompt_matches(target_dir, search, cancel): + files.append({ + "filename": os.path.basename(fpath), + "subfolder": _relative_subfolder(os.path.dirname(fpath), output_dir), + "type": "output", + "mtime": stat.st_mtime, + "size": stat.st_size + }) + else: + # Instant single-directory listing without reading files from disk + with os.scandir(target_dir) as entries: + for entry in entries: + if entry.name.startswith("."): + continue + try: + if entry.is_dir(follow_symlinks=False): + folders.append({ + "name": entry.name, + "subfolder": os.path.relpath(entry.path, output_dir).replace("\\", "/"), + "mtime": entry.stat().st_mtime + }) + elif entry.is_file(follow_symlinks=False): + if os.path.splitext(entry.name)[1].lower() in _GALLERY_EXTS: + stat = entry.stat() + files.append({ + "filename": entry.name, + "subfolder": _relative_subfolder(os.path.dirname(entry.path), output_dir), + "type": "output", + "mtime": stat.st_mtime, + "size": stat.st_size + }) + except (PermissionError, FileNotFoundError): + continue + + folders.sort(key=lambda x: x["name"].lower()) + files.sort(key=lambda x: x["mtime"], reverse=True) + + return 200, { + "subfolder": subfolder.replace("\\", "/"), + "folders": folders, + "files": files + } @PromptServer.instance.routes.get("/comfy-sidebar/gallery/list") @@ -185,92 +228,17 @@ async def gallery_list_handler(request): try: subfolder = request.query.get("subfolder", "").strip().strip("/\\") search = request.query.get("search", "").strip().lower() - output_dir = folder_paths.get_output_directory() - target_dir = os.path.realpath(os.path.join(output_dir, subfolder)) if subfolder else os.path.realpath(output_dir) - - if not is_path_safe(output_dir, target_dir): - return web.json_response({"error": "Unsafe path"}, status=403) - - if not os.path.exists(target_dir) or not os.path.isdir(target_dir): - return web.json_response({"subfolder": subfolder, "folders": [], "files": []}) - - folders = [] - files = [] - - supported_exts = { - ".png", ".jpg", ".jpeg", ".webp", ".gif", ".avif", - ".mp4", ".webm", ".mov", ".mkv", - ".mp3", ".wav", ".ogg", ".flac", ".m4a", ".aac", - ".glb", ".gltf", ".obj", ".fbx", ".stl", - ".json" - } - - if search: - # Recursive search in current directory and all subfolders - for root, dirs, fnames in os.walk(target_dir): - dirs[:] = [d for d in dirs if not d.startswith(".")] - for fname in fnames: - if fname.startswith("."): - continue - ext = os.path.splitext(fname)[1].lower() - if ext in supported_exts: - fpath = os.path.join(root, fname) - matched = search in fname.lower() - if not matched and ext == ".png": - matched = check_png_prompt_match(fpath, search) - - if matched: - rel_sub = os.path.relpath(root, output_dir).replace("\\", "/") - if rel_sub == ".": - rel_sub = "" - stat = os.stat(fpath) - files.append({ - "filename": fname, - "subfolder": rel_sub, - "type": "output", - "mtime": stat.st_mtime, - "size": stat.st_size - }) - else: - # Instant single-directory listing without reading files from disk - with os.scandir(target_dir) as entries: - for entry in entries: - if entry.name.startswith("."): - continue - try: - if entry.is_dir(follow_symlinks=False): - rel_sub = os.path.relpath(entry.path, output_dir).replace("\\", "/") - folders.append({ - "name": entry.name, - "subfolder": rel_sub, - "mtime": entry.stat().st_mtime - }) - elif entry.is_file(follow_symlinks=False): - ext = os.path.splitext(entry.name)[1].lower() - if ext in supported_exts: - rel_sub = os.path.relpath(os.path.dirname(entry.path), output_dir).replace("\\", "/") - if rel_sub == ".": - rel_sub = "" - stat = entry.stat() - files.append({ - "filename": entry.name, - "subfolder": rel_sub, - "type": "output", - "mtime": stat.st_mtime, - "size": stat.st_size - }) - except (PermissionError, FileNotFoundError): - continue - - folders.sort(key=lambda x: x["name"].lower()) - files.sort(key=lambda x: x["mtime"], reverse=True) - - return web.json_response({ - "subfolder": subfolder.replace("\\", "/"), - "folders": folders, - "files": files - }) + # Scanning can touch thousands of files: keep it off the event loop so + # ComfyUI's server (websocket, queue updates) stays responsive. + cancel = threading.Event() + loop = asyncio.get_running_loop() + try: + status, payload = await loop.run_in_executor(None, _list_gallery, subfolder, search, cancel) + except asyncio.CancelledError: + cancel.set() # client went away / sent a newer search: stop scanning + raise + return web.json_response(payload, status=status) except Exception as e: return web.json_response({"error": str(e)}, status=500) diff --git a/js/locales/ar/main.json b/js/locales/ar/main.json index d2daab7..93275a9 100644 --- a/js/locales/ar/main.json +++ b/js/locales/ar/main.json @@ -113,6 +113,6 @@ "Failed to load gallery from disk.": "فشل تحميل المعرض من القرص.", "Sort: Newest First": "الترتيب: الأحدث أولاً", "Sort: Oldest First": "الترتيب: الأقدم أولاً", - "Search Gallery (Filename or Prompt)": "البحث في المعرض (اسم الملف أو الموجه)", - "Search filenames or prompt metadata...": "البحث في أسماء الملفات أو بيانات الموجه..." + "Search Gallery (Prompt Text)": "البحث في المعرض (نص الموجه)", + "Search prompts...": "البحث في الموجهات..." } \ No newline at end of file diff --git a/js/locales/en/main.json b/js/locales/en/main.json index e51bfb0..d13167b 100644 --- a/js/locales/en/main.json +++ b/js/locales/en/main.json @@ -113,6 +113,6 @@ "Failed to load gallery from disk.": "Failed to load gallery from disk.", "Sort: Newest First": "Sort: Newest First", "Sort: Oldest First": "Sort: Oldest First", - "Search Gallery (Filename or Prompt)": "Search Gallery (Filename or Prompt)", - "Search filenames or prompt metadata...": "Search filenames or prompt metadata..." + "Search Gallery (Prompt Text)": "Search Gallery (Prompt Text)", + "Search prompts...": "Search prompts..." } \ No newline at end of file diff --git a/js/locales/es/main.json b/js/locales/es/main.json index 2731a9f..94ec67b 100644 --- a/js/locales/es/main.json +++ b/js/locales/es/main.json @@ -113,6 +113,6 @@ "Failed to load gallery from disk.": "Error al cargar la galería desde el disco.", "Sort: Newest First": "Ordenar: más recientes primero", "Sort: Oldest First": "Ordenar: más antiguos primero", - "Search Gallery (Filename or Prompt)": "Buscar en galería (nombre de archivo o prompt)", - "Search filenames or prompt metadata...": "Buscar nombres de archivo o metadatos de prompt..." + "Search Gallery (Prompt Text)": "Buscar en galería (texto del prompt)", + "Search prompts...": "Buscar prompts..." } \ No newline at end of file diff --git a/js/locales/fa/main.json b/js/locales/fa/main.json index 2b18326..cf004c8 100644 --- a/js/locales/fa/main.json +++ b/js/locales/fa/main.json @@ -113,6 +113,6 @@ "Failed to load gallery from disk.": "بارگذاری گالری از دیسک ناموفق بود.", "Sort: Newest First": "مرتب‌سازی: جدیدترین‌ها در ابتدا", "Sort: Oldest First": "مرتب‌سازی: قدیمی‌ترین‌ها در ابتدا", - "Search Gallery (Filename or Prompt)": "جستجو در گالری (نام فایل یا پرامپت)", - "Search filenames or prompt metadata...": "جستجوی نام فایل‌ها یا متادیتای پرامپت..." + "Search Gallery (Prompt Text)": "جستجو در گالری (متن پرامپت)", + "Search prompts...": "جستجوی پرامپت‌ها..." } \ No newline at end of file diff --git a/js/locales/fr/main.json b/js/locales/fr/main.json index f8a607d..22a9919 100644 --- a/js/locales/fr/main.json +++ b/js/locales/fr/main.json @@ -113,6 +113,6 @@ "Failed to load gallery from disk.": "Échec du chargement de la galerie depuis le disque.", "Sort: Newest First": "Trier : Plus récents d'abord", "Sort: Oldest First": "Trier : Plus anciens d'abord", - "Search Gallery (Filename or Prompt)": "Rechercher dans la galerie (nom de fichier ou prompt)", - "Search filenames or prompt metadata...": "Rechercher des noms de fichiers ou métadonnées de prompt..." + "Search Gallery (Prompt Text)": "Rechercher dans la galerie (texte du prompt)", + "Search prompts...": "Rechercher des prompts..." } \ No newline at end of file diff --git a/js/locales/he/main.json b/js/locales/he/main.json index e2896d4..9c8f9f5 100644 --- a/js/locales/he/main.json +++ b/js/locales/he/main.json @@ -113,6 +113,6 @@ "Failed to load gallery from disk.": "טעינת הגלריה מהכונן נכשלה.", "Sort: Newest First": "מיון: החדשים ביותר קודם", "Sort: Oldest First": "מיון: הישנים ביותר קודם", - "Search Gallery (Filename or Prompt)": "חיפוש בגלריה (שם קובץ או פרומפט)", - "Search filenames or prompt metadata...": "חיפוש שמות קבצים או מטא-דאטה של פרומפט..." + "Search Gallery (Prompt Text)": "חיפוש בגלריה (טקסט הפרומפט)", + "Search prompts...": "חיפוש פרומפטים..." } \ No newline at end of file diff --git a/js/locales/it/main.json b/js/locales/it/main.json index 77b3b15..accd7f0 100644 --- a/js/locales/it/main.json +++ b/js/locales/it/main.json @@ -113,6 +113,6 @@ "Failed to load gallery from disk.": "Impossibile caricare la galleria dal disco.", "Sort: Newest First": "Ordina: Più recenti prima", "Sort: Oldest First": "Ordina: Più vecchi prima", - "Search Gallery (Filename or Prompt)": "Cerca nella galleria (nome file o prompt)", - "Search filenames or prompt metadata...": "Cerca nomi di file o metadati del prompt..." + "Search Gallery (Prompt Text)": "Cerca nella galleria (testo del prompt)", + "Search prompts...": "Cerca nei prompt..." } \ No newline at end of file diff --git a/js/locales/ja/main.json b/js/locales/ja/main.json index a0e7539..49bc466 100644 --- a/js/locales/ja/main.json +++ b/js/locales/ja/main.json @@ -113,6 +113,6 @@ "Failed to load gallery from disk.": "ディスクからギャラリーを読み込めませんでした。", "Sort: Newest First": "並び替え：新しい順", "Sort: Oldest First": "並び替え：古い順", - "Search Gallery (Filename or Prompt)": "ギャラリーを検索 (ファイル名またはプロンプト)", - "Search filenames or prompt metadata...": "ファイル名またはプロンプトを検索..." + "Search Gallery (Prompt Text)": "ギャラリーを検索 (プロンプトのテキスト)", + "Search prompts...": "プロンプトを検索..." } \ No newline at end of file diff --git a/js/locales/ko/main.json b/js/locales/ko/main.json index 2e0d9ef..a207b16 100644 --- a/js/locales/ko/main.json +++ b/js/locales/ko/main.json @@ -113,6 +113,6 @@ "Failed to load gallery from disk.": "디스크에서 갤러리를 불러오지 못했습니다.", "Sort: Newest First": "정렬: 최신순", "Sort: Oldest First": "정렬: 오래된순", - "Search Gallery (Filename or Prompt)": "갤러리 검색 (파일명 또는 프롬프트)", - "Search filenames or prompt metadata...": "파일명 또는 프롬프트 메타데이터 검색..." + "Search Gallery (Prompt Text)": "갤러리 검색 (프롬프트 텍스트)", + "Search prompts...": "프롬프트 검색..." } \ No newline at end of file diff --git a/js/locales/pt-br/main.json b/js/locales/pt-br/main.json index 78d0dd2..369e5d6 100644 --- a/js/locales/pt-br/main.json +++ b/js/locales/pt-br/main.json @@ -113,6 +113,6 @@ "Failed to load gallery from disk.": "Falha ao carregar a galeria do disco.", "Sort: Newest First": "Ordenar: Mais recentes primeiro", "Sort: Oldest First": "Ordenar: Mais antigos primeiro", - "Search Gallery (Filename or Prompt)": "Pesquisar na galeria (nome do arquivo ou prompt)", - "Search filenames or prompt metadata...": "Pesquisar nomes de arquivos ou metadados de prompt..." + "Search Gallery (Prompt Text)": "Pesquisar na galeria (texto do prompt)", + "Search prompts...": "Pesquisar prompts..." } \ No newline at end of file diff --git a/js/locales/ru/main.json b/js/locales/ru/main.json index 1359dd4..7853c75 100644 --- a/js/locales/ru/main.json +++ b/js/locales/ru/main.json @@ -113,6 +113,6 @@ "Failed to load gallery from disk.": "Не удалось загрузить галерею с диска.", "Sort: Newest First": "Сортировка: сначала новые", "Sort: Oldest First": "Сортировка: сначала старые", - "Search Gallery (Filename or Prompt)": "Поиск по галерее (имя файла или промпт)", - "Search filenames or prompt metadata...": "Поиск по именам файлов или промптам..." + "Search Gallery (Prompt Text)": "Поиск по галерее (текст промпта)", + "Search prompts...": "Поиск по промптам..." } \ No newline at end of file diff --git a/js/locales/tr/main.json b/js/locales/tr/main.json index f94c7a7..9e7943a 100644 --- a/js/locales/tr/main.json +++ b/js/locales/tr/main.json @@ -113,6 +113,6 @@ "Failed to load gallery from disk.": "Galeri diskten yüklenemedi.", "Sort: Newest First": "Sıralama: Önce En Yeniler", "Sort: Oldest First": "Sıralama: Önce En Eskiler", - "Search Gallery (Filename or Prompt)": "Galeride Ara (Dosya Adı veya Prompt)", - "Search filenames or prompt metadata...": "Dosya adlarını veya prompt meta verilerini ara..." + "Search Gallery (Prompt Text)": "Galeride Ara (Prompt Metni)", + "Search prompts...": "Prompt ara..." } \ No newline at end of file diff --git a/js/locales/zh-tw/main.json b/js/locales/zh-tw/main.json index 8670eb2..b077b21 100644 --- a/js/locales/zh-tw/main.json +++ b/js/locales/zh-tw/main.json @@ -113,6 +113,6 @@ "Failed to load gallery from disk.": "從磁碟載入圖庫失敗。", "Sort: Newest First": "排序：最新優先", "Sort: Oldest First": "排序：最舊優先", - "Search Gallery (Filename or Prompt)": "搜尋圖庫 (檔案名稱或提示詞)", - "Search filenames or prompt metadata...": "搜尋檔案名稱或提示詞元資料..." + "Search Gallery (Prompt Text)": "搜尋圖庫 (提示詞文字)", + "Search prompts...": "搜尋提示詞..." } \ No newline at end of file diff --git a/js/locales/zh/main.json b/js/locales/zh/main.json index 8274774..299d96c 100644 --- a/js/locales/zh/main.json +++ b/js/locales/zh/main.json @@ -113,6 +113,6 @@ "Failed to load gallery from disk.": "从磁盘加载图库失败。", "Sort: Newest First": "排序：最新优先", "Sort: Oldest First": "排序：最早优先", - "Search Gallery (Filename or Prompt)": "搜索图库 (文件名或提示词)", - "Search filenames or prompt metadata...": "搜索文件名或提示词元数据..." + "Search Gallery (Prompt Text)": "搜索图库 (提示词文本)", + "Search prompts...": "搜索提示词..." } \ No newline at end of file diff --git a/js/ui/galleryView.js b/js/ui/galleryView.js index 93215b5..dd9bf05 100644 --- a/js/ui/galleryView.js +++ b/js/ui/galleryView.js @@ -115,14 +115,23 @@ function handleHeaderScroll(scrollEl) { } } +let galleryFetchController = null; + +// Returns false only when the request was superseded by a newer one (aborted), +// so callers can skip rendering stale results. async function fetchGalleryData(subfolder = "", searchQuery = "") { + // A newer request (e.g. the next keystroke in the search box) cancels the in-flight one; + // the server stops scanning as soon as the client disconnects. + if (galleryFetchController) galleryFetchController.abort(); + const controller = new AbortController(); + galleryFetchController = controller; try { isLoading = true; const normalized = subfolder.replace(/\\/g, "/"); const queryParams = new URLSearchParams({ subfolder: normalized }); if (searchQuery) queryParams.set("search", searchQuery); - const res = await fetch(`/comfy-sidebar/gallery/list?${queryParams.toString()}`); + const res = await fetch(`/comfy-sidebar/gallery/list?${queryParams.toString()}`, { signal: controller.signal }); if (!res.ok) throw new Error(`HTTP ${res.status}`); const data = await res.json(); cachedFolders = data.folders || []; @@ -131,11 +140,16 @@ async function fetchGalleryData(subfolder = "", searchQuery = "") { applySorting(); displayedFilesCount = 0; } catch (e) { + if (e?.name === "AbortError") return false; console.error("[ComfySidebar] Failed to load gallery", e); showNotification(t("Failed to load gallery from disk."), "error"); } finally { - isLoading = false; + if (galleryFetchController === controller) { + galleryFetchController = null; + isLoading = false; + } } + return true; } function applySorting() { @@ -497,7 +511,7 @@ export function setupGalleryView() { const searchIcon = document.createElement("span"); searchIcon.className = "pi pi-search comfy-sidebar-icon-btn"; - searchIcon.title = t("Search Gallery (Filename or Prompt)"); + searchIcon.title = t("Search Gallery (Prompt Text)"); const title = document.createElement("h3"); title.textContent = t("Gallery"); @@ -566,7 +580,7 @@ export function setupGalleryView() { sIcon.className = "pi pi-search"; Object.assign(sIcon.style, { fontSize: "11px", opacity: "0.6", marginRight: "6px", transition: "transform 0.2s" }); const sInput = document.createElement("input"); - Object.assign(sInput, { type: "text", placeholder: t("Search filenames or prompt metadata...") }); + Object.assign(sInput, { type: "text", placeholder: t("Search prompts...") }); Object.assign(sInput.style, { flex: "1", background: "transparent", border: "none", outline: "none", color: "var(--comfy-input-color, var(--fg-color, #eee))", fontSize: "11px", padding: "0" @@ -596,6 +610,7 @@ export function setupGalleryView() { }; let searchDebounce = null; + let searchSeq = 0; const closeSearch = () => { const hadFilter = Boolean(currentSearchQuery || sInput.value.trim()); sInput.value = ""; @@ -626,12 +641,15 @@ export function setupGalleryView() { sIcon.className = "pi pi-spin pi-spinner"; + const seq = ++searchSeq; searchDebounce = setTimeout(async () => { + let applied = false; try { - await fetchGalleryData(currentSubfolder, currentSearchQuery); - renderGallery(); + applied = (await fetchGalleryData(currentSubfolder, currentSearchQuery)) !== false; + if (applied) renderGallery(); } finally { - sIcon.className = "pi pi-search"; + // Only the latest search may stop the spinner + if (seq === searchSeq) sIcon.className = "pi pi-search"; } }, 350); }; diff --git a/prompt_search.py b/prompt_search.py new file mode 100644 index 0000000..bfc0c3b --- /dev/null +++ b/prompt_search.py @@ -0,0 +1,218 @@ +"""Find the text prompts embedded in ComfyUI PNG outputs. + +This module has no ComfyUI imports so it can be unit-tested on its own +(``python -m unittest discover tests``). + +How it works +------------ +* Only the PNG ``prompt`` text chunk is read. Other chunks (notably the often + multi-megabyte ``workflow`` chunk) are skipped with a seek, never parsed. +* From the API-format prompt graph we collect only the strings that actually + feed a text-encoder node (``CLIPTextEncode`` and friends), following links to + upstream string/primitive/concat nodes. Model names, sampler names, file + names and other string widgets are ignored. +* Extracted text is cached per file (keyed by mtime + size), so repeated + searches only pay for files that are new or changed. +""" + +import json +import os +import threading +import zlib +from collections import OrderedDict + +_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" + +# Input names that can carry prompt text (substring match, lower-case). +# Deliberately does not match plain "clip", "model", "vae", "image", etc. +_PROMPT_KEY_HINTS = ( + "text", "prompt", "positive", "negative", "caption", + "string", "value", "t5xxl", "clip_l", "clip_g", +) +_MAX_LINK_DEPTH = 6 +_MAX_PROMPT_CHUNK_BYTES = 16 * 1024 * 1024 +_MAX_TEXT_CHARS_PER_FILE = 50_000 + +_CACHE_MAX_ENTRIES = 20_000 +_cache = OrderedDict() # path -> ((mtime_ns, size), lowercased prompt text) +_cache_lock = threading.Lock() + + +# --------------------------------------------------------------------------- # +# Prompt graph -> text +# --------------------------------------------------------------------------- # + +def _is_prompt_key(name): + lowered = str(name).lower() + return any(hint in lowered for hint in _PROMPT_KEY_HINTS) + + +def _is_encoder(class_type): + lowered = str(class_type or "").lower() + return "textencode" in lowered or "prompt" in lowered + + +def _is_link(value): + # API-format link: [source_node_id, output_slot] + return isinstance(value, list) and len(value) == 2 and isinstance(value[1], int) + + +def extract_prompt_texts(prompt): + """Return the prompt strings found in an API-format ComfyUI prompt dict.""" + if not isinstance(prompt, dict): + return [] + + texts = [] + seen = set() + + def visit(node_id, depth): + key = str(node_id) + if key in seen or depth > _MAX_LINK_DEPTH: + return + seen.add(key) + node = prompt.get(key) + inputs = node.get("inputs") if isinstance(node, dict) else None + if not isinstance(inputs, dict): + return + for name, value in inputs.items(): + if not _is_prompt_key(name): + continue + if isinstance(value, str): + if value.strip(): + texts.append(value) + elif _is_link(value): + visit(value[0], depth + 1) + + starts = [ + node_id for node_id, node in prompt.items() + if isinstance(node, dict) and _is_encoder(node.get("class_type")) + ] + if not starts: + # Unusual graph with no recognisable text-encoder node: fall back to + # every node, still restricted to prompt-like input names. + starts = list(prompt) + for node_id in starts: + visit(node_id, 0) + return texts + + +# --------------------------------------------------------------------------- # +# PNG chunk reading +# --------------------------------------------------------------------------- # + +def _inflate(data, limit=_MAX_PROMPT_CHUNK_BYTES): + # Bounded decompress: refuses to expand beyond `limit` bytes. + return zlib.decompressobj().decompress(data, limit) + + +def _decode_text_chunk(chunk_type, data, nul): + """Decode tEXt / zTXt / iTXt chunk data (keyword already located at `nul`).""" + if chunk_type == b"tEXt": + return data[nul + 1:].decode("latin-1") + + if chunk_type == b"zTXt": + if data[nul + 1] != 0: # only compression method 0 (zlib) exists + return None + return _inflate(data[nul + 2:]).decode("latin-1") + + # iTXt: keyword\0 compression_flag compression_method language\0 translated\0 text + flag = data[nul + 1] + rest = data[nul + 3:] + lang_end = rest.index(b"\x00") + translated_end = rest.index(b"\x00", lang_end + 1) + body = rest[translated_end + 1:] + if flag == 1: + body = _inflate(body) + return body.decode("utf-8", errors="replace") + + +def read_png_prompt(fileobj): + """Return the parsed ``prompt`` JSON from an open PNG file, or None.""" + if fileobj.read(8) != _PNG_SIGNATURE: + return None + while True: + header = fileobj.read(8) + if len(header) < 8: + return None + length = int.from_bytes(header[:4], "big") + chunk_type = header[4:] + if chunk_type in (b"IDAT", b"IEND"): + return None # ComfyUI writes its text chunks before the image data + + if chunk_type in (b"tEXt", b"zTXt", b"iTXt"): + head = fileobj.read(min(length, 80)) # keyword is at most 79 bytes + nul = head.find(b"\x00") + if nul != -1 and head[:nul] == b"prompt": + if length > _MAX_PROMPT_CHUNK_BYTES: + return None + data = head + fileobj.read(length - len(head)) + text = _decode_text_chunk(chunk_type, data, nul) + return json.loads(text) if text else None + fileobj.seek(length - len(head) + 4, 1) # rest of data + CRC + else: + fileobj.seek(length + 4, 1) + + +# --------------------------------------------------------------------------- # +# Cached per-file text + directory scan +# --------------------------------------------------------------------------- # + +def get_prompt_search_text(path): + """Lower-cased prompt text of a PNG ('' if none). Cached by mtime + size.""" + try: + st = os.stat(path) + except OSError: + return "" + signature = (st.st_mtime_ns, st.st_size) + + with _cache_lock: + hit = _cache.get(path) + if hit and hit[0] == signature: + _cache.move_to_end(path) + return hit[1] + + text = "" + try: + with open(path, "rb") as f: + prompt = read_png_prompt(f) + if prompt is not None: + text = "\n".join(extract_prompt_texts(prompt)).lower()[:_MAX_TEXT_CHARS_PER_FILE] + except Exception: + # Corrupt / truncated / non-ComfyUI PNG: treat as "no prompt". + text = "" + + with _cache_lock: + _cache[path] = (signature, text) + _cache.move_to_end(path) + while len(_cache) > _CACHE_MAX_ENTRIES: + _cache.popitem(last=False) + return text + + +def clear_cache(): + with _cache_lock: + _cache.clear() + + +def iter_prompt_matches(target_dir, needle, cancel=None): + """Yield ``(path, os.stat_result)`` for PNGs under `target_dir` whose prompt + text contains `needle` (case-insensitive). Hidden files/folders are skipped. + + `cancel` may be a ``threading.Event``; the scan stops soon after it is set. + """ + needle = needle.lower() + if not needle: + return + for root, dirs, fnames in os.walk(target_dir): + dirs[:] = [d for d in dirs if not d.startswith(".")] + for fname in fnames: + if cancel is not None and cancel.is_set(): + return + if fname.startswith(".") or not fname.lower().endswith(".png"): + continue + path = os.path.join(root, fname) + if needle in get_prompt_search_text(path): + try: + yield path, os.stat(path) + except OSError: + continue diff --git a/tests/test_prompt_search.py b/tests/test_prompt_search.py new file mode 100644 index 0000000..9b90afb --- /dev/null +++ b/tests/test_prompt_search.py @@ -0,0 +1,180 @@ +import json +import os +import struct +import sys +import tempfile +import threading +import unittest +import zlib + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import prompt_search as ps # noqa: E402 + + +def chunk(ctype, data): + return struct.pack(">I", len(data)) + ctype + data + struct.pack(">I", zlib.crc32(ctype + data) & 0xFFFFFFFF) + + +def png_bytes(text_chunks=(), after_idat=()): + """Minimal valid 1x1 PNG. text_chunks go before IDAT, after_idat after it.""" + ihdr = chunk(b"IHDR", struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0)) + idat = chunk(b"IDAT", zlib.compress(b"\x00\x00\x00\x00")) + return (ps._PNG_SIGNATURE + ihdr + b"".join(text_chunks) + idat + + b"".join(after_idat) + chunk(b"IEND", b"")) + + +def text_chunk(key, value): + return chunk(b"tEXt", key.encode("latin-1") + b"\x00" + value.encode("latin-1")) + + +def itxt_chunk(key, value, compressed=False): + body = value.encode("utf-8") + if compressed: + body = zlib.compress(body) + return chunk(b"iTXt", key.encode() + b"\x00" + bytes([1 if compressed else 0, 0]) + b"\x00\x00" + body) + + +def ztxt_chunk(key, value): + return chunk(b"zTXt", key.encode() + b"\x00\x00" + zlib.compress(value.encode("latin-1"))) + + +def comfy_prompt(positive="a red fox", negative="blurry"): + return { + "3": {"class_type": "KSampler", "inputs": { + "seed": 1, "steps": 20, "cfg": 8, "sampler_name": "euler", "scheduler": "normal", + "denoise": 1, "model": ["4", 0], "positive": ["6", 0], "negative": ["7", 0], "latent_image": ["5", 0]}}, + "4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "decoycheckpoint.safetensors"}}, + "5": {"class_type": "EmptyLatentImage", "inputs": {"width": 512, "height": 512, "batch_size": 1}}, + "6": {"class_type": "CLIPTextEncode", "inputs": {"text": positive, "clip": ["4", 1]}}, + "7": {"class_type": "CLIPTextEncode", "inputs": {"text": negative, "clip": ["4", 1]}}, + "9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "decoyprefix", "images": ["8", 0]}}, + } + + +class ExtractTests(unittest.TestCase): + def test_only_encoder_text(self): + texts = ps.extract_prompt_texts(comfy_prompt("a red fox", "blurry")) + self.assertEqual(sorted(texts), ["a red fox", "blurry"]) + + def test_ignores_model_sampler_and_filename_strings(self): + joined = "\n".join(ps.extract_prompt_texts(comfy_prompt())).lower() + for decoy in ("euler", "decoycheckpoint", "decoyprefix", "normal"): + self.assertNotIn(decoy, joined) + + def test_follows_links_to_string_nodes_and_concat(self): + prompt = comfy_prompt() + prompt["6"]["inputs"]["text"] = ["10", 0] + prompt["10"] = {"class_type": "StringConcatenate", "inputs": { + "text_a": ["11", 0], "text_b": "wearing a hat", "delimiter": " ZZDELIM "}} + prompt["11"] = {"class_type": "PrimitiveStringMultiline", "inputs": {"value": "a snowy owl"}} + texts = ps.extract_prompt_texts(prompt) + self.assertIn("a snowy owl", texts) + self.assertIn("wearing a hat", texts) + self.assertNotIn(" ZZDELIM ", texts) + + def test_does_not_follow_clip_link_to_loader(self): + prompt = comfy_prompt() + prompt["4"]["inputs"]["clip_name"] = "decoyclipmodel.safetensors" + self.assertNotIn("decoyclipmodel", "\n".join(ps.extract_prompt_texts(prompt))) + + def test_link_cycle_terminates(self): + prompt = { + "1": {"class_type": "CLIPTextEncode", "inputs": {"text": ["2", 0]}}, + "2": {"class_type": "Whatever", "inputs": {"text": ["1", 0], "value": "loop text"}}, + } + self.assertEqual(ps.extract_prompt_texts(prompt), ["loop text"]) + + def test_fallback_without_encoder_nodes(self): + prompt = {"1": {"class_type": "easy positive", "inputs": {"positive": "fallback words", "mode": "fixed"}}} + self.assertEqual(ps.extract_prompt_texts(prompt), ["fallback words"]) + + def test_garbage_input(self): + self.assertEqual(ps.extract_prompt_texts(None), []) + self.assertEqual(ps.extract_prompt_texts({"1": "x", "2": {"inputs": 5}}), []) + + +class FileTests(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.dir = self.tmp.name + ps.clear_cache() + + def tearDown(self): + self.tmp.cleanup() + + def write(self, rel, data): + path = os.path.join(self.dir, rel) + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "wb") as f: + f.write(data) + return path + + def search(self, needle, **kw): + return sorted(os.path.relpath(p, self.dir).replace("\\", "/") + for p, _ in ps.iter_prompt_matches(self.dir, needle, **kw)) + + def comfy_png(self, positive="a red fox", workflow_note="workflowonlyword"): + return png_bytes([ + text_chunk("prompt", json.dumps(comfy_prompt(positive))), + text_chunk("workflow", json.dumps({"nodes": [{"inputs": [], "note": workflow_note}]})), + ]) + + def test_matches_prompt_not_workflow_not_filename(self): + self.write("sub/ComfyUI_00001_.png", self.comfy_png("A Red Fox in snow")) + self.assertEqual(self.search("red fox"), ["sub/ComfyUI_00001_.png"]) # case-insensitive, recursive + self.assertEqual(self.search("workflowonlyword"), []) # workflow chunk ignored + self.assertEqual(self.search("comfyui_00001"), []) # filename ignored + self.assertEqual(self.search("euler"), []) # non-prompt metadata ignored + + def test_non_ascii_json_escaped_text_chunk(self): + # ComfyUI uses json.dumps (ensure_ascii=True) -> latin-1-safe tEXt chunk + self.write("a.png", self.comfy_png("Красная лиса")) + self.assertEqual(self.search("красная лиса"), ["a.png"]) + + def test_itxt_and_ztxt_variants(self): + raw = json.dumps(comfy_prompt("Серый волк"), ensure_ascii=False) + self.write("i.png", png_bytes([itxt_chunk("prompt", raw)])) + self.write("ic.png", png_bytes([itxt_chunk("prompt", raw, compressed=True)])) + self.write("z.png", png_bytes([ztxt_chunk("prompt", json.dumps(comfy_prompt("grey wolf")))])) + self.assertEqual(self.search("серый волк"), ["i.png", "ic.png"]) + self.assertEqual(self.search("grey wolf"), ["z.png"]) + + def test_prompt_after_other_chunks_and_idat_boundary(self): + self.write("ok.png", png_bytes([text_chunk("workflow", "{}"), text_chunk("prompt", json.dumps(comfy_prompt("late fox")))])) + self.write("late.png", png_bytes(after_idat=[text_chunk("prompt", json.dumps(comfy_prompt("too late fox")))])) + self.assertEqual(self.search("fox"), ["ok.png"]) + + def test_skips_bad_files_hidden_and_non_png(self): + self.write("bad.png", b"not a png at all") + self.write("trunc.png", self.comfy_png()[:40]) + self.write("badjson.png", png_bytes([text_chunk("prompt", "{not json")])) + self.write(".hidden/h.png", self.comfy_png("hidden fox")) + self.write(".dot.png", self.comfy_png("dot fox")) + self.write("a.jpg", self.comfy_png("jpg fox")) + self.write("good.png", self.comfy_png("good fox")) + self.assertEqual(self.search("fox"), ["good.png"]) + + def test_cache_invalidated_on_change(self): + path = self.write("a.png", self.comfy_png("first prompt")) + self.assertEqual(self.search("first"), ["a.png"]) + self.assertEqual(self.search("second"), []) + with open(path, "wb") as f: + f.write(self.comfy_png("second prompt, longer so size differs")) + os.utime(path, ns=(10**18, 10**18)) + self.assertEqual(self.search("second"), ["a.png"]) + self.assertEqual(self.search("first"), []) + + def test_cancel_stops_scan(self): + for i in range(5): + self.write(f"{i}.png", self.comfy_png("fox")) + cancel = threading.Event() + cancel.set() + self.assertEqual(self.search("fox", cancel=cancel), []) + + def test_empty_needle(self): + self.write("a.png", self.comfy_png()) + self.assertEqual(self.search(""), []) + + +if __name__ == "__main__": + unittest.main()