feat(scan): enrich image pairs without guessing or duplicate links (#193)

This commit is contained in:
Vito Sansevero
2026-10-17 06:52:53 -07:00
parent 53596eaeb5
commit 9b4f1ee085
8 changed files with 255 additions and 26 deletions
+31 -3
View File
@@ -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],
+55 -11
View File
@@ -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 # '<db>.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"
+49
View File
@@ -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);
}
});
+28 -6
View File
@@ -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)]
+19
View File
@@ -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)
+57
View File
@@ -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")
+13 -3
View File
@@ -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"],
};
+3 -3
View File
@@ -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;
}
}