From 9b4f1ee085066efd5895cce7121970538b575d6e Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sat, 17 Oct 2026 06:47:10 -0700 Subject: [PATCH] feat(scan): enrich image pairs without guessing or duplicate links (#193) --- database/operations.py | 34 ++++++++++++++++-- py/api/admin.py | 66 ++++++++++++++++++++++++++++------ tests/js/pair-metadata.test.js | 49 +++++++++++++++++++++++++ tests/test_api_admin.py | 34 ++++++++++++++---- tests/test_pair_ingestion.py | 19 ++++++++++ tests/test_pair_scan.py | 57 +++++++++++++++++++++++++++++ web/js/comfy-metadata.js | 16 +++++++-- web/js/gallery.js | 6 ++-- 8 files changed, 255 insertions(+), 26 deletions(-) create mode 100644 tests/js/pair-metadata.test.js create mode 100644 tests/test_pair_scan.py diff --git a/database/operations.py b/database/operations.py index 7a9a8bf..9424f1e 100644 --- a/database/operations.py +++ b/database/operations.py @@ -1995,7 +1995,7 @@ class PromptDatabase(PairOperations): cursor = conn.execute( """SELECT gi.prompt_id, p.text, p.category, p.rating, p.notes, gi.workflow_data, gi.prompt_metadata, gi.generation_time, - gi.parameters + gi.parameters, gi.id, gi.image_path, gi.file_path FROM generated_images gi JOIN prompts p ON gi.prompt_id = p.id WHERE gi.image_path = ? OR gi.file_path = ? @@ -2007,9 +2007,35 @@ class PromptDatabase(PairOperations): f"%{escape_like(os.path.basename(image_path))}", ), ) - row = cursor.fetchone() - if not row: + rows = cursor.fetchall() + exact = [ + row + for row in rows + if row[10] == image_path or row[11] == normalize_image_path(image_path) + ] + matches = exact or rows + if len(matches) != 1: return None + row = matches[0] + associations = self.get_image_prompt_pairs(row[9]) + pair_ids = {item["pair_id"] for item in associations} + pair = ( + self.get_prompt_pair(next(iter(pair_ids))) + if len(pair_ids) == 1 + else None + ) + contexts = { + ( + item["provenance"].get("positive", {}).get("text"), + item["provenance"].get("negative", {}).get("text"), + ) + for item in associations + } + context = None + if len(contexts) == 1: + positive, negative = next(iter(contexts)) + if isinstance(positive, str) and isinstance(negative, str): + context = {"positive_text": positive, "negative_text": negative} tags = self._get_prompt_tags(conn, row[0]) try: workflow = json.loads(row[5]) if row[5] else None @@ -2020,6 +2046,8 @@ class PromptDatabase(PairOperations): except (json.JSONDecodeError, TypeError): prompt_meta = None return { + "pair": pair, + "pair_context": context, "prompt_id": row[0], "text": row[1], "category": row[2], diff --git a/py/api/admin.py b/py/api/admin.py index d13c034..5bb065b 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -26,8 +26,12 @@ except ImportError: try: from ...database.operations import PromptDatabase + from ...database.pair_ingestion import ingest_resolved_pair + from ...utils.prompt_pair_resolver import resolve_prompt_pair except ImportError: from database.operations import PromptDatabase + from database.pair_ingestion import ingest_resolved_pair + from utils.prompt_pair_resolver import resolve_prompt_pair try: @@ -40,6 +44,20 @@ VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".m4v", ".wmv") SCAN_BATCH_SIZE = 50 # Files per executor call during the output scan MAX_SAFETY_BACKUPS = 10 # '.backup_*' files kept next to the live database MAX_SCAN_FILES = 50_000 # Upper bound on media files one scan will look at +SCAN_COUNT_KEYS = ( + "processed", + "found", + "added", + "linked", + "enriched", + "unchanged", + "ambiguous", + "unresolved", + "conflict", + "failed", + "pairs_added", + "pair_links_added", +) MAX_SCAN_DEPTH = 12 # Directory levels below an output root that are scanned # Files the duplicate cleaner may delete: everything the scan reports, plus # the common single-f TIFF spelling. Anything else inside the output @@ -1351,7 +1369,7 @@ class AdminRoutesMixin: media_files = await self._run_in_executor( self._collect_output_media_sync, output_dirs ) - counts = {"processed": 0, "found": 0, "added": 0, "linked": 0} + counts = dict.fromkeys(SCAN_COUNT_KEYS, 0) if media_files: yield self._scan_progress( 5, f"Found {len(media_files)} media files to process...", counts @@ -1405,7 +1423,7 @@ class AdminRoutesMixin: "progress": progress, "status": status, "processed": counts.get("processed", 0), - "found": counts.get("found", 0), + **{key: counts.get(key, 0) for key in SCAN_COUNT_KEYS}, } def _collect_output_media_sync(self, output_dirs): @@ -1431,25 +1449,49 @@ class AdminRoutesMixin: def _ingest_scan_batch_sync(self, batch_results): """Store the prompts found in a metadata batch (blocking DB work).""" - counts = {"processed": 0, "found": 0, "added": 0, "linked": 0} + counts = dict.fromkeys(SCAN_COUNT_KEYS, 0) for media_file, metadata in batch_results: counts["processed"] += 1 try: - outcome = self._ingest_scanned_file(media_file, metadata) + outcome = self._ingest_scanned_file(media_file, metadata, detailed=True) except Exception as e: self.logger.error(f"Error processing {media_file.name}: {e}") + counts["failed"] += 1 continue if outcome is None: continue counts["found"] += 1 - if outcome in ("added", "linked"): - counts[outcome] += 1 + status = outcome["status"] if isinstance(outcome, dict) else outcome + if status in counts: + counts[status] += 1 + if isinstance(outcome, dict): + counts["pairs_added"] += outcome.get("pairs_added", 0) + counts["pair_links_added"] += outcome.get("links_added", 0) return counts - def _ingest_scanned_file(self, media_file, metadata): - """Save or link one scanned file; returns 'added', 'linked', 'found' or None.""" + def _ingest_scanned_file(self, media_file, metadata, detailed=False): + """Enrich exact pairs; retain legacy positive-only import for absent pairs.""" if not metadata: return None + canonical = dict(metadata) + if isinstance(canonical.get("prompt"), str): + try: + canonical["prompt"] = json.loads(canonical["prompt"]) + except (ValueError, TypeError): + pass + if not canonical.get("prompt") and not canonical.get("parameters"): + return None + resolution = resolve_prompt_pair(canonical) + if resolution["state"] == "resolved": + result = ingest_resolved_pair( + self.db, str(media_file), canonical, resolution + ) + return result if detailed else result["status"] + if resolution["state"] in ("ambiguous", "unresolved"): + if canonical.get("prompt") or not isinstance( + canonical.get("parameters"), str + ): + return resolution["state"] parsed = self._parse_comfyui_prompt(metadata) if not (parsed.get("prompt") or parsed.get("parameters")): return None @@ -1465,13 +1507,15 @@ class AdminRoutesMixin: existing = self.db.get_prompt_by_hash(prompt_hash) if existing: try: - self.db.link_image_to_prompt(existing["id"], str(media_file)) + image_id = self.db.link_image_to_prompt( + existing["id"], str(media_file), metadata=canonical + ) except Exception as e: self.logger.error( f"Failed to link {media_file.name} to existing prompt: {e}" ) return "found" - return "linked" + return "linked" if image_id else "unchanged" prompt_id = self.db.save_prompt( prompt_text, @@ -1484,7 +1528,7 @@ class AdminRoutesMixin: if not prompt_id: return "found" try: - self.db.link_image_to_prompt(prompt_id, str(media_file)) + self.db.link_image_to_prompt(prompt_id, str(media_file), metadata=canonical) except Exception as e: self.logger.error(f"Failed to link {media_file.name} to new prompt: {e}") return "added" diff --git a/tests/js/pair-metadata.test.js b/tests/js/pair-metadata.test.js new file mode 100644 index 0000000..8958bca --- /dev/null +++ b/tests/js/pair-metadata.test.js @@ -0,0 +1,49 @@ +const test = require('node:test'); +const assert = require('node:assert/strict'); +const fs = require('node:fs'); +const vm = require('node:vm'); +const path = require('node:path'); +const ImageHelpers = require('../../web/js/image-helpers.js'); +const {extractGenerationParams} = require('../../web/js/comfy-metadata.js'); + +function graph(positive = 0, negative = 1) { + return { + 1:{class_type:'PromptManagerPair',inputs:{text:'cat',negative_text:'blur',prepend_text:'trigger',append_text:'detail'}}, + 2:{class_type:'KSampler',inputs:{positive:['1',positive],negative:['1',negative],model:['3',0]}}, + 3:{class_type:'CheckpointLoaderSimple',inputs:{ckpt_name:'m.safetensors'}}, + }; +} + +test('paired conditioning uses actual output roles and composes only positive', () => { + const metadata = extractGenerationParams({prompt:graph()}); + assert.equal(metadata.positivePrompt,'trigger cat detail'); + assert.equal(metadata.negativePrompt,'blur'); + const swapped = extractGenerationParams({prompt:graph(1,0)}); + assert.equal(swapped.positivePrompt,'blur'); + assert.equal(swapped.negativePrompt,'trigger cat detail'); +}); + +test('paired strings passed through encoders retain the connected output slot', () => { + const prompt=graph(); + prompt[4]={class_type:'CLIPTextEncode',inputs:{text:['1',2]}}; + prompt[5]={class_type:'CLIPTextEncode',inputs:{text:['1',3]}}; + prompt[2].inputs.positive=['4',0];prompt[2].inputs.negative=['5',0]; + const metadata=extractGenerationParams({prompt}); + assert.equal(metadata.positivePrompt,'trigger cat detail'); + assert.equal(metadata.negativePrompt,'blur'); +}); + +test('recorded image pair overrides unrelated PNG fallback, including an empty negative', async () => { + for (const negative of ['blur','']) { + const context=vm.createContext({document:{addEventListener(){}},window:{},console,ImageHelpers, + ApiClient:{fetchUrl:async()=>({ok:true,json:async()=>({success:true,prompt:{text:'raw cat',pair:{negative_text:negative},pair_context:{positive_text:'trigger cat',negative_text:negative}}})})}, + }); + const Manager=vm.runInContext(fs.readFileSync(path.join(__dirname,'../../web/js/gallery.js'),'utf8')+'\nGalleryManager;',context); + const manager=Object.create(Manager.prototype); + manager.extractImageMetadata=async()=>({positivePrompt:'wrong',negativePrompt:'wrong'}); + manager.displayMetadata=data=>{manager.last=data;}; + await manager.loadImageMetadata('/prompt_manager/images/serve/example.png',{}); + assert.equal(manager.last.positivePrompt,'trigger cat'); + assert.equal(manager.last.negativePrompt,negative); + } +}); diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index c4f242e..870e093 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -476,10 +476,13 @@ class TestScansRunOffTheEventLoop(AdminAPITestCase): ).text() self.assertIn('"added": 1', first) - self.assertIn('"linked": 1', second) + self.assertIn('"unchanged": 1', second) self.assertIn('"added": 0', second) prompts = self.api.db.get_recent_prompts(limit=10)["prompts"] - self.assertEqual([p["text"] for p in prompts], ["a red square"]) + self.assertEqual({p["text"] for p in prompts}, {"a red square", "bad"}) + pairs = self.api.db.list_prompt_pairs()["pairs"] + self.assertEqual(len(pairs), 1) + self.assertEqual(pairs[0]["negative_text"], "bad") async def test_output_scan_without_directories_reports_error(self): GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.outside_dir)] @@ -1556,7 +1559,7 @@ class TestOutputScanInternals(AdminAPITestCase): "found", ) self.assertEqual(self.api._ingest_scanned_file(f, meta), "added") - self.assertEqual(self.api._ingest_scanned_file(f, meta), "linked") + self.assertEqual(self.api._ingest_scanned_file(f, meta), "unchanged") def test_ingest_link_failures_are_tolerated(self): f = Path("gen.png") @@ -1589,7 +1592,23 @@ class TestOutputScanInternals(AdminAPITestCase): counts = self.api._ingest_scan_batch_sync( [(Path("a.png"), meta), (Path("b.png"), {})] ) - self.assertEqual(counts, {"processed": 2, "found": 0, "added": 0, "linked": 0}) + self.assertEqual( + counts, + { + "processed": 2, + "found": 0, + "added": 0, + "linked": 0, + "enriched": 0, + "unchanged": 0, + "ambiguous": 0, + "unresolved": 0, + "conflict": 0, + "failed": 1, + "pairs_added": 0, + "pair_links_added": 0, + }, + ) async def test_scan_internal_error_streams_error_event(self): GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)] @@ -1637,7 +1656,10 @@ class TestOutputScanIsABackgroundJob(AdminAPITestCase): await asyncio.wait_for(self.api._scan_job.task, 10) prompts = self.api.db.get_recent_prompts(limit=10)["prompts"] - self.assertEqual([p["text"] for p in prompts], ["a red square"]) + self.assertEqual({p["text"] for p in prompts}, {"a red square", "bad"}) + pairs = self.api.db.list_prompt_pairs()["pairs"] + self.assertEqual(len(pairs), 1) + self.assertEqual(pairs[0]["negative_text"], "bad") self.assertEqual(self.api._scan_job.last_event["type"], "complete") self.assertFalse(self.api._scan_job.running) @@ -1704,7 +1726,7 @@ class TestOutputScanIsABackgroundJob(AdminAPITestCase): body = await (await self.client.request("POST", "/prompt_manager/scan")).text() self.assertEqual(len(calls), 2) - self.assertIn('"linked": 1', body) + self.assertIn('"unchanged": 1', body) async def test_failed_scan_is_not_left_running(self): GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)] diff --git a/tests/test_pair_ingestion.py b/tests/test_pair_ingestion.py index 7e24094..952df4f 100644 --- a/tests/test_pair_ingestion.py +++ b/tests/test_pair_ingestion.py @@ -121,3 +121,22 @@ class TestPairIngestion(unittest.TestCase): self.assertEqual( association["provenance"]["positive"]["text"], "prefix cat suffix" ) + + def test_image_lookup_exposes_its_recorded_pair_not_latest_negative(self): + value = resolved("trigger cat", "blur") + value["positive"]["storage_text"] = "cat" + first = ingest_resolved_pair(self.db, self.path, {}, value) + ingest_resolved_pair(self.db, self.path + "2", {}, resolved("cat", "noise")) + info = self.db.get_image_prompt_info(self.path) + self.assertEqual(info["pair"]["id"], first["pair_id"]) + self.assertEqual(info["pair"]["negative_text"], "blur") + self.assertEqual(info["pair_context"]["positive_text"], "trigger cat") + self.assertEqual(info["pair_context"]["negative_text"], "blur") + + def test_overwritten_file_with_different_negative_reports_conflict(self): + ingest_resolved_pair(self.db, self.path, {}, resolved()) + result = ingest_resolved_pair( + self.db, self.path, {}, resolved(negative="noise") + ) + self.assertEqual(result["status"], "conflict") + self.assertEqual(self.db.list_prompt_pairs()["total"], 1) diff --git a/tests/test_pair_scan.py b/tests/test_pair_scan.py new file mode 100644 index 0000000..fbeb9ac --- /dev/null +++ b/tests/test_pair_scan.py @@ -0,0 +1,57 @@ +"""Metadata scanning enriches known pairs and conservatively leaves conflicts.""" + +import json +from pathlib import Path + +from tests.test_api_admin import AdminAPITestCase + + +def graph(positive="cat", negative="blur"): + return { + "1": {"class_type": "CLIPTextEncode", "inputs": {"text": positive}}, + "2": {"class_type": "CLIPTextEncode", "inputs": {"text": negative}}, + "3": { + "class_type": "KSampler", + "inputs": {"positive": ["1", 0], "negative": ["2", 0]}, + }, + } + + +class TestPairScan(AdminAPITestCase): + def test_pair_scan_repeated_is_unchanged_and_empty_negative_is_supported(self): + path = self.output_dir / "pair.png" + meta = {"prompt": json.dumps(graph(negative=""))} + self.assertEqual(self.api._ingest_scanned_file(path, meta), "added") + self.assertEqual(self.api._ingest_scanned_file(path, meta), "unchanged") + pairs = self.api.db.list_prompt_pairs()["pairs"] + self.assertEqual(len(pairs), 1) + self.assertEqual(pairs[0]["negative_text"], "") + self.assertEqual(pairs[0]["run_count"], 0) + + def test_ambiguous_sampler_scan_does_not_save_a_guess(self): + g = graph() + g["4"] = {**g["3"]} + meta = {"prompt": json.dumps(g)} + self.assertEqual( + self.api._ingest_scanned_file(self.output_dir / "multi.png", meta), + "ambiguous", + ) + self.assertEqual(self.api.db.list_prompt_pairs()["total"], 0) + self.assertEqual(self.api.db.get_recent_prompts()["prompts"], []) + + def test_batch_exposes_pair_counts_and_actual_unchanged_status(self): + meta = {"parameters": "cat\nNegative prompt: blur\nSteps: 20, Sampler: Euler"} + batch = [(Path(self.output_dir / "a.png"), meta)] + first = self.api._ingest_scan_batch_sync(batch) + self.assertEqual(first["pairs_added"], 1) + self.assertEqual(first["pair_links_added"], 1) + second = self.api._ingest_scan_batch_sync(batch) + self.assertEqual(second["unchanged"], 1) + self.assertEqual(second["linked"], 0) + self.assertEqual(second["pairs_added"], 0) + + def test_legacy_single_prompt_second_pass_not_reported_as_new_link(self): + meta = {"prompt": json.dumps({"1": graph()["1"]})} + path = self.output_dir / "old.png" + self.assertEqual(self.api._ingest_scanned_file(path, meta), "added") + self.assertEqual(self.api._ingest_scanned_file(path, meta), "unchanged") diff --git a/web/js/comfy-metadata.js b/web/js/comfy-metadata.js index 8e35fc1..234f66a 100644 --- a/web/js/comfy-metadata.js +++ b/web/js/comfy-metadata.js @@ -63,7 +63,7 @@ const widgets = new Map(); for (const n of workflowNodes(workflow)) { const text = widgetsOf(n)[0]; - if (/promptmanager/i.test(n.type || "") && typeof text === "string") { + if (n.type !== "PromptManagerPair" && /promptmanager/i.test(n.type || "") && typeof text === "string") { widgets.set(String(n.id), text); } } @@ -94,8 +94,9 @@ if (!isLink(link) || seen.size >= MAX_VISITS) return []; const id = String(link[0]); const node = ctx.graph[id]; - if (!node || !node.inputs || seen.has(id)) return []; - seen.add(id); + const key = `${id}:${link[1]}`; + if (!node || !node.inputs || seen.has(key)) return []; + seen.add(key); return [id, node]; } @@ -132,6 +133,10 @@ if (typeof value === "string") return value; const [id, node] = enter(ctx, value, seen); if (!node) return ""; + if (node.class_type === "PromptManagerPair") { + if (value[1] === 3) return resolveString(ctx, node.inputs.negative_text, seen); + return value[1] === 2 ? composeNodeText(ctx, id, node, seen) : ""; + } if ("text" in node.inputs) return composeNodeText(ctx, id, node, seen); const delimiter = typeof node.inputs.delimiter === "string" ? node.inputs.delimiter : " "; return Object.entries(node.inputs) @@ -165,6 +170,10 @@ function traceConditioning(ctx, link, role, seen = new Set()) { const [id, node] = enter(ctx, link, seen); if (!node) return ""; + if (node.class_type === "PromptManagerPair") { + if (link[1] === 1) return resolveString(ctx, node.inputs.negative_text, seen); + return link[1] === 0 ? composeNodeText(ctx, id, node, seen) : ""; + } if ("text" in node.inputs) return composeNodeText(ctx, id, node, seen); // Pass-through nodes (ControlNet apply, conditioning combine/set area, ...) const { inputs } = node; @@ -237,6 +246,7 @@ PromptManagerText: PROMPT_MANAGER_WIDGETS, PromptManagerFocused: PROMPT_MANAGER_WIDGETS, PromptManagerSplit: PROMPT_MANAGER_WIDGETS, + PromptManagerPair: ["text", "negative_text", "category", "tags", "search_text", "prepend_text", "append_text"], CheckpointLoader: ["config_name", "ckpt_name"], UNETLoader: ["unet_name", "weight_dtype"], }; diff --git a/web/js/gallery.js b/web/js/gallery.js index c94e736..4b9563d 100644 --- a/web/js/gallery.js +++ b/web/js/gallery.js @@ -1221,8 +1221,8 @@ if (databasePrompt) { // Create metadata object from database prompt metadata = { - positivePrompt: databasePrompt.text, - negativePrompt: 'No negative prompt found', // Database doesn't separate positive/negative + positivePrompt: databasePrompt.pair_context?.positive_text ?? databasePrompt.text, + negativePrompt: databasePrompt.pair_context?.negative_text ?? databasePrompt.pair?.negative_text ?? 'No negative prompt found', checkpoint: 'Unknown', steps: 'Unknown', cfgScale: 'Unknown', @@ -1248,7 +1248,7 @@ metadata.sampler = pngMetadata.sampler; metadata.seed = pngMetadata.seed; // Use negative prompt from PNG if available - if (pngMetadata.negativePrompt && pngMetadata.negativePrompt !== 'No negative prompt found') { + if (!databasePrompt.pair && pngMetadata.negativePrompt && pngMetadata.negativePrompt !== 'No negative prompt found') { metadata.negativePrompt = pngMetadata.negativePrompt; } }