diff --git a/database/pair_ingestion.py b/database/pair_ingestion.py index 567bc7e..cad3641 100644 --- a/database/pair_ingestion.py +++ b/database/pair_ingestion.py @@ -51,14 +51,21 @@ def ingest_resolved_pair( ) if any(row["negative_hash"] != expected for row in prior): return {"status": "conflict", "pairs_added": 0, "links_added": 0} - before = conn.execute("SELECT COUNT(*) FROM prompt_pairs").fetchone()[0] + prior_pair = conn.execute( + "SELECT pp.id FROM prompt_pairs pp " + "JOIN prompts p ON p.id=pp.positive_prompt_id " + "LEFT JOIN prompts n ON n.id=pp.negative_prompt_id " + "WHERE p.hash=? AND n.hash IS ?", + ( + generate_prompt_hash(positive_text), + generate_prompt_hash(negative_text) if negative_text.strip() else None, + ), + ).fetchone() pair_id = db._save_pair(conn, positive_text, negative_text) pair = conn.execute( "SELECT * FROM prompt_pairs WHERE id=?", (pair_id,) ).fetchone() - pairs_added = ( - conn.execute("SELECT COUNT(*) FROM prompt_pairs").fetchone()[0] - before - ) + pairs_added = int(prior_pair is None) file_size, width, height, image_format = _coerce_file_info(metadata) snapshots = { column: json.dumps(metadata[key], sort_keys=True, allow_nan=False) diff --git a/py/api/admin.py b/py/api/admin.py index 5bb065b..35b6552 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -377,6 +377,10 @@ class AdminRoutesMixin: async def scan_images_route(request): return await self.scan_images(request) + @routes.post("/prompt_manager/scan/cancel") + async def cancel_scan_route(request): + return await self.cancel_scan(request) + @routes.get("/prompt_manager/scan/status") async def scan_status_route(request): return await self.scan_status(request) @@ -1298,6 +1302,14 @@ class AdminRoutesMixin: job = self._start_scan_job() return await self._stream_scan_job(request, job) + async def cancel_scan(self, request): + """Stop between file transactions; an in-flight extraction can finish.""" + job = self._scan_job + requested = bool(job and job.running) + if requested: + job.cancelled = True + return web.json_response({"success": True, "cancellation_requested": requested}) + async def scan_status(self, request): """Whether an output scan is running, and its latest progress event.""" job = self._scan_job @@ -1382,7 +1394,9 @@ class AdminRoutesMixin: f"found={counts['found']}, new_prompts_added={counts['added']}, " f"images_linked_to_existing={counts['linked']}" ) - yield {"type": "complete", **counts, "directories": dir_names} + job = self._scan_job + final_type = "cancelled" if job and job.cancelled else "complete" + yield {"type": final_type, **counts, "directories": dir_names} except Exception: self.logger.exception("Scan error") @@ -1398,6 +1412,8 @@ class AdminRoutesMixin: """ total = len(media_files) for batch_start in range(0, total, SCAN_BATCH_SIZE): + if self._scan_job and self._scan_job.cancelled: + return batch = media_files[batch_start : batch_start + SCAN_BATCH_SIZE] batch_results = await self._run_in_executor( self._extract_batch_metadata_sync, batch @@ -1451,6 +1467,8 @@ class AdminRoutesMixin: """Store the prompts found in a metadata batch (blocking DB work).""" counts = dict.fromkeys(SCAN_COUNT_KEYS, 0) for media_file, metadata in batch_results: + if self._scan_job and self._scan_job.cancelled: + break counts["processed"] += 1 try: outcome = self._ingest_scanned_file(media_file, metadata, detailed=True) diff --git a/py/api/scan_job.py b/py/api/scan_job.py index d3eb9ce..e630470 100644 --- a/py/api/scan_job.py +++ b/py/api/scan_job.py @@ -16,6 +16,7 @@ class ScanJob: def __init__(self): self.last_event = None self.finished = False + self.cancelled = False self.task = None self._subscribers = set() diff --git a/tests/test_pair_image_linking.py b/tests/test_pair_image_linking.py new file mode 100644 index 0000000..350d055 --- /dev/null +++ b/tests/test_pair_image_linking.py @@ -0,0 +1,55 @@ +"""Live metadata annotation uses the same conservative ingestion as scan.""" + +import logging +import tempfile +import unittest +from pathlib import Path + +from database.operations import PromptDatabase +from utils.hashing import generate_prompt_hash +from utils.image_monitor import ImageGenerationHandler + + +def metadata(positive="cat", negative="blur"): + return { + "prompt": { + "1": {"class_type": "PromptManager", "inputs": {"text": positive}}, + "2": {"class_type": "CLIPTextEncode", "inputs": {"text": negative}}, + "3": { + "class_type": "KSampler", + "inputs": {"positive": ["1", 0], "negative": ["2", 0]}, + }, + } + } + + +class TestPairImageLinking(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.db = PromptDatabase(str(Path(self.temp.name) / "db.sqlite")) + self.path = str(Path(self.temp.name) / "image.png") + self.pid = self.db.save_prompt("cat", prompt_hash=generate_prompt_hash("cat")) + self.handler = ImageGenerationHandler.__new__(ImageGenerationHandler) + self.handler.db_manager = self.db + self.handler.logger = logging.getLogger("test.pairlink") + + def tearDown(self): + self.db.close_all() + self.temp.cleanup() + + def test_exact_metadata_creates_pair_link_without_counting_a_run(self): + self.handler.link_image_to_prompt(self.path, {"id": self.pid}, metadata()) + self.handler.link_image_to_prompt(self.path, {"id": self.pid}, metadata()) + pairs = self.db.list_prompt_pairs()["pairs"] + self.assertEqual(len(pairs), 1) + self.assertEqual(pairs[0]["image_count"], 1) + self.assertEqual(pairs[0]["run_count"], 0) + + def test_stale_context_or_ambiguous_metadata_does_not_create_pair(self): + meta = metadata(positive="dog") + self.handler.link_image_to_prompt(self.path, {"id": self.pid}, meta) + self.assertEqual(self.db.list_prompt_pairs()["total"], 0) + meta = metadata() + meta["prompt"]["4"] = meta["prompt"]["3"] + self.handler.link_image_to_prompt(self.path + "2", {"id": self.pid}, meta) + self.assertEqual(self.db.list_prompt_pairs()["total"], 0) diff --git a/tests/test_pair_scan.py b/tests/test_pair_scan.py index fbeb9ac..939a909 100644 --- a/tests/test_pair_scan.py +++ b/tests/test_pair_scan.py @@ -55,3 +55,29 @@ class TestPairScan(AdminAPITestCase): 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") + + async def test_cancel_finishes_at_safe_boundary_and_restart_is_idempotent(self): + from py.api.scan_job import ScanJob + + job = ScanJob() + self.api._scan_job = job + response = await self.client.post("/prompt_manager/scan/cancel") + self.assertEqual(response.status, 200) + self.assertTrue((await response.json())["cancellation_requested"]) + counts = self.api._ingest_scan_batch_sync( + [(self.output_dir / "cancel.png", {"prompt": json.dumps(graph())})] + ) + self.assertEqual(counts["processed"], 0) + self.assertEqual(self.api.db.list_prompt_pairs()["total"], 0) + job.finish() + self.api._scan_job = ScanJob() + first = self.api._ingest_scan_batch_sync( + [(self.output_dir / "cancel.png", {"prompt": json.dumps(graph())})] + ) + self.assertEqual(first["pairs_added"], 1) + self.assertEqual( + self.api._ingest_scan_batch_sync( + [(self.output_dir / "cancel.png", {"prompt": json.dumps(graph())})] + )["unchanged"], + 1, + ) diff --git a/utils/image_monitor.py b/utils/image_monitor.py index 5a8167b..69b2b8f 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -35,6 +35,12 @@ from .hashing import generate_prompt_hash from .metadata_extractor import ComfyUIMetadataExtractor from .logging_config import get_logger from .prompt_graph import resolve_text, run_prompt_nodes +from .prompt_pair_resolver import resolve_prompt_pair + +try: + from ..database.pair_ingestion import ingest_resolved_pair +except ImportError: + from database.pair_ingestion import ingest_resolved_pair # Upper bound on images waiting to be processed (a burst beyond this is skipped) MAX_PENDING_IMAGES = 1000 @@ -529,6 +535,17 @@ class ImageGenerationHandler(FileSystemEventHandler): etc.) """ try: + resolution = resolve_prompt_pair(metadata) + if resolution["state"] == "resolved": + positive = resolution["positive"] + expected = self.db_manager.get_prompt_by_hash( + generate_prompt_hash(positive.get("storage_text", positive["text"])) + ) + if expected and expected["id"] == prompt_context["id"]: + ingest_resolved_pair( + self.db_manager, image_path, metadata, resolution + ) + return image_id = self.db_manager.link_image_to_prompt( prompt_id=prompt_context["id"], image_path=image_path, metadata=metadata )