diff --git a/py/api/admin.py b/py/api/admin.py index 0f3eb88..f5f3b0c 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -27,10 +27,22 @@ except ImportError: try: from ...database.operations import PromptDatabase from ...database.pair_ingestion import ingest_resolved_pair + from ...utils.generation_params import ( + PLACEHOLDERS, + GenerationParamsError, + extract_generation_params, + parse_js_json, + ) 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.generation_params import ( + PLACEHOLDERS, + GenerationParamsError, + extract_generation_params, + parse_js_json, + ) from utils.prompt_pair_resolver import resolve_prompt_pair @@ -1511,6 +1523,21 @@ class AdminRoutesMixin: counts["pair_links_added"] += outcome.get("links_added", 0) return counts + @staticmethod + def _graph_prompt_text(metadata): + """The positive prompt the browser parser traces from the graphs, or None.""" + records = {} + for key in ("prompt", "workflow"): + value = metadata.get(key) + records[key] = parse_js_json(value) if isinstance(value, str) else value + try: + positive = extract_generation_params(records)["positivePrompt"] + except GenerationParamsError: + return None + if not isinstance(positive, str) or positive == PLACEHOLDERS["positivePrompt"]: + return None + return positive.strip() 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: @@ -1531,23 +1558,32 @@ class AdminRoutesMixin: self.db, str(media_file), canonical, resolution ) return result if detailed else result["status"] - if resolution["state"] == "ambiguous": - return "ambiguous" - if resolution["state"] == "unresolved": - if canonical.get("prompt") or not isinstance( - canonical.get("parameters"), str - ): + if canonical.get("prompt") and resolution["state"] in ( + "ambiguous", + "unresolved", + ): + # The resolver could not decide: link the prompt the inspector and the + # legacy pages show for this file (the browser parser, in Python). + # Graphs without a recognized sampler keep the legacy extractor below. + prompt_text = self._graph_prompt_text(metadata) + if not prompt_text: return resolution["state"] - parsed = self._parse_comfyui_prompt(metadata) - if not (parsed.get("prompt") or parsed.get("parameters")): - return None - - prompt_text = self._extract_readable_prompt(parsed) - if prompt_text is not None and not isinstance(prompt_text, str): - prompt_text = str(prompt_text) - if not (prompt_text and prompt_text.strip()): - return "found" - prompt_text = prompt_text.strip() + elif resolution["state"] == "ambiguous": + return "ambiguous" + elif resolution["state"] == "unresolved" and not isinstance( + canonical.get("parameters"), str + ): + return resolution["state"] + else: + parsed = self._parse_comfyui_prompt(metadata) + if not (parsed.get("prompt") or parsed.get("parameters")): + return None + prompt_text = self._extract_readable_prompt(parsed) + if prompt_text is not None and not isinstance(prompt_text, str): + prompt_text = str(prompt_text) + if not (prompt_text and prompt_text.strip()): + return "found" + prompt_text = prompt_text.strip() prompt_hash = generate_prompt_hash(prompt_text) existing = self.db.get_prompt_by_hash(prompt_hash) diff --git a/tests/test_ingestion_resolver_profile.py b/tests/test_ingestion_resolver_profile.py index bbc717c..a222849 100644 --- a/tests/test_ingestion_resolver_profile.py +++ b/tests/test_ingestion_resolver_profile.py @@ -8,6 +8,7 @@ leaves it unresolved, while the inspection profile resolves it. import logging import unittest +from pathlib import Path from unittest.mock import Mock, patch from utils import prompt_pair_resolver as resolver @@ -72,12 +73,19 @@ class TestIngestionKeepsDefaultProfile(unittest.TestCase): mixin = admin.AdminRoutesMixin.__new__(admin.AdminRoutesMixin) mixin.db = Mock() + mixin.db.get_prompt_by_hash.return_value = None + mixin.db.save_prompt.return_value = 7 spy = _Spy() with patch.object(admin, "resolve_prompt_pair", spy): - status = mixin._ingest_scanned_file("/out/f03.png", _f03_metadata()) - self.assertEqual(status, "unresolved") + status = mixin._ingest_scanned_file(Path("/out/f03.png"), _f03_metadata()) spy.assert_default_and_gated(self) - mixin.db.assert_not_called() + # The resolver leaves f03 unresolved; the scan links the prompt the + # browser parser traces, as the inspector shows it. + self.assertEqual(status, "added") + self.assertEqual( + mixin.db.save_prompt.call_args.args[0], "a red fox in snow, detailed fur" + ) + mixin.db.link_image_to_prompt.assert_called_once() if __name__ == "__main__": diff --git a/tests/test_pair_scan.py b/tests/test_pair_scan.py index c8d055d..945d19c 100644 --- a/tests/test_pair_scan.py +++ b/tests/test_pair_scan.py @@ -17,6 +17,14 @@ def graph(positive="cat", negative="blur"): } +def sampled_graph(positive="cat", negative="blur"): + """A KSampler with a model link, as real workflows have.""" + g = graph(positive, negative) + g["4"] = {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "m"}} + g["3"]["inputs"]["model"] = ["4", 0] + return g + + class TestPairScan(AdminAPITestCase): def test_pair_scan_repeated_is_unchanged_and_empty_negative_is_supported(self): path = self.output_dir / "pair.png" @@ -28,16 +36,66 @@ class TestPairScan(AdminAPITestCase): 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", - ) + def _scanned(self, name, prompt_graph): + meta = {"prompt": json.dumps(prompt_graph)} + return self.api._ingest_scanned_file(self.output_dir / name, meta) + + def _texts(self): + return [p["text"] for p in self.api.db.get_recent_prompts()["prompts"]] + + def test_samplers_sharing_one_prompt_link_that_prompt(self): + g = sampled_graph() + g["5"] = {**g["3"]} + self.assertEqual(self._scanned("multi.png", g), "added") + self.assertEqual(self._texts(), ["cat"]) + + def test_samplers_with_different_prompts_link_the_first_sampler_prompt(self): + # The prompt the inspector and the legacy panels show for this file. + g = sampled_graph() + g["6"] = {"class_type": "CLIPTextEncode", "inputs": {"text": "dog"}} + g["12"] = {**g["3"], "inputs": {**g["3"]["inputs"], "positive": ["6", 0]}} + self.assertEqual(self._scanned("different.png", g), "added") + self.assertEqual(self._texts(), ["cat"]) + + def test_an_unsupported_pass_through_still_links_the_traced_prompt(self): + g = sampled_graph() + g["7"] = { + "class_type": "MyConditioningTweak", + "inputs": {"conditioning": ["1", 0]}, + } + g["3"]["inputs"]["positive"] = ["7", 0] + g["5"] = {**g["3"]} + self.assertEqual(self._scanned("tweak.png", g), "added") + self.assertEqual(self._texts(), ["cat"]) + + def test_a_graph_without_a_traceable_prompt_saves_nothing(self): + g = sampled_graph() + g["5"] = {**g["3"]} + for sampler in ("3", "5"): + g[sampler]["inputs"]["positive"] = ["9", 0] + g["9"] = { + "class_type": "BasicGuider", + "inputs": {"width": 2048, "height": 2048}, + } + self.assertEqual(self._scanned("untraceable.png", g), "ambiguous") self.assertEqual(self.api.db.list_prompt_pairs()["total"], 0) - self.assertEqual(self.api.db.get_recent_prompts()["prompts"], []) + self.assertEqual(self._texts(), []) + + def test_numeric_widgets_are_never_stored_as_a_prompt(self): + # The pre-redesign scan stored "2048" from a resolution widget here. + g = { + "1": {"class_type": "CLIPTextEncode", "inputs": {"text": ["8", 0]}}, + "8": {"class_type": "DisplayText", "inputs": {"text_size": 2048}}, + "3": { + "class_type": "SamplerCustomAdvanced", + "inputs": {"guider": ["4", 0], "sigmas": ["5", 0]}, + }, + "4": {"class_type": "BasicGuider", "inputs": {"conditioning": ["1", 0]}}, + } + self.assertIn( + self._scanned("numbers.png", g), ("ambiguous", "unresolved", "found") + ) + self.assertEqual(self._texts(), []) def test_batch_exposes_pair_counts_and_actual_unchanged_status(self): meta = {"parameters": "cat\nNegative prompt: blur\nSteps: 20, Sampler: Euler"}