fix(ingestion): link an image only to a prompt that provably produced it (#344)
* fix(ingestion): link an image only to a prompt that provably produced it An LLM fed a PromptManagerPair's prepend_text while its text editor was empty. The node refused to save (an empty positive raised and logged a warning), and the image monitor then linked the image to a stale tracker entry or to the newest prompt in the database. - Nodes: an empty text editor saves and tracks nothing (prepend and append text are never the stored prompt), logs at info, and marks the node's run key empty so no earlier entry can be used for its images. - Tracker entries carry a run key: a digest of the prompt node and its whole upstream graph, the identity ComfyUI's cache uses. The monitor computes the same key from the image's embedded graph and uses a queued or snapshot entry only when the keys match (a snapshot also needs a known run with one entry, so a batch's last item is never applied to every image). - The recency fallback is gone; anything uncertain stays unlinked and is logged at info. - Recorded execution history attributes a file when exactly one run reported it, the file was written during that run, and the run's positive contexts name one stored prompt. The monitor and the output scan use it as their last certain source. - The scan and pair ingestion treat a resolved pair with an empty stored positive as unresolved instead of raising. - tools/review_image_links.py reviews selected links by these rules; dry run by default, --apply changes only named ids after a backup. Closes #342 * test(ingestion): real-host attribution case and the #342 testing contract The managed host disables the image monitor, so the owned fixture pack gains a route that starts and stops the plugin's monitor singleton for the attribution case only. The case runs in both renderers: a typed Pair image links, a Pair image whose positive is only linked prepend text stays unlinked with nothing stored, an unknown node feeding text links through its run key, and its cached re-run links again. * fix(ingestion): take a queued entry only from the run the image appeared in Dropping a key's older queue entries whenever the node ran again let a later queued run, executing during the monitor's processing delay, take the earlier image's entry. The snapshot taken when the image appears now names the key's latest run, and the queue is popped by key and run: same-key entries of earlier runs ahead of it are stale, entries of later runs wait for their own images. Also keeps the attribution case out of the stock real-host suite. * test(ingestion): host lookups degrade to no proof * fix(ingestion): address the #344 review (history re-check, keys, repair) - Images without an embedded API graph: the scan now asks recorded history before skipping a file with no prompt metadata, and the monitor gives a keyless unlinked image one history-only re-check when its run's outputs are recorded (a seal notification from the outputs write; at most 1000 waiting images, each for an hour). - Run keys: an image is unkeyable when any positive node has no key, so a two-pass graph is never linked through the other pass; links follow ComfyUI's is_link exactly. - A failed save marks the node's run key like an empty editor, so an older entry of that key can't stand in for this run. - A placeholder is the latest entry: the image stays unlinked without falling through to history, the same rule as a popped placeholder. - An index on execution_run_outputs(filename, subfolder, root_key, path_style) for the history lookup, added by migration and accepted by the output schema check (any other index is still refused). - Repair tool: unlink removes only the link record and only while the image is still linked to the prompt it was judged on; every run prints what --apply does. - Real-host case proves the re-run's PromptManager node was cached. * test(ingestion): real-host bare-output re-check and the review-round contract A metadata-disabled save, held open behind the fixture gate after SaveImage, waits unlinked while its run is open and links once the run's outputs are recorded. The fixture's monitor route lists the waiting files. TESTING.md covers the history re-check, the read index and what the repair tool's --apply does. * fix(ingestion): batch history lookups instead of adding an index An index on execution_run_outputs would make the library unopenable by older builds, whose schema check refuses any extra index there. The migration and the schema-check allowance are reverted, so the history schema is exactly main's. stored_run_prompts instead looks up a batch of files with one parameterized filename IN (...) query per chunk of at most 500 names and matches in memory; the scan computes it once per scan batch. * fix(ingestion): address the #344 re-review (seal race, linked roots) - A sealed history re-check that can't be scheduled yet (the image's own pass still pending, or the pending set full) is kept due and scheduled when a pending image finishes, instead of being dropped. - History matches a file through ComfyUI's spelling of its output folder first, then with both sides resolved, so a symlinked or junction output folder still attributes the scan's resolved paths; waiting files are keyed by their resolved path. - The scan makes its batch history lookup only when a file needs it. - set_current_prompt pushes to the queue last, so marking a failed execution never adds a second queue slot. - A keyless image outside the output folder doesn't take a waiting slot.
This commit is contained in:
Vito
@@ -1,5 +1,6 @@
|
||||
"""Deterministic nodes and queue controls installed only in an owned test host."""
|
||||
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
|
||||
@@ -255,3 +256,50 @@ async def cache_mode(request: web.Request) -> web.Response:
|
||||
raise web.HTTPBadRequest()
|
||||
_CACHE_MODE = mode
|
||||
return web.json_response({"configured": True})
|
||||
|
||||
|
||||
def _plugin_module(suffix: str):
|
||||
"""A loaded PromptManager module (ComfyUI names the package by its path)."""
|
||||
import sys
|
||||
|
||||
for name, module in list(sys.modules.items()):
|
||||
if name.endswith("promptmanager." + suffix):
|
||||
return module
|
||||
raise web.HTTPServiceUnavailable(text=f"PromptManager {suffix} not loaded")
|
||||
|
||||
|
||||
@PromptServer.instance.routes.post("/pm_execution_fixture/monitor")
|
||||
async def image_monitor(request: web.Request) -> web.Response:
|
||||
"""Start or stop PromptManager's image monitor on the host output folder.
|
||||
|
||||
The managed host's config disables monitoring; only the attribution cases
|
||||
(#342) turn it on, and they turn it off again.
|
||||
"""
|
||||
enabled = (await request.json())["enabled"]
|
||||
if type(enabled) is not bool:
|
||||
raise web.HTTPBadRequest()
|
||||
import folder_paths
|
||||
|
||||
monitor_module = _plugin_module("utils.image_monitor")
|
||||
config = _plugin_module("py.config").GalleryConfig
|
||||
monitor = monitor_module._monitor_instance
|
||||
if monitor is None:
|
||||
database = _plugin_module("database.operations").PromptDatabase()
|
||||
tracker = _plugin_module("utils.prompt_tracker").get_prompt_tracker(database)
|
||||
monitor = monitor_module.get_image_monitor(database, tracker)
|
||||
if enabled:
|
||||
config.MONITORING_ENABLED = True
|
||||
monitor.start_monitoring([folder_paths.get_output_directory()])
|
||||
else:
|
||||
monitor.stop_monitoring()
|
||||
config.MONITORING_ENABLED = False
|
||||
return web.json_response({"monitoring": monitor.is_monitoring})
|
||||
|
||||
|
||||
@PromptServer.instance.routes.get("/pm_execution_fixture/monitor")
|
||||
async def image_monitor_state(request: web.Request) -> web.Response:
|
||||
"""File names of keyless images waiting for their run's recorded outputs."""
|
||||
monitor = _plugin_module("utils.image_monitor")._monitor_instance
|
||||
handler = getattr(monitor, "handler", None)
|
||||
waiting = handler.awaiting_history() if handler is not None else []
|
||||
return web.json_response({"awaiting": [os.path.basename(p) for p in waiting]})
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
// #342 on the real pinned host: the image monitor links an output only to a prompt
|
||||
// that provably produced it. An unknown node (PMHistoryCached) stands in for the
|
||||
// maintainer's LLM; KSampler is the fixture's deterministic stand-in.
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import { execFile } from 'node:child_process';
|
||||
import { promisify } from 'node:util';
|
||||
import { fileURLToPath } from 'node:url';
|
||||
import { expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs';
|
||||
import { test } from './execution-fixture.mjs';
|
||||
|
||||
const run = promisify(execFile);
|
||||
const reader = fileURLToPath(new URL('./read_execution_history.py', import.meta.url));
|
||||
|
||||
async function readLinks(databasePath, filenames, texts = []) {
|
||||
const { stdout } = await run(process.env.PM_TEST_PYTHON || 'python',
|
||||
[reader, databasePath, '--links', ...filenames, '--texts', ...texts], { timeout: 5000, maxBuffer: 1024 * 1024 });
|
||||
return JSON.parse(stdout);
|
||||
}
|
||||
|
||||
// kind: the positive node; text: a typed editor or 'llm' (linked from the unknown
|
||||
// node); prepend: 'llm' links the unknown node into prepend_text.
|
||||
function graph({ kind = 'PromptManagerPair', text = '', prepend = '', llm = '', seed = 1, prefix }) {
|
||||
const linked = value => value === 'llm' ? ['11', 0] : value;
|
||||
const nodes = {
|
||||
'10': { class_type: 'PMHistoryInputs', inputs: {} },
|
||||
'20': { class_type: kind, inputs: { text: linked(text), clip: ['10', 0], append_text: '',
|
||||
...(prepend ? { prepend_text: linked(prepend) } : {}),
|
||||
...(kind === 'PromptManagerPair' ? { negative_text: '' } : {}) } },
|
||||
'30': { class_type: 'CLIPTextEncode', inputs: { text: '', clip: ['10', 0] } },
|
||||
'40': { class_type: 'KSampler', inputs: { model: ['10', 1], latent_image: ['10', 2],
|
||||
positive: ['20', 0], negative: kind === 'PromptManagerPair' ? ['20', 1] : ['30', 0],
|
||||
seed, steps: 1, cfg: 1, sampler_name: 'euler', scheduler: 'normal', denoise: 1 } },
|
||||
'60': { class_type: 'PMHistoryImages', inputs: { value: 64, batch: 1, latent: ['40', 0] } },
|
||||
'70': { class_type: 'SaveImage', inputs: { images: ['60', 0], filename_prefix: prefix } },
|
||||
};
|
||||
if (llm) nodes['11'] = { class_type: 'PMHistoryCached', inputs: { value: llm } };
|
||||
return nodes;
|
||||
}
|
||||
|
||||
async function saved(request, prompt, history = {}) {
|
||||
const id = randomUUID();
|
||||
const response = await request.post('/prompt', { data: { prompt, prompt_id: id } });
|
||||
expect(response.ok()).toBe(true);
|
||||
let item;
|
||||
await expect.poll(async () => {
|
||||
item = (await (await request.get(`/history/${id}`)).json())[id];
|
||||
return item?.status?.messages?.some(([kind]) => kind === 'execution_success') || false;
|
||||
}, { timeout: 15000 }).toBe(true);
|
||||
const [image] = item.outputs['70'].images;
|
||||
expect(image.type).toBe('output');
|
||||
history.item = item;
|
||||
return image.filename;
|
||||
}
|
||||
|
||||
async function linkedTo(execution, filename, expected) {
|
||||
await expect.poll(async () => (await readLinks(execution.databasePath, [filename])).links[filename],
|
||||
{ timeout: 20000, message: `${filename} is linked by the image monitor` }).toBe(expected);
|
||||
}
|
||||
|
||||
async function monitoring(execution, enabled) {
|
||||
const response = await execution.request.post('/pm_execution_fixture/monitor', { data: { enabled } });
|
||||
expect(response.ok()).toBe(true);
|
||||
expect((await response.json()).monitoring).toBe(enabled);
|
||||
}
|
||||
|
||||
async function cacheMode(execution, mode) {
|
||||
expect((await execution.request.post('/pm_execution_fixture/cache', { data: { mode } })).ok()).toBe(true);
|
||||
}
|
||||
|
||||
test.beforeEach(async ({ execution }) => {
|
||||
expect((await execution.request.post('/pm_execution_fixture/metadata', { data: { disabled: false } })).ok()).toBe(true);
|
||||
await cacheMode(execution, 'CLASSIC'); // keeps the re-run's PromptManager node cached
|
||||
await monitoring(execution, true);
|
||||
});
|
||||
test.afterEach(async ({ execution }) => {
|
||||
await monitoring(execution, false);
|
||||
await cacheMode(execution, 'RAM_PRESSURE');
|
||||
});
|
||||
|
||||
test('an output whose positive is only linked prepend text stays unlinked, never given an earlier prompt', async ({ execution }, testInfo) => {
|
||||
const tag = `${testInfo.project.name} ${randomUUID()}`;
|
||||
const earlier = `a woman holding a coffee cup ${tag}`;
|
||||
const prepended = `llm prepend landscape ${tag}`, llmText = `llm text landscape ${tag}`;
|
||||
|
||||
// A typed prompt runs first: before #342 the next unidentifiable image took it.
|
||||
const first = await saved(execution.request, graph({ text: earlier, prefix: 'misslink/earlier' }));
|
||||
await linkedTo(execution, first, earlier);
|
||||
|
||||
// The maintainer's workflow: empty editor, the generated text arrives in prepend.
|
||||
const orphan = await saved(execution.request,
|
||||
graph({ prepend: 'llm', llm: prepended, prefix: 'misslink/prepend' }));
|
||||
|
||||
// An LLM feeding the editor itself is stored and links through its run key; the
|
||||
// monitor works in creation order, so the prepend image was processed first.
|
||||
const llmGraph = graph({ kind: 'PromptManager', text: 'llm', llm: llmText, prefix: 'misslink/text' });
|
||||
const proven = await saved(execution.request, llmGraph);
|
||||
await linkedTo(execution, proven, llmText);
|
||||
const state = await readLinks(execution.databasePath, [orphan], [prepended]);
|
||||
expect(state).toEqual({ links: { [orphan]: null }, stored: { [prepended]: false } });
|
||||
|
||||
// Only the sampler seed changes: the PromptManager node is served from the cache
|
||||
// without running (the host says so), so no queue entry exists for this run and
|
||||
// only the snapshot's run match can link it.
|
||||
const rerunHistory = {};
|
||||
const rerun = await saved(execution.request, { ...llmGraph, '40': { ...llmGraph['40'],
|
||||
inputs: { ...llmGraph['40'].inputs, seed: 2 } } }, rerunHistory);
|
||||
const cached = rerunHistory.item.status.messages.find(([kind]) => kind === 'execution_cached');
|
||||
expect(cached?.[1].nodes).toEqual(expect.arrayContaining(['11', '20']));
|
||||
await linkedTo(execution, rerun, llmText);
|
||||
});
|
||||
|
||||
test('an output saved without metadata links once its recorded run names a stored prompt', async ({ execution }, testInfo) => {
|
||||
const llmText = `bare output landscape ${testInfo.project.name} ${randomUUID()}`;
|
||||
const request = execution.request;
|
||||
expect((await request.post('/pm_execution_fixture/metadata', { data: { disabled: true } })).ok()).toBe(true);
|
||||
try {
|
||||
// No embedded graph: no run key can prove the tracker entry, so only the
|
||||
// run's recorded history can, and only once the run's outputs are sealed.
|
||||
// The gate after SaveImage holds the run open while the monitor looks.
|
||||
const prompt = graph({ kind: 'PromptManager', text: 'llm', llm: llmText, prefix: 'misslink/bare' });
|
||||
prompt['80'] = { class_type: 'PMHistoryImageOutput', inputs: { images: ['70', 0], mode: 'gate' } };
|
||||
expect((await request.post('/pm_execution_fixture/gate/hold')).ok()).toBe(true);
|
||||
const id = randomUUID();
|
||||
try {
|
||||
expect((await request.post('/prompt', { data: { prompt, prompt_id: id } })).ok()).toBe(true);
|
||||
let waiting = [];
|
||||
await expect.poll(async () => {
|
||||
waiting = (await (await request.get('/pm_execution_fixture/monitor')).json()).awaiting;
|
||||
return waiting.filter(name => name.startsWith('bare_')).length;
|
||||
}, { timeout: 20000, message: 'the monitor processed the bare output while its run was open' }).toBe(1);
|
||||
const [bare] = waiting.filter(name => name.startsWith('bare_'));
|
||||
expect((await readLinks(execution.databasePath, [bare])).links[bare]).toBeNull();
|
||||
expect((await request.post('/pm_execution_fixture/gate/release')).ok()).toBe(true);
|
||||
await linkedTo(execution, bare, llmText);
|
||||
expect((await (await request.get('/pm_execution_fixture/monitor')).json()).awaiting).not.toContain(bare);
|
||||
} finally {
|
||||
expect((await request.post('/pm_execution_fixture/gate/release')).ok()).toBe(true);
|
||||
}
|
||||
} finally {
|
||||
expect((await request.post('/pm_execution_fixture/metadata', { data: { disabled: false } })).ok()).toBe(true);
|
||||
}
|
||||
});
|
||||
@@ -1,6 +1,7 @@
|
||||
"""Bounded read-only SQLite inspection for the owned execution fixture."""
|
||||
|
||||
import argparse
|
||||
from contextlib import closing
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sqlite3
|
||||
@@ -89,14 +90,48 @@ def read_runs(database: str, job_id: str, session_id: str | None = None) -> dict
|
||||
}
|
||||
|
||||
|
||||
def read_links(database: str, filenames: list, texts: list) -> dict:
|
||||
"""Each output file's linked prompt text (None: unlinked), and which texts
|
||||
are stored as prompts (#342)."""
|
||||
uri = Path(database).resolve().as_uri() + "?mode=ro"
|
||||
with closing(sqlite3.connect(uri, uri=True, timeout=0.25)) as connection:
|
||||
connection.execute("PRAGMA query_only=ON")
|
||||
links = {}
|
||||
for name in filenames[:64]:
|
||||
rows = connection.execute(
|
||||
"SELECT p.text FROM generated_images g JOIN prompts p "
|
||||
"ON p.id = g.prompt_id WHERE g.filename = ? LIMIT 2",
|
||||
(name,),
|
||||
).fetchall()
|
||||
values = [row[0] for row in rows]
|
||||
links[name] = values[0] if len(values) == 1 else None
|
||||
if len(values) > 1:
|
||||
links[name] = "<several links>"
|
||||
stored = {
|
||||
text: connection.execute(
|
||||
"SELECT 1 FROM prompts WHERE text = ? LIMIT 1", (text,)
|
||||
).fetchone()
|
||||
is not None
|
||||
for text in texts[:64]
|
||||
}
|
||||
return {"links": links, "stored": stored}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("database")
|
||||
parser.add_argument("job_id")
|
||||
parser.add_argument("job_id", nargs="?")
|
||||
parser.add_argument("--session-id")
|
||||
parser.add_argument("--links", nargs="*", default=None)
|
||||
parser.add_argument("--texts", nargs="*", default=[])
|
||||
arguments = parser.parse_args()
|
||||
print(
|
||||
json.dumps(
|
||||
read_runs(arguments.database, arguments.job_id, arguments.session_id)
|
||||
if arguments.links is not None:
|
||||
print(
|
||||
json.dumps(read_links(arguments.database, arguments.links, arguments.texts))
|
||||
)
|
||||
else:
|
||||
print(
|
||||
json.dumps(
|
||||
read_runs(arguments.database, arguments.job_id, arguments.session_id)
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
@@ -2466,6 +2466,7 @@ class TestActivityJobs(AdminAPITestCase):
|
||||
ambiguous=None,
|
||||
untraced=None,
|
||||
records=None,
|
||||
history=None,
|
||||
):
|
||||
key = Path(media_file).relative_to(self.output_dir).as_posix()
|
||||
outcome = names[key]
|
||||
|
||||
+17
-13
@@ -12,6 +12,7 @@ from database.operations import PromptDatabase
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.image_monitor import ImageGenerationHandler
|
||||
from utils.prompt_tracker import PromptTracker
|
||||
from utils.run_keys import prompt_node_key
|
||||
|
||||
LOADER = {
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
@@ -103,8 +104,11 @@ class TestMetadataLinking(LinkingTestCase):
|
||||
batch_graph = graph()
|
||||
batch_graph["134"]["inputs"]["text"] = ["7", 0]
|
||||
batch_graph["7"] = {"class_type": "PromptSearchList", "inputs": {}}
|
||||
self.tracker.set_current_prompt("item one", {"prompt_id": first})
|
||||
self.tracker.set_current_prompt("item two", {"prompt_id": second})
|
||||
key = prompt_node_key(batch_graph, "134")
|
||||
for text, prompt_id in (("item one", first), ("item two", second)):
|
||||
self.tracker.set_current_prompt(
|
||||
text, {"prompt_id": prompt_id}, run_key=key, run_id="run-1"
|
||||
)
|
||||
|
||||
self.assertEqual(self._process({"prompt": batch_graph}), first)
|
||||
self.assertEqual(self._process({"prompt": batch_graph}), second)
|
||||
@@ -122,9 +126,10 @@ class TestMetadataLinking(LinkingTestCase):
|
||||
self.assertEqual(
|
||||
self.handler._find_prompt_from_metadata(metadata), ("unknown", None)
|
||||
)
|
||||
latest = self._save("most recent prompt") # recency fallback picks this
|
||||
self.assertEqual(self._process(metadata), latest)
|
||||
self.assertNotEqual(latest, negative)
|
||||
self._save("most recent prompt") # never picked by recency (#342)
|
||||
self.assertIsNone(self._process(metadata))
|
||||
self.assertEqual(self.linked, [])
|
||||
self.assertIsNotNone(negative)
|
||||
|
||||
def test_workflow_fallback_still_used_without_an_api_graph(self):
|
||||
positive = self._save("from workflow only")
|
||||
@@ -200,12 +205,14 @@ class TestForeignWorkflows(LinkingTestCase):
|
||||
self.handler._find_prompt_from_metadata(metadata), ("foreign", None)
|
||||
)
|
||||
|
||||
def test_missing_metadata_is_unknown_and_uses_the_queue(self):
|
||||
def test_missing_metadata_is_unknown_and_left_unlinked(self):
|
||||
first = self._queue_batch_item("item one")
|
||||
self.assertEqual(
|
||||
self.handler._find_prompt_from_metadata(None), ("unknown", None)
|
||||
)
|
||||
self.assertEqual(self._process(None), first)
|
||||
# Nothing proves which run made it (#342): unlinked, queue untouched
|
||||
self.assertIsNone(self._process(None))
|
||||
self.assertEqual(self.tracker.pop_next_prompt()["id"], first)
|
||||
|
||||
def test_tri_state_for_linked_unknown_and_foreign(self):
|
||||
positive = self._save("a cat")
|
||||
@@ -224,14 +231,11 @@ class TestForeignWorkflows(LinkingTestCase):
|
||||
("unknown", None),
|
||||
)
|
||||
|
||||
def test_unknown_image_without_any_context_falls_back_to_the_latest_prompt(self):
|
||||
def test_unknown_image_is_never_given_the_latest_prompt(self):
|
||||
self._save("older")
|
||||
latest = self._save("latest")
|
||||
self.assertEqual(self._process(None), latest)
|
||||
|
||||
def test_recent_prompt_fallback_handles_an_empty_database(self):
|
||||
self.assertIsNone(self.handler._get_fallback_prompt())
|
||||
self._save("latest")
|
||||
self.assertIsNone(self._process(None))
|
||||
self.assertIsNone(self._process({"prompt": graph("not saved")}))
|
||||
|
||||
def test_metadata_lookup_errors_are_unknown(self):
|
||||
self._save("a cat")
|
||||
|
||||
@@ -326,7 +326,7 @@ class TestHandlerScheduling(unittest.TestCase):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.tmp.cleanup)
|
||||
self.tracker = MagicMock()
|
||||
self.tracker.get_current_prompt.return_value = {"id": 7}
|
||||
self.tracker.keyed_snapshot.return_value = {"key": ({"id": 7}, "run", 1)}
|
||||
self.handler = im_mod.ImageGenerationHandler(MagicMock(), self.tracker)
|
||||
self.handler.processing_delay = 0
|
||||
self.handler.settle_interval = 0.01
|
||||
@@ -354,7 +354,14 @@ class TestHandlerScheduling(unittest.TestCase):
|
||||
self.handler.on_created(FakeEvent(os.path.join(self.tmp.name, "c.txt")))
|
||||
self.handler.on_created(FakeEvent(self.tmp.name, is_directory=True))
|
||||
self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.assertEqual(self.processed, [(path, {"id": 7})])
|
||||
self.assertEqual(self.processed, [(path, {"key": ({"id": 7}, "run", 1)})])
|
||||
|
||||
def test_a_failed_snapshot_still_schedules_the_image(self):
|
||||
self.tracker.keyed_snapshot.side_effect = RuntimeError("tracker gone")
|
||||
path = self._image("d.png")
|
||||
self.handler.on_created(FakeEvent(path))
|
||||
self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.assertEqual(self.processed, [(path, None)])
|
||||
|
||||
def test_images_are_processed_in_creation_order(self):
|
||||
paths = [self._image(f"{i}.png") for i in range(5)]
|
||||
|
||||
@@ -0,0 +1,662 @@
|
||||
"""#342: an image is linked only to a prompt that provably produced it.
|
||||
|
||||
The maintainer's run: an LLM writes the prompt into a PromptManagerPair's
|
||||
``prepend_text`` while its ``text`` editor is empty. Prepend and append text are
|
||||
never stored as the prompt (product decision), so nothing is saved, and the
|
||||
image must stay unlinked instead of taking the newest prompt in the database or
|
||||
a stale tracker entry.
|
||||
"""
|
||||
|
||||
import copy
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from prompt_manager import PromptManager
|
||||
from prompt_manager_pair import PromptManagerPair
|
||||
from prompt_manager_text import PromptManagerText
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.image_monitor import ImageGenerationHandler
|
||||
from utils.prompt_tracker import PromptTracker
|
||||
from utils.run_keys import image_run_keys, prompt_node_key
|
||||
from utils.usage_tracking import (
|
||||
PendingFirstUse,
|
||||
queued_pair_runs,
|
||||
queued_prompt_runs,
|
||||
)
|
||||
|
||||
LLM_TEXT = "Photorealistic aerial landscape of misty mountains at dawn"
|
||||
EVIDENCE_KEY = "comfyassets_prompt_pairs"
|
||||
|
||||
|
||||
def issue_graph(text="", seed=1):
|
||||
"""The workflow from #342 (subgraph ids kept), PromptManagerPair positive."""
|
||||
return {
|
||||
"18:35": {
|
||||
"class_type": "PromptManagerPair",
|
||||
"inputs": {
|
||||
"text": text,
|
||||
"negative_text": "",
|
||||
"prepend_text": ["18:13", 0],
|
||||
"append_text": "",
|
||||
"category": "",
|
||||
"tags": "",
|
||||
"clip": ["4", 1],
|
||||
},
|
||||
},
|
||||
"18:13": {"class_type": "PreviewAny", "inputs": {"source": ["18:16", 0]}},
|
||||
"18:16": {
|
||||
"class_type": "ComfySwitchNode",
|
||||
"inputs": {
|
||||
"switch": True,
|
||||
"on_false": ["18:17", 0],
|
||||
"on_true": ["18:11", 0],
|
||||
},
|
||||
},
|
||||
"18:11": {
|
||||
"class_type": "TextGenerate",
|
||||
"inputs": {"prompt": ["18:17", 0], "seed": 7},
|
||||
},
|
||||
"18:17": {
|
||||
"class_type": "PrimitiveStringMultiline",
|
||||
"inputs": {"value": "Describe an aerial mountain landscape"},
|
||||
},
|
||||
"18:5": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": "blurry", "clip": ["4", 1]},
|
||||
},
|
||||
"4": {
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
"inputs": {"ckpt_name": "moss.safetensors"},
|
||||
},
|
||||
"3": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"model": ["4", 0],
|
||||
"positive": ["18:35", 0],
|
||||
"negative": ["18:5", 0],
|
||||
"latent_image": ["6", 0],
|
||||
"seed": seed,
|
||||
},
|
||||
},
|
||||
"6": {"class_type": "EmptyLatentImage", "inputs": {"batch_size": 1}},
|
||||
}
|
||||
|
||||
|
||||
def llm_text_graph(seed=1):
|
||||
"""An LLM feeding a PromptManager's ``text``: stored, but not in the metadata."""
|
||||
graph = issue_graph(seed=seed)
|
||||
graph["18:35"] = {
|
||||
"class_type": "PromptManager",
|
||||
"inputs": {"text": ["18:13", 0], "clip": ["4", 1]},
|
||||
}
|
||||
return graph
|
||||
|
||||
|
||||
def pair_evidence(text="", prepend=LLM_TEXT):
|
||||
"""The execution record PromptManagerPair embeds in the saved image."""
|
||||
return {
|
||||
"version": 1,
|
||||
"nodes": {
|
||||
"18:35": {
|
||||
"records": [
|
||||
{
|
||||
"text": text,
|
||||
"negative_text": "",
|
||||
"positive_text": " ".join(p for p in (prepend, text) if p),
|
||||
"negative_effective_text": "",
|
||||
"prepend_text": prepend,
|
||||
"append_text": "",
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def make_tracker(db):
|
||||
"""A real tracker without its cleanup thread or config import."""
|
||||
with mock.patch.object(PromptTracker, "__init__", lambda self, db: None):
|
||||
tracker = PromptTracker.__new__(PromptTracker)
|
||||
tracker.logger = logging.getLogger("test.prompt_attribution.tracker")
|
||||
tracker.db_manager = db
|
||||
tracker._local = threading.local()
|
||||
tracker.active_prompts = {}
|
||||
tracker.lock = threading.Lock()
|
||||
tracker._prompt_queue = []
|
||||
tracker._queue_lock = threading.Lock()
|
||||
tracker.cleanup_interval = 300
|
||||
tracker.prompt_timeout = 600
|
||||
return tracker
|
||||
|
||||
|
||||
class DatabaseCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.tmp.cleanup)
|
||||
self.addCleanup(assert_closed, self, self.tmp.name)
|
||||
self.db = PromptDatabase(os.path.join(self.tmp.name, "prompts.db"))
|
||||
self.addCleanup(self.db.close_all)
|
||||
self.tracker = make_tracker(self.db)
|
||||
|
||||
def save(self, text):
|
||||
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
||||
|
||||
|
||||
class NodeCase(DatabaseCase):
|
||||
def node(self, cls):
|
||||
node = cls.__new__(cls)
|
||||
node.db = self.db
|
||||
node.logger = logging.getLogger("test.prompt_attribution.node")
|
||||
node.prompt_tracker = self.tracker
|
||||
node.comfyui_integration = mock.Mock()
|
||||
node._inject_lora_trigger_words = lambda text: text
|
||||
node.pending_first_use = PendingFirstUse()
|
||||
node.pending_pair_first_use = PendingFirstUse()
|
||||
return node
|
||||
|
||||
def clip(self):
|
||||
clip = mock.Mock()
|
||||
clip.tokenize.side_effect = lambda text: text
|
||||
clip.encode_from_tokens_scheduled.side_effect = lambda t: ("cond", t)
|
||||
return clip
|
||||
|
||||
def count_rows(self, table):
|
||||
conn = self.db.model.get_connection()
|
||||
return conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]
|
||||
|
||||
|
||||
class TestEmptyEditorSavesNothing(NodeCase):
|
||||
"""Prepend/append are never the stored prompt; an empty editor saves nothing."""
|
||||
|
||||
def test_pair_with_only_prepend_text_saves_nothing_and_logs_info(self):
|
||||
graph = issue_graph()
|
||||
with self.assertLogs("test.prompt_attribution.node", "INFO") as logs:
|
||||
result = self.node(PromptManagerPair).encode_prompt(
|
||||
self.clip(),
|
||||
"",
|
||||
negative_text="",
|
||||
prepend_text=LLM_TEXT,
|
||||
prompt=graph,
|
||||
unique_id="18:35",
|
||||
)
|
||||
self.assertEqual(result[2], LLM_TEXT, "the composed text is still encoded")
|
||||
self.assertEqual(self.count_rows("prompts"), 0)
|
||||
self.assertEqual(self.count_rows("prompt_pairs"), 0)
|
||||
self.assertFalse(
|
||||
any(record.levelno >= logging.WARNING for record in logs.records),
|
||||
[record.getMessage() for record in logs.records],
|
||||
)
|
||||
self.assertTrue(any("Prompt text is empty" in line for line in logs.output))
|
||||
|
||||
def test_prompt_manager_and_text_nodes_with_only_prepend_save_nothing(self):
|
||||
graph = issue_graph()
|
||||
graph["18:35"] = {
|
||||
"class_type": "PromptManager",
|
||||
"inputs": {"text": "", "prepend_text": ["18:13", 0], "clip": ["4", 1]},
|
||||
}
|
||||
self.node(PromptManager).encode_prompt(
|
||||
self.clip(), "", prepend_text=LLM_TEXT, prompt=graph, unique_id="18:35"
|
||||
)
|
||||
graph["18:35"]["class_type"] = "PromptManagerText"
|
||||
self.node(PromptManagerText).process_text(
|
||||
"", prepend_text=LLM_TEXT, prompt=graph, unique_id="18:35"
|
||||
)
|
||||
self.assertEqual(self.count_rows("prompts"), 0)
|
||||
|
||||
def test_empty_editor_marks_its_run_key_so_no_older_entry_is_used(self):
|
||||
graph = issue_graph()
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
older = self.save("a woman holding a coffee cup")
|
||||
self.tracker.set_current_prompt(
|
||||
"a woman holding a coffee cup",
|
||||
{"prompt_id": older},
|
||||
push_to_queue=True,
|
||||
run_key=key,
|
||||
run_id="run-1",
|
||||
)
|
||||
self.node(PromptManagerPair).encode_prompt(
|
||||
self.clip(), "", prepend_text=LLM_TEXT, prompt=graph, unique_id="18:35"
|
||||
)
|
||||
latest, _run, _count = self.tracker.keyed_snapshot()[key]
|
||||
self.assertTrue(latest.get("placeholder"))
|
||||
self.assertIsNone(self.tracker.pop_for_key(key, "another run"))
|
||||
popped = self.tracker.pop_for_key(key)
|
||||
self.assertEqual(popped["id"], older, "an earlier queued item stays first")
|
||||
self.assertTrue(self.tracker.pop_for_key(key).get("placeholder"))
|
||||
|
||||
def test_empty_negative_node_is_not_marked(self):
|
||||
graph = issue_graph(text="a typed positive")
|
||||
graph["18:5"] = {
|
||||
"class_type": "PromptManager",
|
||||
"inputs": {"text": "", "clip": ["4", 1]},
|
||||
}
|
||||
self.node(PromptManager).encode_prompt(
|
||||
self.clip(), "", prompt=graph, unique_id="18:5"
|
||||
)
|
||||
self.assertEqual(self.tracker.keyed_snapshot(), {})
|
||||
|
||||
def test_queue_hook_counts_nothing_for_an_empty_editor(self):
|
||||
graph = issue_graph()
|
||||
self.assertEqual(queued_prompt_runs(graph), [])
|
||||
self.assertEqual(queued_pair_runs(graph), [])
|
||||
|
||||
def test_typed_text_with_prepend_still_stores_only_the_editor_text(self):
|
||||
graph = issue_graph(text="mountains")
|
||||
self.node(PromptManagerPair).encode_prompt(
|
||||
self.clip(),
|
||||
"mountains",
|
||||
prepend_text=LLM_TEXT,
|
||||
prompt=graph,
|
||||
unique_id="18:35",
|
||||
)
|
||||
self.assertIsNotNone(
|
||||
self.db.get_prompt_by_hash(generate_prompt_hash("mountains"))
|
||||
)
|
||||
self.assertIsNone(self.db.get_prompt_by_hash(generate_prompt_hash(LLM_TEXT)))
|
||||
context, _run, _count = self.tracker.keyed_snapshot()[
|
||||
prompt_node_key(graph, "18:35")
|
||||
]
|
||||
self.assertEqual(context["text"], f"{LLM_TEXT} mountains")
|
||||
|
||||
|
||||
class MonitorCase(DatabaseCase):
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
fd, self.image = tempfile.mkstemp(suffix=".png", dir=self.tmp.name)
|
||||
os.close(fd)
|
||||
self.handler = ImageGenerationHandler(self.db, self.tracker)
|
||||
self.linked = []
|
||||
self.handler.link_image_to_prompt = lambda path, prompt, meta: (
|
||||
self.linked.append(prompt["id"])
|
||||
)
|
||||
|
||||
def process(self, metadata, snapshot=None):
|
||||
before = len(self.linked)
|
||||
with mock.patch.object(
|
||||
self.handler.metadata_extractor, "extract_metadata", return_value=metadata
|
||||
):
|
||||
self.handler.process_new_image(self.image, prompt_snapshot=snapshot)
|
||||
return self.linked[-1] if len(self.linked) > before else None
|
||||
|
||||
|
||||
class TestIssueReproduction(MonitorCase):
|
||||
def test_llm_prepend_image_is_left_unlinked_not_given_the_newest_prompt(self):
|
||||
"""Failing before #342: linked to the newest prompt in the database."""
|
||||
self.save("a woman holding a coffee cup")
|
||||
metadata = {"prompt": issue_graph(), EVIDENCE_KEY: pair_evidence()}
|
||||
with self.assertLogs("prompt_manager.image_monitor", "INFO") as logs:
|
||||
self.assertIsNone(self.process(metadata))
|
||||
self.assertTrue(any("left unlinked" in line for line in logs.output))
|
||||
|
||||
def test_a_stale_entry_of_another_workflow_is_never_used(self):
|
||||
"""Failing before #342: the live tracker's entry was taken."""
|
||||
stale = self.save("a woman holding a coffee cup")
|
||||
other = issue_graph(text="a woman holding a coffee cup")
|
||||
self.tracker.set_current_prompt(
|
||||
"a woman holding a coffee cup",
|
||||
{"prompt_id": stale},
|
||||
push_to_queue=True,
|
||||
run_key=prompt_node_key(other, "18:35"),
|
||||
run_id="run-0",
|
||||
)
|
||||
metadata = {"prompt": issue_graph(), EVIDENCE_KEY: pair_evidence()}
|
||||
self.assertIsNone(self.process(metadata, self.tracker.keyed_snapshot()))
|
||||
self.assertIsNone(self.process(metadata))
|
||||
self.assertEqual(len(self.tracker._prompt_queue), 1, "not consumed")
|
||||
|
||||
def test_an_entry_without_a_run_key_is_never_used(self):
|
||||
stale = self.save("previous run")
|
||||
self.tracker.set_current_prompt("previous run", {"prompt_id": stale})
|
||||
self.assertIsNone(self.process({"prompt": llm_text_graph()}))
|
||||
|
||||
def test_an_image_without_metadata_is_left_unlinked(self):
|
||||
"""Failing before #342: popped the queue, else took the newest prompt."""
|
||||
self.save("latest")
|
||||
self.assertIsNone(self.process(None))
|
||||
|
||||
def test_the_recency_fallback_is_gone(self):
|
||||
self.assertFalse(hasattr(self.handler, "_get_fallback_prompt"))
|
||||
|
||||
|
||||
class TestProvenRuns(MonitorCase):
|
||||
"""Entries the image's own graph proves still link (LLM into ``text``)."""
|
||||
|
||||
def track(self, graph, text, run_id="run-1", push=True):
|
||||
prompt_id = (
|
||||
self.save(text)
|
||||
if not self.db.get_prompt_by_hash(generate_prompt_hash(text))
|
||||
else self.db.get_prompt_by_hash(generate_prompt_hash(text))["id"]
|
||||
)
|
||||
self.tracker.set_current_prompt(
|
||||
text,
|
||||
{"prompt_id": prompt_id},
|
||||
push_to_queue=push,
|
||||
run_key=prompt_node_key(graph, "18:35"),
|
||||
run_id=run_id,
|
||||
)
|
||||
return prompt_id
|
||||
|
||||
def test_queued_entry_of_the_same_run_key_links(self):
|
||||
graph = llm_text_graph()
|
||||
prompt_id = self.track(graph, LLM_TEXT)
|
||||
self.assertEqual(self.process({"prompt": graph}), prompt_id)
|
||||
|
||||
def test_cached_node_rerun_links_through_the_snapshot(self):
|
||||
graph = llm_text_graph(seed=1)
|
||||
prompt_id = self.track(graph, LLM_TEXT)
|
||||
self.assertEqual(self.process({"prompt": graph}), prompt_id)
|
||||
# Only the sampler seed changed: the node was cached and did not run,
|
||||
# but its upstream (the run key) is unchanged.
|
||||
rerun = llm_text_graph(seed=2)
|
||||
self.assertEqual(
|
||||
prompt_node_key(rerun, "18:35"), prompt_node_key(graph, "18:35")
|
||||
)
|
||||
self.assertEqual(
|
||||
self.process({"prompt": rerun}, self.tracker.keyed_snapshot()), prompt_id
|
||||
)
|
||||
|
||||
def test_snapshot_needs_a_known_run(self):
|
||||
graph = llm_text_graph()
|
||||
self.track(graph, LLM_TEXT, run_id=None, push=False)
|
||||
self.assertIsNone(self.process({"prompt": graph}))
|
||||
|
||||
def test_batch_items_link_in_order_and_the_latest_alone_proves_nothing(self):
|
||||
graph = llm_text_graph()
|
||||
first = self.track(graph, "item one")
|
||||
second = self.track(graph, "item two")
|
||||
self.assertEqual(self.process({"prompt": graph}), first)
|
||||
self.assertEqual(self.process({"prompt": graph}), second)
|
||||
# Queue drained: two entries in one run, so the latest is not proof
|
||||
self.assertIsNone(self.process({"prompt": graph}))
|
||||
|
||||
def test_an_empty_batch_item_keeps_the_queue_aligned(self):
|
||||
graph = llm_text_graph()
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
first = self.track(graph, "item one")
|
||||
self.tracker.mark_empty(key, "run-1", push_to_queue=True)
|
||||
third = self.track(graph, "item three")
|
||||
self.assertEqual(self.process({"prompt": graph}), first)
|
||||
self.assertIsNone(self.process({"prompt": graph}))
|
||||
self.assertEqual(self.process({"prompt": graph}), third)
|
||||
|
||||
def test_a_newer_run_drops_stale_entries_of_the_same_key(self):
|
||||
graph = llm_text_graph()
|
||||
# run-1 saved no image, so its entry was never taken
|
||||
self.track(graph, "earlier run text", run_id="run-1")
|
||||
newer = self.track(graph, "newer run text", run_id="run-2")
|
||||
self.assertEqual(self.process({"prompt": graph}), newer)
|
||||
self.assertEqual(self.tracker._prompt_queue, [])
|
||||
|
||||
def test_a_later_run_queued_before_processing_keeps_its_own_entry(self):
|
||||
# The image appears (snapshot), the next queued run executes the same
|
||||
# node again before the monitor's delay is over, then the image is
|
||||
# processed: it takes its own run's entry, the later run keeps its own.
|
||||
graph = llm_text_graph()
|
||||
own = self.track(graph, "first run text", run_id="run-1")
|
||||
snapshot = self.tracker.keyed_snapshot()
|
||||
later = self.track(graph, "second run text", run_id="run-2")
|
||||
self.assertEqual(self.process({"prompt": graph}, snapshot), own)
|
||||
self.assertEqual(self.process({"prompt": graph}), later)
|
||||
|
||||
def test_no_entry_when_the_image_appeared_means_no_queue_match(self):
|
||||
graph = llm_text_graph()
|
||||
empty_snapshot = self.tracker.keyed_snapshot()
|
||||
self.track(graph, "queued after the image appeared")
|
||||
self.assertIsNone(self.process({"prompt": graph}, empty_snapshot))
|
||||
self.assertEqual(len(self.tracker._prompt_queue), 1)
|
||||
|
||||
def test_hosts_without_run_ids_still_link_queued_entries_in_order(self):
|
||||
graph = llm_text_graph()
|
||||
first = self.track(graph, "item one", run_id=None)
|
||||
second = self.track(graph, "item two", run_id=None)
|
||||
self.assertEqual(self.process({"prompt": graph}), first)
|
||||
self.assertEqual(self.process({"prompt": graph}), second)
|
||||
self.assertIsNone(self.process({"prompt": graph}))
|
||||
|
||||
def test_other_keys_stay_queued_for_their_own_images(self):
|
||||
mine, other = llm_text_graph(), copy.deepcopy(llm_text_graph())
|
||||
other["18:17"]["inputs"]["value"] = "Describe a beach"
|
||||
theirs = self.track(other, "a beach")
|
||||
prompt_id = self.track(mine, LLM_TEXT)
|
||||
self.assertEqual(self.process({"prompt": mine}), prompt_id)
|
||||
self.assertEqual(self.process({"prompt": other}), theirs)
|
||||
|
||||
def test_two_positive_prompt_nodes_are_not_guessed(self):
|
||||
graph = llm_text_graph()
|
||||
self.track(graph, LLM_TEXT)
|
||||
graph["30"] = {
|
||||
"class_type": "PromptManager",
|
||||
"inputs": {"text": ["18:13", 0], "clip": ["4", 1]},
|
||||
}
|
||||
graph["31"] = dict(graph["3"])
|
||||
graph["31"]["inputs"] = {**graph["3"]["inputs"], "positive": ["30", 0]}
|
||||
self.assertEqual(len(image_run_keys({"prompt": graph})), 2)
|
||||
self.assertIsNone(self.process({"prompt": graph}))
|
||||
|
||||
def test_metadata_still_wins_and_leaves_the_queue_alone(self):
|
||||
typed = self.save("a typed prompt")
|
||||
graph = issue_graph(text="a typed prompt")
|
||||
self.tracker.set_current_prompt(
|
||||
"x",
|
||||
{"prompt_id": typed},
|
||||
run_key=prompt_node_key(graph, "18:35"),
|
||||
run_id="r",
|
||||
)
|
||||
self.assertEqual(self.process({"prompt": graph}), typed)
|
||||
self.assertEqual(len(self.tracker._prompt_queue), 1)
|
||||
|
||||
|
||||
class TestRunKeys(unittest.TestCase):
|
||||
def test_key_covers_the_upstream_and_ignores_downstream_and_titles(self):
|
||||
graph = issue_graph()
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
changed = copy.deepcopy(graph)
|
||||
changed["3"]["inputs"]["seed"] = 99
|
||||
changed["18:35"]["_meta"] = {"title": "renamed"}
|
||||
self.assertEqual(prompt_node_key(changed, "18:35"), key)
|
||||
changed["18:11"]["inputs"]["seed"] = 8
|
||||
self.assertNotEqual(prompt_node_key(changed, "18:35"), key)
|
||||
|
||||
def test_unknowable_keys_are_none(self):
|
||||
graph = issue_graph()
|
||||
self.assertIsNone(prompt_node_key(None, "1"))
|
||||
self.assertIsNone(prompt_node_key(graph, None))
|
||||
self.assertIsNone(prompt_node_key(graph, "missing"))
|
||||
broken = copy.deepcopy(graph)
|
||||
broken["18:13"]["inputs"]["source"] = ["gone", 0]
|
||||
self.assertIsNone(prompt_node_key(broken, "18:35"))
|
||||
broken = copy.deepcopy(graph)
|
||||
broken["18:17"]["inputs"]["value"] = object()
|
||||
self.assertIsNone(prompt_node_key(broken, "18:35"))
|
||||
|
||||
def test_image_keys_read_string_graphs_and_reject_garbage(self):
|
||||
import json
|
||||
|
||||
graph = issue_graph()
|
||||
self.assertEqual(
|
||||
image_run_keys({"prompt": json.dumps(graph)}),
|
||||
{"18:35": prompt_node_key(graph, "18:35")},
|
||||
)
|
||||
self.assertEqual(image_run_keys({"prompt": "{not json"}), {})
|
||||
self.assertEqual(image_run_keys(None), {})
|
||||
self.assertEqual(image_run_keys({"prompt": []}), {})
|
||||
|
||||
def test_cyclic_and_huge_closures_have_no_key(self):
|
||||
from utils.prompt_graph import MAX_VISITS
|
||||
|
||||
cyclic = {
|
||||
"1": {"class_type": "A", "inputs": {"x": ["2", 0]}},
|
||||
"2": {"class_type": "B", "inputs": {"x": ["1", 0]}},
|
||||
}
|
||||
self.assertIsNotNone(prompt_node_key(cyclic, "1"))
|
||||
chain = {
|
||||
str(i): {"class_type": "N", "inputs": {"x": [str(i + 1), 0]}}
|
||||
for i in range(MAX_VISITS + 1)
|
||||
}
|
||||
chain[str(MAX_VISITS + 1)] = {"class_type": "N", "inputs": {}}
|
||||
self.assertIsNone(prompt_node_key(chain, "0"))
|
||||
|
||||
|
||||
class TestReviewFindings(MonitorCase, NodeCase):
|
||||
"""Final-review findings M1-M4 (#344)."""
|
||||
|
||||
def test_m1_a_positive_node_without_a_key_makes_the_image_unkeyable(self):
|
||||
graph = llm_text_graph()
|
||||
graph["30"] = {
|
||||
"class_type": "PromptManager",
|
||||
"inputs": {"text": ["missing", 0], "clip": ["4", 1]},
|
||||
}
|
||||
graph["31"] = {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {**graph["3"]["inputs"], "positive": ["30", 0]},
|
||||
}
|
||||
self.assertIsNone(prompt_node_key(graph, "30"))
|
||||
self.assertEqual(image_run_keys({"prompt": graph}), {})
|
||||
prompt_id = self.save(LLM_TEXT)
|
||||
self.tracker.set_current_prompt(
|
||||
LLM_TEXT,
|
||||
{"prompt_id": prompt_id},
|
||||
run_key=prompt_node_key(graph, "18:35"),
|
||||
run_id="run-1",
|
||||
)
|
||||
self.assertIsNone(self.process({"prompt": graph}))
|
||||
|
||||
def test_m2_links_are_exactly_comfyui_is_link(self):
|
||||
from utils.run_keys import _links
|
||||
|
||||
inputs = {
|
||||
"string_slot_int": ["7", 0],
|
||||
"string_slot_float": ["7", 0.0],
|
||||
"string_slot_bool": ["7", True], # bool is an int to ComfyUI too
|
||||
"int_source": [7, 0],
|
||||
"tuple": ("7", 0),
|
||||
"three": ["7", 0, 1],
|
||||
"slot_text": ["7", "0"],
|
||||
}
|
||||
self.assertEqual(sorted(_links(inputs)), ["7", "7", "7"])
|
||||
widget = issue_graph()
|
||||
widget["18:17"]["inputs"]["size"] = [512, 768] # a value, not a link
|
||||
self.assertIsNotNone(prompt_node_key(widget, "18:35"))
|
||||
floated = issue_graph()
|
||||
floated["18:35"]["inputs"]["prepend_text"] = ["18:13", 0.0]
|
||||
changed = copy.deepcopy(floated)
|
||||
changed["18:17"]["inputs"]["value"] = "another instruction"
|
||||
self.assertNotEqual(
|
||||
prompt_node_key(floated, "18:35"), prompt_node_key(changed, "18:35")
|
||||
)
|
||||
|
||||
def _failing_save(self, node):
|
||||
def explode(*args, **kwargs):
|
||||
raise RuntimeError("database is locked")
|
||||
|
||||
node.db = mock.Mock(wraps=self.db)
|
||||
for name in ("save_prompt", "save_prompt_pair", "get_prompt_by_hash"):
|
||||
setattr(node.db, name, mock.Mock(side_effect=explode))
|
||||
return node
|
||||
|
||||
def test_m3_a_failed_save_takes_over_its_run_key(self):
|
||||
graph = llm_text_graph()
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
older = self.save("previous run text")
|
||||
self.tracker.set_current_prompt(
|
||||
"previous run text", {"prompt_id": older}, run_key=key, run_id="run-1"
|
||||
)
|
||||
node = self._failing_save(self.node(PromptManager))
|
||||
with mock.patch("prompt_manager_base.executing_run_id", return_value="run-2"):
|
||||
node.encode_prompt(
|
||||
self.clip(), "new run text", prompt=graph, unique_id="18:35"
|
||||
)
|
||||
latest, run_id, _count = self.tracker.keyed_snapshot()[key]
|
||||
self.assertTrue(latest.get("placeholder"))
|
||||
self.assertEqual(run_id, "run-2")
|
||||
self.assertIsNone(self.process({"prompt": graph}))
|
||||
|
||||
def test_m3_a_failed_pair_save_takes_over_its_run_key(self):
|
||||
graph = issue_graph(text="typed")
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
node = self._failing_save(self.node(PromptManagerPair))
|
||||
with self.assertLogs("test.prompt_attribution.node", "WARNING"):
|
||||
node.encode_prompt(self.clip(), "typed", prompt=graph, unique_id="18:35")
|
||||
self.assertTrue(self.tracker.keyed_snapshot()[key][0].get("placeholder"))
|
||||
|
||||
def test_m3_a_failed_text_node_save_takes_over_its_run_key(self):
|
||||
graph = issue_graph()
|
||||
graph["18:35"] = {
|
||||
"class_type": "PromptManagerText",
|
||||
"inputs": {"text": "typed"},
|
||||
}
|
||||
graph["40"] = {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": ["18:35", 0], "clip": ["4", 1]},
|
||||
}
|
||||
graph["3"]["inputs"]["positive"] = ["40", 0]
|
||||
node = self._failing_save(self.node(PromptManagerText))
|
||||
with mock.patch.object(
|
||||
node, "_track_prompt_execution", side_effect=RuntimeError("boom")
|
||||
):
|
||||
node.process_text("typed", prompt=graph, unique_id="18:35")
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
self.assertTrue(self.tracker.keyed_snapshot()[key][0].get("placeholder"))
|
||||
|
||||
def test_m4_a_placeholder_never_falls_through_to_history(self):
|
||||
graph = llm_text_graph()
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
self.tracker.mark_empty(key, "run-1", push_to_queue=False)
|
||||
with mock.patch.object(
|
||||
self.handler, "_history_prompt", return_value=99
|
||||
) as history:
|
||||
self.assertIsNone(self.process({"prompt": graph})) # snapshot form
|
||||
self.tracker.mark_empty(key, "run-2", push_to_queue=True)
|
||||
self.assertIsNone(self.process({"prompt": graph})) # popped form
|
||||
history.assert_not_called()
|
||||
|
||||
|
||||
class TestRereviewFindings(MonitorCase, NodeCase):
|
||||
"""Re-review findings R4 and R5 (#344)."""
|
||||
|
||||
def test_r4_a_logging_failure_after_tracking_never_doubles_a_queue_slot(self):
|
||||
graph = llm_text_graph()
|
||||
key = prompt_node_key(graph, "18:35")
|
||||
node = self.node(PromptManager)
|
||||
logger = mock.Mock(wraps=self.tracker.logger)
|
||||
logger.debug.side_effect = RuntimeError("log handler broke")
|
||||
self.tracker.logger = logger
|
||||
node.encode_prompt(self.clip(), LLM_TEXT, prompt=graph, unique_id="18:35")
|
||||
queued = [ctx for ctx in self.tracker._prompt_queue if ctx["run_key"] == key]
|
||||
self.assertEqual(len(queued), 1)
|
||||
|
||||
def test_r5_a_keyless_image_outside_the_output_folder_does_not_wait(self):
|
||||
outside = tempfile.mkdtemp()
|
||||
self.addCleanup(lambda: __import__("shutil").rmtree(outside, True))
|
||||
stray = os.path.join(outside, "copied.png")
|
||||
with open(stray, "wb") as handle:
|
||||
handle.write(b"png")
|
||||
with mock.patch(
|
||||
"utils.image_monitor.host_output_directory", return_value=self.tmp.name
|
||||
):
|
||||
with mock.patch.object(
|
||||
self.handler.metadata_extractor, "extract_metadata", return_value=None
|
||||
):
|
||||
self.handler.process_new_image(stray)
|
||||
self.handler.process_new_image(self.image) # inside: waits
|
||||
self.assertEqual(
|
||||
self.handler.awaiting_history(),
|
||||
[os.path.normcase(os.path.realpath(self.image))],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,255 @@
|
||||
"""tools/review_image_links.py: a dry run by default, changes only named images."""
|
||||
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import unittest
|
||||
import unittest.mock
|
||||
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from tests.test_api_images import ImageAPITestCase
|
||||
from tests.test_prompt_attribution import issue_graph, pair_evidence
|
||||
from tools import review_image_links as tool
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
|
||||
def png(path, prompt=None, evidence=None):
|
||||
info = PngInfo()
|
||||
if prompt is not None:
|
||||
info.add_text("prompt", json.dumps(prompt))
|
||||
if evidence is not None:
|
||||
info.add_text("comfyassets_prompt_pairs", json.dumps(evidence))
|
||||
Image.new("RGB", (4, 4)).save(path, pnginfo=info)
|
||||
|
||||
|
||||
class ToolFixture(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.mkdtemp()
|
||||
self.addCleanup(shutil.rmtree, self.tmp, True)
|
||||
self.addCleanup(assert_closed, self, self.tmp)
|
||||
self.output = os.path.join(self.tmp, "output")
|
||||
os.makedirs(self.output)
|
||||
self.db_path = os.path.join(self.tmp, "prompts.db")
|
||||
db = PromptDatabase(self.db_path)
|
||||
try:
|
||||
save = lambda text: db.save_prompt( # noqa: E731
|
||||
text=text, prompt_hash=generate_prompt_hash(text)
|
||||
)
|
||||
self.coffee = save("a woman holding a coffee cup")
|
||||
self.typed = save("snowy peaks")
|
||||
# The six #342 images: only prepend text, linked by the old fallback
|
||||
self.guessed = []
|
||||
for n in range(9, 12):
|
||||
path = os.path.join(self.output, f"moss_{n:05}_.png")
|
||||
png(path, issue_graph(), pair_evidence())
|
||||
self.guessed.append(
|
||||
db.link_image_to_prompt(self.coffee, path, metadata={})
|
||||
)
|
||||
# A typed image linked to the wrong prompt, and a correct one
|
||||
wrong = os.path.join(self.output, "typed_wrong.png")
|
||||
png(wrong, issue_graph(text="snowy peaks"))
|
||||
self.wrong = db.link_image_to_prompt(self.coffee, wrong, metadata={})
|
||||
right = os.path.join(self.output, "typed_right.png")
|
||||
png(right, issue_graph(text="snowy peaks"))
|
||||
self.right = db.link_image_to_prompt(self.typed, right, metadata={})
|
||||
gone = os.path.join(self.output, "gone.png")
|
||||
self.gone = db.link_image_to_prompt(self.coffee, gone, metadata={})
|
||||
finally:
|
||||
db.close_all()
|
||||
|
||||
def run_tool(self, *args):
|
||||
out = io.StringIO()
|
||||
code = tool.main([self.db_path, "--output-dir", self.output, *args], out=out)
|
||||
return code, out.getvalue()
|
||||
|
||||
def links(self):
|
||||
db = PromptDatabase(self.db_path)
|
||||
try:
|
||||
rows = db.model.get_connection().execute(
|
||||
"SELECT id, prompt_id FROM generated_images"
|
||||
)
|
||||
return {row[0]: row[1] for row in rows}
|
||||
finally:
|
||||
db.close_all()
|
||||
|
||||
|
||||
class ReviewToolTests(ToolFixture):
|
||||
def test_dry_run_judges_and_changes_nothing(self):
|
||||
before = self.links()
|
||||
ids = f"{self.guessed[0]}-{self.gone}"
|
||||
code, out = self.run_tool("--ids", ids, "--json")
|
||||
self.assertEqual(code, 0)
|
||||
verdicts = {v["image_id"]: v for v in json.loads(out.split("--apply, for")[0])}
|
||||
for image_id in self.guessed:
|
||||
self.assertEqual(verdicts[image_id]["verdict"], "unlink")
|
||||
self.assertEqual(verdicts[image_id]["sources"], {})
|
||||
self.assertEqual(verdicts[self.wrong]["verdict"], "relink")
|
||||
self.assertEqual(verdicts[self.wrong]["target"], self.typed)
|
||||
self.assertEqual(verdicts[self.right]["verdict"], "keep")
|
||||
self.assertEqual(verdicts[self.gone]["verdict"], "missing")
|
||||
self.assertIn("Dry run: nothing was changed.", out)
|
||||
self.assertEqual(self.links(), before)
|
||||
self.assertEqual(
|
||||
sorted(os.listdir(self.tmp)), ["output", "prompts.db"], "no side files"
|
||||
)
|
||||
|
||||
def test_prompt_selection_and_text_report(self):
|
||||
code, out = self.run_tool("--prompt", str(self.coffee))
|
||||
self.assertEqual(code, 0)
|
||||
self.assertIn("unlink: 3", out)
|
||||
self.assertIn("relink: 1", out)
|
||||
self.assertIn("missing: 1", out)
|
||||
self.assertIn(f"-> {self.typed}", out)
|
||||
|
||||
def test_apply_backs_up_then_changes_only_named_images(self):
|
||||
named = f"{self.guessed[0]}-{self.guessed[1]},{self.wrong}"
|
||||
code, out = self.run_tool(
|
||||
"--ids", named, "--prompt", str(self.coffee), "--apply"
|
||||
)
|
||||
self.assertEqual(code, 0)
|
||||
self.assertIn("Applied: 1 relinked, 2 unlinked, 0 skipped", out)
|
||||
links = self.links()
|
||||
self.assertNotIn(self.guessed[0], links)
|
||||
self.assertNotIn(self.guessed[1], links)
|
||||
self.assertEqual(links[self.guessed[2]], self.coffee, "not named: kept")
|
||||
self.assertEqual(links[self.wrong], self.typed)
|
||||
self.assertEqual(links[self.gone], self.coffee, "missing files are kept")
|
||||
backups = [n for n in os.listdir(self.tmp) if ".pre-relink_" in n]
|
||||
self.assertEqual(len(backups), 1)
|
||||
restored = PromptDatabase(os.path.join(self.tmp, backups[0]))
|
||||
try:
|
||||
count = (
|
||||
restored.model.get_connection()
|
||||
.execute("SELECT COUNT(*) FROM generated_images")
|
||||
.fetchone()[0]
|
||||
)
|
||||
finally:
|
||||
restored.close_all()
|
||||
self.assertEqual(count, 6)
|
||||
|
||||
def test_arguments_are_validated(self):
|
||||
for args in (
|
||||
[],
|
||||
["--ids", "5-2"],
|
||||
["--ids", "x"],
|
||||
["--prompt", "1", "--apply"],
|
||||
):
|
||||
with self.subTest(args=args), self.assertRaises(SystemExit):
|
||||
with unittest.mock.patch("sys.stderr", io.StringIO()):
|
||||
self.run_tool(*args)
|
||||
with (
|
||||
self.assertRaises(SystemExit),
|
||||
unittest.mock.patch("sys.stderr", io.StringIO()),
|
||||
):
|
||||
tool.main(
|
||||
[os.path.join(self.tmp, "none.db"), "--output-dir", "x", "--ids", "1"]
|
||||
)
|
||||
self.assertEqual(tool.parse_ids(" 3, 1-2 ,"), [1, 2, 3])
|
||||
|
||||
def test_conflicting_sources_are_left_to_the_maintainer(self):
|
||||
db = PromptDatabase(self.db_path)
|
||||
try:
|
||||
handler = tool.ImageGenerationHandler(db, None)
|
||||
row = {
|
||||
"id": 1,
|
||||
"prompt_id": self.coffee,
|
||||
"image_path": os.path.join(self.output, "typed_right.png"),
|
||||
}
|
||||
with unittest.mock.patch.object(
|
||||
tool, "certain_prompts", return_value={"scan": 1, "history": 2}
|
||||
):
|
||||
verdict = tool.judge(db, handler, row, self.output)
|
||||
finally:
|
||||
db.close_all()
|
||||
self.assertEqual(verdict["verdict"], "conflict")
|
||||
|
||||
|
||||
class ApplyGuardTests(ToolFixture):
|
||||
"""M6: unlink clears only the link, and only the link it was judged on."""
|
||||
|
||||
def test_unlink_skips_a_row_linked_elsewhere_since_the_review(self):
|
||||
db = PromptDatabase(self.db_path)
|
||||
try:
|
||||
verdicts = tool.review(db, [self.guessed[0]], None, self.output)
|
||||
self.assertEqual(verdicts[0]["verdict"], "unlink")
|
||||
db.model.get_connection().execute(
|
||||
"UPDATE generated_images SET prompt_id=? WHERE id=?",
|
||||
(self.typed, self.guessed[0]),
|
||||
)
|
||||
db.model.get_connection().commit()
|
||||
result = tool.apply(db, verdicts, [self.guessed[0]])
|
||||
finally:
|
||||
db.close_all()
|
||||
self.assertEqual(result["unlinked"], 0)
|
||||
self.assertEqual(self.links()[self.guessed[0]], self.typed)
|
||||
|
||||
def test_unlink_keeps_the_file_and_its_favorite(self):
|
||||
path = os.path.join(self.output, "moss_00009_.png")
|
||||
db = PromptDatabase(self.db_path)
|
||||
try:
|
||||
db.set_image_favorite("a" * 64, "moss_00009_.png")
|
||||
finally:
|
||||
db.close_all()
|
||||
self.run_tool("--ids", str(self.guessed[0]), "--apply")
|
||||
self.assertTrue(os.path.isfile(path))
|
||||
db = PromptDatabase(self.db_path)
|
||||
try:
|
||||
self.assertEqual(db.get_root_favorites("a" * 64), {"moss_00009_.png"})
|
||||
summary = db.get_image_link_summaries([path]).get(path) or {}
|
||||
finally:
|
||||
db.close_all()
|
||||
self.assertNotEqual(summary.get("status"), "linked")
|
||||
|
||||
def test_dry_run_states_what_apply_does(self):
|
||||
_code, out = self.run_tool("--ids", str(self.guessed[0]))
|
||||
self.assertIn("removes only the image's link record", out)
|
||||
self.assertIn("file stays on disk and in the gallery as unlinked", out)
|
||||
|
||||
|
||||
class GalleryAfterUnlinkTests(ImageAPITestCase):
|
||||
"""M6: after --apply the file is still listed, as an unlinked image."""
|
||||
|
||||
async def test_unlinked_file_stays_in_the_workspace_gallery(self):
|
||||
db = self.api.db
|
||||
prompt = db.save_prompt(
|
||||
text="a woman holding a coffee cup",
|
||||
prompt_hash=generate_prompt_hash("a woman holding a coffee cup"),
|
||||
)
|
||||
path = self.output_dir / "moss_00009_.png"
|
||||
png(str(path), issue_graph(), pair_evidence())
|
||||
image_id = db.link_image_to_prompt(prompt, str(path), metadata={})
|
||||
out = io.StringIO()
|
||||
tool.main(
|
||||
[
|
||||
db.model.db_path,
|
||||
"--output-dir",
|
||||
str(self.output_dir),
|
||||
"--ids",
|
||||
str(image_id),
|
||||
"--apply",
|
||||
],
|
||||
out=out,
|
||||
)
|
||||
self.assertIn("Applied: 0 relinked, 1 unlinked", out.getvalue())
|
||||
response = await self.client.get(
|
||||
"/prompt_manager/images/output", params={"mode": "workspace"}
|
||||
)
|
||||
self.assertEqual(response.status, 200)
|
||||
images = {i["filename"]: i for i in (await response.json())["images"]}
|
||||
self.assertIn("moss_00009_.png", images)
|
||||
link = images["moss_00009_.png"].get("linked") or {}
|
||||
self.assertNotEqual(link.get("status"), "linked")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,674 @@
|
||||
"""#342: recorded history attributes a file only to one certain, stored prompt."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
from database.run_attribution import stored_run_prompt
|
||||
from tests import test_execution_history_database as history_tests
|
||||
from tests.test_api_admin import AdminAPITestCase
|
||||
from tests.test_execution_context_database import snapshot
|
||||
from tests.test_prompt_attribution import issue_graph, llm_text_graph, make_tracker
|
||||
from utils.execution_contexts import ContextBundle, RecipeContext
|
||||
from utils.execution_outputs import (
|
||||
OutputBundle,
|
||||
OutputReference,
|
||||
snapshot_output_roots,
|
||||
)
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.image_monitor import ImageGenerationHandler
|
||||
|
||||
STYLE = "windows" if os.name == "nt" else "posix"
|
||||
LLM_TEXT = "Photorealistic aerial landscape of misty mountains at dawn"
|
||||
|
||||
|
||||
def now_ms():
|
||||
return time.time_ns() // 1_000_000
|
||||
|
||||
|
||||
class RecordedRuns:
|
||||
"""Write runs into execution history the way the capture pipeline does."""
|
||||
|
||||
start = history_tests.HistoryDatabaseTests.start
|
||||
finish = history_tests.HistoryDatabaseTests.finish
|
||||
|
||||
def record_run(self, db, output_root, files, positives, item=0, window=None):
|
||||
started, finished = window or (now_ms() - 5000, now_ms() + 5000)
|
||||
start = self.start(db, item, started)
|
||||
db.write_execution_start(start)
|
||||
finish = self.finish(start, time=finished)
|
||||
db.write_execution_finish(finish)
|
||||
db.write_execution_contexts(
|
||||
start.run_id,
|
||||
start.storage_epoch,
|
||||
ContextBundle(
|
||||
tuple(
|
||||
RecipeContext(str(4 + i), snapshot(text), snapshot(""), "inferred")
|
||||
for i, text in enumerate(positives)
|
||||
),
|
||||
"complete",
|
||||
"captured",
|
||||
),
|
||||
)
|
||||
root = snapshot_output_roots(STYLE, output_root, output_root).output_key
|
||||
references = tuple(
|
||||
OutputReference("9", i, "output", sub, name, STYLE, root)
|
||||
for i, (sub, name) in enumerate(files)
|
||||
)
|
||||
db.write_execution_outputs_for_finish(
|
||||
finish, OutputBundle(references, "complete", "reported")
|
||||
)
|
||||
return start.run_id
|
||||
|
||||
|
||||
class TestStoredRunPrompt(RecordedRuns, unittest.TestCase):
|
||||
setUp = history_tests.HistoryDatabaseTests.setUp
|
||||
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
||||
database = history_tests.HistoryDatabaseTests.database
|
||||
|
||||
def setUp(self): # noqa: F811 - the shared temp dir and database list first
|
||||
history_tests.HistoryDatabaseTests.setUp(self)
|
||||
self.db = self.database()
|
||||
self.output = str(Path(self.temp.name) / "output")
|
||||
os.makedirs(os.path.join(self.output, "sub"))
|
||||
self.image = os.path.join(self.output, "sub", "img_00001_.png")
|
||||
Path(self.image).write_bytes(b"png")
|
||||
|
||||
def save(self, text):
|
||||
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
||||
|
||||
def attribute(self, path=None):
|
||||
return stored_run_prompt(self.db, path or self.image, self.output)
|
||||
|
||||
def test_one_run_naming_one_stored_prompt_attributes_the_file(self):
|
||||
prompt_id = self.save("an llm prompt")
|
||||
self.record_run(
|
||||
self.db, self.output, [("sub", "img_00001_.png")], ["an llm prompt"]
|
||||
)
|
||||
self.assertEqual(self.attribute(), prompt_id)
|
||||
|
||||
def test_text_that_is_not_stored_attributes_nothing(self):
|
||||
"""The #342 run: raw editor text was empty, nothing was stored."""
|
||||
self.save("a woman holding a coffee cup")
|
||||
self.record_run(self.db, self.output, [("sub", "img_00001_.png")], [""])
|
||||
self.assertIsNone(self.attribute())
|
||||
|
||||
def test_two_runs_reporting_the_same_file_attribute_nothing(self):
|
||||
self.save("first")
|
||||
self.save("second")
|
||||
self.record_run(self.db, self.output, [("sub", "img_00001_.png")], ["first"])
|
||||
self.record_run(
|
||||
self.db, self.output, [("sub", "img_00001_.png")], ["second"], item=1
|
||||
)
|
||||
self.assertIsNone(self.attribute())
|
||||
|
||||
def test_a_file_written_outside_the_run_attributes_nothing(self):
|
||||
self.save("an llm prompt")
|
||||
self.record_run(
|
||||
self.db,
|
||||
self.output,
|
||||
[("sub", "img_00001_.png")],
|
||||
["an llm prompt"],
|
||||
window=(1_000_000, 2_000_000),
|
||||
)
|
||||
self.assertIsNone(self.attribute())
|
||||
|
||||
def test_contexts_naming_different_prompts_attribute_nothing(self):
|
||||
self.save("left")
|
||||
self.save("right")
|
||||
self.record_run(
|
||||
self.db, self.output, [("sub", "img_00001_.png")], ["left", "right"]
|
||||
)
|
||||
self.assertIsNone(self.attribute())
|
||||
|
||||
def test_files_outside_the_output_root_or_unrecorded_attribute_nothing(self):
|
||||
self.save("an llm prompt")
|
||||
self.record_run(
|
||||
self.db, self.output, [("sub", "img_00001_.png")], ["an llm prompt"]
|
||||
)
|
||||
outside = os.path.join(self.temp.name, "img_00001_.png")
|
||||
Path(outside).write_bytes(b"png")
|
||||
self.assertIsNone(self.attribute(outside))
|
||||
other = os.path.join(self.output, "sub", "other.png")
|
||||
Path(other).write_bytes(b"png")
|
||||
self.assertIsNone(self.attribute(other))
|
||||
self.assertIsNone(stored_run_prompt(self.db, self.image, None))
|
||||
self.assertIsNone(stored_run_prompt(self.db, self.image, "relative/root"))
|
||||
|
||||
def test_a_library_without_history_tables_attributes_nothing(self):
|
||||
broken = mock.Mock()
|
||||
broken.model.get_connection.return_value.execute.side_effect = (
|
||||
sqlite3.OperationalError("no such table")
|
||||
)
|
||||
self.assertIsNone(stored_run_prompt(broken, self.image, self.output))
|
||||
|
||||
def test_the_monitor_links_from_history_when_nothing_else_is_certain(self):
|
||||
prompt_id = self.save(LLM_TEXT)
|
||||
self.record_run(self.db, self.output, [("sub", "img_00001_.png")], [LLM_TEXT])
|
||||
handler = ImageGenerationHandler(self.db, make_tracker(self.db))
|
||||
linked = []
|
||||
handler.link_image_to_prompt = lambda path, prompt, meta: linked.append(
|
||||
prompt["id"]
|
||||
)
|
||||
with (
|
||||
mock.patch(
|
||||
"utils.image_monitor.host_output_directory", return_value=self.output
|
||||
),
|
||||
mock.patch.object(
|
||||
handler.metadata_extractor,
|
||||
"extract_metadata",
|
||||
return_value={"prompt": llm_text_graph()},
|
||||
),
|
||||
):
|
||||
handler.process_new_image(self.image)
|
||||
self.assertEqual(linked, [prompt_id])
|
||||
|
||||
|
||||
class TestScanUsesHistory(RecordedRuns, AdminAPITestCase):
|
||||
"""The output scan applies the same rules as the monitor."""
|
||||
|
||||
def linked_prompt(self, path):
|
||||
row = (
|
||||
self.api.db.model.get_connection()
|
||||
.execute(
|
||||
"SELECT prompt_id FROM generated_images WHERE image_path=?",
|
||||
(str(path),),
|
||||
)
|
||||
.fetchone()
|
||||
)
|
||||
return row[0] if row else None
|
||||
|
||||
def write(self, name):
|
||||
path = self.output_dir / name
|
||||
path.write_bytes(b"png")
|
||||
return path
|
||||
|
||||
def test_an_llm_text_file_links_to_the_prompt_its_run_stored(self):
|
||||
db = self.api.db
|
||||
prompt_id = db.save_prompt(
|
||||
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
||||
)
|
||||
path = self.write("llm_00001_.png")
|
||||
self.record_run(db, str(self.output_dir), [("", path.name)], [LLM_TEXT])
|
||||
meta = {"prompt": json.dumps(llm_text_graph())}
|
||||
self.assertEqual(self.api._ingest_scanned_file(path, meta), "linked")
|
||||
self.assertEqual(self.linked_prompt(path), prompt_id)
|
||||
|
||||
def test_a_metadata_less_file_links_through_its_recorded_run(self):
|
||||
"""I1: --disable-metadata outputs carry no prompt; history still proves one."""
|
||||
db = self.api.db
|
||||
prompt_id = db.save_prompt(
|
||||
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
||||
)
|
||||
path = self.write("bare_00001_.png")
|
||||
self.record_run(db, str(self.output_dir), [("", path.name)], [LLM_TEXT])
|
||||
self.assertEqual(self.api._ingest_scanned_file(path, {}), "linked")
|
||||
self.assertEqual(self.linked_prompt(path), prompt_id)
|
||||
self.assertEqual(self.api._ingest_scanned_file(path, None), "unchanged")
|
||||
detailed = self.api._ingest_scanned_file(path, {}, detailed=True)
|
||||
self.assertEqual(detailed["status"], "unchanged")
|
||||
self.assertTrue(detailed["title"])
|
||||
|
||||
def test_a_metadata_less_file_without_a_recorded_run_is_still_skipped(self):
|
||||
path = self.write("bare_00002_.png")
|
||||
self.assertIsNone(self.api._ingest_scanned_file(path, {}))
|
||||
self.assertIsNone(self.api._ingest_scanned_file(path, {"workflow": {}}))
|
||||
self.assertIsNone(self.linked_prompt(path))
|
||||
|
||||
def test_the_issue_file_stays_unlinked_even_with_its_pair_record(self):
|
||||
"""Resolved with an empty stored positive: before #342 the scan raised."""
|
||||
db = self.api.db
|
||||
db.save_prompt(
|
||||
text="a woman holding a coffee cup",
|
||||
prompt_hash=generate_prompt_hash("a woman holding a coffee cup"),
|
||||
)
|
||||
path = self.write("moss_00009_.png")
|
||||
self.record_run(db, str(self.output_dir), [("", path.name)], [""])
|
||||
evidence = {
|
||||
"version": 1,
|
||||
"nodes": {
|
||||
"18:35": {
|
||||
"records": [
|
||||
{
|
||||
"text": "",
|
||||
"negative_text": "",
|
||||
"positive_text": LLM_TEXT,
|
||||
"negative_effective_text": "",
|
||||
"prepend_text": LLM_TEXT,
|
||||
"append_text": "",
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
}
|
||||
meta = {
|
||||
"prompt": json.dumps(issue_graph()),
|
||||
"comfyassets_prompt_pairs": json.dumps(evidence),
|
||||
}
|
||||
self.assertEqual(self.api._ingest_scanned_file(path, meta), "unresolved")
|
||||
self.assertIsNone(self.linked_prompt(path))
|
||||
self.assertEqual(db.list_prompt_pairs()["pairs"], [])
|
||||
|
||||
|
||||
class TestHostHelpers(unittest.TestCase):
|
||||
"""The host lookups degrade to "no proof", never to an error."""
|
||||
|
||||
def test_host_output_directory(self):
|
||||
import sys
|
||||
import types
|
||||
|
||||
from database import run_attribution
|
||||
|
||||
cases = (
|
||||
(types.SimpleNamespace(get_output_directory=lambda: "/out"), "/out"),
|
||||
(types.SimpleNamespace(get_output_directory=lambda: ""), None),
|
||||
(types.SimpleNamespace(get_output_directory=lambda: 1 / 0), None),
|
||||
)
|
||||
original = sys.modules.get("folder_paths")
|
||||
try:
|
||||
for module, expected in cases:
|
||||
sys.modules["folder_paths"] = module
|
||||
self.assertEqual(run_attribution.host_output_directory(), expected)
|
||||
sys.modules["folder_paths"] = None # import fails: outside ComfyUI
|
||||
self.assertIsNone(run_attribution.host_output_directory())
|
||||
finally:
|
||||
if original is None:
|
||||
sys.modules.pop("folder_paths", None)
|
||||
else:
|
||||
sys.modules["folder_paths"] = original
|
||||
|
||||
def test_run_window_needs_integer_times_and_a_readable_file(self):
|
||||
from database.run_attribution import _written_during
|
||||
|
||||
self.assertFalse(_written_during(__file__, None, 5))
|
||||
self.assertFalse(_written_during(os.path.join("no", "such.png"), 0, 5))
|
||||
|
||||
def test_executing_run_id(self):
|
||||
import sys
|
||||
import types
|
||||
|
||||
from utils.run_keys import executing_run_id
|
||||
|
||||
def host(context):
|
||||
def get_executing_context():
|
||||
if isinstance(context, Exception):
|
||||
raise context
|
||||
return context
|
||||
|
||||
return types.SimpleNamespace(get_executing_context=get_executing_context)
|
||||
|
||||
original = sys.modules.get("comfy_execution.utils")
|
||||
try:
|
||||
for context, expected in (
|
||||
(types.SimpleNamespace(prompt_id="run-1"), "run-1"),
|
||||
(types.SimpleNamespace(prompt_id=""), None),
|
||||
(None, None),
|
||||
(RuntimeError("no context"), None),
|
||||
):
|
||||
sys.modules["comfy_execution.utils"] = host(context)
|
||||
self.assertEqual(executing_run_id(), expected)
|
||||
sys.modules["comfy_execution.utils"] = None
|
||||
self.assertIsNone(executing_run_id())
|
||||
finally:
|
||||
if original is None:
|
||||
sys.modules.pop("comfy_execution.utils", None)
|
||||
else:
|
||||
sys.modules["comfy_execution.utils"] = original
|
||||
|
||||
|
||||
class TestMonitorHistoryRecheck(RecordedRuns, unittest.TestCase):
|
||||
"""I1: a keyless image gets one history-only re-check when its run is sealed."""
|
||||
|
||||
def setUp(self):
|
||||
history_tests.HistoryDatabaseTests.setUp(self)
|
||||
self.db = history_tests.HistoryDatabaseTests.database(self)
|
||||
self.output = str(Path(self.temp.name) / "output")
|
||||
os.makedirs(self.output)
|
||||
self.handler = ImageGenerationHandler(self.db, make_tracker(self.db))
|
||||
self.handler.processing_delay = 0
|
||||
self.handler.settle_interval = 0.01
|
||||
self.addCleanup(self.handler.shutdown)
|
||||
self.linked = []
|
||||
self.handler.link_image_to_prompt = lambda path, prompt, meta: (
|
||||
self.linked.append((os.path.basename(path), prompt["id"]))
|
||||
)
|
||||
patches = (
|
||||
mock.patch(
|
||||
"utils.image_monitor.host_output_directory", return_value=self.output
|
||||
),
|
||||
mock.patch.object(
|
||||
self.handler.metadata_extractor, "extract_metadata", return_value=None
|
||||
),
|
||||
)
|
||||
for patch in patches:
|
||||
patch.start()
|
||||
self.addCleanup(patch.stop)
|
||||
|
||||
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
||||
|
||||
def image(self, name):
|
||||
path = os.path.join(self.output, name)
|
||||
Path(path).write_bytes(b"png")
|
||||
return path
|
||||
|
||||
def save(self, text):
|
||||
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
||||
|
||||
def wait_for(self, predicate, timeout=5.0):
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
if predicate():
|
||||
return True
|
||||
time.sleep(0.01)
|
||||
return predicate()
|
||||
|
||||
def test_the_run_ending_after_processing_links_the_image_once(self):
|
||||
prompt_id = self.save(LLM_TEXT)
|
||||
path = self.image("late_00001_.png")
|
||||
self.handler.process_new_image(path) # the run is still executing
|
||||
self.assertEqual(self.linked, [])
|
||||
self.assertIn(
|
||||
os.path.normcase(os.path.realpath(path)), self.handler.awaiting_history()
|
||||
)
|
||||
self.record_run(self.db, self.output, [("", "late_00001_.png")], [LLM_TEXT])
|
||||
self.assertTrue(self.wait_for(lambda: self.linked))
|
||||
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.assertEqual(self.linked, [("late_00001_.png", prompt_id)])
|
||||
self.assertEqual(self.handler.awaiting_history(), [])
|
||||
# Another run reporting the same name later is not a second re-check
|
||||
self.record_run(
|
||||
self.db, self.output, [("", "late_00001_.png")], [LLM_TEXT], item=1
|
||||
)
|
||||
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.assertEqual(len(self.linked), 1)
|
||||
|
||||
def test_a_seal_during_the_first_pass_is_not_lost(self):
|
||||
"""R1: the run seals while the image's own pass is still pending."""
|
||||
prompt_id = self.save(LLM_TEXT)
|
||||
path = self.image("race_00001_.png")
|
||||
real = self.handler._history_prompt
|
||||
sealed = []
|
||||
|
||||
def history(image_path):
|
||||
result = real(image_path) # before the run's outputs commit: None
|
||||
if not sealed:
|
||||
sealed.append(True)
|
||||
# The writer commits and notifies now; the path is still pending
|
||||
self.record_run(
|
||||
self.db, self.output, [("", "race_00001_.png")], [LLM_TEXT]
|
||||
)
|
||||
return result
|
||||
|
||||
self.handler._history_prompt = history
|
||||
self.assertTrue(self.handler.schedule(path))
|
||||
self.assertTrue(self.wait_for(lambda: self.linked))
|
||||
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.assertEqual(self.linked, [("race_00001_.png", prompt_id)])
|
||||
|
||||
def test_a_seal_while_the_pending_set_is_full_is_rechecked_later(self):
|
||||
from utils import image_monitor
|
||||
|
||||
prompt_id = self.save(LLM_TEXT)
|
||||
path = self.image("full_00001_.png")
|
||||
self.handler.process_new_image(path) # waiting, unlinked
|
||||
blocker = self.image("blocker.png")
|
||||
release = threading.Event()
|
||||
original = self.handler.process_new_image
|
||||
|
||||
def slow(image_path, prompt_snapshot=None):
|
||||
if image_path == blocker:
|
||||
release.wait(5)
|
||||
return original(image_path, prompt_snapshot=prompt_snapshot)
|
||||
|
||||
self.handler.process_new_image = slow
|
||||
with mock.patch.object(image_monitor, "MAX_PENDING_IMAGES", 1):
|
||||
self.assertTrue(self.handler.schedule(blocker))
|
||||
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 1))
|
||||
self.record_run(self.db, self.output, [("", "full_00001_.png")], [LLM_TEXT])
|
||||
self.assertEqual(self.linked, [], "not scheduled while the set is full")
|
||||
release.set()
|
||||
self.assertTrue(self.wait_for(lambda: self.linked))
|
||||
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.assertEqual(self.linked, [("full_00001_.png", prompt_id)])
|
||||
|
||||
def test_a_run_naming_nothing_stored_leaves_the_image_unlinked(self):
|
||||
path = self.image("empty_00001_.png")
|
||||
self.handler.process_new_image(path)
|
||||
self.record_run(self.db, self.output, [("", "empty_00001_.png")], [""])
|
||||
self.assertTrue(self.wait_for(lambda: not self.handler.awaiting_history()))
|
||||
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.assertEqual(self.linked, [])
|
||||
|
||||
def test_history_found_at_once_leaves_nothing_awaiting(self):
|
||||
prompt_id = self.save(LLM_TEXT)
|
||||
self.record_run(self.db, self.output, [("", "done_00001_.png")], [LLM_TEXT])
|
||||
self.handler.process_new_image(self.image("done_00001_.png"))
|
||||
self.assertEqual(self.linked, [("done_00001_.png", prompt_id)])
|
||||
self.assertEqual(self.handler.awaiting_history(), [])
|
||||
|
||||
def test_keyed_images_do_not_wait_for_history(self):
|
||||
with mock.patch.object(
|
||||
self.handler.metadata_extractor,
|
||||
"extract_metadata",
|
||||
return_value={"prompt": llm_text_graph()},
|
||||
):
|
||||
self.handler.process_new_image(self.image("keyed_00001_.png"))
|
||||
self.assertEqual(self.handler.awaiting_history(), [])
|
||||
|
||||
def test_the_waiting_set_is_bounded_and_expires(self):
|
||||
from utils import image_monitor
|
||||
|
||||
with mock.patch.object(image_monitor, "MAX_AWAITING_HISTORY", 2):
|
||||
for n in range(3):
|
||||
self.handler.process_new_image(self.image(f"b_{n}.png"))
|
||||
self.assertEqual(
|
||||
[os.path.basename(p) for p in self.handler.awaiting_history()],
|
||||
["b_1.png", "b_2.png"],
|
||||
)
|
||||
later = time.monotonic() + image_monitor.AWAIT_HISTORY_SECONDS + 1
|
||||
with mock.patch.object(image_monitor.time, "monotonic", return_value=later):
|
||||
self.assertEqual(self.handler.awaiting_history(), [])
|
||||
|
||||
def test_shutdown_forgets_waiting_images_and_a_dead_handler_is_dropped(self):
|
||||
import gc
|
||||
|
||||
from database import run_attribution
|
||||
|
||||
self.handler.process_new_image(self.image("gone_00001_.png"))
|
||||
self.handler.shutdown()
|
||||
self.assertEqual(self.handler.awaiting_history(), [])
|
||||
other = ImageGenerationHandler(self.db, make_tracker(self.db))
|
||||
with mock.patch.object(
|
||||
other.metadata_extractor, "extract_metadata", return_value=None
|
||||
):
|
||||
other.process_new_image(self.image("gone_00002_.png"))
|
||||
del other
|
||||
gc.collect()
|
||||
run_attribution.notify_outputs_sealed("run", []) # no dead callback raises
|
||||
|
||||
|
||||
class TestBoundedLookup(RecordedRuns, unittest.TestCase):
|
||||
"""M5: history lookups are batched, and the history schema is unchanged."""
|
||||
|
||||
setUp = history_tests.HistoryDatabaseTests.setUp
|
||||
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
||||
database = history_tests.HistoryDatabaseTests.database
|
||||
|
||||
def output_lookups(self, db, action):
|
||||
statements = []
|
||||
conn = db.model.get_connection()
|
||||
conn.set_trace_callback(statements.append)
|
||||
try:
|
||||
result = action()
|
||||
finally:
|
||||
conn.set_trace_callback(None)
|
||||
return result, [s for s in statements if "FROM execution_run_outputs" in s]
|
||||
|
||||
def test_one_query_per_chunk_not_per_file_and_same_answers(self):
|
||||
from database import run_attribution
|
||||
|
||||
db = self.database()
|
||||
output = str(Path(self.temp.name) / "output")
|
||||
os.makedirs(output)
|
||||
prompt_id = db.save_prompt(
|
||||
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
||||
)
|
||||
names = [f"img_{n:05}_.png" for n in range(25)]
|
||||
paths = [os.path.join(output, name) for name in names]
|
||||
for path in paths:
|
||||
Path(path).write_bytes(b"png")
|
||||
self.record_run(db, output, [("", n) for n in names[:10]], [LLM_TEXT])
|
||||
self.record_run(db, output, [("", n) for n in names[5:10]], [""], item=1)
|
||||
with mock.patch.object(run_attribution, "LOOKUP_CHUNK", 10):
|
||||
result, lookups = self.output_lookups(
|
||||
db, lambda: run_attribution.stored_run_prompts(db, paths, output)
|
||||
)
|
||||
self.assertEqual(len(lookups), 3) # 25 names in chunks of 10
|
||||
expected = {path: None for path in paths}
|
||||
expected.update({path: prompt_id for path in paths[:5]})
|
||||
self.assertEqual(result, expected)
|
||||
for path in paths[:12]:
|
||||
self.assertEqual(
|
||||
run_attribution.stored_run_prompt(db, path, output), expected[path]
|
||||
)
|
||||
|
||||
def test_the_history_schema_gains_no_index(self):
|
||||
from database.execution_history_schema import history_schema_error
|
||||
|
||||
db = self.database()
|
||||
conn = db.model.get_connection()
|
||||
indexes = [
|
||||
row[1]
|
||||
for row in conn.execute("PRAGMA index_list(execution_run_outputs)")
|
||||
if row[3] != "pk"
|
||||
]
|
||||
self.assertEqual(indexes, [])
|
||||
self.assertIsNone(history_schema_error(conn, allow_missing=False))
|
||||
|
||||
|
||||
class TestScanBatchLookup(RecordedRuns, AdminAPITestCase):
|
||||
"""M5: a scan batch asks history once per chunk of names, not per file."""
|
||||
|
||||
def test_a_batch_its_metadata_links_never_asks_history(self):
|
||||
"""R3: the batch lookup is made only when a file needs it."""
|
||||
files = []
|
||||
for n in range(3):
|
||||
path = self.output_dir / f"typed_{n:05}_.png"
|
||||
path.write_bytes(b"png")
|
||||
files.append(path)
|
||||
meta = {"prompt": json.dumps(issue_graph(text="typed mountains"))}
|
||||
statements = []
|
||||
conn = self.api.db.model.get_connection()
|
||||
conn.set_trace_callback(statements.append)
|
||||
try:
|
||||
counts = self.api._ingest_scan_batch_sync([(f, meta) for f in files])
|
||||
finally:
|
||||
conn.set_trace_callback(None)
|
||||
self.assertEqual(counts["added"] + counts["linked"], 3, counts)
|
||||
lookups = [s for s in statements if "FROM execution_run_outputs" in s]
|
||||
self.assertEqual(lookups, [])
|
||||
|
||||
def test_metadata_less_batch_links_with_one_history_lookup(self):
|
||||
db = self.api.db
|
||||
prompt_id = db.save_prompt(
|
||||
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
||||
)
|
||||
files = []
|
||||
for n in range(6):
|
||||
path = self.output_dir / f"bare_{n:05}_.png"
|
||||
path.write_bytes(b"png")
|
||||
files.append(path)
|
||||
self.record_run(
|
||||
db, str(self.output_dir), [("", f.name) for f in files[:4]], [LLM_TEXT]
|
||||
)
|
||||
statements = []
|
||||
conn = db.model.get_connection()
|
||||
conn.set_trace_callback(statements.append)
|
||||
try:
|
||||
counts = self.api._ingest_scan_batch_sync([(f, {}) for f in files])
|
||||
finally:
|
||||
conn.set_trace_callback(None)
|
||||
lookups = [s for s in statements if "FROM execution_run_outputs" in s]
|
||||
self.assertEqual(len(lookups), 1)
|
||||
self.assertEqual(counts["linked"], 4)
|
||||
linked = {
|
||||
row[0]
|
||||
for row in conn.execute(
|
||||
"SELECT filename FROM generated_images WHERE prompt_id=?",
|
||||
(prompt_id,),
|
||||
)
|
||||
}
|
||||
self.assertEqual(linked, {f.name for f in files[:4]})
|
||||
|
||||
|
||||
class TestLinkedOutputRoot(RecordedRuns, unittest.TestCase):
|
||||
"""R2: ComfyUI's output folder reached through a symlink or junction."""
|
||||
|
||||
setUp = history_tests.HistoryDatabaseTests.setUp
|
||||
close_databases = history_tests.HistoryDatabaseTests.close_databases
|
||||
database = history_tests.HistoryDatabaseTests.database
|
||||
|
||||
def linked_root(self):
|
||||
self.real = str(Path(self.temp.name) / "big_drive" / "output")
|
||||
os.makedirs(os.path.join(self.real, "sub"))
|
||||
link = str(Path(self.temp.name) / "output_link")
|
||||
try:
|
||||
os.symlink(self.real, link, target_is_directory=True)
|
||||
except (OSError, NotImplementedError) as error:
|
||||
self.skipTest(f"symlinks unavailable: {error}")
|
||||
return link
|
||||
|
||||
def test_either_spelling_of_the_file_is_attributed(self):
|
||||
db = self.database()
|
||||
link = self.linked_root()
|
||||
prompt_id = db.save_prompt(
|
||||
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
||||
)
|
||||
real_file = os.path.join(self.real, "sub", "img_00001_.png")
|
||||
Path(real_file).write_bytes(b"png")
|
||||
# History keeps the host's spelling of its output folder (the link)
|
||||
self.record_run(db, link, [("sub", "img_00001_.png")], [LLM_TEXT])
|
||||
link_file = os.path.join(link, "sub", "img_00001_.png")
|
||||
self.assertEqual(stored_run_prompt(db, real_file, link), prompt_id)
|
||||
self.assertEqual(stored_run_prompt(db, link_file, link), prompt_id)
|
||||
# The real folder named as the root is another root key: not attributed
|
||||
self.assertIsNone(stored_run_prompt(db, real_file, self.real))
|
||||
|
||||
def test_the_monitor_rechecks_a_file_seen_through_the_real_path(self):
|
||||
db = self.database()
|
||||
link = self.linked_root()
|
||||
prompt_id = db.save_prompt(
|
||||
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
|
||||
)
|
||||
real_file = os.path.join(self.real, "late_00001_.png")
|
||||
Path(real_file).write_bytes(b"png")
|
||||
handler = ImageGenerationHandler(db, make_tracker(db))
|
||||
handler.processing_delay = 0
|
||||
handler.settle_interval = 0.01
|
||||
self.addCleanup(handler.shutdown)
|
||||
linked = []
|
||||
handler.link_image_to_prompt = lambda path, prompt, meta: linked.append(
|
||||
prompt["id"]
|
||||
)
|
||||
with (
|
||||
mock.patch("utils.image_monitor.host_output_directory", return_value=link),
|
||||
mock.patch.object(
|
||||
handler.metadata_extractor, "extract_metadata", return_value=None
|
||||
),
|
||||
):
|
||||
handler.process_new_image(real_file)
|
||||
self.assertEqual(linked, [])
|
||||
self.record_run(db, link, [("", "late_00001_.png")], [LLM_TEXT])
|
||||
deadline = time.monotonic() + 5
|
||||
while not linked and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
self.assertEqual(linked, [prompt_id])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -477,11 +477,20 @@ class TestPendingFirstUse(unittest.TestCase):
|
||||
class FakeTracker:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
self.runs = []
|
||||
|
||||
def set_current_prompt(self, prompt_text, additional_data=None, push_to_queue=True):
|
||||
def set_current_prompt(
|
||||
self, prompt_text, additional_data=None, push_to_queue=True, **run
|
||||
):
|
||||
self.calls.append({"text": prompt_text, "push_to_queue": push_to_queue})
|
||||
self.runs.append(run)
|
||||
return "exec-1"
|
||||
|
||||
def mark_empty(self, run_key=None, run_id=None, push_to_queue=False):
|
||||
self.calls.append(
|
||||
{"empty": True, "run_key": run_key, "push_to_queue": push_to_queue}
|
||||
)
|
||||
|
||||
|
||||
class TestNodeRoles(DbTestCase):
|
||||
def _node(self):
|
||||
|
||||
Vendored
+4
-1
@@ -29,15 +29,18 @@ test('execution script discovers only the new matrix and preserves all 78 origin
|
||||
assert.doesNotMatch(stock, /execution-history/);
|
||||
assert.doesNotMatch(stock, /execution-outputs/);
|
||||
assert.doesNotMatch(stock, /library-location/); // it starts its own fresh host (#249 S6b)
|
||||
assert.doesNotMatch(stock, /image-attribution/); // needs the owned fixture pack (#342)
|
||||
const executions = list('tools/vendor/playwright.execution.config.mjs');
|
||||
assert.match(executions, /\[classic-start\]/);
|
||||
assert.match(executions, /\[nodes2-start\]/);
|
||||
assert.match(executions, /Total: 86 tests in 4 files/); // + the fresh-host library location case, + #343 linked text
|
||||
// + the fresh-host library location case, + #342 image attribution, + #343 linked text
|
||||
assert.match(executions, /Total: 90 tests in 5 files/);
|
||||
for (const project of ['classic-start', 'nodes2-start']) {
|
||||
const cases = executions.split('\n').filter(line => line.includes(`[${project}]`));
|
||||
assert.equal(cases.filter(line => line.includes('execution-history.spec.mjs')).length, 15);
|
||||
assert.equal(cases.filter(line => line.includes('execution-outputs.spec.mjs')).length, 8);
|
||||
assert.equal(cases.filter(line => line.includes('library-location.spec.mjs')).length, 1);
|
||||
assert.equal(cases.filter(line => line.includes('image-attribution.spec.mjs')).length, 2);
|
||||
assert.equal(cases.filter(line => line.includes('execution-text-link.spec.mjs')).length, 19);
|
||||
}
|
||||
assert.doesNotMatch(executions, /prompt-manager.spec/);
|
||||
|
||||
Reference in New Issue
Block a user