Merge pull request #4 from m0rtus59/feat/prompt-only-search
Improve gallery prompt search while retaining filename search
This commit is contained in:
m0rtus59
@@ -18,7 +18,7 @@ A compact and powerful layout overhaul for ComfyUI that makes generation results
|
||||
|
||||
<img src="assets/compare.png" alt="ComfySidebar compare" width="100%">
|
||||
|
||||
* **[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, filename and prompt metadata search across subfolders, 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.
|
||||
|
||||
|
||||
+112
-121
@@ -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,104 @@ 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 filename + prompt-text search. Prompt metadata is only read for PNGs.
|
||||
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_set():
|
||||
return 200, {
|
||||
"subfolder": subfolder.replace("\\", "/"),
|
||||
"folders": folders,
|
||||
"files": files
|
||||
}
|
||||
if fname.startswith("."):
|
||||
continue
|
||||
ext = os.path.splitext(fname)[1].lower()
|
||||
if ext not in _GALLERY_EXTS:
|
||||
continue
|
||||
|
||||
fpath = os.path.join(root, fname)
|
||||
matched = search in fname.lower()
|
||||
if not matched and ext == ".png":
|
||||
matched = search in prompt_search.get_prompt_search_text(fpath)
|
||||
if matched:
|
||||
try:
|
||||
stat = os.stat(fpath)
|
||||
except (PermissionError, FileNotFoundError):
|
||||
continue
|
||||
files.append({
|
||||
"filename": fname,
|
||||
"subfolder": _relative_subfolder(root, 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 +251,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)
|
||||
|
||||
|
||||
+23
-5
@@ -125,14 +125,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 || [];
|
||||
@@ -141,11 +150,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() {
|
||||
@@ -606,6 +620,7 @@ export function setupGalleryView() {
|
||||
};
|
||||
|
||||
let searchDebounce = null;
|
||||
let searchSeq = 0;
|
||||
const closeSearch = () => {
|
||||
const hadFilter = Boolean(currentSearchQuery || sInput.value.trim());
|
||||
sInput.value = "";
|
||||
@@ -636,12 +651,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);
|
||||
};
|
||||
|
||||
@@ -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
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user