Compare commits

..
Author SHA1 Message Date
Vito Sansevero 112d66f4e1 Merge remote-tracking branch 'origin/main' into feat/331-collections-pr1
Main's #332 PR 2 (#406) admits Images by each file's effective category
(own, else inherited) and passes an own-category map through the listing;
this branch replaced the favorites set with collection or favorite members.
The admission now returns both (members, own), and a new test covers a
collection filter together with an image's own category.
2026-10-28 10:56:53 -07:00
Vito cdb771d722 feat(images): filter, inspect and count by each image's effective category (#332, PR 2) (#406)
* feat(images): store images' own categories in the library database (#332)

One new table, image_categories, and its category index, created empty
by a forward-only migration that reads and writes only its own objects.
Startup checks the structure; restore also checks the values. The store
sets, clears and reads own categories by file key (the deepest output
source holding the file and the path below it), and copies a removed
duplicate's category to a kept copy that has none.

* feat(images): duplicate cleanup keeps images' own categories (#332)

The kept copy's own category wins; a kept copy without one receives the
removed copy's, keyed by the deepest output source holding each file.
The removed copy's row is never deleted, so a give-back or a failed
delete loses nothing. The job log reports how many were kept.

* feat(images): set, remove and list images' own categories (#332)

PUT /prompt_manager/images/category sets the own category of up to 500
listed files, DELETE removes it (listed or not), and
GET /prompt_manager/images/categories lists the values whose file is
listed. Targets are the Images listing identity of an output source,
checked by media_target (shared with Collections). The routes never
touch prompts, recipes or links; POST /prompt_manager/bulk/category is
unchanged.

* feat(vocabulary): count and rename image categories (#332)

Tags & categories lists every category value with image_count, the
stored image_categories rows holding it. Values only images hold follow
every prompt-backed value (first_pid at SQLite's INTEGER maximum), so
keyset cursors cross the boundary. /library/options keeps its own query
and never lists image-only values.

Manage category counts image rows (impact.images), treats an image-held
target as existing, renames image rows in the same transaction, allows
renaming an image-only value, and hashes the sorted image rowids into
the impact token so a change between preview and apply is refused with
vocabulary-conflict. Every image-category statement reads
idx_image_categories_category.

* feat(images): the category filter matches each image's effective category (#332)

An image's own category wins; without one it inherits its link's: the
recipe's category when set, else the positive prompt's (MC4, one rule
with the inspector chip). Ambiguous and unlinked images inherit nothing,
so an unlinked image with its own category is now listed. The link read
no longer filters by category in SQL; the page filter, select all N
matching and saved views compare the effective value. A failed
own-category read answers the existing listing-failed 500.

* fix(images): key symlinked output folders and keep key work off the event loop (#332)

Review fixes: a source matches by its configured spelling too, so a
detected output folder that is a symlink keys its files (routes and
duplicate cleanup). File keys are computed in the executor. A path no
source holds is a plain 403, a lone surrogate a 400. The install repairs
a table that lacks its category index instead of refusing to open.

* feat(inspector): show an image's own category in the inspector and History (#332)

The inspector's library facts gain category_scope "image": an image's
own category overrides the inherited one for any link status (linked,
ambiguous, none). Images reads it for output sources only; History reads
it once per inspection inside its snapshot, keyed by the same file key
from the symlink-resolved output folder. A failed read hides the chip
rather than show a possibly wrong value. The chip's hidden prefix names
the origin ("Image category: "); the visible origin text comes in PR 3.

* perf(images): keep the category filter's link read narrow; record D141, D142, D145 (#332)

Under a category the link read keeps only links whose raw recipe or
positive category is the value (Python still decides the inherited
one), the links of images whose own category is the value are read by
path, and plain reads skip the inherited-category work. The own-category
map is keyed by normalized path. TESTING.md records D141, D142 and D145
and amends D13, D36 and D39; the filter budget seeds 2,000 own
categories.

* feat(images): set or remove a category on all N matching images (#332)

The category routes also accept the Images listing filter of "Select
all N matching" (#392) instead of explicit items, resolved now and
capped at 500 images like an explicit selection.

* fix(images): Windows path keys for own categories; plan check for the category read (#332)

Review fixes: on Windows the own-category map keys each file with
normalize_image_path (abspath drops trailing dots and spaces), and the
plan test also covers the category's narrowing clause.

* test(vendor): count the #332 host case in the execution matrix

image-attribution.spec.mjs gains the image category route case, so the
execution config lists 92 tests (3 attribution cases per project).

* test(legacy): record the image category routes in the legacy contract snapshot (#332)

#332 adds PUT/DELETE /prompt_manager/images/category and
GET /prompt_manager/images/categories; the legacy pages' routes are
unchanged.
2026-10-28 17:19:11 +00:00
Vito 55fa2c4870 fix(runs): count cached PromptManager nodes fed by a runtime text link on every job (#347) (#423)
* fix(usage): count cached re-runs of runtime-linked prompt text (#347)

A PromptManager-family node whose text comes from a runtime link
(PreviewAny, an LLM, a list node) counted only the job that executed it;
later jobs served from ComfyUI's output cache counted nothing.

Each execution now records a receipt (digests of its STRING outputs and
hashes of the text it stores) per run key and queue item. On
execution_cached the receipt whose digests equal the cache entry's text is
replayed once for the job, with the node's role and what the queue hook
already saw recomputed from the job's own graph. The executing node and the
replay count through one shared own_counts. Unmatched cache entries count
nothing, as before; a node that executes after its replay in the same job
skips its own count. No schema change, no backfill.

* fix(usage): never replay a node twice in one job; dedupe receipts once complete (#347)

A repeated execution_cached id or message for the same job now counts the
node once. Receipts replace equal older ones only when their job finishes
(or the next job starts), so a list still recording [A] of [A, B] no longer
drops an older [A] receipt. Node-level tests check the receipt each
execution records.

* test(comfyui): count #347's cached-run cases in the execution suite pin (#347)

The pin listed 19 execution-text-link cases per project; #347 adds five node
kinds and eight one-renderer cases, so the suite lists 32 per project and 118
in total with #332's image-attribution case from main.
2026-10-28 16:37:24 +00:00
Vito Sansevero 6325d33831 Merge remote-tracking branch 'origin/main' into feat/331-collections-pr1 2026-10-28 09:23:15 -07:00
Vito 11422d4584 perf(images): read prompt links in one SQLite step (#354) (#418)
* perf(images): read prompt links in one SQLite step (#354)

get_image_link_matches streamed one sqlite3_step per linked row, and
Python's sqlite3 releases and retakes the GIL around every step. While
ComfyUI generates, each retake waits on the busy thread, so the
first-page link read of the workspace Images page stretched from tens of
milliseconds to seconds.

SQLite now returns the unchanged filter SELECT's rows as one JSON value
(json_group_array), decoded with json.loads: a constant two or three GIL
hand-backs instead of one per row. Nothing is cached. Any
OperationalError from the wrapped statement, or a BLOB file_path or
rating, makes that read stream as before with one warning per process; a
missing JSON function stops the attempt for the rest of the process.

tests/test_database_link_matches.py proves the result equals the
streamed read on every filter combination and on NULL ratings, quoted,
backslashed, non-ASCII and multi-line paths, ambiguous and pair links,
an empty database and both fallbacks; the structural budget test also
checks the wrapped query plan.

* fix(images): keep the maintenance link reader, refine the one-step fallback (#354)

The one-step reader was named _link_rows, which shadowed
MaintenanceOperations._link_rows on PromptDatabase and broke the Missing
links preview and apply; it is now _one_step_link_rows, and a test pins
that the maintenance reader is not shadowed.

An interrupted read (Move library retiring the connection) is raised as
before instead of falling back and spending the one warning; a JSON text
over SQLite's length limit (DataError) streams like the other value
failures. The equivalence test compares value types per path rather than
by row order, and a test checks that the read leaves no WAL snapshot
pinned.

* docs(testing): name every one-step link read fallback (#354)
2026-10-28 16:00:15 +00:00
Vito Sansevero 8aebfd4aaa Merge remote-tracking branch 'origin/main' into feat/331-collections-pr1
# Conflicts:
#	database/duplicate_copies.py
#	database/models.py
#	py/api/duplicate_files.py
#	tests/test_duplicate_image_categories.py
2026-10-28 08:52:10 -07:00
Vito b5a56bee0b feat(images): store images' own categories and set them by route (#332, PR 1) (#401)
* feat(images): store images' own categories in the library database (#332)

One new table, image_categories, and its category index, created empty
by a forward-only migration that reads and writes only its own objects.
Startup checks the structure; restore also checks the values. The store
sets, clears and reads own categories by file key (the deepest output
source holding the file and the path below it), and copies a removed
duplicate's category to a kept copy that has none.

* feat(images): duplicate cleanup keeps images' own categories (#332)

The kept copy's own category wins; a kept copy without one receives the
removed copy's, keyed by the deepest output source holding each file.
The removed copy's row is never deleted, so a give-back or a failed
delete loses nothing. The job log reports how many were kept.

* feat(images): set, remove and list images' own categories (#332)

PUT /prompt_manager/images/category sets the own category of up to 500
listed files, DELETE removes it (listed or not), and
GET /prompt_manager/images/categories lists the values whose file is
listed. Targets are the Images listing identity of an output source,
checked by media_target (shared with Collections). The routes never
touch prompts, recipes or links; POST /prompt_manager/bulk/category is
unchanged.

* fix(images): key symlinked output folders and keep key work off the event loop (#332)

Review fixes: a source matches by its configured spelling too, so a
detected output folder that is a symlink keys its files (routes and
duplicate cleanup). File keys are computed in the executor. A path no
source holds is a plain 403, a lone surrogate a 400. The install repairs
a table that lacks its category index instead of refusing to open.

* feat(images): set or remove a category on all N matching images (#332)

The category routes also accept the Images listing filter of "Select
all N matching" (#392) instead of explicit items, resolved now and
capped at 500 images like an explicit selection.

* test(vendor): count the #332 host case in the execution matrix

image-attribution.spec.mjs gains the image category route case, so the
execution config lists 92 tests (3 attribution cases per project).

* test(legacy): record the image category routes in the legacy contract snapshot (#332)

#332 adds PUT/DELETE /prompt_manager/images/category and
GET /prompt_manager/images/categories; the legacy pages' routes are
unchanged.
2026-10-28 15:29:41 +00:00
Vito 95586e6395 refactor(gallery): open System for gallery settings, maintenance and tagging models (#330, PR 3) (#408)
* test(legacy): snapshot the routes and the settings response

Pins every registered route and the key paths of GET /prompt_manager/settings
while System S13 (#330) moves the legacy pages' tools to native System. Recorded
on the unchanged main; cutover (#334) removes it.

* test(legacy): prove each removed legacy control has native coverage

TESTING.md gains the "Legacy System UI removed by S13" table (no rows yet).
legacy-coverage.test.mjs fails when a removal cites a capability that is not
covered, copies another native test, lists an id or method twice, or leaves
the id or method on its page. legacy-system-links.spec.mjs holds the System
link table and guards the kept admin flows.

* refactor(admin): open System for Diagnostics and logs

The admin page's Diagnostics and View Logs buttons become links that open
System › Diagnostics & logs in a new tab. The legacy Diagnostics and System Logs
dialogs and their methods go; TESTING.md lists them against their covered
capabilities (rows B1–B7).

* refactor(admin): open System for settings and keep page options

Settings opens System › Overview in a new tab. The settings dialog becomes Page
options (opened from the list header) with only the Infinite Scrolling default,
and its Save posts {"infinite_scroll": …} alone, so a stale legacy form can no
longer overwrite scan folders or other values edited in System (#330 S13-9).
Scan folders, monitoring status, auto-hide, display mode, worker threads,
LoraManager and the LoRA import leave the page (rows B8–B14).

* refactor(admin): open System for maintenance

Maintenance opens System › Maintenance in a new tab, where each operation runs
as a reviewed Activity job. The one-click legacy Database Maintenance dialog
and its methods go (rows B15–B18); formatFileSize stays for the image modal.

* refactor(admin): open System for backup and restore

Backup DB and Restore DB open System › Backup & export in a new tab, where a
restore is staged, reviewed and run as an Activity job. The one-click Restore
dialog and its methods go (rows B19–B20); the scan dialog's quick backup stays.
The removed browser tests' data checks live on at the route: a new round trip
test restores every prompt with its tags and image links.

* test(legacy): find a removed modal id passed to showModal and settle the links spec

The removal check also finds an id passed as a string to a ...Modal() call and no
longer reads data-id= as id=. The no-request test waits for the page's own
start-up reads to finish, the sort check compares as the server sorts, and the
kept-flow test checks the added prompt shows in the list.

* test(admin): open Page options after the saved default is read

Drops the leftover LoraManager section comment above the Page options save.

* refactor(gallery): open System for settings and maintenance

The gallery's Settings opens System › Appearance in a new tab. Its settings
dialog becomes Page options (a new header button) with only the page's own
settings: thumbnail startup check, lazy loading, items per page, grid columns,
infinite scrolling, debug mode and API timeout. Save and Reset keep every other
stored key, so the Appearance values System saves survive (S13-13). Generate
Thumbnails, Clear cache, Rescan and Find duplicates leave the page; Page options
links System › Maintenance and Sources & scans (rows B21–B28).

* refactor(autotag): manage tagging models in System

In both auto-tag dialogs a model that is not downloaded shows a Manage in
System link to System › Auto-tagging (a new tab) instead of a Download button,
and the gallery's Unload Model becomes the same link. The legacy download
dialog and the download, cancel and unload methods go (rows B29–B31); the batch
and review stay. A test hook marks a fake model loaded.

* docs(system): name System as the home for settings and maintenance

README and TESTING.md send settings, maintenance, backups, diagnostics, the
gallery's tools and auto-tag model downloads to their System section and name
each legacy page's Page options.

* fix(gallery): keep the default view in Page options and save viewer choices alone

System › Appearance's Default layout is the workspace's own key, so the legacy
gallery's Default View Mode stays a page option instead of losing its only
control. The viewer's mute and loop buttons remember just that choice instead of
saving the Page options dialog's unfilled inputs over the stored page options.
The pending-thumbnails toast names System › Sources & scans.

* test(legacy): check each System link sends no request in its own test

The no-request check runs inside each link's test, so a slow page no longer
times out the whole table, and it ignores lazy thumbnail loads the layout
starts on its own.

* test(legacy): a lazy thumbnail load is not a System link's request

* fix(gallery): let the header buttons wrap on a narrow screen

With Page options beside Settings the header no longer fits a phone's width;
it wrapped off the page and covered the Page options dialog in the touch runs.

* feat(gallery): remove Default View Mode from Page options (#330)

System > Appearance's Default layout is the one default-layout control, as
the S13 design's G5 says. The legacy gallery's Page options no longer show
or save a default view and Reset no longer resets it. A stored legacy
defaultViewMode still picks the view the gallery opens in (grid otherwise),
and Save and Reset keep that key for the workspace-preferences migration.
2026-10-28 13:31:23 +00:00
Vito 5583159497 feat(inspector): use recorded run history when the graph alone can't name the prompt (#386) (#396)
* feat(inspector): use recorded run history when the graph alone can't name the prompt (#386)

When a file's own graph leaves the prompt open, both inspect routes ask
recorded run history about the file (one run reported it and it was written
during that run). The save node that run reported it from anchors the walk
like a file-name match, so dataset branches that share one saver prefix
show their own prompt, labelled "Recorded run". Still open at one sampler
whose PromptManager text was recorded at execution (a language model's
output), the recorded text shows as "Execution recorded". Results the file
resolves alone never change; a failed lookup judges the file alone.

* fix(inspector): recorded texts follow history's saver; conflicts anchor nothing (#386)

Review follow-ups: the texts PromptManager recorded are looked up at the
sampler history's saver leads to, never at a sampler the unanchored walk
named; a file name naming a different saver than history anchors nothing;
outputs and texts are read in one transaction; "Recorded run" ranks below
every workflow-inferred source; unexpected lookup errors log at error level.

* fix(history): look up recorded runs by history's own spelling of the file (#386)

The history route inspects a case-normalized target on Windows, while
history matches file names exactly, so the occurrence's reference names
the file.
2026-10-28 13:29:29 +00:00
Vito 3f085d0724 ci: full matrix once per PR in the merge queue; Windows/macOS non-blocking (#420) (#421)
* ci: run the full matrix once per PR, in the merge queue (#420)

Pull requests now run a fast subset: Python 3.12 on ubuntu-latest, the
coverage gate, black/flake8, frontend modules, bandit and the Chromium/touch
frontend coverage shards plus merge. The merge queue, pushes to main and
manual runs keep the full matrix: Linux/Windows/macOS x Python 3.12-3.14,
real ComfyUI compatibility and the Firefox/WebKit page workflows.

A dropped matrix leg never reports a check, so the new required "all tests"
gate requires every test.yml job to succeed outside pull requests and real
ComfyUI to be skipped only on pull requests; viewerjs does the same for the
browser shards. branch-protection.json drops the eight legs a pull request
no longer runs and adds "all tests".

* docs(ci): ruleset edit keeps every rule and integration id; find merge-queue runs

* ci: move Windows and macOS out of the blocking path (#420)

The maintainer deferred Windows and macOS to the end of the redesign.
The merge queue now runs Linux on Python 3.12-3.14, real ComfyUI, the
frontend coverage shards and the Firefox/WebKit page workflows; pull
requests keep the Linux 3.12 subset. The six Windows/macOS legs move to
the new portability.yml, which runs on push to main, nightly and on
demand, keeps the 27-minute limit and is never a required check. Red
runs there are fixed at their cause before the cutover (#334).

* docs(ci): required checks live in the ruleset and the classic protection (#420)
2026-10-28 12:44:06 +00:00
Vito Sansevero 1701e927f7 test(duplicates): a removed copy with a collection and a category moves both (#331, #332) 2026-10-28 04:35:36 -07:00
Vito Sansevero 894b48964c Merge remote-tracking branch 'origin/feat/332-image-categories-pr1' into feat/331-collections-pr1
Keep both features: a duplicate removal moves collection memberships and
copies the image category; both migrations and schema checks run; the
legacy route snapshot now records the six collection routes beside the
image category routes.
2026-10-28 04:34:24 -07:00
Vito e769845417 refactor(admin): open System for settings, maintenance, backup and restore (#330, PR 2) (#407)
* test(legacy): snapshot the routes and the settings response

Pins every registered route and the key paths of GET /prompt_manager/settings
while System S13 (#330) moves the legacy pages' tools to native System. Recorded
on the unchanged main; cutover (#334) removes it.

* test(legacy): prove each removed legacy control has native coverage

TESTING.md gains the "Legacy System UI removed by S13" table (no rows yet).
legacy-coverage.test.mjs fails when a removal cites a capability that is not
covered, copies another native test, lists an id or method twice, or leaves
the id or method on its page. legacy-system-links.spec.mjs holds the System
link table and guards the kept admin flows.

* refactor(admin): open System for Diagnostics and logs

The admin page's Diagnostics and View Logs buttons become links that open
System › Diagnostics & logs in a new tab. The legacy Diagnostics and System Logs
dialogs and their methods go; TESTING.md lists them against their covered
capabilities (rows B1–B7).

* refactor(admin): open System for settings and keep page options

Settings opens System › Overview in a new tab. The settings dialog becomes Page
options (opened from the list header) with only the Infinite Scrolling default,
and its Save posts {"infinite_scroll": …} alone, so a stale legacy form can no
longer overwrite scan folders or other values edited in System (#330 S13-9).
Scan folders, monitoring status, auto-hide, display mode, worker threads,
LoraManager and the LoRA import leave the page (rows B8–B14).

* refactor(admin): open System for maintenance

Maintenance opens System › Maintenance in a new tab, where each operation runs
as a reviewed Activity job. The one-click legacy Database Maintenance dialog
and its methods go (rows B15–B18); formatFileSize stays for the image modal.

* refactor(admin): open System for backup and restore

Backup DB and Restore DB open System › Backup & export in a new tab, where a
restore is staged, reviewed and run as an Activity job. The one-click Restore
dialog and its methods go (rows B19–B20); the scan dialog's quick backup stays.
The removed browser tests' data checks live on at the route: a new round trip
test restores every prompt with its tags and image links.

* test(legacy): find a removed modal id passed to showModal and settle the links spec

The removal check also finds an id passed as a string to a ...Modal() call and no
longer reads data-id= as id=. The no-request test waits for the page's own
start-up reads to finish, the sort check compares as the server sorts, and the
kept-flow test checks the added prompt shows in the list.

* test(admin): open Page options after the saved default is read

Drops the leftover LoraManager section comment above the Page options save.

* test(legacy): a lazy thumbnail load is not a System link's request
2026-10-28 11:25:20 +00:00
Vito Sansevero 463bda0e38 Merge remote-tracking branch 'origin/main' into feat/331-collections-pr1 2026-10-28 04:19:13 -07:00
Vito Sansevero 74e7cdaa62 test(legacy): record the image category routes in the legacy contract snapshot (#332)
#332 adds PUT/DELETE /prompt_manager/images/category and
GET /prompt_manager/images/categories; the legacy pages' routes are
unchanged.
2026-10-28 03:26:23 -07:00
Vito Sansevero 12c27ae2b0 Merge remote-tracking branch 'origin/main' into feat/332-image-categories-pr1 2026-10-28 03:26:14 -07:00
Vito 9a90c7677f fix(exit): stop running jobs, downloads and encoders at Ctrl+C instead of waiting for them (#376) (#417)
* fix(exit): stop running jobs before Python joins the executor threads (#376)

ComfyUI returns from main.py after Ctrl+C without closing its event loop,
and Python joins every executor thread before any atexit callback runs, so
a scan batch, thumbnail batch or model download in flight kept ComfyUI
alive until it ended (on Windows each ffmpeg up to its 30 s timeout).

The package now registers a threading._register_atexit hook (after
concurrent.futures', so it runs before that join) that sets an exit signal
and kills the running encoder processes. Job.cancelled follows the signal,
so every check a job already makes between files stops it at exit;
map_parallel starts no further item, the scan walk stops between folders and
a tagging model download stops before its next chunk. Without the internal
hook exit behaves as before.

* fix(exit): abort Xet downloads and stop LoRA staging at exit (#376)

Review follow-ups: Xet transfers report progress from native code and ignore
the chunk checkpoint, so the exit hook calls huggingface_hub's own
abort_xet_session when it is loaded; an aborted download logs as stopped.
LoRA import staging reads the exit signal itself, so the legacy LoRA scan
stops too. The test's fake encoder outlives the ceiling by less, and
TESTING.md states the download limits of older huggingface_hub.
2026-10-28 10:15:28 +00:00
Vito Sansevero 94089d452f Merge remote-tracking branch 'origin/main' into feat/331-collections-pr1 2026-10-27 20:01:17 -07:00
Vito Sansevero 718d4a5cf1 test(browser): install the slideshow clock before wrapping setTimeout
The slideshow fixture wrapped window.setTimeout before page.clock.install(),
so the viewer's timer ran on real time and clock.runFor() did not drive it;
a late WebKit timer failed 'slideshow exposes exact state labels' on CI.
2026-10-27 20:01:13 -07:00
Vito Sansevero 1110ccbd33 test(vendor): count the #332 host case in the execution matrix
image-attribution.spec.mjs gains the image category route case, so the
execution config lists 92 tests (3 attribution cases per project).
2026-10-27 19:57:13 -07:00
Vito Sansevero 20b6514882 feat(images): set or remove a category on all N matching images (#332)
The category routes also accept the Images listing filter of "Select
all N matching" (#392) instead of explicit items, resolved now and
capped at 500 images like an explicit selection.
2026-10-27 18:47:58 -07:00
Vito Sansevero e1ae60d5eb fix(collections): address review findings 2026-10-27 18:37:30 -07:00
Vito Sansevero 0de5c0a4e7 fix(images): key symlinked output folders and keep key work off the event loop (#332)
Review fixes: a source matches by its configured spelling too, so a
detected output folder that is a symlink keys its files (routes and
duplicate cleanup). File keys are computed in the executor. A path no
source holds is a plain 403, a lone surrogate a 400. The install repairs
a table that lacks its category index instead of refusing to open.
2026-10-27 18:31:43 -07:00
Vito Sansevero df6f538283 feat(images): set, remove and list images' own categories (#332)
PUT /prompt_manager/images/category sets the own category of up to 500
listed files, DELETE removes it (listed or not), and
GET /prompt_manager/images/categories lists the values whose file is
listed. Targets are the Images listing identity of an output source,
checked by media_target (shared with Collections). The routes never
touch prompts, recipes or links; POST /prompt_manager/bulk/category is
unchanged.
2026-10-27 18:21:35 -07:00
Vito Sansevero 0237fba77a feat(images): duplicate cleanup keeps images' own categories (#332)
The kept copy's own category wins; a kept copy without one receives the
removed copy's, keyed by the deepest output source holding each file.
The removed copy's row is never deleted, so a give-back or a failed
delete loses nothing. The job log reports how many were kept.
2026-10-27 18:17:28 -07:00
Vito Sansevero cd573ccf1a feat(images): store images' own categories in the library database (#332)
One new table, image_categories, and its category index, created empty
by a forward-only migration that reads and writes only its own objects.
Startup checks the structure; restore also checks the values. The store
sets, clears and reads own categories by file key (the deepest output
source holding the file and the path below it), and copies a removed
duplicate's category to a kept copy that has none.
2026-10-27 18:15:47 -07:00
Vito Sansevero 1db6088b2f Merge remote-tracking branch 'origin/main' into feat/331-collections-pr1 2026-10-27 18:11:51 -07:00
Vito Sansevero cb874d7826 feat(collections): add collection routes and the Images collection filter 2026-10-27 18:02:29 -07:00
Vito Sansevero b543b7db89 feat(collections): add, remove and rename collections and keep members through duplicate cleanup 2026-10-27 17:46:12 -07:00
Vito Sansevero 2455ab21dd feat(collections): store collections in the library database 2026-10-27 17:43:08 -07:00
96 changed files with 10951 additions and 2917 deletions
+29 -38
View File
@@ -56,7 +56,7 @@ Both nodes include the complete PromptManager feature set:
- **📊 Analytics**: Track prompt usage patterns and effectiveness over time
- **🔬 Workflow Analysis**: Extract and analyze ComfyUI workflow data from PNG images
- **📋 Metadata Viewer**: Standalone tool for analyzing ComfyUI-generated images
- **🛠️ System Management**: Built-in diagnostics, backup/restore, and maintenance tools
- **🛠️ System Management**: Built-in diagnostics, backup/restore, and maintenance tools in System
- **🏷️ AI AutoTag**: Automatically tag your image collection using WD14 or JoyCaption vision models
- **🔗 LoRA Manager Integration**: Import LoRA metadata and preview images from [ComfyUI-Lora-Manager](https://github.com/willmiao/ComfyUI-Lora-Manager)
@@ -206,14 +206,11 @@ match Classic and Focused.
- **Responsive Design**: Works on desktop, tablet, and mobile devices
- **Real-time Search**: Instant search results as you type
- **Bulk Operations**: Edit multiple prompts simultaneously with batch tools
- **Settings Panel**: Configure behavior and display options
- **System Management**: Complete toolset for database operations:
- **Settings and System tools live in System**: the dashboard's **Settings** opens System › Overview, **Maintenance** opens System › Maintenance, **Backup DB** and **Restore DB** open System › Backup & export, and **Diagnostics** and **View Logs** open System › Diagnostics & logs, each in a new tab. System reviews every destructive operation before it runs it as an Activity job.
- **Page options**: the list header's **Page options** keeps the prompt list's Infinite Scrolling default
- **Dashboard tools**:
- **🔍 Metadata Viewer**: Launch standalone PNG analysis tool
- **📸 Scan Images**: Bulk scan output directories for ComfyUI images
- **💾 Backup/Restore**: Database backup and recovery operations
- **🔧 Maintenance**: Database optimization and cleanup tools including orphaned prompt removal
- **📋 Logs**: Comprehensive logging system with download capability
- **🔍 Diagnostics**: System health checks and troubleshooting
- **📸 Scan Images**: Bulk scan output directories for ComfyUI images (with a quick backup first)
- **Export Tools**: Download prompts and metadata in various formats
- **Progress Tracking**: Real-time progress indicators for long operations
- **Error Handling**: Graceful error display with recovery suggestions
@@ -333,14 +330,11 @@ The comprehensive gallery system provides complete access to all your ComfyUI ou
#### **⚙️ Gallery Settings**
Configure your viewing experience in the settings panel:
The gallery's **Settings** opens System › Appearance in a new tab: image quality, details on cards, collapsed viewer details, video autoplay/mute/loop, metadata loading and file paths are set there and apply to this page too.
- **Display Settings**: Default view, items per page, grid columns, media info, infinite scrolling
- **Performance Settings**: Image quality, lazy loading
- **Video Settings**: Autoplay, mute by default, loop videos
- **Thumbnail Management**: Generate optimized thumbnails for faster loading (runs on the configured worker threads; a run that stops at the per-run cap tells you how many remain)
- **Find Duplicate Images**: Hashes every file on the worker threads with a progress bar; close the window and it keeps running, results wait in a pill
- **Rescan Output Folder**: Drops the server-side listing cache and rewalks the folders
The gallery's **Page options** keep the page's own settings: the thumbnail check at startup, lazy loading, items per page, grid columns, infinite scrolling, debug mode and API timeout. Reset puts back only these.
Thumbnail generation, the server thumbnail and browser cache clears, refreshing the file listing and finding duplicate images are in System › Sources & scans and System › Maintenance (Page options links both); each runs as an Activity job you can follow after closing the page.
#### **🖼️ Image Viewing Experience**
@@ -398,7 +392,7 @@ Import existing ComfyUI images into your database:
- Any errors or issues encountered
5. **Access imported data** through the normal gallery interface
The scan reads PNG text chunks and A1111-style parameters, and, when `ffprobe` is installed, the prompt ComfyUI's video savers embed in MP4/MOV/WebM files. Files are read on the number of **worker threads** set in Settings → Performance (default: half your cores), which makes a large library several times faster.
The scan reads PNG text chunks and A1111-style parameters, and, when `ffprobe` is installed, the prompt ComfyUI's video savers embed in MP4/MOV/WebM files. Files are read on the number of **worker threads** set in System › Sources & scans (default: half your cores), which makes a large library several times faster.
### 🏷️ AI AutoTag
@@ -414,7 +408,7 @@ Automatically tag your entire image collection using AI vision models:
1. **AutoTag (Batch Mode)**: Tag your entire collection automatically
- Click **"🏷️ AutoTag"** in the admin dashboard
- Choose your model type (FP16 or GGUF)
- Choose your model type (FP16 or GGUF). A model that is not downloaded yet shows **Manage in System**: download (or unload) it in System › Auto-tagging, then reopen the dialog
- Select how to handle already-tagged images:
- **Skip images with existing tags**: Preserve your meticulous manual tagging work
- **Re-tag all images**: Overwrite existing tags with fresh AI analysis
@@ -445,7 +439,7 @@ Adjust the system prompt to match your tagging style:
#### **Requirements**
- Sufficient VRAM for the chosen model (FP16 requires more, GGUF is lighter)
- Models are downloaded automatically on first use
- Models are downloaded in System › Auto-tagging
- Feature requests and improvements are welcome!
### 📁 Folder Filter
@@ -470,18 +464,18 @@ If you use [ComfyUI-Lora-Manager](https://github.com/willmiao/ComfyUI-Lora-Manag
#### **Enabling the Integration**
1. Open the admin dashboard and click **Settings**
2. Scroll to the **Integrations** section — PromptManager will auto-detect if LoRA Manager is installed
1. Open System › Integrations (the admin dashboard's **Settings** opens System)
2. PromptManager auto-detects whether LoRA Manager is installed
![LoRA Enabled](images/pm-lora-enabled.png)
3. Toggle the **LoraManager** switch to enable
4. Optionally enable **Auto-inject trigger words** to automatically append LoRA trigger words when `<lora:name:weight>` is detected in your prompts
5. Click **Save Settings** (requires a ComfyUI restart to take effect)
5. Click **Save preferences** (requires a ComfyUI restart to take effect)
#### **Importing LoRA Data**
Click **Import LoRA Data** in the settings panel to start the import. A progress popup shows real-time status as each LoRA is processed:
Click **Review import** in System › Integrations; Start runs the import as an Activity job:
- Reads LoRA Manager's own settings the way LoRA Manager does: from the folder in `LORA_MANAGER_SETTINGS_DIR`, from the plugin's `settings.json` in portable mode, or from your user config folder, and only the keys that say where its files are
- Finds each LoRA's metadata in both storage modes: beside the model (the default) or in the central sidecar folder, and imports only LoRAs whose model file still exists in a LoRA folder (the review says how many metadata files it skipped)
@@ -502,7 +496,7 @@ A CivitAI API key is **required if your library contains NSFW LoRAs** — CivitA
To add your key:
1. Get your API key from [civitai.com/user/account](https://civitai.com/user/account)
2. Paste it into the **CivitAI API Key** field in the Integrations settings
2. Paste it into the **CivitAI API key** field in System › Integrations
3. Save and re-import to download any previously skipped images
The key is sent only over HTTPS to CivitAI's own hosts. CivitAI serves images through a redirect to its storage host; PromptManager follows such a redirect only over HTTPS to a host on CivitAI's own `civitai.com`, `civitai.red` or `civitai.green` domains, at most three hops, and never sends the key with it. When downloads fail, the import's status line says how many; when none could be downloaded, it starts with a warning. The server log has one warning per failed image with the reason (HTTP status, refused redirect, size cap, unreadable image, timeout or network error).
@@ -517,23 +511,20 @@ The comprehensive web interface provides:
- **Filter Options**: Filter by category, tags, rating, and date
- **Bulk Operations**: Select multiple prompts for batch editing
- **Export Tools**: Download your prompt collection
- **Settings Panel**: Configure auto-save and display options
- **Diagnostics**: System health checks and troubleshooting
- **Settings**: System (new tab) for settings, maintenance, backups and diagnostics; **Page options** for the page's own display settings
### 🔧 Database Maintenance
Keep your database optimized with built-in maintenance tools:
Keep your database optimized in System › Maintenance (the admin dashboard's **Maintenance** opens it in a new tab):
1. **Open the admin dashboard** and click **"🔧 Maintenance"**
2. **Select maintenance operations**:
- **Remove Duplicates**: Find and merge duplicate prompts while preserving images
- **Vacuum Database**: Optimize database file size and performance
- **Clean Orphaned Images**: Remove image records for missing files
- **🚮 Prune Orphaned Prompts**: Remove prompts that have no linked images
- **Check Consistency**: Verify database integrity and relationships
- **Database Statistics**: Generate detailed database reports
3. **Run selected operations** and monitor progress in real-time
4. **Review results** showing detailed information about each operation
- **Duplicate prompt text**: merge duplicate prompts while preserving images
- **Optimize database**: reclaim unused SQLite space
- **Missing-file links**: remove image records for missing files
- **Unused prompt components**: remove prompts nothing uses
- **Run health checks**: verify database integrity and relationships (the Health report)
- **Database size and prompt counts**: System › Diagnostics & logs (Database) and the Overview's Statistics
Each operation opens a review that shows what it will change and the safety backup it takes, then runs as an Activity job.
### Database Location
@@ -835,8 +826,8 @@ Before upgrading or changing this setting:
1. Identify the database your current installation actually uses and confirm its
library contents. Check both the configured path and the prior fallback
`<extension>/prompts.db`; keep track of which existing file you intend to use.
2. Preserve a consistent backup of the active database with the dashboard's
Backup/Restore backup action or SQLite's online backup API. Do not copy only
2. Preserve a consistent backup of the active database with
System › Backup & export's Back up or SQLite's online backup API. Do not copy only
the main file of a live SQLite database: committed data may still be in its
`-wal` file. For a filesystem copy, stop ComfyUI cleanly first and preserve any
remaining `-wal` file together with the main database.
+253 -19
View File
@@ -322,6 +322,26 @@ from the tests' package unload (`cleanup_package`), and it removes its own
ComfyUI-Manager restart through `os.execv`, a closed Windows console, or
ComfyUI Desktop stopping its server. SQLite's WAL journal covers those.
Jobs running at exit (#376): before any `atexit` callback, Python joins every
executor thread (`concurrent.futures` registers that join with
`threading._register_atexit`). The package registers its own callback there
(`utils/exit_signal.install`, after `concurrent.futures`, so it runs first):
it sets the exit signal and kills the running encoder processes
(`shutdown_encoders`, the ffmpeg that never sees Ctrl+C on Windows). The call in
flight then stops at its next safe point: `Job.cancelled` is True, so every
check a job already makes between files stops it, `map_parallel` starts no
further item (`ExitRequested`), the scan walk stops between folders, LoRA
import staging stops before its next sidecar (also for the legacy scan) and a
tagging model download stops before its next chunk where huggingface_hub takes
`tqdm_class` (1.x removes its temp file; 0.x `hf_hub_download` does not take it
and finishes the file). Xet transfers ignore that checkpoint, so the hook also
calls hub's own `abort_xet_session` when huggingface_hub is loaded. The
trigger-word cache load (small JSON sidecars) has no checkpoint. No check sits
inside a library write. That hook is
CPython-internal; without it (or when it is refused) `install` returns False and
exit waits for the call in flight, as before. `tests/test_exit_cancel.py` runs a
child process with three such calls running and requires its exit within 15 s.
Stop contract: `PromptTracker.shutdown(timeout)`,
`ImageGenerationHandler.shutdown(timeout)` (restartable),
`ImageGenerationHandler.close(timeout)` (final: `schedule()` refuses and a
@@ -371,6 +391,10 @@ A spec performance budget gets two tests:
`test_filter_budget` (Images filters, spec D12: p95 < 150 ms at 50,000 files
and 20,000 links).
A hot SQLite read that returns many rows aggregates them into one step
(`json_group_array`) instead of stepping once per row, so it does not wait for
the GIL once per row while ComfyUI generates (#354, `get_image_link_matches`).
The `perf budget` job in `.github/workflows/test.yml` is the reference
measurement: `ubuntu-latest`, one Python, no coverage, only the wall-clock tests.
It runs on every push to `main` and on `workflow_dispatch`, never on pull
@@ -869,6 +893,25 @@ unattributed stage, when it is a plain first stage: not a detailer, no prompt ov
set, and no listed sampler in its own latent ancestry (else it stays unattributed) (`tests/test_resolver_saver_anchor.py`, shapes in
`tests/saver_anchor_graphs.py` from the maintainer's `gpu_unlock_bench`, dataset and
ANIMA runs; browser parity through the `anchor-*` cases, which carry `__file_name`).
Recorded run history (#386): both inspect routes (Images and a history occurrence)
ask `database.run_attribution.recorded_output` inside the file's stat window. It
answers only for a file that exactly one recorded run reported and that was written
during that run (the #342 rules): the one output node that run reported it from (none
when several did) and, per sampler anchor, the effective texts PromptManager recorded
at execution (a role only inferred from the graph is none). Only when the file alone
leaves the prompt open (`ambiguous` or `unresolved`) the inspector resolves again
anchored on that node (`recorded_saver_id`, walked like a file-name match; a file
name naming a different saver contradicts history and anchors nothing), and keeps
that result only when history's saver chose its samplers (`recorded_anchor`) and it
names the positive (resolved, or a guider with no negative); its `graph-inferred`
sides read `recorded-run` ("Recorded run" badge, ranked below every "Workflow
inferred" source). Still open at the one sampler that saver leads to, with that
sampler's positive text recorded (a language model's output feeding PromptManager),
the recorded texts show as `execution`. Outputs and texts are read in one
transaction. A result the file resolves alone never changes; a lookup failure is
logged and judges the file alone; local Inspect and ingestion are unchanged
(`tests/test_recorded_saver_anchor.py`, `tests/test_run_attribution.py`,
`tests/test_api_history_outputs.py`).
`tools/inspection_corpus.py --root DIR --out JSONL [--workers N] [--baseline JSONL]
[--summary TXT] [--code TREE]` measures a media folder read-only against the #240
@@ -1301,7 +1344,7 @@ How the earlier workspace Images filters moved (old → new). The legacy `/galle
| Folder | "Refine this view" → Folder (D14); shown in the scope chip |
| Media (All media / Images / Videos) | "Refine this view" → Media (D14) |
| Sort: Newest first / Oldest first / Filename A–Z / Filename Z–A | toolbar sort: Recently added / Oldest first / Name A–Z / Name Z–A, plus Highest rated (D13, D14) |
| (none) | toolbar category and "Refine this view" → Tags, through the linked positive (D13, D14) |
| (none) | toolbar category (the effective category, D141) and "Refine this view" → Tags, through the linked positive (D13, D14) |
| Refresh | Import → "Scan an output folder" → Rescan (D15) |
| Grid / List / Masonry, Comfortable / Large / Compact | the same choices in the concept filter row |
| Card image → side panel; preview icon → viewer | image → viewer; Info → side panel (D19) |
@@ -1384,7 +1427,8 @@ Recorded deviations on this page (spec §8):
`overscroll-behavior: contain`.
- **D35**: counts pluralize ("1 positive prompt", "1 image", "1 more image loaded").
- **D36**: the category list is the Library's, sorted alphabetically (case-insensitive)
by the page, whatever order the server sends.
by the page, whatever order the server sends (amended by D141: Images adds its
image categories).
- **D37**: production loads 40 images per page (the concept 12), so the continuation
and toast use 40; after Escape production returns focus to the opening image.
- **D44** (#249 PR E): Filters reads "Filters · N" while N of Source, Folder, Media and
@@ -1470,15 +1514,46 @@ Recorded deviations on this page (spec §8):
- **D39** (ruled 2026-10-21: fallback): a recipe-linked image shows its recipe's own
category when set, otherwise the linked positive prompt's, and its recipe's own tags
when it has any, otherwise the positive's; each field falls back on its own and the
two tag sets are never merged. Set category (#259) writes the positive's category and
the Images category and Tags filters match the positive's, so the panel now agrees
with them whenever the recipe adds nothing. The inspection response names each
two tag sets are never merged. The Images Tags filter matches the positive's tags.
The category filter matches the same inherited category as the panel (D142, #332), and
an image's own category overrides both (D141, D146). The inspection response names each
field's source (`library.category_scope`, `library.tags_scope`: `pair`, `positive` or
null); the category chip's visually hidden prefix (D21) says "Recipe category: " or
"Positive prompt category: " from it, and tag chips are Tags filter buttons only when
they are the positive's (#260 D2). History's media inspector uses the same rule. The
concept shows the prompt's category and tags for every image and the concept
fixture's recipes have neither, so the inspector comparison needs no mapping for it.
- **D141** (#332, amends D13 and D36): the Images category filter matches each image's
effective category: its own category (Images Set category, PR 3) when it has one,
else its inherited one (D142), else none. An image whose own category differs never
matches its prompt's category, and an unlinked or ambiguous image with an own
category is listed under it (the members-only gate now applies to the Tags filter
only). "Select all N matching" and saved views use the same value. The caption's
positive-prompt count still counts the linked positives of the images shown. With a
category set the listing reads the own categories once per source (a primary-key
range); a failed read is the existing 500 `listing-failed`, never an unfiltered page.
The link read then keeps only links whose raw recipe or positive category is the
value (a superset of those that inherit it; Python decides), and reads the links of
files whose own category is the value by path, so their rating and positive count.
`test_filter_budget` seeds 2,000 own categories on its 50,000-file fixture.
The option list adds the image categories of listed files
(`GET /prompt_manager/images/categories`, PR 3).
- **D142** (#332 MC4, amends D39's filter note): the inherited category is one rule,
`inherited_category` in `utils/inspection_response.py`: a linked recipe's own
category when it is not blank, else the linked positive's. Before #332 the filter
matched the positive's even when the chip showed the recipe's. The link reads select
the raw recipe and positive categories and never decide inheritance in SQL
(`tests/test_api_workspace_image_categories.py` checks the filter and the chip agree
on blank values stored as spaces, tabs, U+00A0 and BLOB). The concept has no recipe
categories, so this is behavior only.
- **D145** (#332, DC4): History's media inspector shows an image's own category with
the chip rule of the side panel and viewer. Only the inspection request reads it, once,
inside the History snapshot: the file is the symlink-resolved output folder plus the
receipt's subfolder and filename, keyed with `image_category_key` over the configured
output sources, as Images keys it. A temp output or a file no source holds has none;
a failed read hides the chip (DC5). The outputs list and media requests read nothing
new. The chip's hidden prefix reads "Image category: " for an own category; the
visible origin text and its comparison are D146 (PR 3).
- **D42**: the concept never shows a toast and a page alert together. Production keeps
the concept's fixed toast (centred, 22px from the bottom, 12px at 700px and below)
and, when an Import → Rescan toast leaves the Images discovery alert
@@ -2514,6 +2589,14 @@ production hub `#/overview` is unchanged (D28); its System card opens
the pressed section button and disposes the previous section's scope, so its reads,
listeners and dialogs never outlive it.
System is the home for settings and maintenance (#330 System S13). The legacy pages keep
their prompt list, tags, Scan Images, auto-tag batch and review and gallery browsing; their
Settings, Diagnostics, View Logs, Maintenance, Backup DB, Restore DB, the gallery's
thumbnail, cache, rescan and duplicate tools and the auto-tag model downloads open the
System section that owns them in a new tab (`legacy-system-links.spec.mjs`). Each page
keeps only its own settings in Page options. What left the legacy pages is listed in
"Legacy System UI removed by S13".
Drafts (`system-state.mjs`) live for the page session (Q10). Save writes the dirty
browser fields, then the host fields (`POST /prompt_manager/settings` with
`expected_revision`, #220), then the runtime fields. A failed domain keeps its drafts and
@@ -4128,6 +4211,9 @@ auto-tag / review", which the Images selection bar owns, not a System section. A
is not `covered` names none (`—`). Test is the exact title of a test that runs, then its
file. Anchors are line numbers on `main` at this PR.
S13 (#330) removed the legacy copy of the `covered` rows the legacy pages linked to System;
"Legacy System UI removed by S13" below lists every removed element against its row.
`tests/vendor/legacy-coverage.test.mjs` fails on a row with any of these problems:
- no owner, or a status outside the three;
- a malformed Native control cell;
@@ -4208,7 +4294,11 @@ opens `workspace#/system/<id>` in a new tab) and removes the legacy copy of a ca
only when its row above is `covered`. One row per removed capability: Capability is the
exact Capability cell of its coverage row, Page is `admin` or `gallery`, Removed ids and
Removed methods list what went (`—` for none), and Native test copies that row's Test
cell. An id that stays on a System link is not listed, only its handler.
cell. An id that stays on a System link is not listed, only its handler. The gallery's
Default View Mode left Page options too: System › Appearance's Default layout is the one
default-layout control. Nothing writes the legacy `defaultViewMode` any more; a stored
value still picks the view the legacy gallery opens in (grid otherwise), and Page options
Save and Reset keep it.
`tests/vendor/legacy-coverage.test.mjs` fails when a row cites a capability that is not
`covered` or copies another test, when an id or method is listed twice for one page, or
@@ -4227,6 +4317,31 @@ legacy page workflows), not by this check.
| Log files list and download | admin | `#downloadLogsBtn`, `#logFilesList` | `loadLogFiles()`, `downloadLogs()`, `downloadLogFile()` | each listed file downloads by its returned name and a link has no Download (workspace-system-diagnostics.spec.mjs) |
| Clear/truncate logs | admin | `#clearLogsBtn` | `clearLogs()` | Clear all logs states both counts, confirms first and reads everything again (workspace-system-diagnostics.spec.mjs) |
| Set log level (incl. CRITICAL), console logging | admin | `#updateLogConfigBtn`, `#setLogLevel`, `#consoleLogging` | `loadLogConfig()`, `updateLogConfig()` | runtime level (including CRITICAL) and console logging save and say they apply until the server restarts (workspace-system-diagnostics.spec.mjs) |
| Image scan directories add/remove | admin | `#scanPathsList`, `#addScanPathBtn` | `renderScanPaths()`, `addScanPath()`, `removeScanPath()`, `_collectScanPaths()` | Add source validates on the server, then marks the host draft dirty; Save persists (workspace-system-sources.spec.mjs) |
| Currently monitoring list + Refresh | admin | `#refreshMonitoringStatus`, `#monitoringStatus` | `updateMonitoringStatus()` | watching only with a live observer; restart pending otherwise (workspace-system-sources.spec.mjs) |
| Search results auto-hide `result_timeout` | admin | `#resultTimeout` | — | host fields round-trip through /settings with revision (workspace-system.spec.mjs) |
| Web UI display mode | admin | `#webuiDisplayMode` | — | host fields round-trip through /settings with revision (workspace-system.spec.mjs) |
| Worker threads | admin | `#workerThreads`, `#workerThreadsValue` | `renderWorkerThreads()` | worker range max is cpu_count (workspace-system-sources.spec.mjs) |
| LoraManager badge, enable, path, trigger words, CivitAI key, clear key | admin | `#loraDetectionBadge`, `#loraEnabled`, `#loraSettings`, `#loraManagerPath`, `#loraTriggerWords`, `#civitaiApiKey`, `#clearCivitaiKey` | `detectLoraManager()`, `_bindLoraEvents()`, `saveLoraSettings()` | LoRA rows show detected / enabled / cache-loaded separately (workspace-system-integrations.spec.mjs) |
| LoRA data import with progress | admin | `#loraImportModal`, `#loraImportBtn`, `#loraImportStatus`, `#loraImportPercent`, `#loraImportBar`, `#loraImportProcessed`, `#loraImportImported` | `runLoraImport()` | Review import shows replacement, blocked items and network effects with the backup box; Start opens the job (D107) (workspace-system-integrations.spec.mjs) |
| Maintenance cleanup_duplicates, vacuum, cleanup_orphaned_images, prune_orphaned_prompts | admin | `#maintenanceModal`, `#cleanup_duplicates`, `#vacuum`, `#cleanup_orphaned_images`, `#prune_orphaned_prompts`, `#runMaintenanceBtn`, `#maintenanceResults` | `showMaintenanceModal()`, `closeMaintenance()`, `runMaintenance()`, `renderMaintenanceResults()` | Apply reviewed operation closes, toasts, opens the job in Activity (workspace-system-maintenance.spec.mjs) |
| Maintenance check_hash_duplicates, check_consistency | admin | `#check_hash_duplicates`, `#check_consistency` | — | Run health checks opens Health report and updates metric 4 (workspace-system.spec.mjs) |
| Maintenance statistics (database info) | admin | `#statistics` | — | meta-grid shows real database and source status, the pinned Viewer.js version and the host renderers (workspace-system-diagnostics.spec.mjs) |
| Maintenance select all / clear all / run | admin | `#selectAllMaintenanceBtn`, `#clearAllMaintenanceBtn` | `selectAllMaintenance()`, `clearAllMaintenance()` | plan dialog shows the server's scope, protected, backup policy and media lines with the backup box (D100); Cancel writes nothing (workspace-system-maintenance.spec.mjs) |
| Database backup | admin | — | `backupDatabase()` | Back up downloads through the api client and revokes the object URL (D103) (workspace-system-backups.spec.mjs) |
| Restore from file | admin | `#restoreModal`, `#restoreFileInput`, `#confirmRestore`, `#cancelRestore` | `showRestoreModal()`, `confirmRestore()` | Review restore stages the chosen file and shows the real summary with the backup box (D105) (workspace-system-backups.spec.mjs) |
| Gallery thumbnail generation | gallery | `#generateThumbnailsBtn`, `#thumbnailStatus`, `#thumbnailProgress`, `#progressText`, `#progressPercent`, `#cancelThumbnailsBtn`, `#progressBar`, `#progressDetails`, `#progressETA` | `generateThumbnails()`, `generateThumbnailsWithProgress()` | Start closes the dialog, toasts, and opens the job in Activity (workspace-system-sources.spec.mjs) |
| Default view mode, image quality | gallery | `#defaultViewModeSelect`, `#imageQualitySelect` | — | defaults apply to a new visit but not to the current view or a saved view (workspace-system.spec.mjs) |
| Image info on cards; sidebar collapsed | gallery | `#showImageInfoToggle`, `#sidebarCollapsedByDefaultToggle` | — | added rows reach their consumers (workspace-system.spec.mjs) |
| Video autoplay, mute, loop | gallery | `#videoAutoplayToggle`, `#videoMuteToggle`, `#videoLoopToggle` | — | video defaults reach the viewer (workspace-system.spec.mjs) |
| Auto-load metadata, cache metadata, show file paths | gallery | `#autoLoadMetadataToggle`, `#cacheMetadataToggle`, `#showFilePathsToggle` | — | metadata off shows an explicit load and loads once on request (workspace-system.spec.mjs) |
| Browser cache clear | gallery | `#clearCacheBtn` | `clearCache()` | Metadata cache reviews and clears only owned browser entries when there are some (D102) (workspace-system-maintenance.spec.mjs) |
| Clear server thumbnails | gallery | — | — | Thumbnail cache opens the S5 dialog (workspace-system-maintenance.spec.mjs) |
| Refresh file listing | gallery | `#rescanFolderBtn` | `rescanFolder()` | Refresh files calls /gallery/rescan and toasts its counts (workspace-system-sources.spec.mjs) |
| Duplicate image scan, review, delete | gallery | `#scanDuplicatesBtn`, `#dupScanPill`, `#dupScanPillText`, `#duplicatesModal`, `#closeDuplicatesBtn`, `#duplicatesScanStatus`, `#dupProgressText`, `#dupProgressBar`, `#duplicatesContent`, `#duplicatesList`, `#duplicatesFooter`, `#duplicatesCount`, `#selectAllBtn`, `#removeDuplicatesBtn` | `scanDuplicates()`, `attachDuplicateStream()`, `resumeDuplicateScan()`, `applyDuplicateEvent()`, `renderDuplicateProgress()`, `setDuplicateButton()`, `showDuplicatePill()`, `hideDuplicatePill()`, `openDuplicatesModal()`, `showDuplicatesModal()`, `hideDuplicatesModal()`, `displayDuplicates()`, `selectAllDuplicates()`, `removeDuplicates()` | Find duplicates starts a job; its view links Review groups; Keep radios build a duplicate_files plan (D101) (workspace-system-maintenance.spec.mjs) |
| Auto-tag model download | admin | `#autoTagDownloadModal`, `#downloadModelName`, `#downloadStatus`, `#downloadProgressPercent`, `#downloadProgressBar`, `#cancelDownloadBtn` | `downloadModel()`, `cancelDownload()` | Download starts the job and Activity shows it without Cancel (workspace-system-tagging.spec.mjs) |
| Auto-tag model download | gallery | `#autoTagDownloadModal`, `#downloadModelName`, `#downloadStatus`, `#downloadProgressPercent`, `#downloadProgressBar`, `#cancelDownloadBtn` | `downloadModel()`, `cancelDownload()` | Download starts the job and Activity shows it without Cancel (workspace-system-tagging.spec.mjs) |
| Unload model | gallery | — | `unloadModel()` | Unload is enabled only while a model is loaded, and a held gate's refusal names the operation (workspace-system-tagging.spec.mjs) |
### Browser cache ownership (#222)
@@ -4245,17 +4360,15 @@ page writes a key that is neither a preference nor inventoried, writes one it ca
CacheStorage or IndexedDB; `the storage scan reports every way of writing browser
storage` pins those patterns. It skips `web/lib` (third-party builds) and the pinned
Toolkit build (`web/js/shared/toolkit`, never hand-edited, verified by `npm run check`;
its storage belongs to the Toolkit repository and a new pin is reviewed on its own). The
gallery's server thumbnail clear stays its own request and runs first: refused or
unreachable, it changes nothing and says so (the old code toasted success when the
request failed); refused browser storage is reported, never hidden
(`tests/js/legacy-refusals.test.js`; the browser case `the gallery cache clear keeps
preferences and other tools caches` in `workflows.spec.mjs`).
its storage belongs to the Toolkit repository and a new pin is reviewed on its own). Since
System S13 (#330) the legacy gallery no longer offers Clear cache: System › Maintenance's
Metadata cache and Thumbnail cache reviews own both parts (removal table above), and
`cache-ownership.js` serves the native review.
The legacy duplicates dialog's View Full now names its source (`?root=<index>`, with the
path URL-encoded) when the gallery's output folder is a listed source; before, the serve
route opened the first source holding a file at that path
(`test_view_full_names_the_gallery_folder_source`).
The duplicates dialog's View Full named its source (`?root=<index>`, with the path
URL-encoded) when the gallery's output folder is a listed source
(`test_view_full_names_the_gallery_folder_source`); S13 removed the legacy dialog, and the
route keeps that behaviour for System › Maintenance's duplicate review.
### Library location contract (#249 System S6b)
@@ -4547,8 +4660,18 @@ Filters (`mode=workspace`; anything invalid is 400 "Invalid gallery filters"):
rating 5 to 1, then linked but unrated, then unlinked and ambiguous; ties break
newest first. A signed cursor covers every filter and the sort, so reusing it
with a change is 400.
- One link read (`get_image_link_matches`, streamed and bounded to the listed
paths) serves the filters, the rating sort and the count.
- One link read (`get_image_link_matches`) serves the filters, the rating sort
and the count. SQLite returns its rows as one JSON value in one step (#354):
Python's `sqlite3` releases and retakes the GIL around every step, and while
ComfyUI generates each retake can wait on the busy thread, so a read that
stepped once per linked row took seconds. The result keeps only the listed
paths; peak memory grows with the database's matching linked rows (about
530 bytes each, freed on return). When the JSON step fails (a SQLite without
JSON functions, a BLOB value, a JSON text over SQLite's length limit) the read
streams the rows as before and logs one warning per process; a missing JSON
function stops the attempt for the rest of the process. An interrupted read
(Move library) raises as before. `tests/test_database_link_matches.py` proves both reads
equal on every filter and on odd values.
`positive_count` is the number of distinct linked positives over the full matching
set, not the page:
@@ -4599,6 +4722,70 @@ concept fixture stars the concept's favorites (each other prompt's first image).
restore. Old backups without the table restore, and backup plus restore keeps
WAL-resident favorites and closes every handle (`tests/test_database_backup.py`).
Image categories (#332, store and routes): an image file may carry its own
category; without one it shows the category its link supplies.
- The table is `image_categories(root_key, relative_path, category, updated_at)`
with primary key `(root_key, relative_path)` and the index
`idx_image_categories_category`. `migrate_image_categories` creates both empty
inside a savepoint, with no backfill. It reads only `sqlite_master` and its own
table's PRAGMAs, and writes only its two objects. `tests/test_image_categories.py`
checks this with an SQLite authorizer on a 23,000-link library, and checks that
`prompts` and `prompt_pairs` `(id, category, updated_at)` are byte-for-byte equal
afterwards.
- The **file key** is `image_category_key(sources, absolute path)` in
`py/api/images.py`: the deepest output source that contains the path (compared
lexically, never resolving the path, so a symlinked subfolder keeps its listed
spelling) and the path below it, keyed with the favorites helpers. A source
matches by its real path or its configured spelling (a detected output folder
can be a symlink). Without nested sources it equals a favorite's key. Every write and read uses it. A path
that no source contains has no key, so it has no own category.
- `image_categories_schema_error` checks the structure at startup (a value saved
under older rules never stops the library opening; a table without its index,
which this code never leaves, is repaired by the install). Restore also checks the
values: text columns, keys within the favorites bound, `category` 1 to 100
characters with no character below U+0020 and no leading or trailing space, and
`updated_at` of at most 64 bytes. Backups made before #332 restore without the
table; a migrated restore stage must have it (`require_library_state`).
- `PUT /prompt_manager/images/category` `{items, category}` answers
`{success, category, updated, unchanged, not_found}`.
`DELETE /prompt_manager/images/category` `{items}` answers `{success, cleared}`.
`GET /prompt_manager/images/categories` answers `{success, categories}`: the
sorted distinct own values whose file is in some source's listing.
- `items` are 1 to 500 Images targets `{root_index, root_scope, path}`, checked by
`py/api/media_target.py` (shared with Collections, #331). Refusals:
- 400: a malformed body, a field other than `items` and `category`, a
`root_index` outside the output sources (a LoRA folder), or a category that is
blank after stripping, over 100 characters, or holds a control character;
- 403: traversal, a drive, a backslash, a non-media name ("Only media files can
have a category") or a key over the stored bound;
- 409 `source-changed`: a stale `root_scope`;
- 413: a body over 1 MiB.
Any invalid item fails the whole request before a write. Targets that fold to one
key are written and counted once. Instead of `items` a request may send `filter`,
the Images listing's own query ("Select all N matching", #392, keys as in
`parse_listing_filter`), resolved now: over 500 matches is a 400
`too-many-images` ("Select at most 500 images."), a renumbered source a 409
`source-changed`.
- PUT writes only files the source's listing holds and counts the rest in
`not_found`; when none is listed it answers 404 `files-not-found` and writes
nothing. A file that already has the value counts as `unchanged` and keeps its
`updated_at`. DELETE accepts any well-formed target, so a missing file's
category can still be removed. An unexpected failure is a logged 500 ("The
category could not be changed." / "Image categories could not be read.") without
the exception text. The routes never read or write `prompts`, `prompt_pairs` or
`generated_images`; `POST /prompt_manager/bulk/category` is unchanged.
- Duplicate cleanup (System › Maintenance): the kept copy's own category wins. A
kept copy without one receives the removed copy's (`COPY_SQL`, part of the
cleanup's query-plan check). The removed copy's row is never deleted, so a
give-back (the kept copy changed) and a failed delete lose nothing. The job log
says "Kept N image categories on the kept copies".
`tests/test_duplicate_image_categories.py` covers nested sources in both orders.
- A moved, renamed or deleted file's row stays and is hidden; a file that later
appears at the same path shows it. The real-host check is the last test in
`tests/comfyui/image-attribution.spec.mjs`: PUT, list, DELETE on a generated
output.
Image annotations schema (#249, `database/annotation_schema.py`): the group
`media_annotation_metadata` (one row: schema version, `image_index_revision`),
`media_annotation_files` (one private locator per physical file,
@@ -5128,7 +5315,7 @@ which used to save `null` under a link.
Primitive or through `PreviewAny`, for PromptManager, PromptManagerText, Focused, Split
and the Pair (positive, negative, both, neither linked, and a Primitive negative). A
wildcard Primitive on PromptManager (queued as written) and Focused (one choice) saves
and counts what was queued. The linked text is saved, counted once, recorded as
and counts what was queued. The linked text is saved, counted once per job, recorded as
the run's raw positive/negative, and after a scan the run's image links to it. Text
only known at run time (`PreviewAny`) cannot be read from the image; the case only
requires that it is never linked to another prompt (its attribution from the recorded
@@ -5142,6 +5329,39 @@ PM_COMFY_ROOT=... PM_TEST_PYTHON=... npm run --prefix tools/vendor test:comfyui
PM_COMFY_ROOT=... PM_TEST_PYTHON=... npm run --prefix tools/vendor test:executions -- execution-text-link
```
## Cached runtime-linked runs (#347)
Text only known at run time (a `PreviewAny` or LLM link, a list node) is counted by the
node when it executes; the queue hook cannot read it. When ComfyUI serves such a node from
its output cache, `utils/cached_runs.py` counts the job instead: every execution records
a receipt (digests of its STRING outputs and hashes of the text it stores, per run key and
queue item), and on `execution_cached` the receipt whose digests equal the cache entry's
text is replayed once. The role (positive, negative, which Pair outputs are used) and what
the hook already saw are recomputed from the job's own graph, and the node and replay
decide through the same `own_counts`.
- Fail closed: no receipt in this process (after a restart, an external cache provider),
changed cached text, more than 64 list values, a node the job's outputs do not need, or
a host without the queue/executor interfaces: nothing is counted, as before #347.
- A cached node counts at job start like the hook: a job failing later still counts it.
A node that executes after its replay in the same job (RAM-pressure freed the entry)
skips its own count but still consumes the pending first-use catch-up.
- No backfill: earlier cached jobs keep their counts. The history page's "tracked runs"
can stay below the Library count for these prompts (#403).
- Tests: `tests/test_cached_runs.py` (whole jobs with the real nodes, a real database,
fake host caches); `tests/comfyui/execution-text-link.spec.mjs` "#347" cases on the
managed host: PreviewAny into each node kind, 3 jobs changing only the sampler seed,
count 3 in both renderers; in one renderer the Classic, LRU, RAM-pressure and no-cache
modes, A, B, A from the owned `PMCycledText` under LRU, typed text, a downstream
failure, a "run selected output" job that does not need the node, and a rewire to the
negative input. The RAM-pressure eviction path cannot be forced on a real host and is
covered by the Python tests.
```sh
python3 -m unittest discover -s tests -p "test_cached_runs.py"
PM_COMFY_ROOT=... PM_TEST_PYTHON=... npm run --prefix tools/vendor test:executions -- execution-text-link --grep "#347"
```
## Durable execution ledger
`npm run --prefix tools/vendor test:executions` runs 52 real execution cases
@@ -5561,6 +5781,20 @@ rotates the epoch, so a pre-restore `expected_epoch` is refused with `storage-ch
Owner rows and their timestamps never change on a tag command; a category apply only
touches rows carrying the exact source.
Image categories (#332 §7, MC9, MC10, DC7): each `GET categories` row also carries
`image_count`, the stored `image_categories` rows holding that exact value (files
missing on disk included). A value prompts and images share appears once, ordered by
its oldest positive; a value only images hold appears once after every prompt-backed
value (`first_pid` 9223372036854775807 in the keyset cursor), ordered by value, with
`prompt_count` 0. `/library/options` keeps its own query and never lists image-only
values. Category preview and apply report `impact.images` (rows holding the source),
the target exists when an image row holds it, an image-only value can be renamed, the
apply renames the image rows in the same transaction (file keys and rowids kept, so a
merge cannot collide), and the impact token hashes the sorted image rowids, so an image
row gaining or losing the value between preview and apply is 409
`vocabulary-conflict` with nothing written. Every image-category statement reads
`idx_image_categories_category` (`test_category_reads_use_the_image_category_index`).
Two legacy routes change additively: `GET /tags` now lists unused registry tags, and
`/tags/suggest` without context appends unused names at count 0. Unused tags become
deletable only through the native Delete tag (vocabulary UI PR), not on the legacy
+58
View File
@@ -195,6 +195,26 @@ except Exception as e:
f"Failed to register usage hook: {e}"
)
# A cached node counts what its last execution of the same text counted (#347).
try:
from .py.config import server_instance
from .utils.cached_runs import register_cached_runs
register_cached_runs(server_instance, _db_factory)
except Exception as e:
try:
from .utils.logging_config import get_logger
get_logger("prompt_manager.init").error(
f"Failed to register cached-run counting: {e}"
)
except Exception:
import logging
logging.getLogger("prompt_manager.init").error(
f"Failed to register cached-run counting: {e}"
)
def _history_db_factory() -> "PromptDatabase":
"""Initialize history on its worker even if monitoring initialization failed."""
@@ -355,6 +375,44 @@ async def _stop_background_work_on_cleanup(app):
_stop_background_work()
def _kill_encoders_at_exit():
"""Run by the exit hook before Python joins the executor threads (#376):
a JPEG XL or ffmpeg encoder (no console on Windows, so Ctrl+C never
reaches it) is killed instead of being waited for."""
from .py.api.thumbnail_encode import shutdown_encoders
shutdown_encoders()
def _abort_downloads_at_exit():
"""Run by the exit hook (#376): huggingface_hub's Xet transfers ignore
the chunk checkpoint, so they are aborted the way hub aborts on Ctrl+C."""
import sys
if "huggingface_hub" not in sys.modules:
return # no download ever ran: import nothing at exit
from .py.autotag import abort_hub_downloads
abort_hub_downloads()
try:
# Before the executor join: running jobs stop at their next file, item or
# chunk (#376). The teardown below still runs from atexit afterwards.
from .utils import exit_signal as _exit_signal
_exit_signal.install(_kill_encoders_at_exit)
_exit_signal.install(_abort_downloads_at_exit)
except Exception as e:
try:
from .utils.logging_config import get_logger
get_logger("prompt_manager.init").error(
f"Failed to register the exit signal: {e}"
)
except Exception:
pass
try:
atexit.register(_stop_background_work)
from .py.config import server_instance as _teardown_server
+42 -5
View File
@@ -9,8 +9,12 @@ IMMEDIATE`` transaction, so no row ever names a deleted file:
already links the kept copy: then its saved-pair links join that record
and it is deleted;
- its favorite;
- its collection memberships (#331), joined onto the kept copy's as their
union under the same keys a favorite moves by;
- its image tags (``media_image_tags`` of its byte record), joined onto the
kept copy's byte record as their union, in order;
- its own image category (#332), copied to the kept copy when that copy has
none; the removed copy's row is kept, so a give-back loses nothing;
- its pending Needs review items, settled: only an item still ``pending`` at
revision 0 with no decision row is removed (``review._SETTLE_ITEM``), so a
decided item is never touched.
@@ -22,6 +26,8 @@ import os
from .annotation_schema import BYTES_TABLE, FILES_TABLE
from .annotations import _BUMP_INDEX
from .image_categories import COPY_SQL as COPY_CATEGORY_SQL
from .image_categories import copy_image_category
from .models import normalize_image_path, utc_now_iso
from .review import _BUMP_REVIEW_INDEX, _PENDING_OF_LOCATOR, _SETTLE_ITEM
from .review_changes import notify_review_change
@@ -50,6 +56,15 @@ INSERT_FAVORITE_SQL = "INSERT OR IGNORE INTO image_favorites VALUES (?, ?, ?)"
DELETE_FAVORITE_SQL = (
"DELETE FROM image_favorites WHERE root_key = ? AND relative_path = ?"
)
COPY_MEMBERSHIPS_SQL = (
"INSERT OR IGNORE INTO image_collection_items "
"(collection_id, root_key, relative_path, added_at) "
"SELECT collection_id, ?, ?, added_at FROM image_collection_items "
"WHERE root_key = ? AND relative_path = ?"
)
DELETE_MEMBERSHIPS_SQL = (
"DELETE FROM image_collection_items WHERE root_key = ? AND relative_path = ?"
)
COPY_BYTE_SQL = (
f"SELECT b.id FROM {FILES_TABLE} f JOIN {BYTES_TABLE} b " # nosec B608
"ON b.file_id=f.id WHERE f.locator=? AND b.sha256=?"
@@ -73,9 +88,12 @@ FILE_LINK_SQL = (
DELETE_LINK_SQL,
FAVORITE_SQL,
DELETE_FAVORITE_SQL,
COPY_MEMBERSHIPS_SQL,
DELETE_MEMBERSHIPS_SQL,
COPY_BYTE_SQL,
COPY_TAG_COUNT_SQL,
PENDING_COUNT_SQL,
COPY_CATEGORY_SQL,
)
@@ -135,12 +153,17 @@ class DuplicateCopyOperations:
by), ``targets`` (``{removed spelling: kept spelling}``, the spelling
a moved record takes; else ``kept_path``), ``favorites`` (``(root key,
removed key, kept root key, kept key)`` per listed root holding the
file), ``removed_locator``,
``kept_locator``, ``sha256`` and the kept copy's ``kept_size`` and
``kept_mtime_ns``. Returns ``{moved, merged, favorites, tags,
review}``.
file; they move its favorite and its collection memberships),
``removed_locator``, ``kept_locator``, ``sha256`` and the kept copy's
``kept_size`` and ``kept_mtime_ns``; ``categories``: ``(removed file
key, kept file key)`` pairs of the image categories (#332). Returns
``{moved, merged, favorites, tags, review, collections, categories}``
(``collections``: the memberships inserted on the kept copy;
``categories``: the image categories copied to the kept copy).
"""
counts = dict.fromkeys(("moved", "merged", "favorites", "tags", "review"), 0)
counts = dict.fromkeys(
("moved", "merged", "favorites", "tags", "review", "collections"), 0
)
kept_keys = [normalize_image_path(path) for path in move["kept"]]
with self._pair_transaction() as conn:
for spelled in dict.fromkeys(move["removed"]):
@@ -173,7 +196,12 @@ class DuplicateCopyOperations:
counts["moved"] += 1
for favorite in move.get("favorites", ()):
counts["favorites"] += _move_favorite(conn, *favorite)
counts["collections"] += _move_memberships(conn, *favorite)
counts["tags"] = self._move_copy_tags(conn, move)
counts["categories"] = sum(
copy_image_category(conn, removed, kept)
for removed, kept in move.get("categories", ())
)
for (item_id,) in conn.execute(
_PENDING_OF_LOCATOR, (move["removed_locator"],)
).fetchall():
@@ -213,3 +241,12 @@ def _move_favorite(conn, root_key, removed_key, kept_root, kept_key):
conn.execute(INSERT_FAVORITE_SQL, (kept_root, kept_key, row[0]))
conn.execute(DELETE_FAVORITE_SQL, (root_key, removed_key))
return 1
def _move_memberships(conn, root_key, removed_key, kept_root, kept_key):
"""Copy a removed copy's memberships to the kept copy; the rows inserted."""
inserted = conn.execute(
COPY_MEMBERSHIPS_SQL, (kept_root, kept_key, root_key, removed_key)
).rowcount
conn.execute(DELETE_MEMBERSHIPS_SQL, (root_key, removed_key))
return inserted
+249
View File
@@ -0,0 +1,249 @@
"""Image categories (#332): an image file's own category.
An image shows its own category when it has one, otherwise the category its
link supplies (its recipe's, else its positive prompt's); this table holds
only the own values. A row is keyed by the file key: the deepest configured
output source that contains the file (its real, case-normalized path) and the
path below it, case-folded where the OS folds case
(``py/api/images.py`` ``image_category_key``). Without nested sources the key
equals a favorite's key. Nothing follows a moved file: its row stays and is
hidden until a file appears at the same path again.
Writes validate strictly (the route strips the value and applies
``validate_category``). Startup checks the structure only, so a value saved
under older rules never stops the library opening; restore also checks the
values. Optional in old backups, required once a database is current.
"""
import re
import sqlite3
from typing import Iterable
TABLE = "image_categories"
INDEX = "idx_image_categories_category"
MAX_KEY_LENGTH = 4096
MAX_CATEGORY_LENGTH = 100
_DDL = (
f"CREATE TABLE {TABLE} (root_key TEXT NOT NULL, relative_path TEXT NOT NULL, "
"category TEXT NOT NULL, updated_at TEXT NOT NULL, "
"PRIMARY KEY (root_key, relative_path))"
)
_INDEX_DDL = f"CREATE INDEX {INDEX} ON {TABLE}(category)"
_AUTOINDEX = f"sqlite_autoindex_{TABLE}_1"
_COLUMNS = (
("root_key", "TEXT", 1, 1),
("relative_path", "TEXT", 1, 2),
("category", "TEXT", 1, 0),
("updated_at", "TEXT", 1, 0),
)
# A category may not hold a character below U+0020 (validate_category).
_CONTROL_GLOB = "*[" + "".join(chr(code) for code in range(1, 32)) + "]*"
CATEGORY_OF_SQL = (
f"SELECT category FROM {TABLE} WHERE root_key = ? AND relative_path = ?"
)
CATEGORIES_OF_ROOT_SQL = (
f"SELECT relative_path, category FROM {TABLE} WHERE root_key = ?"
)
UPSERT_SQL = (
f"INSERT INTO {TABLE} (root_key, relative_path, category, updated_at) "
"VALUES (?, ?, ?, ?) ON CONFLICT(root_key, relative_path) "
"DO UPDATE SET category = excluded.category, updated_at = excluded.updated_at"
)
DELETE_SQL = f"DELETE FROM {TABLE} WHERE root_key = ? AND relative_path = ?"
# The duplicate cleanup's copy (#332 MC11): the removed copy's row is copied
# to the kept copy only when the kept copy has none; it is never deleted.
COPY_SQL = (
f"INSERT OR IGNORE INTO {TABLE} (root_key, relative_path, category, updated_at) "
f"SELECT ?, ?, category, updated_at FROM {TABLE} "
"WHERE root_key = ? AND relative_path = ?"
)
STATEMENTS = (CATEGORY_OF_SQL, CATEGORIES_OF_ROOT_SQL, UPSERT_SQL, DELETE_SQL, COPY_SQL)
def _compact(sql: str) -> str:
return re.sub(r"\s+", "", sql or "").lower()
def _structure_error(
conn: sqlite3.Connection, allow_missing: bool, allow_missing_index: bool = False
) -> str | None:
objects = conn.execute(
"SELECT type,name,sql FROM sqlite_master WHERE name=? "
"OR (type IN ('index','trigger') AND tbl_name=?)",
(TABLE, TABLE),
).fetchall()
if not objects:
return None if allow_missing else "Missing image categories schema"
tables = [sql for kind, _name, sql in objects if kind == "table"]
others = {(kind, name) for kind, name, _sql in objects if kind != "table"}
expected = {("index", _AUTOINDEX), ("index", INDEX)}
# Before its install, a table without the category index is repaired by
# CREATE INDEX IF NOT EXISTS rather than refused.
if allow_missing_index and others == {("index", _AUTOINDEX)}:
expected = others
if len(tables) != 1 or others != expected:
return "Image categories name is reserved by an unsupported schema object"
if _compact(tables[0]) != _compact(_DDL):
return "Image categories has an invalid table definition"
index_sql = next((sql for kind, name, sql in objects if name == INDEX), None)
if ("index", INDEX) in expected and _compact(index_sql) != _compact(_INDEX_DDL):
return "Image categories has an invalid index"
autoindex_sql = next(sql for kind, name, sql in objects if name == _AUTOINDEX)
if autoindex_sql is not None:
return "Image categories has an invalid index"
columns = tuple(
(row[1], row[2].upper(), row[3], row[5])
for row in conn.execute(f"PRAGMA table_info({TABLE})")
)
if columns != _COLUMNS:
return "Image categories has invalid columns"
return None
def _values_error(conn: sqlite3.Connection) -> str | None:
# Bytes bound characters from above (UTF-8 is at most 4 bytes each).
limit = MAX_KEY_LENGTH * 4
if conn.execute(
f"SELECT 1 FROM {TABLE} WHERE typeof(root_key) != 'text' " # nosec B608
"OR typeof(relative_path) != 'text' OR typeof(category) != 'text' "
"OR typeof(updated_at) != 'text' "
f"OR length(CAST(root_key AS BLOB)) > {limit} "
f"OR length(CAST(relative_path AS BLOB)) > {limit} "
"OR length(CAST(updated_at AS BLOB)) > 64 "
"OR instr(CAST(category AS BLOB), x'00') > 0 "
f"OR length(category) < 1 OR length(category) > {MAX_CATEGORY_LENGTH} "
"OR category != TRIM(category) OR category GLOB ? LIMIT 1",
(_CONTROL_GLOB,),
).fetchone():
return "Image categories contains invalid values"
for root, path in conn.execute(
f"SELECT CAST(root_key AS BLOB), CAST(relative_path AS BLOB) " # nosec B608
f"FROM {TABLE} WHERE length(CAST(root_key AS BLOB)) > {MAX_KEY_LENGTH} "
f"OR length(CAST(relative_path AS BLOB)) > {MAX_KEY_LENGTH}"
):
for raw in (root, path):
if len(bytes(raw).decode("utf-8", "replace")) > MAX_KEY_LENGTH:
return "Image categories contains invalid values"
return None
def image_categories_schema_error(
conn: sqlite3.Connection,
*,
allow_missing: bool = True,
check_values: bool = True,
) -> str | None:
"""Reject an unsupported image categories schema (and, optionally, values).
Reads only. A missing table is acceptable for backups that predate it
unless ``allow_missing`` is false. ``check_values`` also refuses a stored
value the write path could never have written.
"""
error = _structure_error(conn, allow_missing)
if error is not None or not check_values:
return error
exists = conn.execute(
"SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (TABLE,)
).fetchone()
return _values_error(conn) if exists else None
def migrate_image_categories(conn: sqlite3.Connection) -> None:
"""Install the table and its index atomically; structure is checked
before and after. Reads only ``sqlite_master`` and this table's PRAGMAs
and writes only its own two objects (no backfill)."""
conn.execute("SAVEPOINT image_categories_schema")
try:
error = _structure_error(conn, True, allow_missing_index=True)
if error:
raise sqlite3.DatabaseError(error)
conn.execute(_DDL.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS", 1))
conn.execute(
_INDEX_DDL.replace("CREATE INDEX", "CREATE INDEX IF NOT EXISTS", 1)
)
error = image_categories_schema_error(
conn, allow_missing=False, check_values=False
)
if error:
raise sqlite3.DatabaseError(error)
conn.execute("RELEASE image_categories_schema")
except BaseException:
conn.execute("ROLLBACK TO image_categories_schema")
conn.execute("RELEASE image_categories_schema")
raise
def copy_image_category(conn: sqlite3.Connection, removed_key, kept_key) -> int:
"""Give the kept copy the removed copy's category when it has none.
Runs inside the caller's transaction; never deletes the removed copy's
row, so a give-back or a failed delete loses nothing. Returns 1 when a
row was inserted, else 0.
"""
if tuple(removed_key) == tuple(kept_key):
return 0
return max(conn.execute(COPY_SQL, (*kept_key, *removed_key)).rowcount, 0)
class ImageCategoryOperations:
"""PromptDatabase mixin: set, clear and read images' own categories.
Never reads or writes ``prompts``, ``prompt_pairs`` or
``generated_images``.
"""
def image_category(self, root_key: str, relative_key: str) -> str | None:
"""One file's own category, or None."""
row = (
self.model.get_connection()
.execute(CATEGORY_OF_SQL, (root_key, relative_key))
.fetchone()
)
return row[0] if row else None
def own_categories_of(self, root_keys: Iterable[str]) -> dict:
"""``{(root key, relative key): category}`` under the given sources.
One primary-key range read per source.
"""
found = {}
conn = self.model.get_connection()
for root_key in dict.fromkeys(root_keys):
for relative, category in conn.execute(CATEGORIES_OF_ROOT_SQL, (root_key,)):
found[(root_key, relative)] = category
return found
def set_image_categories(self, file_keys, category: str) -> dict:
"""Set *category* on every file key, in one transaction.
``file_keys``: ``(root key, relative key)`` pairs from any number of
sources. Returns ``{updated, unchanged}``; a file that already had
exactly this value is ``unchanged`` and keeps its ``updated_at``.
"""
from .models import utc_now_iso # models imports this module's schema
keys = list(dict.fromkeys(tuple(key) for key in file_keys))
counts = {"updated": 0, "unchanged": 0}
now = utc_now_iso()
with self._pair_transaction() as conn:
for key in keys:
row = conn.execute(CATEGORY_OF_SQL, key).fetchone()
if row is not None and row[0] == category:
counts["unchanged"] += 1
continue
conn.execute(UPSERT_SQL, (*key, category, now))
counts["updated"] += 1
return counts
def clear_image_categories(self, file_keys) -> int:
"""Remove the own category of every file key (one transaction).
Returns the number of rows removed.
"""
keys = list(dict.fromkeys(tuple(key) for key in file_keys))
removed = 0
with self._pair_transaction() as conn:
for key in keys:
removed += conn.execute(DELETE_SQL, key).rowcount
return removed
+490
View File
@@ -0,0 +1,490 @@
"""Collections (#331): named, hand-picked groups of image files.
A member names its file exactly as a favorite does (``image_favorites.py``):
the source folder's real, case-normalized path (``root_key``) and the path
below it (``relative_path``), both server-side only. A member whose file is
moved, renamed or deleted outside PromptManager is kept and simply not listed;
nothing follows a file to a new place.
Names are trimmed with each run of whitespace collapsed to one space, 1 to 100
characters without control characters; ``name_key`` (NFC of the case-folded
NFC name) is UNIQUE, so names are unique ignoring case and Unicode form.
The two tables install like ``image_favorites`` and ``saved_views``: optional
in old backups, required once a database is current, created empty with no
backfill. Startup checks only the structure, so a value saved under older
rules never stops the library from opening; restore checks the values too,
but never recomputes ``name_key`` (case tables differ between Python
versions).
"""
import re
import sqlite3
import unicodedata
from typing import Optional
from .image_favorites import MAX_KEY_LENGTH
COLLECTIONS_TABLE = "image_collections"
ITEMS_TABLE = "image_collection_items"
FILE_INDEX = "idx_image_collection_items_file"
TABLES = (COLLECTIONS_TABLE, ITEMS_TABLE)
MAX_NAME_LENGTH = 100
_TIMESTAMP_BYTES = 64
# A name key's byte bound: case folding expands a character to at most
# three, each at most 4 UTF-8 bytes.
_MAX_KEY_BYTES = MAX_NAME_LENGTH * 3 * 4
_WHITESPACE = re.compile(r"\s+")
_TABLE_DDL = {
COLLECTIONS_TABLE: (
f"CREATE TABLE {COLLECTIONS_TABLE} (id INTEGER PRIMARY KEY AUTOINCREMENT, "
"name TEXT NOT NULL, name_key TEXT NOT NULL UNIQUE, "
"created_at TEXT NOT NULL)"
),
ITEMS_TABLE: (
f"CREATE TABLE {ITEMS_TABLE} (collection_id INTEGER NOT NULL REFERENCES "
f"{COLLECTIONS_TABLE}(id) ON DELETE CASCADE, root_key TEXT NOT NULL, "
"relative_path TEXT NOT NULL, added_at TEXT NOT NULL, "
"PRIMARY KEY (collection_id, root_key, relative_path))"
),
}
_INDEX_DDL = f"CREATE INDEX {FILE_INDEX} ON {ITEMS_TABLE}(root_key, relative_path)"
# (name, declared type, not null, default, primary key position)
_COLUMNS = {
COLLECTIONS_TABLE: (
("id", "INTEGER", 0, None, 1),
("name", "TEXT", 1, None, 0),
("name_key", "TEXT", 1, None, 0),
("created_at", "TEXT", 1, None, 0),
),
ITEMS_TABLE: (
("collection_id", "INTEGER", 1, None, 1),
("root_key", "TEXT", 1, None, 2),
("relative_path", "TEXT", 1, None, 3),
("added_at", "TEXT", 1, None, 0),
),
}
# (column, parent table, parent column, on delete)
_FOREIGN_KEYS = {
COLLECTIONS_TABLE: set(),
ITEMS_TABLE: {("collection_id", COLLECTIONS_TABLE, "id", "CASCADE")},
}
class CollectionInvalid(ValueError):
"""A name, id or target breaks the write rules (a 400)."""
def _has_control(text: str) -> bool:
"""Any Unicode control character (category Cc: C0, DEL and C1)."""
return any(unicodedata.category(char) == "Cc" for char in text)
def normalize_collection_name(name: object) -> str:
"""The stored spelling of *name*: trimmed, whitespace collapsed, checked.
Raises ``CollectionInvalid`` unless the result is 1 to 100 characters
without control characters.
"""
if not isinstance(name, str):
raise CollectionInvalid("Name must be a string")
collapsed = _WHITESPACE.sub(" ", name).strip()
if not 1 <= len(collapsed) <= MAX_NAME_LENGTH:
raise CollectionInvalid(f"Name must be 1 to {MAX_NAME_LENGTH} characters")
if _has_control(collapsed):
raise CollectionInvalid("Name contains control characters")
return collapsed
def collection_name_key(name: str) -> str:
"""The uniqueness key of a name: NFC of the case-folded NFC spelling."""
return unicodedata.normalize("NFC", unicodedata.normalize("NFC", name).casefold())
def _compact(sql: str) -> str:
return re.sub(r"\s+", "", sql).lower()
def _expected_catalog() -> set:
objects = {("table", name, _compact(sql)) for name, sql in _TABLE_DDL.items()}
objects.add(("index", FILE_INDEX, _compact(_INDEX_DDL)))
objects |= {("index", f"sqlite_autoindex_{table}_1", "") for table in TABLES}
return objects
def _catalog(conn: sqlite3.Connection) -> list:
# sqlite_sequence (AUTOINCREMENT) is shared with other tables and ignored.
return conn.execute(
"SELECT type,name,sql FROM sqlite_master "
"WHERE name IN (?,?,?) OR tbl_name IN (?,?)",
(*TABLES, FILE_INDEX, *TABLES),
).fetchall()
def _structure_error(conn: sqlite3.Connection) -> Optional[str]:
for table, expected in _COLUMNS.items():
columns = tuple(
(row[1], row[2].upper(), row[3], row[4], row[5])
for row in conn.execute(f"PRAGMA table_info({table})")
)
if columns != expected:
return f"Image collections table {table} has invalid columns"
foreign_keys = {
(row[3], row[2], row[4], row[6])
for row in conn.execute(f"PRAGMA foreign_key_list({table})")
}
if foreign_keys != _FOREIGN_KEYS[table]:
return f"Image collections table {table} has invalid foreign keys"
return None
def _decoded_too_long(conn, sql: str, limit: int) -> bool:
"""Whether a byte-bounded candidate row decodes to over *limit* characters."""
for row in conn.execute(sql):
for raw in row:
if len(bytes(raw).decode("utf-8", "replace")) > limit:
return True
return False
def _value_error(conn: sqlite3.Connection) -> Optional[str]:
# Only this module's constants are interpolated into the SQL below. Bytes
# bound characters from above (UTF-8 is at most 4 bytes each); the few
# candidates over the character bound are measured decoded in Python.
if conn.execute(
f"SELECT 1 FROM {COLLECTIONS_TABLE} WHERE " # nosec B608
"typeof(id) != 'integer' OR typeof(name) != 'text' "
"OR typeof(name_key) != 'text' OR typeof(created_at) != 'text' "
f"OR length(CAST(name AS BLOB)) > {MAX_NAME_LENGTH * 4} "
"OR length(CAST(name_key AS BLOB)) = 0 "
f"OR length(CAST(name_key AS BLOB)) > {_MAX_KEY_BYTES} "
f"OR length(CAST(created_at AS BLOB)) > {_TIMESTAMP_BYTES} LIMIT 1"
).fetchone():
return "Image collections contains invalid values"
# The key is not recomputed: Unicode case tables differ between Python
# versions, and a key written by another version must stay restorable.
# Its UNIQUE constraint keeps names distinct; writes compute it.
for (name,) in conn.execute(f"SELECT name FROM {COLLECTIONS_TABLE}"): # nosec B608
if (
not 1 <= len(name) <= MAX_NAME_LENGTH
or _has_control(name)
or name != name.strip()
):
return "Image collections contains invalid names"
limit = MAX_KEY_LENGTH * 4
if conn.execute(
f"SELECT 1 FROM {ITEMS_TABLE} WHERE " # nosec B608
"typeof(collection_id) != 'integer' OR typeof(root_key) != 'text' "
"OR typeof(relative_path) != 'text' OR typeof(added_at) != 'text' "
f"OR length(CAST(root_key AS BLOB)) > {limit} "
f"OR length(CAST(relative_path AS BLOB)) > {limit} "
f"OR length(CAST(added_at AS BLOB)) > {_TIMESTAMP_BYTES} LIMIT 1"
).fetchone() or _decoded_too_long(
conn,
f"SELECT CAST(root_key AS BLOB), CAST(relative_path AS BLOB) " # nosec B608
f"FROM {ITEMS_TABLE} WHERE length(CAST(root_key AS BLOB)) > {MAX_KEY_LENGTH} "
f"OR length(CAST(relative_path AS BLOB)) > {MAX_KEY_LENGTH}",
MAX_KEY_LENGTH,
):
return "Image collections contains invalid members"
if conn.execute(f"PRAGMA foreign_key_check({ITEMS_TABLE})").fetchone():
return "Image collections contains members of a missing collection"
return None
def image_collections_schema_error(
conn: sqlite3.Connection,
*,
allow_missing: bool = True,
check_values: bool = True,
) -> Optional[str]:
"""Reject partial or unsupported collections schemas and invalid rows.
Reads only: nothing is ever repaired. A database without either table is
acceptable for backups that predate them unless ``allow_missing`` is
false. ``check_values=False`` (live startup) checks the structure only.
"""
objects = _catalog(conn)
if not objects:
return None if allow_missing else "Missing image collections schema"
present = {name for kind, name, _sql in objects if kind == "table"}
if present != set(TABLES):
return "Incomplete image collections schema"
actual = {(kind, name, _compact(sql or "")) for kind, name, sql in objects}
if actual != _expected_catalog():
return "Image collections has an invalid schema definition"
error = _structure_error(conn)
if error or not check_values:
return error
return _value_error(conn)
def migrate_image_collections(conn: sqlite3.Connection) -> None:
"""Install both tables atomically; validate before and after creating them."""
conn.execute("SAVEPOINT image_collections_schema")
try:
error = image_collections_schema_error(conn, check_values=False)
if error:
raise sqlite3.DatabaseError(error)
for sql in _TABLE_DDL.values():
conn.execute(sql.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS", 1))
conn.execute(_INDEX_DDL.replace("INDEX", "INDEX IF NOT EXISTS", 1))
error = image_collections_schema_error(
conn, allow_missing=False, check_values=False
)
if error:
raise sqlite3.DatabaseError(error)
conn.execute("RELEASE image_collections_schema")
except BaseException:
conn.execute("ROLLBACK TO image_collections_schema")
conn.execute("RELEASE image_collections_schema")
raise
class CollectionMissing(LookupError):
"""No collection has the id (a 404 ``collection-not-found``)."""
class CollectionNameTaken(ValueError):
"""Another collection already has the name's key (a 409 ``name-taken``)."""
def __init__(self, collection_id: int):
super().__init__("Another collection has this name")
self.collection_id = collection_id
class CollectionFilesMissing(LookupError):
"""A new name with nothing to insert (a 404 ``files-not-found``)."""
LIST_SQL = f"SELECT id, name FROM {COLLECTIONS_TABLE} ORDER BY name_key, id"
GET_SQL = f"SELECT id, name FROM {COLLECTIONS_TABLE} WHERE id = ?"
BY_KEY_SQL = f"SELECT id, name FROM {COLLECTIONS_TABLE} WHERE name_key = ?"
MEMBERS_SQL = (
f"SELECT relative_path FROM {ITEMS_TABLE} "
"WHERE collection_id = ? AND root_key = ?"
)
ROOT_MEMBERS_SQL = (
f"SELECT collection_id, relative_path FROM {ITEMS_TABLE} WHERE root_key = ?"
)
COUNTS_SQL = (
f"SELECT c.id, (SELECT COUNT(*) FROM {ITEMS_TABLE} i "
f"WHERE i.collection_id = c.id) FROM {COLLECTIONS_TABLE} c"
)
FILE_SQL = (
f"SELECT collection_id FROM {ITEMS_TABLE} "
"WHERE root_key = ? AND relative_path = ? ORDER BY collection_id"
)
# Every read, for the EXPLAIN QUERY PLAN check: none scans the items table and
# only the two whole-list reads (LIST_SQL, COUNTS_SQL) scan the collections.
READ_SQL = (
LIST_SQL,
GET_SQL,
BY_KEY_SQL,
MEMBERS_SQL,
ROOT_MEMBERS_SQL,
COUNTS_SQL,
FILE_SQL,
)
_INSERT_ITEM_SQL = f"INSERT OR IGNORE INTO {ITEMS_TABLE} VALUES (?, ?, ?, ?)"
_DELETE_ITEM_SQL = (
f"DELETE FROM {ITEMS_TABLE} "
"WHERE collection_id = ? AND root_key = ? AND relative_path = ?"
)
def _collection_id(value: object) -> int:
if type(value) is not int or not 1 <= value <= 999_999_999:
raise CollectionInvalid("Collection id must be a positive integer")
return value
def _items(items: object) -> list:
"""Distinct ``(root key, relative key)`` pairs, in order."""
if not isinstance(items, (list, tuple)):
raise CollectionInvalid("Items must be a list")
found = []
for item in items:
if (
not isinstance(item, (list, tuple))
or len(item) != 2
or not all(
isinstance(part, str) and 1 <= len(part) <= MAX_KEY_LENGTH
for part in item
)
):
raise CollectionInvalid("Each item is a root key and a relative key")
found.append((item[0], item[1]))
return list(dict.fromkeys(found))
def _row(row) -> dict:
return {"id": row[0], "name": row[1]}
class ImageCollectionOperations:
"""PromptDatabase mixin: collections and their members.
Each write runs in its own ``BEGIN IMMEDIATE`` transaction
(``_pair_transaction``), so two tabs creating one name are serialized.
"""
def list_collections(self) -> list[dict]:
"""Every collection ``{id, name}`` in name-key order, then id."""
rows = self.model.get_connection().execute(LIST_SQL)
return [_row(row) for row in rows]
def get_collection(self, collection_id: int) -> Optional[dict]:
"""``{id, name}`` of one collection, or None."""
row = (
self.model.get_connection()
.execute(GET_SQL, (_collection_id(collection_id),))
.fetchone()
)
return _row(row) if row else None
def collection_exists(self, collection_id: int) -> bool:
return self.get_collection(collection_id) is not None
def collection_members(self, collection_id: int, root_key: str) -> set[str]:
"""The relative keys stored for one collection under one source."""
rows = self.model.get_connection().execute(
MEMBERS_SQL, (_collection_id(collection_id), root_key)
)
return {path for (path,) in rows}
def root_collection_members(self, root_key: str) -> dict[int, set[str]]:
"""``{collection id: relative keys}`` stored under one source."""
found: dict[int, set[str]] = {}
for collection_id, path in self.model.get_connection().execute(
ROOT_MEMBERS_SQL, (root_key,)
):
found.setdefault(collection_id, set()).add(path)
return found
def collection_count_snapshot(self, root_keys: list) -> tuple[dict, dict]:
"""``({root key: {id: relative keys}}, {id: stored rows})`` read in
one read transaction, so a write between the reads cannot skew them.
"""
conn = self.model.get_connection()
if conn.in_transaction:
raise RuntimeError("Collection counts require their own transaction")
conn.execute("BEGIN")
try:
members = {}
for root_key in dict.fromkeys(root_keys):
found: dict[int, set[str]] = {}
for collection_id, path in conn.execute(ROOT_MEMBERS_SQL, (root_key,)):
found.setdefault(collection_id, set()).add(path)
members[root_key] = found
stored = {cid: count for cid, count in conn.execute(COUNTS_SQL)}
finally:
conn.rollback() # it only read
return members, stored
def collection_counts(self) -> dict[int, int]:
"""``{collection id: stored rows}`` for every collection."""
rows = self.model.get_connection().execute(COUNTS_SQL)
return {collection_id: count for collection_id, count in rows}
def file_collections(self, root_key: str, relative_path: str) -> list[int]:
"""Ids of the collections holding one file."""
rows = self.model.get_connection().execute(FILE_SQL, (root_key, relative_path))
return [collection_id for (collection_id,) in rows]
def add_to_collection(
self, *, items: object, collection_id: object = None, name: object = None
) -> dict:
"""Add files to an existing id or a name; create the name when needed.
A name whose key exists uses that collection; otherwise a collection
is created only when at least one item is inserted. Returns
``{collection, created, added, already}``.
"""
from .models import utc_now_iso # models imports this module's schema
if (collection_id is None) == (name is None):
raise CollectionInvalid("Give either a collection id or a name")
pairs = _items(items)
if name is not None:
spelled = normalize_collection_name(name)
else:
collection_id = _collection_id(collection_id)
now = utc_now_iso()
with self._pair_transaction() as conn:
created = False
if name is not None:
row = conn.execute(
BY_KEY_SQL, (collection_name_key(spelled),)
).fetchone()
if row is None:
if not pairs:
raise CollectionFilesMissing("Nothing to add")
cursor = conn.execute(
f"INSERT INTO {COLLECTIONS_TABLE} " # nosec B608
"(name, name_key, created_at) VALUES (?, ?, ?)",
(spelled, collection_name_key(spelled), now),
)
row, created = (cursor.lastrowid, spelled), True
else:
row = conn.execute(GET_SQL, (collection_id,)).fetchone()
if row is None:
raise CollectionMissing("No such collection")
collection = _row(row)
added = sum(
conn.execute(
_INSERT_ITEM_SQL, (collection["id"], root, path, now)
).rowcount
for root, path in pairs
)
return {
"collection": collection,
"created": created,
"added": added,
"already": len(pairs) - added,
}
def remove_from_collection(self, collection_id: object, items: object) -> dict:
"""Remove files from one collection; ``{collection, removed}``."""
collection_id = _collection_id(collection_id)
pairs = _items(items)
with self._pair_transaction() as conn:
row = conn.execute(GET_SQL, (collection_id,)).fetchone()
if row is None:
raise CollectionMissing("No such collection")
removed = sum(
conn.execute(_DELETE_ITEM_SQL, (collection_id, root, path)).rowcount
for root, path in pairs
)
return {"collection": _row(row), "removed": removed}
def rename_collection(self, collection_id: object, name: object) -> dict:
"""The renamed ``{id, name}``; a case variant of its own name is fine."""
collection_id = _collection_id(collection_id)
spelled = normalize_collection_name(name)
key = collection_name_key(spelled)
with self._pair_transaction() as conn:
if conn.execute(GET_SQL, (collection_id,)).fetchone() is None:
raise CollectionMissing("No such collection")
taken = conn.execute(BY_KEY_SQL, (key,)).fetchone()
if taken is not None and taken[0] != collection_id:
raise CollectionNameTaken(taken[0])
conn.execute(
f"UPDATE {COLLECTIONS_TABLE} SET name = ?, name_key = ? " # nosec B608
"WHERE id = ?",
(spelled, key, collection_id),
)
return {"id": collection_id, "name": spelled}
def delete_collection(self, collection_id: object) -> bool:
"""Delete one collection and its members (never a file). True."""
collection_id = _collection_id(collection_id)
with self._pair_transaction() as conn:
cursor = conn.execute(
f"DELETE FROM {COLLECTIONS_TABLE} WHERE id = ?", # nosec B608
(collection_id,),
)
if cursor.rowcount == 0:
raise CollectionMissing("No such collection")
return True
+21 -2
View File
@@ -22,6 +22,10 @@ from .managed_connection import (
from . import restore_stage
from .library_schema import library_schema_error, migrate_library_state
from .image_favorites import image_favorites_schema_error, migrate_image_favorites
from .image_collections import (
image_collections_schema_error,
migrate_image_collections,
)
from .annotation_schema import annotation_schema_error, migrate_annotation_state
from .autotag_schema import (
autotag_schema_error,
@@ -34,6 +38,7 @@ from .review_schema import (
review_without_annotations_error,
)
from .saved_views import migrate_saved_views, saved_views_schema_error
from .image_categories import image_categories_schema_error, migrate_image_categories
from .workflow_documents import (
DOCUMENT_ID_COLUMNS,
migrate_workflow_documents,
@@ -337,6 +342,7 @@ class PromptModel:
migrate_library_state(conn)
migrate_image_favorites(conn)
migrate_saved_views(conn)
migrate_image_collections(conn)
# A review group without its annotation group is refused before the
# annotation install could make it look valid (catalog only; the rows
# are validated once, in migrate_review_state).
@@ -350,6 +356,8 @@ class PromptModel:
migrate_annotation_state(conn)
migrate_review_state(conn)
migrate_autotag_state(conn)
# Image categories (#332): independent of every table above.
migrate_image_categories(conn)
# Last: the older generated_images rebuilds run only before file_path.
migrate_workflow_documents(conn)
@@ -1155,8 +1163,9 @@ class PromptModel:
Args:
path: File to inspect
require_library_state: Also require the tables every current
database has (Library state, image favorites, saved views,
image annotations and the review queue), as a migrated
database has (Library state, image favorites, image
categories, saved views, collections, image annotations and
the review queue), as a migrated
restore stage must; old backups may lack them.
documents: Read every workflow document and refuse a damaged
one. A restore checks its upload without it and repairs
@@ -1220,11 +1229,21 @@ class PromptModel:
)
if favorites_error:
return False, favorites_error
categories_error = image_categories_schema_error(
conn, allow_missing=not require_library_state
)
if categories_error:
return False, categories_error
views_error = saved_views_schema_error(
conn, allow_missing=not require_library_state
)
if views_error:
return False, views_error
collections_error = image_collections_schema_error(
conn, allow_missing=not require_library_state
)
if collections_error:
return False, collections_error
annotation_error = annotation_schema_error(
conn, allow_missing=not require_library_state
)
+188 -41
View File
@@ -7,12 +7,18 @@ import json
import math
import datetime
import os
from typing import Optional, List, Dict, Any, Set, Tuple
from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple
try:
from ..utils.inspection_response import inherited_category
except ImportError:
from utils.inspection_response import inherited_category
from .models import PromptModel, normalize_image_path, utc_now_iso
from .pairs import PairDependencyError, PairOperations # noqa: F401 (re-exported)
from .library import LibraryOperations, titled
from .image_favorites import ImageFavoriteOperations
from .image_collections import ImageCollectionOperations
from .annotations import AnnotationOperations
from .review import ReviewOperations
from .review_decisions import ReviewDecisionOperations
@@ -27,6 +33,7 @@ from .health import HealthOperations, iter_consistency_issues
from .duplicate_copies import DuplicateCopyOperations
from .workflow_storage import WorkflowStorageOperations
from .category_import import CategoryImportOperations, EmptyReplacement # noqa: F401
from .image_categories import ImageCategoryOperations
from .workflow_documents import (
DocumentReader,
documents_of_rows,
@@ -237,12 +244,17 @@ LEFT JOIN prompt_library_state pls ON pls.prompt_id = pp.positive_prompt_id
WHERE gi.image_path IN ({marks}) OR gi.file_path IN ({marks}){fallback}"""
# One image row describes a card only when no other row claims the same file
# (as in get_image_link_summaries, several rows are ambiguous). The pair scope
# and the linked positive follow _LINK_SUMMARY_SQL; category, rating, tags and
# the searchable text and titles come from that positive (spec D12/D13). Every
# and the linked positive follow _LINK_SUMMARY_SQL; rating, tags and the
# searchable text and titles come from that positive (spec D12/D13). The raw
# recipe and positive categories are read for the inherited category, which
# Python decides (``inherited_category``, #332 MC4); SQL never filters by
# category. Every
# writer since the 3.2.4 migration stores file_path, so it is the path key.
# The prompt-level conditions are uncorrelated IN sets SQLite builds once, and
# the CASE evaluates the per-row ambiguity probe only for admitted rows.
_LINK_FILTER_SQL = """SELECT gi.file_path, p.rating, {hit}, p.id
_LINK_FILTER_SQL = """SELECT gi.file_path AS file_path, p.rating AS rating,
{hit} AS matched, p.id AS positive, pp.id IS NOT NULL AS has_pair,
pp.category AS pair_category, p.category AS category
FROM generated_images gi
LEFT JOIN prompt_pairs pp ON pp.id = (
SELECT CASE WHEN COUNT(DISTINCT ip.pair_id) = 1 THEN MIN(ip.pair_id) END
@@ -251,10 +263,31 @@ JOIN prompts p ON p.id = COALESCE(pp.positive_prompt_id, gi.prompt_id)
WHERE gi.file_path IS NOT NULL AND CASE WHEN {admit} THEN NOT EXISTS (
SELECT 1 FROM generated_images other
WHERE other.file_path = gi.file_path AND other.id <> gi.id) ELSE 0 END"""
# _LINK_FILTER_SQL's rows as one JSON array, built by SQLite in one step (#354).
# Python's sqlite3 releases and retakes the GIL around every step, so a read
# that steps once per linked row waits on a busy ComfyUI thread thousands of
# times. A BLOB value (no writer stores one) makes the value NULL (or
# json_array refuses it), and the read streams instead. The raw categories are
# put in only for a category read (#332), so plain reads stay as small as #354
# measured them.
_LINK_MATCHES_JSON_SQL = """SELECT CASE
WHEN max({blob}) THEN NULL
ELSE json_group_array(json_array({columns})) END
FROM ({select})"""
_LINK_PLAIN_COLUMNS = ("file_path", "rating", "matched", "positive")
_LINK_CATEGORY_COLUMNS = _LINK_PLAIN_COLUMNS + ("has_pair", "pair_category", "category")
# The columns a BLOB can reach: matched, positive and has_pair are computed.
_LINK_STORED_COLUMNS = ("file_path", "rating", "pair_category", "category")
# Set when this SQLite has no JSON functions: every later read streams.
_one_step_unsupported = False
# One warning per process when a read falls back to streaming.
_one_step_warned = False
# The linked positive of each given path (one unambiguous link, pair scope as
# in _LINK_FILTER_SQL) with its category; {hit} as in link_filter_query. Used
# for the few files whose current image tags meet a Tags filter (#249 B3).
_LINK_FACTS_SQL = """SELECT gi.file_path, p.rating, {hit}, p.id, p.category
# in _LINK_FILTER_SQL) with the raw categories of its inherited category; {hit}
# as in link_filter_query. Used for the few files whose current image tags meet
# a Tags filter (#249 B3).
_LINK_FACTS_SQL = """SELECT gi.file_path, p.rating, {hit}, p.id,
pp.id IS NOT NULL, pp.category, p.category
FROM generated_images gi
LEFT JOIN prompt_pairs pp ON pp.id = (
SELECT CASE WHEN COUNT(DISTINCT ip.pair_id) = 1 THEN MIN(ip.pair_id) END
@@ -283,8 +316,24 @@ _LINK_FILTER_TAG = (
)
# One unambiguously linked image's facts for the Images filters, a plain tuple
# (hit, rating, positive id, inherited category): built per linked row, so no
# NamedTuple. The category is the inherited one (#332 MC4: the recipe's when
# set, else the positive's), never the image's own category.
LinkMatch = Tuple[bool, Optional[int], int, Any]
def _link_match(rating, matched, positive, has_pair, pair_category, category):
inherited, _scope = inherited_category(category, pair_category, bool(has_pair))
return (bool(matched), rating, positive, inherited)
def link_filter_query(
q: str, category: Optional[str], tags: List[str], *, every_link: bool = False
q: str,
tags: List[str],
*,
every_link: bool = False,
category: Optional[str] = None,
) -> Tuple[str, Dict[str, Any]]:
"""The workspace filter SELECT and its named parameters (all bound).
@@ -298,7 +347,9 @@ def link_filter_query(
params["pattern"] = f"%{escape_like(q)}%"
admitted = []
if category is not None:
admitted.append("p.category = :category")
# Only narrows: an inherited category is one of the two raw values;
# which one is decided in Python (inherited_category, #332 MC4).
admitted.append("(p.category = :category OR pp.category = :category)")
params["category"] = category
for index, tag in enumerate(tags):
admitted.append(_LINK_FILTER_TAG.format(index=index))
@@ -309,6 +360,50 @@ def link_filter_query(
return sql, params
def link_matches_json_query(sql: str, categories: bool = False) -> str:
"""*sql* from link_filter_query, wrapped to return its rows as one JSON value.
With *categories* each row also carries the raw values the inherited
category is decided from; otherwise only path, rating, hit and positive.
"""
columns = _LINK_CATEGORY_COLUMNS if categories else _LINK_PLAIN_COLUMNS
blob = " OR ".join(
f"typeof({name}) = 'blob'" for name in _LINK_STORED_COLUMNS if name in columns
)
return _LINK_MATCHES_JSON_SQL.format(
blob=blob, columns=", ".join(columns), select=sql
)
def _link_matches_from_rows(
rows: Iterable[Sequence[Any]], categories: bool, wanted: Optional[Set[str]]
) -> Dict[str, LinkMatch]:
"""LinkMatch per wanted path from link rows (one-step or streamed).
Without *categories* the inherited category is None, as the caller then
never compares it.
"""
matches: Dict[str, LinkMatch] = {}
for path, rating, matched, positive, *raw in rows:
if wanted is None or path in wanted:
if categories:
matches[path] = _link_match(rating, matched, positive, *raw[:3])
else:
matches[path] = (bool(matched), rating, positive, None)
return matches
def _streamed_link_matches(
conn: sqlite3.Connection,
sql: str,
params: Dict[str, Any],
wanted: Optional[Set[str]],
categories: bool = False,
) -> Dict[str, LinkMatch]:
"""The link read stepped once per row (the fallback of the one-step read)."""
return _link_matches_from_rows(conn.execute(sql, params), categories, wanted)
_ASCII_FOLD = str.maketrans("ABCDEFGHIJKLMNOPQRSTUVWXYZ", "abcdefghijklmnopqrstuvwxyz")
@@ -359,6 +454,7 @@ class PromptDatabase(
ExecutionOutputOperations,
HistoryReadOperations,
ImageFavoriteOperations,
ImageCollectionOperations,
AnnotationOperations,
ReviewOperations,
ReviewDecisionOperations,
@@ -368,6 +464,7 @@ class PromptDatabase(
DuplicateCopyOperations,
WorkflowStorageOperations,
CategoryImportOperations,
ImageCategoryOperations,
):
"""Database operations class for managing prompts."""
@@ -2541,38 +2638,94 @@ class PromptDatabase(
def get_image_link_matches(
self,
q: str,
category: Optional[str],
tags: List[str],
*,
every_link: bool = False,
wanted: Optional[Set[str]] = None,
) -> Dict[str, Tuple[bool, Optional[int], int]]:
"""Map normalized paths of unambiguously linked images to
(hit, rating, positive id).
category: Optional[str] = None,
) -> Dict[str, LinkMatch]:
"""Map normalized paths of unambiguously linked images to a LinkMatch
(hit, rating, positive id, inherited category).
Rows are limited to images whose linked positive has *category* and
every tag in *tags* (exact names). ``hit`` says whether *q* occurs in
Rows are limited to images whose linked positive has every tag in
*tags* (exact names). The category filter is applied by the caller to
the effective category (#332). ``hit`` says whether *q* occurs in
the positive's text, Library title or a tag name; a pair's own title
names its negative variant and is not searched.
Without category or tags only hits are returned, unless *every_link*
Without tags only hits are returned, unless *every_link*
asks for every linked image (the rating sort and the positive count).
The positive id is the card's linked positive. One SELECT whose rows
are streamed; with *wanted* (the listing's normalized paths) only those
paths are kept, so memory is bounded by the listing, not the database.
The positive id is the card's linked positive. With *category* only
links whose raw recipe or positive category is that value are read (a
superset of the links that inherit it; the caller compares the
inherited value) and the inherited category is decided per row;
otherwise it is None, which keeps the plain reads cheap.
One SELECT, returned as one JSON value (#354): SQLite steps once, so
the read takes the GIL back a constant two or three times instead of
once per linked row while ComfyUI generates. Peak memory therefore
grows with the database's matching linked rows (about 530 bytes each,
more with the raw categories of a category read, freed on return); the
result is bounded by *wanted* (the listing's normalized paths) when
given. Where the JSON step fails (no JSON functions, a BLOB value, a
JSON text over SQLite's length limit) the rows are streamed as before.
"""
sql, params = link_filter_query(q, category, tags, every_link=every_link)
matches: Dict[str, Tuple[bool, Optional[int], int]] = {}
sql, params = link_filter_query(
q, tags, every_link=every_link, category=category
)
categories = category is not None
with self.model.get_connection() as conn:
for path, rating, matched, positive in conn.execute(sql, params):
if wanted is None or path in wanted:
matches[path] = (bool(matched), rating, positive)
return matches
rows = None
if not _one_step_unsupported:
rows = self._one_step_link_rows(conn, sql, params, categories)
if rows is None:
return _streamed_link_matches(conn, sql, params, wanted, categories)
return _link_matches_from_rows(rows, categories, wanted)
def _one_step_link_rows(
self,
conn: sqlite3.Connection,
sql: str,
params: Dict[str, Any],
categories: bool = False,
) -> Optional[List[List[Any]]]:
"""*sql*'s rows from one JSON value, or None when this read must stream.
An interrupted read (Move library retiring the connection) is raised
as before, never retried row by row.
"""
global _one_step_unsupported, _one_step_warned
error: Optional[str] = None
try:
cursor = conn.execute(link_matches_json_query(sql, categories), params)
try:
((document,),) = cursor.fetchall()
finally:
cursor.close() # resets the statement: its snapshot ends now
except (sqlite3.OperationalError, sqlite3.DataError) as failure:
# DataError: the JSON text is longer than SQLITE_LIMIT_LENGTH.
error = str(failure)
if error == "interrupted":
raise
if error.startswith("no such function"):
_one_step_unsupported = True
else:
if document is not None:
return json.loads(document)
error = "a BLOB value"
if not _one_step_warned:
_one_step_warned = True
self.logger.warning(
"Reading image links row by row (slower while ComfyUI "
"generates); the one-step read failed: %s",
error,
)
return None
def get_image_link_facts(
self, paths: List[str], q: str = ""
) -> Dict[str, Tuple[bool, Optional[int], int, Optional[str], Tuple[str, ...]]]:
) -> Dict[str, Tuple[bool, Optional[int], int, Any, Tuple[str, ...]]]:
"""Map each normalized path with one unambiguous link to its positive's
(hit, rating, positive id, category, tag names).
(hit, rating, positive id, inherited category, tag names).
The link is resolved as in get_image_link_matches, and ``hit`` says
whether *q* occurs in the positive's text, Library title or a tag
@@ -2583,7 +2736,7 @@ class PromptDatabase(
if q:
hit = _LINK_FILTER_HIT
params["pattern"] = f"%{escape_like(q)}%"
rows: Dict[str, Tuple[bool, Optional[int], int, Optional[str]]] = {}
rows: Dict[str, LinkMatch] = {}
names: Dict[int, List[str]] = {}
unique = list(dict.fromkeys(paths))
with self.model.get_connection() as conn:
@@ -2592,10 +2745,8 @@ class PromptDatabase(
bound = {f"path{i}": path for i, path in enumerate(chunk)}
marks = ",".join(f":{key}" for key in bound)
sql = _LINK_FACTS_SQL.format(hit=hit, marks=marks)
for path, rating, matched, positive, category in conn.execute(
sql, {**params, **bound}
):
rows[path] = (bool(matched), rating, positive, category)
for path, *row in conn.execute(sql, {**params, **bound}):
rows[path] = _link_match(*row)
positives = sorted({row[2] for row in rows.values()})
for start in range(0, len(positives), _FACTS_CHUNK):
chunk = positives[start : start + _FACTS_CHUNK]
@@ -2608,20 +2759,16 @@ class PromptDatabase(
path: (*row, tuple(names.get(row[2], ()))) for path, row in rows.items()
}
def get_image_paths_matching(
self, q: str, category: Optional[str], tags: List[str]
) -> Set[str]:
"""Normalized paths whose linked positive matches *q*, *category* and *tags*.
def get_image_paths_matching(self, q: str, tags: List[str]) -> Set[str]:
"""Normalized paths whose linked positive matches *q* and *tags*.
See get_image_link_matches; an empty *q* matches every linked image
the category and tags admit.
the tags admit.
"""
return {
path
for path, (hit, _rating, _positive) in self.get_image_link_matches(
q, category, tags
).items()
if hit or not q
for path, match in self.get_image_link_matches(q, tags).items()
if match[0] or not q
}
def get_prompt_id_for_image(self, image_path: str) -> Optional[int]:
+8 -4
View File
@@ -59,6 +59,7 @@ SUMMARY_KEYS = (
"image_links",
"runs",
"saved_views",
"collections",
)
@@ -144,7 +145,8 @@ def library_summary(conn: sqlite3.Connection) -> Dict[str, int]:
``positives`` are the Library's prompts (``library_counts``'s
``positive_prompts``); ``pairs`` the active saved pairs and
``archived_pairs`` the archived ones; ``image_links`` the indexed image
rows; ``runs`` the recorded executions.
rows; ``runs`` the recorded executions; ``collections`` the image
collections (#331).
"""
from .library import _ELIGIBLE
@@ -169,6 +171,7 @@ def library_summary(conn: sqlite3.Connection) -> Dict[str, int]:
"image_links": _count(conn, "generated_images"),
"runs": _count(conn, "execution_runs"),
"saved_views": _count(conn, "saved_views"),
"collections": _count(conn, "image_collections"),
}
@@ -180,6 +183,7 @@ _FINGERPRINT_SQL = (
("prompt_pairs", "SELECT MAX(id), MAX(updated_at) FROM prompt_pairs"),
("execution_runs", "SELECT COUNT(*), MAX(rowid) FROM execution_runs"),
("saved_views", "SELECT MAX(rowid) FROM saved_views"),
("image_collections", "SELECT MAX(rowid) FROM image_collections"),
)
@@ -190,9 +194,9 @@ def live_fingerprint(conn: sqlite3.Connection):
tables a restore replaces, or when a prompt, pair, run or saved view is
the newest edited (counts plus the newest row and edit of each table);
the apply refuses with ``plan-changed`` then. It does not see every
edit (a tag or favorite change on an older row); those are kept in the
restore's safety backup. A read transaction the caller already has is
used and left open; one opened here is ended here.
edit (a tag, favorite or collection membership change on an older row);
those are kept in the restore's safety backup. A read transaction the
caller already has is used and left open; one opened here is ended here.
"""
import json
+150 -26
View File
@@ -15,6 +15,10 @@ only when all of this holds:
Anything else returns None: the caller leaves the file unlinked. Read-only.
``recorded_output`` applies the first two rules and returns, for the
inspector (#386), the one output node (save node) that run reported the file
from and the texts PromptManager recorded at execution per sampler.
Lookups are batched (``stored_run_prompts``): one parameterized query per
chunk of at most ``LOOKUP_CHUNK`` distinct file names, matched in memory. The
history schema is left exactly as it is (no extra index), so older builds can
@@ -26,6 +30,7 @@ import os
import sqlite3
import threading
import weakref
from dataclasses import dataclass
from typing import Any, Callable, Dict, Iterable, List, Optional
try:
@@ -45,7 +50,8 @@ LOOKUP_CHUNK = 500
_RUNS_FOR_NAMES = (
"SELECT DISTINCT o.filename, o.subfolder, o.root_key, o.path_style, o.run_id, "
"r.observed_dequeue_at, COALESCE(r.terminal_event_at, r.observed_callback_at) "
"r.observed_dequeue_at, COALESCE(r.terminal_event_at, r.observed_callback_at), "
"o.output_node_id "
"FROM execution_run_outputs o JOIN execution_runs r ON r.run_id = o.run_id "
"WHERE o.folder_type = 'output' AND o.filename IN ({marks})"
)
@@ -138,9 +144,10 @@ def stored_run_prompt(
def _runs_by_location(db: Any, names: List[str]) -> Optional[Dict[tuple, dict]]:
"""``{(filename, subfolder, root_key, style): {run_id: (start, end)}}``.
"""``{(filename, subfolder, root_key, style): {run_id: (start, end, nodes)}}``.
One query per ``LOOKUP_CHUNK`` names; None without history tables.
``nodes`` is the set of output node ids that reported the location in that
run. One query per ``LOOKUP_CHUNK`` names; None without history tables.
"""
found: Dict[tuple, dict] = {}
conn = db.model.get_connection()
@@ -151,12 +158,49 @@ def _runs_by_location(db: Any, names: List[str]) -> Optional[Dict[tuple, dict]]:
rows = conn.execute(query, chunk).fetchall() # nosec B608
except sqlite3.Error:
return None # no history tables (a library without history)
for name, subfolder, root_key, style, run_id, began, ended in rows:
for name, subfolder, root_key, style, run_id, began, ended, node in rows:
runs = found.setdefault((name, subfolder, root_key, style), {})
runs[run_id] = (began, ended)
runs.setdefault(run_id, (began, ended, set()))[2].add(node)
return found
def _recorded_runs(
db: Any,
image_paths: Iterable[str],
output_root: Optional[str],
path_style: Optional[str],
) -> Dict[str, tuple]:
"""``{path: (run_id, output node ids)}`` for each file history attributes.
A file is attributed when exactly one recorded run reported its location
and it was last written while that run executed.
"""
paths = list(dict.fromkeys(path for path in image_paths if path))
if not paths or not output_root:
return {}
style = path_style or ("windows" if os.name == "nt" else "posix")
located = {}
for path in paths:
location = _location(path, output_root, style)
if location is not None:
located[path] = location
if not located:
return {}
names = sorted({filename for _key, _sub, filename in located.values()})
runs = _runs_by_location(db, names)
if not runs:
return {}
attributed = {}
for path, (root_key, subfolder, filename) in located.items():
reported = runs.get((filename, subfolder, root_key, style), {})
if len(reported) != 1:
continue
((run_id, (began, ended, nodes)),) = reported.items()
if _written_during(path, began, ended):
attributed[path] = (run_id, nodes)
return attributed
def _named_prompt(db: Any, run_id: str) -> Optional[int]:
contexts = db.get_execution_contexts(run_id).get("contexts") or []
named = {
@@ -183,35 +227,115 @@ def stored_run_prompts(
"""
paths = list(dict.fromkeys(path for path in image_paths if path))
result: Dict[str, Optional[int]] = dict.fromkeys(paths)
if not paths or not output_root:
return result
style = path_style or ("windows" if os.name == "nt" else "posix")
located = {}
for path in paths:
location = _location(path, output_root, style)
if location is not None:
located[path] = location
if not located:
return result
names = sorted({filename for _key, _sub, filename in located.values()})
runs = _runs_by_location(db, names)
if not runs:
return result
prompts: Dict[str, Optional[int]] = {}
for path, (root_key, subfolder, filename) in located.items():
reported = runs.get((filename, subfolder, root_key, style), {})
if len(reported) != 1:
continue
((run_id, (began, ended)),) = reported.items()
if not _written_during(path, began, ended):
continue
for path, (run_id, _nodes) in _recorded_runs(
db, paths, output_root, path_style
).items():
if run_id not in prompts:
prompts[run_id] = _named_prompt(db, run_id)
result[path] = prompts[run_id]
return result
@dataclass(frozen=True)
class RecordedOutput:
"""What recorded history knows about the run that wrote one file (#386).
``saver_id`` is the one output node the run reported the file from (None
when several did). ``texts`` maps each sampler anchor of the run to the
effective texts PromptManager recorded at execution for its roles
(``{"positive": text or None, "negative": text or None}``); a role whose
text was only inferred from the graph is None.
"""
saver_id: Optional[str]
texts: Dict[str, Dict[str, Optional[str]]]
def _recorded_texts(db: Any, run_id: str) -> Dict[str, Dict[str, Optional[str]]]:
texts = {}
for context in db.get_execution_contexts(run_id).get("contexts") or []:
evidence = context.get("evidence") or {}
roles = {}
for role in ("positive", "negative"):
values = context.get(role) or {}
proof = (evidence.get(role) or {}).get("effective") or {}
value = values.get("effective")
recorded = proof.get("provenance") == "recorded"
roles[role] = value if recorded and isinstance(value, str) else None
texts[context["anchor_id"]] = roles
return texts
def recorded_output(
db: Any,
image_path: str,
output_root: Optional[str],
path_style: Optional[str] = None,
) -> Optional[RecordedOutput]:
"""The recorded run facts for ``image_path``, else None (#386).
The file must be attributed to one run (``_recorded_runs``: one run
reported its location and it was written during that run). Read-only; it
never links anything.
"""
if not image_path:
return None
conn = db.model.get_connection()
own_transaction = not conn.in_transaction
if own_transaction: # outputs and texts from one state of history
conn.execute("BEGIN")
try:
found = _recorded_runs(db, [image_path], output_root, path_style)
if image_path not in found:
return None
run_id, nodes = found[image_path]
saver = next(iter(nodes)) if len(nodes) == 1 else None
return RecordedOutput(saver, _recorded_texts(db, run_id))
finally:
if own_transaction:
conn.rollback()
_LOGGER = logging.getLogger("prompt_manager.run_attribution")
def recorded_output_lookup(
db: Any, *, close: bool = False, reference: Optional[dict] = None
) -> Callable[[Any], Any]:
"""``lookup(path) -> RecordedOutput or None`` for the inspect routes (#386).
Uses ComfyUI's output folder. Any failure (no history, a locked library)
is logged and names no node, so the file is judged from its own metadata.
``close`` closes the calling thread's connection afterwards, for routes
that close theirs after reading. ``reference`` (a history occurrence's
output reference) spells the file as history recorded it: the route's
target is case-normalized on Windows, and history matches names exactly.
"""
def lookup(path: Any) -> Optional[RecordedOutput]:
root = host_output_directory()
if reference is not None:
if reference.get("folder_type") != "output" or not root:
return None
path = os.path.join(
root, reference.get("subfolder") or "", reference.get("filename") or ""
)
try:
return recorded_output(db, str(path), root)
except (sqlite3.Error, OSError, ValueError):
_LOGGER.warning("Recorded run lookup failed", exc_info=True)
return None
except Exception:
_LOGGER.exception("Recorded run lookup failed unexpectedly")
return None
finally:
if close:
db.model.close()
return lookup
_SEALED_LISTENERS: List[weakref.WeakMethod] = []
_SEALED_LOCK = threading.Lock()
+55 -15
View File
@@ -13,7 +13,9 @@ same connection before commit. Owner rows and their timestamps never change.
Category consolidation is previewed in one ``library_snapshot`` and applied in
one ``_pair_transaction`` that recomputes the owner projection and requires the
preview's ``impact_token``; only rows carrying the exact source change.
preview's ``impact_token``; only rows carrying the exact source change. Image
rows holding the source (``image_categories``, #332) are owners too: they are
counted, hashed into the token and renamed with the prompts and pairs.
"""
import sqlite3
@@ -97,16 +99,30 @@ def _validate_tag_id(tag_id: object) -> None:
raise ValueError("Invalid tag ID")
# The /library/options universe: nonblank values on eligible parents or on
# active pairs, grouped exactly (whitespace and case variants stay distinct).
_CATEGORY_GROUPS = (
# Category values in Tags & categories (#332 §7): the prompt-backed values
# (nonblank values on eligible parents or on active pairs, grouped exactly, so
# whitespace and case variants stay distinct; /library/options lists the same
# prompt-backed set with its own query) plus the values only image rows hold.
# A prompt-backed value is ordered by its oldest owning positive; an image-only
# value sorts after all of them (first_pid is SQLite's INTEGER maximum, above
# any prompt id), then by value. ``image_count`` counts stored rows, including
# files missing on disk. No FULL OUTER JOIN, so it runs on SQLite 3.31+.
_IMAGE_ONLY_FIRST_PID = 2**63 - 1
_CATEGORY_ROWS = (
"WITH owners(value,pid) AS ("
"SELECT p.category,p.id FROM prompts p WHERE p.category IS NOT NULL "
"AND TRIM(p.category)!='' AND " + _ELIGIBLE + " "
"UNION SELECT pp.category,pp.positive_prompt_id FROM prompt_pairs pp "
"WHERE pp.archived=0 AND pp.category IS NOT NULL AND TRIM(pp.category)!=''),"
"grouped AS (SELECT value,MIN(pid) AS first_pid,"
"COUNT(DISTINCT pid) AS prompt_count FROM owners GROUP BY value) "
"image_values(value,image_count) AS ("
"SELECT category,COUNT(*) FROM image_categories GROUP BY category),"
"grouped(value,first_pid,prompt_count) AS ("
"SELECT value,MIN(pid),COUNT(DISTINCT pid) FROM owners GROUP BY value "
"UNION ALL SELECT value," + str(_IMAGE_ONLY_FIRST_PID) + ",0 FROM image_values "
"WHERE value NOT IN (SELECT value FROM owners)),"
"counted AS (SELECT g.value,g.first_pid,g.prompt_count,"
"COALESCE(i.image_count,0) AS image_count FROM grouped g "
"LEFT JOIN image_values i ON i.value=g.value) "
)
@@ -117,15 +133,23 @@ _CATEGORY_COMPONENTS = (
"WHERE p.category=? ORDER BY p.id"
)
_CATEGORY_PAIRS = "SELECT id,archived FROM prompt_pairs WHERE category=? ORDER BY id"
# Image rows holding one exact value (#332 DC7), an index search.
_IMAGE_CATEGORY_ROWIDS = (
"SELECT rowid FROM image_categories WHERE category=? ORDER BY rowid"
)
_CATEGORY_EXISTS = (
"SELECT EXISTS(SELECT 1 FROM prompts WHERE category=:v)"
" OR EXISTS(SELECT 1 FROM prompt_pairs WHERE category=:v)"
" OR EXISTS(SELECT 1 FROM image_categories WHERE category=:v)"
)
# Pairs after components, so a failing pair update rolls both back.
# Pairs after components and image rows last, so a failing update rolls every
# earlier one back. An image row is keyed by its file, never by its value, so
# merging into a value other files already hold cannot conflict.
_CATEGORY_UPDATES = (
"UPDATE prompts SET category=?,updated_at=? WHERE category=?",
"UPDATE prompt_pairs SET category=?,updated_at=? WHERE category=?",
"UPDATE image_categories SET category=?,updated_at=? WHERE category=?",
)
@@ -136,8 +160,15 @@ def category_impact_token(
target_exists: bool,
components: list,
pairs: list,
*,
images: list = (),
) -> str:
"""Hash a category projection like ``library_edit_token`` (no secret)."""
"""Hash a category projection like ``library_edit_token`` (no secret).
``images`` are the rowids of the image rows holding ``source`` (#332 DC7).
With none the preimage has no ``images`` key, so a prompt-only token keeps
the value it had before image categories existed.
"""
preimage = {
"epoch": epoch,
"source": source,
@@ -146,6 +177,8 @@ def category_impact_token(
"components": sorted([int(i), bool(flag)] for i, flag in components),
"pairs": sorted([int(i), bool(flag)] for i, flag in pairs),
}
if images:
preimage["images"] = sorted(int(i) for i in images)
return canonical_token(_CATEGORY_DOMAIN, preimage)
@@ -325,7 +358,8 @@ class VocabularyOperations:
pairs = [
[row[0], bool(row[1])] for row in conn.execute(_CATEGORY_PAIRS, (source,))
]
if not components and not pairs:
images = [row[0] for row in conn.execute(_IMAGE_CATEGORY_ROWIDS, (source,))]
if not components and not pairs and not images:
raise LibraryCommandError("vocabulary-conflict", epoch)
target_exists = bool(
conn.execute(_CATEGORY_EXISTS, {"v": target}).fetchone()[0]
@@ -339,9 +373,10 @@ class VocabularyOperations:
"other_prompts": len(components) - eligible,
"negative_variants": len(pairs) - archived,
"archived_negative_variants": archived,
"images": len(images),
},
"impact_token": category_impact_token(
epoch, source, target, target_exists, components, pairs
epoch, source, target, target_exists, components, pairs, images=images
),
}
@@ -365,7 +400,8 @@ class VocabularyOperations:
def apply_vocabulary_category(
self, *, expected_epoch: str, source: str, target: str, impact_token: str
) -> dict:
"""Set the exact ``source`` to ``target`` on every component and pair.
"""Set the exact ``source`` to ``target`` on every component, pair
and image row.
The recomputed projection must still match the preview's token; a
mismatch is ``vocabulary-conflict`` with nothing written.
@@ -422,15 +458,15 @@ class VocabularyOperations:
_limit(limit)
key = _category_after(after)
total = conn.execute(
_CATEGORY_GROUPS + "SELECT COUNT(*) FROM grouped"
_CATEGORY_ROWS + "SELECT COUNT(*) FROM counted"
).fetchone()[0]
where, params = "", []
if key is not None:
where = " WHERE first_pid>? OR (first_pid=? AND value>?)"
params = [key[0], key[0], key[1]]
rows = conn.execute(
_CATEGORY_GROUPS
+ "SELECT value,first_pid,prompt_count FROM grouped"
_CATEGORY_ROWS
+ "SELECT value,first_pid,prompt_count,image_count FROM counted"
+ where
+ " ORDER BY first_pid,value LIMIT ?",
[*params, limit + 1],
@@ -440,7 +476,11 @@ class VocabularyOperations:
return {
"library_epoch": epoch,
"categories": [
{"value": row["value"], "prompt_count": row["prompt_count"]}
{
"value": row["value"],
"prompt_count": row["prompt_count"],
"image_count": row["image_count"],
}
for row in rows
],
"total": total,
+1
View File
@@ -268,6 +268,7 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
record_runtime_context(
type(self).__name__, text, prepend_text, append_text, encoding_text
)
self._record_cached_item(unique_id, (encoding_text,), storage_text)
return result
@classmethod
+58 -3
View File
@@ -26,6 +26,13 @@ try:
from .utils.prompt_graph import resolve_text, run_prompt_nodes
from .utils.run_keys import executing_run_id, prompt_node_key
from .utils.usage_tracking import PENDING_FIRST_USE
from .utils.cached_runs import (
CACHED_RUNS,
Role,
Seen,
own_counts,
stored_facts,
)
except ImportError:
import sys
@@ -38,6 +45,13 @@ except ImportError:
from utils.prompt_graph import resolve_text, run_prompt_nodes
from utils.run_keys import executing_run_id, prompt_node_key
from utils.usage_tracking import PENDING_FIRST_USE
from utils.cached_runs import (
CACHED_RUNS,
Role,
Seen,
own_counts,
stored_facts,
)
class PromptManagerBase:
@@ -71,8 +85,10 @@ class PromptManagerBase:
"""Save the prompt to the SQLite database.
Runs are normally counted by the queue hook (utils/usage_tracking.py), so a
save only counts one when asked to (batch items whose text the hook could not
see) or when the hook queued this prompt before it existed.
save only counts one when asked to (text the hook could not see, unless
this job already counted the node from ComfyUI's cache) or when the hook
queued this prompt before it existed. The two are independent: the
pending catch-up is written even when ``count_run`` is False.
Args:
text: The prompt text
@@ -167,8 +183,19 @@ class PromptManagerBase:
batch_item = (
role["known"] and role["is_run_prompt"] and not role["seen_at_queue"]
)
counts = own_counts(
type(self).__name__,
stored_facts(text),
Role(
frozenset({"text"} if role["known"] and role["is_run_prompt"] else ())
),
Seen(text=role["seen_at_queue"]),
)
prompt_id = self._save_prompt_to_database(
text=text, category=category, tags=tags, count_run=batch_item
text=text,
category=category,
tags=tags,
count_run=bool(counts) and not self._replayed(unique_id),
)
if role["is_run_prompt"] and not prompt_id:
# The save failed: an older entry of this key must not stand in (#342)
@@ -182,6 +209,34 @@ class PromptManagerBase:
)
return prompt_id
def _cached_runs(self) -> Any:
return getattr(self, "cached_runs", CACHED_RUNS)
def _replayed(self, unique_id: Any) -> bool:
"""This job already counted the node from ComfyUI's cache (#347).
A RAM-pressure cache can free the entry after ``execution_cached``, so
the node executes in a job that replayed it: its own count is skipped.
"""
try:
return unique_id is not None and self._cached_runs().replayed(unique_id)
except Exception as error:
self.logger.debug(f"Could not check the cached-run replay: {error}")
return False
def _record_cached_item(
self, unique_id: Any, outputs: tuple, text: Any, negative_text: Any = None
) -> None:
"""Remember this execution's outputs and stored text for cached jobs (#347)."""
if unique_id is None:
return
try:
self._cached_runs().record_item(
unique_id, outputs, stored_facts(text, negative_text)
)
except Exception as error:
self.logger.debug(f"Could not record the cached-run receipt: {error}")
@staticmethod
def _run_identity(prompt_graph: Any, unique_id: Any) -> Dict[str, Any]:
"""The tracker's proof that an image belongs to this execution (#342)."""
+40 -22
View File
@@ -12,6 +12,14 @@ else:
try:
from .prompt_manager import IO, InputTypeDict, PromptManagerFocused
from .utils.cached_runs import (
PAIR,
Role,
Seen,
own_counts,
pair_components,
stored_facts,
)
from .utils.hashing import generate_prompt_hash
from .utils.pair_execution import record_pair_execution
from .utils.prompt_graph import intrinsic_pair_nodes, run_prompt_components
@@ -23,6 +31,14 @@ try:
)
except ImportError:
from prompt_manager import IO, InputTypeDict, PromptManagerFocused
from utils.cached_runs import (
PAIR,
Role,
Seen,
own_counts,
pair_components,
stored_facts,
)
from utils.hashing import generate_prompt_hash
from utils.pair_execution import record_pair_execution
from utils.prompt_graph import intrinsic_pair_nodes, run_prompt_components
@@ -165,6 +181,9 @@ class PromptManagerPair(PromptManagerFocused):
record_runtime_context(
type(self).__name__, negative_text, "", "", negative_text, role="negative"
)
self._record_cached_item(
unique_id, (positive, negative_text), text, negative_text
)
return result
def _save_and_track_pair(
@@ -209,36 +228,35 @@ class PromptManagerPair(PromptManagerFocused):
and generate_prompt_hash(resolved_negative)
== generate_prompt_hash(negative_text)
)
negative_seen = resolved_negative is not None and generate_prompt_hash(
resolved_negative
) == generate_prompt_hash(negative_text)
positive_pending = getattr(self, "pending_first_use", PENDING_FIRST_USE)
pair_pending = getattr(self, "pending_pair_first_use", PENDING_PAIR_FIRST_USE)
counted = set()
for field, component_text, resolved_text, component_id in (
("text", text, resolved_positive, prompt_id),
(
"negative_text",
negative_text,
resolved_negative,
pair["negative_prompt_id"],
),
):
component_hash = generate_prompt_hash(component_text)
if (
field not in used_fields
or not component_id
or component_hash in counted
):
facts = stored_facts(text, negative_text)
role = Role(frozenset(used_fields), pair_used)
# Own counts are skipped when this job already counted the node from
# ComfyUI's cache (#347); the pending first-use catch-up never is.
counts = (
[]
if self._replayed(unique_id)
else own_counts(
PAIR, facts, role, Seen(positive_seen, negative_seen, pair_seen)
)
)
component_ids = {"text": prompt_id, "negative_text": pair["negative_prompt_id"]}
for field, component_hash in pair_components(facts, role):
component_id = component_ids[field]
if not component_id:
continue
counted.add(component_hash)
seen_at_queue = (
resolved_text is not None
and generate_prompt_hash(resolved_text) == component_hash
runs = positive_pending.consume(component_hash) + int(
("prompt", component_hash) in counts
)
runs = positive_pending.consume(component_hash) + int(not seen_at_queue)
if runs:
self.db.record_prompt_use(component_id, times=runs)
if pair_used:
runs = pair_pending.consume(pair_run_key(text, negative_text)) + int(
not pair_seen
("pair", facts.positive_hash, facts.negative_hash) in counts
)
if runs:
self.db.record_prompt_pair_use(pair_id, times=runs)
+1
View File
@@ -239,6 +239,7 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC):
record_runtime_context(
type(self).__name__, text, prepend_text, append_text, final_text
)
self._record_cached_item(unique_id, (final_text,), storage_text)
return result
@classmethod
+6
View File
@@ -48,6 +48,7 @@ from .review import ReviewRoutesMixin
from .review_recording import ReviewRecordingMixin
from .review_events import ReviewEventsMixin
from .saved_views import SavedViewRoutesMixin
from .collections import CollectionRoutesMixin
from .history import HistoryRoutesMixin
from .images import ImageRoutesMixin
from .gallery_listing import WALK_STOPPED, RootListing, WalkOutcome
@@ -72,6 +73,7 @@ from .diagnostics_report import DiagnosticsReportMixin
from .autotag_routes import AutotagRoutesMixin
from .lora_integration import LoraIntegrationMixin
from .lora_import_routes import LoraImportRoutesMixin
from .image_categories import ImageCategoryRoutesMixin
try:
from ...database.operations import PromptDatabase
@@ -314,6 +316,7 @@ class PromptManagerAPI(
ReviewRecordingMixin,
ReviewEventsMixin,
SavedViewRoutesMixin,
CollectionRoutesMixin,
HistoryRoutesMixin,
PromptRoutesMixin,
ImageRoutesMixin,
@@ -336,6 +339,7 @@ class PromptManagerAPI(
AutotagRoutesMixin,
LoraIntegrationMixin,
LoraImportRoutesMixin,
ImageCategoryRoutesMixin,
):
"""REST API handler for PromptManager operations and web interface.
@@ -684,6 +688,7 @@ class PromptManagerAPI(
self._register_review_routes(routes)
self._register_review_event_routes(routes)
self._register_saved_view_routes(routes)
self._register_collection_routes(routes)
self._register_history_routes(routes)
self._register_prompt_routes(routes)
self._register_image_routes(routes)
@@ -700,6 +705,7 @@ class PromptManagerAPI(
self._register_autotag_routes(routes)
self._register_lora_routes(routes)
self._register_lora_import_routes(routes)
self._register_image_category_routes(routes)
# Register gzip compression middleware (once)
global _gzip_registered
+6
View File
@@ -34,6 +34,11 @@ try:
except ImportError:
from utils.parallel import map_parallel
try:
from ...utils import exit_signal
except ImportError:
from utils import exit_signal
try:
from ...database.pair_ingestion import blank_positive, ingest_resolved_pair
from ...database.run_attribution import (
@@ -118,6 +123,7 @@ def _collect_media_files(output_dirs, extensions):
if not os.path.isdir(root):
continue
for current, dirnames, filenames in os.walk(root, followlinks=False):
exit_signal.checkpoint() # #376: stop between directories at exit
depth = 0 if current == root else len(Path(current).relative_to(root).parts)
dirnames[:] = sorted(
d
+8 -36
View File
@@ -27,7 +27,12 @@ from .annotation_suggest import (
validate_options,
)
from .autotag_routes import model_busy_response
from .images import _forbidden, _has_traversal, _unavailable_link, _verified_title
# The target check is shared with the collection routes (#331); both names
# stay importable from here.
from .media_target import MAX_TARGET_PATH # noqa: F401
from .media_target import media_target as _target
from .images import _unavailable_link, _verified_title
from .prompts import bad_request
from .. import model_gate
@@ -40,7 +45,7 @@ try:
)
from ...database.execution_history import subscribe_history_restore
from ...database.library_commands import LibraryCommandError
from ...utils.media_capture import CAPTURE_EXTENSIONS, CaptureError, CaptureStore
from ...utils.media_capture import CaptureError, CaptureStore
except ImportError:
from database.annotations import (
MAX_IMAGE_TAGS,
@@ -50,15 +55,13 @@ except ImportError:
)
from database.execution_history import subscribe_history_restore
from database.library_commands import LibraryCommandError
from utils.media_capture import CAPTURE_EXTENSIONS, CaptureError, CaptureStore
from utils.media_capture import CaptureError, CaptureStore
# Request bodies of the annotation routes are small (a target plus a tag set).
MAX_BODY_BYTES = 64 * 1024
MAX_TARGET_PATH = 1024
CAPTURE_URL = "/prompt_manager/image-annotations/capture/"
_PREVIEW_CHUNK = 64 * 1024
_TARGET_FIELDS = frozenset({"root_index", "root_scope", "path"})
_OBSERVE_FIELDS = frozenset({"target"})
_SAVE_FIELDS = frozenset(
{
@@ -188,37 +191,6 @@ def _capture_error(exc: CaptureError) -> web.Response:
return _coded(exc.code, _MESSAGES.get(exc.code) or str(exc))
def _target(data: object):
"""``((index, scope, relative), None)`` for a target, else ``(None, error)``.
Only the listing identity is accepted: an int root index, the root's
64-hex scope and a bounded root-relative POSIX media path, never an
absolute, drive-qualified, backslash or ``..`` path. Nothing here touches
the filesystem or the database.
"""
if not isinstance(data, dict) or set(data) != _TARGET_FIELDS:
return None, bad_request("target must name root_index, root_scope and path")
index, scope, relative = data["root_index"], data["root_scope"], data["path"]
if type(index) is not int or index < 0:
return None, bad_request("Invalid gallery root")
if not isinstance(scope, str) or _HEX64.fullmatch(scope) is None:
return None, bad_request("Invalid gallery root scope")
if not isinstance(relative, str) or not relative or len(relative) > MAX_TARGET_PATH:
return None, bad_request("path must be a relative media path")
parts = relative.split("/")
if (
"\0" in relative
or "\\" in relative
or relative.startswith("/")
or _has_traversal(relative)
or any(part in ("", ".", "..") or ":" in part for part in parts)
):
return None, _forbidden()
if os.path.splitext(parts[-1])[1].lower() not in CAPTURE_EXTENSIONS:
return None, _forbidden("Only media files can be tagged")
return (index, scope, relative), None
def _expires_at(seconds: float | None) -> str:
moment = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(
seconds=max(0.0, seconds or 0.0)
+350
View File
@@ -0,0 +1,350 @@
"""Collections (#331): named, hand-picked groups of images.
Routes under ``/prompt_manager/collections``. A file is named by its listing
identity ``{root_index, root_scope, path}`` (``media_target.py``) and stored
under the favorites' keys (``_favorite_root_key``, ``_favorite_relative_key``),
which never leave the server. A stale scope (the sources were reordered) is a
409 ``source-changed`` before anything is written. Responses are
``{"success": bool, ...}``; an unexpected failure is a logged 500 with a
generic message, never the exception text.
"""
import hmac
import re
from functools import wraps
from typing import Callable
from aiohttp import web
from .images import (
_favorite_relative_key,
_favorite_root_key,
_forbidden,
_json_body_or_error,
collection_not_found,
)
from .media_target import media_target
from .prompts import bad_request
try:
from ...database.image_collections import (
CollectionFilesMissing,
CollectionInvalid,
CollectionMissing,
CollectionNameTaken,
normalize_collection_name,
)
from ...database.image_favorites import MAX_KEY_LENGTH
except ImportError:
from database.image_collections import (
CollectionFilesMissing,
CollectionInvalid,
CollectionMissing,
CollectionNameTaken,
normalize_collection_name,
)
from database.image_favorites import MAX_KEY_LENGTH
BASE = "/prompt_manager/collections"
# The most files one add or remove names (the bulk tagging cap, MC12).
MAX_ITEMS = 500
MAX_ID = 999_999_999
_PATH_ID = re.compile(r"[1-9][0-9]{0,8}")
_ROOT_INDEX = re.compile(r"[0-9]{1,9}")
_ADD_TARGETS = (frozenset({"collection_id", "items"}), frozenset({"name", "items"}))
_REMOVE_FIELDS = frozenset({"collection_id", "items"})
_RENAME_FIELDS = frozenset({"name"})
_MEDIA_MESSAGE = "Only media files can be added to a collection"
READ_FAILED = "Collections could not be read."
CHANGE_FAILED = "The collection could not be changed."
SOURCES_CHANGED = "The sources changed. Refresh the images and try again."
FILES_NOT_FOUND = (
"None of the selected images could be found. Refresh the images and try again."
)
NAME_TAKEN = "Another collection already has this name."
class _Refused(Exception):
"""A request refused before any write; carries its response."""
def __init__(self, response: web.Response):
super().__init__(response.status)
self.response = response
def _coded(status: int, code: str, message: str, **extra) -> web.Response:
return web.json_response(
{"success": False, "code": code, "error": message, **extra},
status=status,
headers={"Cache-Control": "no-store"},
)
def _collection_response(failure: str) -> Callable:
"""Map the database's refusals to their status; anything else to a 500."""
def decorate(handler: Callable) -> Callable:
@wraps(handler)
async def wrapped(self, request: web.Request) -> web.Response:
try:
return await handler(self, request)
except _Refused as refused:
return refused.response
except CollectionInvalid as exc:
return bad_request(str(exc))
except CollectionMissing:
return collection_not_found()
except CollectionFilesMissing:
return _coded(404, "files-not-found", FILES_NOT_FOUND)
except CollectionNameTaken as exc:
return _coded(
409, "name-taken", NAME_TAKEN, collection_id=exc.collection_id
)
except Exception:
self.logger.exception("Collection request failed")
return web.json_response(
{"success": False, "error": failure}, status=500
)
return wrapped
return decorate
def _path_id(request: web.Request) -> int:
raw = request.match_info["collection_id"]
if not raw.isascii() or _PATH_ID.fullmatch(raw) is None:
raise _Refused(bad_request("Invalid collection id"))
return int(raw)
def _body_id(value: object) -> int:
if type(value) is not int or not 1 <= value <= MAX_ID:
raise _Refused(bad_request("Invalid collection id"))
return value
async def _body(request: web.Request, allowed: tuple) -> dict:
data, error = await _json_body_or_error(request)
if error is not None:
raise _Refused(error)
if frozenset(data) not in allowed:
names = " or ".join(", ".join(sorted(fields)) for fields in allowed)
raise _Refused(bad_request(f"Request must contain exactly {names}"))
return data
class CollectionRoutesMixin:
"""Mixin providing the collection routes (#331)."""
def _register_collection_routes(self, routes: web.RouteTableDef) -> None:
# The fixed paths go first: "/items" must never match "/{collection_id}".
routes.get(BASE)(self.list_collections)
routes.post(BASE + "/items")(self.add_collection_items)
routes.delete(BASE + "/items")(self.remove_collection_items)
routes.get(BASE + "/membership")(self.collection_membership)
routes.patch(BASE + "/{collection_id}")(self.rename_collection)
routes.delete(BASE + "/{collection_id}")(self.delete_collection)
# -- reads -------------------------------------------------------------
@_collection_response(READ_FAILED)
async def list_collections(self, request: web.Request) -> web.Response:
"""Every collection with its listed ``count`` and its ``missing`` rows.
``count`` is the members listed under every configured source (one
primary-key range read per source, matched against its cached
listing); ``missing`` the stored rows beyond it. When a scan or a read
fails both are null for every collection and the names still return.
"""
collections = await self._run_in_executor(self.db.list_collections)
counts = None
if collections:
try:
counts = await self._collection_counts()
except Exception:
self.logger.warning(
"Collection counts could not be read", exc_info=True
)
rows = []
for collection in collections:
listed, stored = (None, None) if counts is None else counts
count = None if listed is None else listed.get(collection["id"], 0)
stored_rows = None if stored is None else stored.get(collection["id"], 0)
rows.append(
{
**collection,
"count": count,
# Clamped: never below zero, whatever the snapshot saw.
"missing": None if count is None else max(0, stored_rows - count),
}
)
return web.json_response({"success": True, "collections": rows})
async def _collection_counts(self) -> tuple[dict, dict]:
"""``({id: listed members}, {id: stored rows})`` over every source.
The scans come from the gallery cache; the root keys (realpath), the
listings and one read snapshot of every member and stored count run
in the executor, so a write between two reads cannot skew them.
"""
roots = list(self._get_all_output_dirs())
root_keys = await self._run_in_executor(
lambda: [_favorite_root_key(root) for root in roots]
)
distinct = {} # one folder configured twice counts once
for root, root_key in zip(roots, root_keys):
distinct.setdefault(root_key, root)
scans = {
root_key: await self._get_gallery_files(root)
for root_key, root in distinct.items()
}
return await self._run_in_executor(self._count_members, distinct, scans)
def _count_members(self, roots: dict, scans: dict) -> tuple[dict, dict]:
"""Blocking: match one snapshot's members against each listing."""
members, stored = self.db.collection_count_snapshot(list(roots))
listed: dict[int, int] = {}
for root_key, root in roots.items():
listing = self._favorite_listing(root, scans[root_key])
for collection_id, keys in members.get(root_key, {}).items():
found = sum(1 for key in keys if key in listing)
listed[collection_id] = listed.get(collection_id, 0) + found
return listed, stored
@_collection_response(READ_FAILED)
async def collection_membership(self, request: web.Request) -> web.Response:
"""GET ``?root_index&root_scope&path``: the ids holding that file."""
query = request.query
if len(query) != 3 or set(query) != {"root_index", "root_scope", "path"}:
# len() counts a repeated key (?path=a&path=b) twice.
return bad_request("Query must name root_index, root_scope and path once")
raw_index = query["root_index"]
if not raw_index.isascii() or _ROOT_INDEX.fullmatch(raw_index) is None:
return bad_request("Invalid gallery root")
roots = list(self._get_all_output_dirs())
target = {
"root_index": int(raw_index),
"root_scope": query["root_scope"],
"path": query["path"],
}
root_key, key = await self._run_in_executor(
self._collection_key, target, roots, {}
)
ids = await self._run_in_executor(self.db.file_collections, root_key, key)
return web.json_response({"success": True, "collection_ids": ids})
# -- writes ------------------------------------------------------------
@_collection_response(CHANGE_FAILED)
async def add_collection_items(self, request: web.Request) -> web.Response:
"""POST ``{collection_id | name, items}``: add the listed files.
Only files the source's listing holds are added; the rest count in
``not_found``. A new name with nothing to add creates nothing (404
``files-not-found``).
"""
data = await _body(request, _ADD_TARGETS)
target = {}
if "name" in data:
target["name"] = normalize_collection_name(data["name"])
else:
target["collection_id"] = _body_id(data["collection_id"])
roots = list(self._get_all_output_dirs())
keys = await self._run_in_executor(self._collection_items, data["items"], roots)
found, not_found = [], 0
listings = {}
for index, root_key, key in keys:
if index not in listings:
root = roots[index]
files = await self._get_gallery_files(root)
listings[index] = await self._run_in_executor(
self._favorite_listing, root, files
)
if key in listings[index]:
found.append((root_key, key))
else:
not_found += 1
result = await self._run_in_executor(
lambda: self.db.add_to_collection(items=found, **target)
)
return web.json_response({"success": True, **result, "not_found": not_found})
@_collection_response(CHANGE_FAILED)
async def remove_collection_items(self, request: web.Request) -> web.Response:
"""DELETE ``{collection_id, items}``: listed or not, every named file."""
data = await _body(request, (_REMOVE_FIELDS,))
collection_id = _body_id(data["collection_id"])
roots = list(self._get_all_output_dirs())
keys = await self._run_in_executor(self._collection_items, data["items"], roots)
result = await self._run_in_executor(
self.db.remove_from_collection,
collection_id,
[(root_key, key) for _index, root_key, key in keys],
)
return web.json_response({"success": True, **result})
@_collection_response(CHANGE_FAILED)
async def rename_collection(self, request: web.Request) -> web.Response:
"""PATCH ``{name}``: a case variant of its own name is allowed."""
collection_id = _path_id(request)
data = await _body(request, (_RENAME_FIELDS,))
name = normalize_collection_name(data["name"])
collection = await self._run_in_executor(
self.db.rename_collection, collection_id, name
)
return web.json_response({"success": True, "collection": collection})
@_collection_response(CHANGE_FAILED)
async def delete_collection(self, request: web.Request) -> web.Response:
"""DELETE: the collection and its members; files are never touched."""
collection_id = _path_id(request)
await self._run_in_executor(self.db.delete_collection, collection_id)
return web.json_response({"success": True, "deleted": True})
# -- request steps -----------------------------------------------------
def _collection_items(self, items: object, roots: list) -> list[tuple]:
"""``(index, root key, relative key)`` per distinct file, in order.
Blocking (realpath per root, once per request): run in the executor.
Every item is checked before anything is written: 400 for a malformed
shape, 403 for an unsafe path, 409 for a stale scope. Two spellings
that fold to one key (Windows) count once.
"""
if not isinstance(items, list) or not 1 <= len(items) <= MAX_ITEMS:
raise _Refused(bad_request(f"items must list 1 to {MAX_ITEMS} images"))
keys, sources = {}, {}
for item in items:
index = item.get("root_index") if isinstance(item, dict) else None
root_key, key = self._collection_key(item, roots, sources)
keys.setdefault((root_key, key), index)
return [(index, root_key, key) for (root_key, key), index in keys.items()]
def _collection_key(
self, item: object, roots: list, sources: dict
) -> tuple[str, str]:
"""The stored keys of one target, or raise its refusal (blocking).
*sources* caches ``{index: (scope, root key)}`` for one request, so
each root's realpath runs once however many items name it.
"""
target, error = media_target(item, media_message=_MEDIA_MESSAGE)
if error is not None:
raise _Refused(error)
index, scope, relative = target
if index >= len(roots):
raise _Refused(_coded(409, "source-changed", SOURCES_CHANGED))
if index not in sources:
sources[index] = (
self._gallery_root_scope(roots[index], index),
_favorite_root_key(roots[index]),
)
current, root_key = sources[index]
if not hmac.compare_digest(scope, current):
raise _Refused(_coded(409, "source-changed", SOURCES_CHANGED))
key = _favorite_relative_key(relative)
if max(len(root_key), len(key)) > MAX_KEY_LENGTH:
# A stored key past the schema bound would fail the next restore.
raise _Refused(_forbidden("Path is too long to add to a collection"))
return root_key, key
+12 -1
View File
@@ -588,7 +588,10 @@ class DuplicateFilesMixin:
job.update(total=len(targets), status_line="Removing duplicate files")
spellings = await self._run_in_executor(self._source_spellings)
removed, failure, keepers = 0, None, {}
totals = dict.fromkeys(("moved", "merged", "favorites", "tags", "review"), 0)
totals = dict.fromkeys(
("moved", "merged", "favorites", "tags", "review", "collections"), 0
)
totals["categories"] = 0 # image categories kept on the kept copies (#332)
for done, target in enumerate(targets, start=1):
if failure is not None:
self._record_target(job, "kept", target, NOT_ATTEMPTED)
@@ -631,10 +634,18 @@ class DuplicateFilesMixin:
job.log_line(
f"Moved {_plural(totals['favorites'], 'favorite')} to the kept copies"
)
if totals["collections"]:
job.log_line(
f"Moved {_plural(totals['collections'], 'collection membership')} "
"to the kept copies"
)
if totals["tags"]:
job.log_line(
f"Added {_plural(totals['tags'], 'image tag')} to the kept copies"
)
if totals.get("categories"):
kept = _plural(totals["categories"], "image category", "image categories")
job.log_line(f"Kept {kept} on the kept copies")
if totals["review"]:
job.log_line(
f"Settled {_plural(totals['review'], 'Needs review item')} "
+26 -2
View File
@@ -18,7 +18,7 @@ Blocking: each lookup stats the file's folders once (cached).
import os
from pathlib import Path
from .images import _favorite_relative_key, _favorite_root_key
from .images import _favorite_relative_key, _favorite_root_key, image_category_key
def _identity(path):
@@ -78,6 +78,15 @@ class Spellings:
)
return spelled
@property
def listed_roots(self):
"""The listed roots (the output sources), in order."""
return [root for root, _identity in self._listed]
def category_key(self, path):
"""*path*'s image category file key (#332), or None outside every source."""
return image_category_key(self.listed_roots, path)
def favorites(self, path):
"""``{listed root: path below it (POSIX)}`` for each listed root holding it."""
path = os.path.abspath(str(path))
@@ -95,7 +104,8 @@ def copy_move(target, kept_stat, spellings):
A record found under a folder spelling moves to the kept copy under
the same spelling when that folder holds it too, else to its real
path. Its favorite moves under every listed root holding the file.
path. Its favorite and its collection memberships move under every
listed root holding the file.
"""
removed = spellings.of(target["abs"])
kept = spellings.of(target["keeper_abs"])
@@ -125,9 +135,23 @@ def copy_move(target, kept_stat, spellings):
"targets": targets,
"kept_path": target["keeper_abs"],
"favorites": favorites,
"categories": _category_pairs(target, spellings),
"removed_locator": target["locator"],
"kept_locator": target["keeper_locator"],
"sha256": target["hash"],
"kept_size": kept_stat.st_size,
"kept_mtime_ns": kept_stat.st_mtime_ns,
}
def _category_pairs(target, spellings):
"""``[(removed file key, kept file key)]`` for the image category (#332).
From the two absolute paths, not the favorites tuples: those follow the
listed root, which differs from the file key when sources nest.
"""
removed = spellings.category_key(target["abs"])
kept = spellings.category_key(target["keeper_abs"])
if removed is None or kept is None:
return []
return [(removed, kept)]
+58 -5
View File
@@ -5,6 +5,7 @@ import base64
import binascii
from functools import wraps
import json
import os
import re
import sqlite3
from typing import Callable
@@ -27,10 +28,12 @@ from .history_current import (
snapshot_current_output_scope,
)
from .inspection_service import CurrentFileChanged, inspect_current_file
from .images import image_category_key
try:
from ...database.execution_history import HistoryEpochChanged
from ...database.history_gallery import read_history_gallery
from ...database.run_attribution import recorded_output_lookup
from ...database.history_reads import (
HistoryCaptureChanged,
OutputPosition,
@@ -38,10 +41,15 @@ try:
WindowPosition,
)
from ...utils.history_read_fence import HistoryReadFence
from ...utils.inspection_response import encode_history_inspection_response
from ...utils.inspection_response import (
OWN_CATEGORY_UNREAD,
apply_own_category,
encode_history_inspection_response,
)
except ImportError:
from database.execution_history import HistoryEpochChanged
from database.history_gallery import read_history_gallery
from database.run_attribution import recorded_output_lookup
from database.history_reads import (
HistoryCaptureChanged,
OutputPosition,
@@ -49,7 +57,11 @@ except ImportError:
WindowPosition,
)
from utils.history_read_fence import HistoryReadFence
from utils.inspection_response import encode_history_inspection_response
from utils.inspection_response import (
OWN_CATEGORY_UNREAD,
apply_own_category,
encode_history_inspection_response,
)
_RUN_MAX_BYTES = 8 * 1024 * 1024
@@ -675,6 +687,32 @@ class HistoryRoutesMixin:
) -> web.Response:
return await self._history_current(request, "inspection")
def _history_own_category(self, reference: dict, scope):
"""The recorded output's own category (#332), inside the snapshot.
The path is the symlink-resolved output folder plus the receipt's
subfolder and filename, keyed by ``image_category_key`` over the
configured output sources as Images keys it. A temp output or a path
no source holds has none. A failed read is OWN_CATEGORY_UNREAD.
"""
if (
reference.get("folder_type") != "output"
or scope.canonical_output is None
or candidate_locator(reference, scope) is None
):
return None
path = os.path.join(
str(scope.canonical_output), reference["subfolder"], reference["filename"]
)
try:
key = image_category_key(self._get_all_output_dirs(), path)
return None if key is None else self.db.image_category(*key)
except Exception:
self.logger.warning(
"History image category could not be read", exc_info=True
)
return OWN_CATEGORY_UNREAD
async def _history_current(
self, request: web.Request, kind: str
) -> web.StreamResponse:
@@ -699,6 +737,7 @@ class HistoryRoutesMixin:
if scope.current_scope != current_scope:
raise _current_stale()
result = db.get_history_occurrence(run_id, node, ordinal)
own = None
if result is not None:
receipt = result["receipt"]
path = candidate_locator(receipt["reference"], scope)
@@ -706,6 +745,10 @@ class HistoryRoutesMixin:
db.model.get_connection(),
(path,) if path is not None else (),
)
if kind == "inspection":
own = self._history_own_category(
receipt["reference"], scope
)
finally:
db.model.close()
if result is None:
@@ -724,9 +767,19 @@ class HistoryRoutesMixin:
resolved.canonical_target,
source_label=resolved.public["source_label"],
saved=saved.get("saved"),
library=saved.get(
"library",
{"status": "unavailable" if saved_unavailable else "none"},
library=apply_own_category(
saved.get(
"library",
{
"status": (
"unavailable" if saved_unavailable else "none"
)
},
),
own,
),
recorded=recorded_output_lookup(
db, close=True, reference=receipt["reference"]
),
)
_reauthorize_current_target(resolved, scope)
+374
View File
@@ -0,0 +1,374 @@
"""Image category routes (#332): an image file's own category.
``PUT /prompt_manager/images/category`` sets the own category of up to 500
listed files; ``DELETE`` removes it (listed or not, so a missing file's
category can still be removed); ``GET /prompt_manager/images/categories``
lists the own-category values of files the sources list. Targets are the
Images listing identity (``{root_index, root_scope, path}``, checked by
``media_target``) of an output source; LoRA folders are not sources here.
Instead of ``items`` a request may name ``filter``, the Images listing's own
query ("Select all N matching", #392), resolved now up to 500 images.
Each file is stored under its file key (``images.image_category_key``).
These routes never read or write prompts, recipes or image links.
"""
import os
from pathlib import Path
from aiohttp import web
from .image_tag_jobs import MAX_JOB_TARGETS, _parse_targets
from .images import (
_favorite_relative_key,
_favorite_root_key,
_forbidden,
_json_body_or_error,
image_category_key,
)
from .prompts import bad_request
try:
from ...database.image_categories import MAX_CATEGORY_LENGTH, MAX_KEY_LENGTH
from ...database.models import normalize_image_path
from ...utils.validators import validate_category
except ImportError:
from database.image_categories import MAX_CATEGORY_LENGTH, MAX_KEY_LENGTH
from database.models import normalize_image_path
from utils.validators import validate_category
CATEGORY_MEDIA_MESSAGE = "Only media files can have a category"
_PUT_FIELDS = frozenset({"category"})
_DELETE_FIELDS = frozenset()
TOO_MANY = f"Select at most {MAX_JOB_TARGETS} images."
SOURCE_CHANGED_MESSAGE = "The sources changed. Refresh the images and try again."
WRITE_FAILED = "The category could not be changed."
READ_FAILED = "Image categories could not be read."
def _coded(status: int, code: str, message: str) -> web.Response:
return web.json_response(
{"success": False, "code": code, "error": message},
status=status,
headers={"Cache-Control": "no-store"},
)
def _category_value(raw) -> tuple[str | None, web.Response | None]:
"""The stripped category, or a 400 (an empty value is a 400: use DELETE)."""
if not isinstance(raw, str):
return None, bad_request("category must be a string")
value = raw.strip()
if not value:
return None, bad_request("category must not be empty")
if len(value) > MAX_CATEGORY_LENGTH:
return None, bad_request(
f"Category cannot exceed {MAX_CATEGORY_LENGTH} characters"
)
try:
value.encode("utf-8")
validate_category(value)
except UnicodeEncodeError:
return None, bad_request("category must be valid text")
except ValueError as exc:
return None, bad_request(str(exc))
return value, None
class _KeyRefused(Exception):
"""A target has no storable file key (outside every source, or too long)."""
def __init__(self, message: str = "Access denied"):
super().__init__(message)
self.message = message
def _file_key(roots, root, relative) -> tuple[str, str]:
"""The target's file key (blocking: resolves the sources)."""
key = image_category_key(roots, Path(root) / relative)
if key is None:
raise _KeyRefused()
if max(len(key[0]), len(key[1])) > MAX_KEY_LENGTH:
raise _KeyRefused("Path is too long")
return key
def _nested_sources(root: Path, sources) -> list[tuple[str, str]]:
"""``[(relative key prefix, source root key)]`` of sources inside *root*.
Deepest first, so the first prefix that matches is the file's source.
"""
base = os.path.join(os.path.normcase(os.path.realpath(str(root))), "")
nested = []
for source in sources:
real = os.path.realpath(str(source))
if not os.path.normcase(real).startswith(base):
continue
relative = os.path.relpath(real, os.path.realpath(str(root)))
nested.append(
(
_favorite_relative_key(relative.replace(os.sep, "/")),
_favorite_root_key(source),
)
)
nested.sort(key=lambda item: item[0].count("/"), reverse=True)
return nested
def _source_prefix(relative_key: str, nested) -> str | None:
"""The prefix of the deepest nested source holding *relative_key*."""
for prefix, _source_key in nested:
if relative_key.startswith(prefix + "/"):
return prefix
return None
class ImageCategoryRoutesMixin:
"""Routes for images' own categories (#332)."""
def _register_image_category_routes(self, routes):
# Fixed paths beside DELETE /prompt_manager/images/{image_id}, as the
# favorite routes are. Registered after it, they still win: aiohttp
# (3.10+) looks fixed paths up before dynamic resources
# (test_delete_category_never_reaches_delete_image).
@routes.put("/prompt_manager/images/category")
async def set_image_category_route(request):
return await self.set_image_category(request)
@routes.delete("/prompt_manager/images/category")
async def clear_image_category_route(request):
return await self.clear_image_category(request)
@routes.get("/prompt_manager/images/categories")
async def list_image_categories_route(request):
return await self.list_image_categories(request)
# -- reading ---------------------------------------------------------------
def _own_category_map(self, sources, listings) -> dict:
"""``{normalized listed path: own category}`` for the listed files (blocking).
One primary-key range read per source; each stored row is matched to
the listing of every selected root that holds its source, so the work
follows the stored rows, not the files. A file under a nested source
is keyed by the deepest one (``image_category_key``). A row whose file
is gone matches nothing. A failed read raises.
"""
sources = list(sources)
rows = self.db.own_categories_of(_favorite_root_key(s) for s in sources)
if not rows:
return {}
by_source: dict[str, dict[str, str]] = {}
for (root_key, relative), category in rows.items():
by_source.setdefault(root_key, {})[relative] = category
found = {}
for _index, root, files in listings:
listing = self._favorite_listing(root, files)
root_key = _favorite_root_key(root)
# normalize_image_path(root / listed) without a call per file: the
# listed spellings are clean relative POSIX paths. Windows abspath
# also drops a component's trailing dots and spaces, so it keeps
# the call there.
base = os.path.join(normalize_image_path(str(root)), "")
def normalized(listed, base=base, root=root):
if os.name == "nt":
return normalize_image_path(str(root / listed))
return base + os.path.normcase(listed.replace("/", os.sep))
nested = [
item for item in _nested_sources(root, sources) if item[1] != root_key
]
for relative, category in by_source.get(root_key, {}).items():
listed = listing.get(relative)
if listed is not None and _source_prefix(relative, nested) is None:
found[normalized(listed)] = category
for prefix, source_key in nested:
for relative, category in by_source.get(source_key, {}).items():
full = f"{prefix}/{relative}"
listed = listing.get(full)
if listed is not None and _source_prefix(full, nested) == prefix:
found[normalized(listed)] = category
return found
async def _source_listings(self, roots) -> list[tuple]:
return [
(index, root, await self._get_gallery_files(root))
for index, root in enumerate(roots)
]
async def list_image_categories(self, request):
"""GET: the own-category values of listed files, sorted."""
try:
roots = list(self._get_all_output_dirs())
listings = await self._source_listings(roots)
found = await self._run_in_executor(self._own_category_map, roots, listings)
except Exception:
self.logger.exception("Image categories read failed")
return _coded(500, "categories-failed", READ_FAILED)
return web.json_response(
{"success": True, "categories": sorted(set(found.values()))},
headers={"Cache-Control": "no-store"},
)
# -- writing ---------------------------------------------------------------
async def _filter_targets(self, raw):
"""``(targets, None)`` for an Images listing filter, else ``(None, error)``.
"Select all N matching" (#392) sends the listing's own query; it is
resolved now, up to 500 images (a 400 ``too-many-images`` above).
"""
from .images import ListingSourceChanged, parse_listing_filter
try:
selection = await self.resolve_listing_filter(
parse_listing_filter(raw), MAX_JOB_TARGETS
)
except ListingSourceChanged:
return None, _coded(409, "source-changed", SOURCE_CHANGED_MESSAGE)
except ValueError as exc:
return None, bad_request(str(exc))
if selection.truncated:
return None, _coded(400, "too-many-images", TOO_MANY)
return tuple(selection.targets), None
async def _category_request(self, request, fields: frozenset):
"""``(data, targets, roots, None)`` or ``(None, None, None, error)``.
The images are ``items`` (1 to 500 targets) or ``filter`` (all images
an Images listing filter admits), besides *fields*.
"""
data, error = await _json_body_or_error(request)
if error is not None:
return None, None, None, error
names = set(data) - fields
if set(data) & fields != fields or names not in ({"items"}, {"filter"}):
return (
None,
None,
None,
bad_request(
"Request must contain "
+ ", ".join(sorted(fields))
+ (" and " if fields else "")
+ "either items or filter"
),
)
if "filter" in data:
targets, error = await self._filter_targets(data["filter"])
if error is not None:
return None, None, None, error
return data, targets, list(self._get_all_output_dirs()), None
items = data["items"]
if not isinstance(items, list) or not 1 <= len(items) <= MAX_JOB_TARGETS:
return (
None,
None,
None,
bad_request(f"items must list 1 to {MAX_JOB_TARGETS} images"),
)
targets, error = _parse_targets(items, media_message=CATEGORY_MEDIA_MESSAGE)
if error is not None:
return None, None, None, error
try:
for _index, _scope, relative in targets:
relative.encode("utf-8")
except UnicodeEncodeError:
return None, None, None, bad_request("path must be valid text")
roots = list(self._get_all_output_dirs())
for index, scope in dict.fromkeys((t[0], t[1]) for t in targets):
if index >= len(roots):
return None, None, None, bad_request("Invalid gallery root")
error = self._gallery_scope_error(
{"root_scope": scope, "root": str(index)}, roots
)
if error is not None:
if error.status == 409:
error = _coded(409, "source-changed", SOURCE_CHANGED_MESSAGE)
return None, None, None, error
return data, targets, roots, None
def _listed_file_keys(self, roots, targets, files) -> tuple[list, list]:
"""``(listed keys, unlisted keys)`` of the targets (blocking).
A target counts as listed when its source's listing holds it; its key
then follows the listed spelling. Targets folding to one key count
once, and a key that is listed through any target is listed.
"""
listings, listed, missing = {}, {}, {}
for index, _scope, relative in targets:
root = roots[index]
if index not in listings:
listings[index] = self._favorite_listing(root, files[index])
spelled = listings[index].get(_favorite_relative_key(relative))
key = _file_key(roots, root, spelled or relative)
(missing if spelled is None else listed)[key] = True
return list(listed), [key for key in missing if key not in listed]
async def set_image_category(self, request):
"""PUT ``{items, category}``: set the own category of listed files."""
data, targets, roots, error = await self._category_request(request, _PUT_FIELDS)
if error is not None:
return error
category, error = _category_value(data["category"])
if error is not None:
return error
try:
files = {}
for index in dict.fromkeys(t[0] for t in targets):
files[index] = await self._get_gallery_files(roots[index])
listed_keys, missing_keys = await self._run_in_executor(
self._listed_file_keys, roots, targets, files
)
except _KeyRefused as refused:
return _forbidden(refused.message)
except Exception:
self.logger.exception("Image category update failed")
return _coded(500, "category-failed", WRITE_FAILED)
try:
if not listed_keys:
return _coded(
404,
"files-not-found",
"None of the selected images could be found.",
)
counts = await self._run_in_executor(
self.db.set_image_categories, listed_keys, category
)
except Exception:
self.logger.exception("Image category update failed")
return _coded(500, "category-failed", WRITE_FAILED)
return web.json_response(
{
"success": True,
"category": category,
"updated": counts["updated"],
"unchanged": counts["unchanged"],
"not_found": len(missing_keys),
},
headers={"Cache-Control": "no-store"},
)
async def clear_image_category(self, request):
"""DELETE ``{items}``: remove the own category, listed or not."""
_data, targets, roots, error = await self._category_request(
request, _DELETE_FIELDS
)
if error is not None:
return error
try:
keys = await self._run_in_executor(
lambda: list(
dict.fromkeys(_file_key(roots, roots[i], r) for i, _s, r in targets)
)
)
cleared = await self._run_in_executor(self.db.clear_image_categories, keys)
except _KeyRefused as refused:
return _forbidden(refused.message)
except Exception:
self.logger.exception("Image category removal failed")
return _coded(500, "category-failed", WRITE_FAILED)
return web.json_response(
{"success": True, "cleared": cleared},
headers={"Cache-Control": "no-store"},
)
+11 -3
View File
@@ -58,7 +58,9 @@ from .annotations import (
from .autotag_routes import BATCH_CANCELLED, MODEL_BUSY_ERROR, model_busy_response
from .images import (
MAX_LISTING_SELECTION,
CollectionMissing,
ListingSourceChanged,
collection_not_found,
_current_fingerprint,
_listing_posix,
_public_path,
@@ -193,13 +195,17 @@ class _Spec:
return scope
def _parse_targets(raw):
"""Unique validated targets in request order, or ``(None, error)``."""
def _parse_targets(raw, **target_options):
"""Unique validated targets in request order, or ``(None, error)``.
``target_options`` reach the target check (the image category routes
pass their own ``media_message``, #332).
"""
if not isinstance(raw, list) or not 1 <= len(raw) <= MAX_JOB_TARGETS:
return None, bad_request(f"targets must list 1 to {MAX_JOB_TARGETS} images")
seen, targets = set(), []
for item in raw:
target, error = _target(item)
target, error = _target(item, **target_options)
if error is not None:
return None, error
if target not in seen:
@@ -685,6 +691,8 @@ class ImageTagJobsMixin:
selection = await self.resolve_listing_filter(spec.filter, cap)
except ListingSourceChanged:
return None, _coded("source-changed")
except CollectionMissing:
return None, collection_not_found()
except ValueError as exc:
return None, bad_request(str(exc))
except Exception:
+248 -59
View File
@@ -18,7 +18,7 @@ import urllib.parse
from bisect import bisect_left, bisect_right, insort
from collections.abc import Mapping
from pathlib import Path
from typing import Callable, NamedTuple
from typing import Callable, Iterable, NamedTuple
from aiohttp import web
from PIL import Image, UnidentifiedImageError
@@ -58,13 +58,17 @@ from .inspection_service import (
)
try:
from ...database.image_collections import CollectionMissing
from ...database.image_favorites import MAX_KEY_LENGTH
from ...database.library import display_title
from ...database.models import normalize_image_path
from ...database.run_attribution import recorded_output_lookup
from ...utils.card_seed import card_seed, header_params, parsed_positive
from ...utils.parallel import map_parallel
from ...utils.json_strict import dumps_strict
from ...utils.inspection_response import (
OWN_CATEGORY_UNREAD,
apply_own_category,
encode_inspection_response,
library_metadata,
)
@@ -78,13 +82,17 @@ try:
reveal_in_file_manager,
)
except ImportError:
from database.image_collections import CollectionMissing
from database.image_favorites import MAX_KEY_LENGTH
from database.library import display_title
from database.models import normalize_image_path
from database.run_attribution import recorded_output_lookup
from utils.card_seed import card_seed, header_params, parsed_positive
from utils.parallel import map_parallel
from utils.json_strict import dumps_strict
from utils.inspection_response import (
OWN_CATEGORY_UNREAD,
apply_own_category,
encode_inspection_response,
library_metadata,
)
@@ -104,6 +112,21 @@ IMAGE_EXTENSIONS = _IMAGE_EXTENSIONS
_CARD_RECORDS = ("prompt", "workflow")
_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
_LISTING_FAILED = "The image list could not be filtered. Check server logs for details."
_COLLECTION_ID = re.compile(r"[1-9][0-9]{0,8}")
COLLECTION_NOT_FOUND = "That collection no longer exists."
def collection_not_found() -> web.Response:
"""The 404 a write or a listing answers for an unknown collection (#331)."""
return web.json_response(
{
"success": False,
"code": "collection-not-found",
"error": COLLECTION_NOT_FOUND,
},
status=404,
headers={"Cache-Control": "no-store"},
)
def _card_header(path: Path) -> tuple:
@@ -563,7 +586,8 @@ class _ImageTagMatch(NamedTuple):
tags: bool # every filter tag is in its image tags or linked positive tags
query: bool # ``q`` is a case-folded substring of one image tag name
link: tuple | None # (hit, rating, positive id) of its unambiguous link
# (hit, rating, positive id, inherited category) of its unambiguous link
link: tuple | None
def _current_fingerprint(path: Path) -> tuple[int, int] | None:
@@ -800,6 +824,53 @@ def _favorite_root_key(root) -> str:
return _favorite_relative_key(os.path.realpath(str(root)))
def _source_depth(source: str) -> int:
return len([part for part in re.split(r"[\\/]+", source) if part])
def deepest_source(sources, absolute_path):
"""``(source, POSIX path below it)`` of the deepest source holding the path.
A source matches by its real path or by its spelling as configured (a
detected output folder can be a symlink), and the deepest real path wins.
The path is compared lexically after ``normpath`` and never resolved, so a
symlinked subfolder keeps its listed spelling. Containment folds case
where the OS does. *source* is the real path. None when no source
contains the path.
"""
path = os.path.normpath(str(absolute_path))
folded = os.path.normcase(path)
best = None
for source in sources:
real = os.path.normpath(os.path.realpath(str(source)))
for spelling in dict.fromkeys((real, os.path.normpath(str(source)))):
prefix = os.path.join(os.path.normcase(spelling), "")
if not folded.startswith(prefix):
continue
if best is None or _source_depth(real) > _source_depth(best[0]):
best = (real, spelling)
break
if best is None:
return None
relative = os.path.relpath(path, best[1])
return best[0], relative.replace(os.sep, "/")
def image_category_key(sources, absolute_path):
"""The file key of an image's own category (#332), or None.
``(root key, relative key)``: the deepest output source that contains the
file and the path below it, keyed with the favorites helpers. Without
nested sources it equals the file's favorite key. A path no source
contains has no key (and so no own category).
"""
found = deepest_source(sources, absolute_path)
if found is None:
return None
source, relative = found
return _favorite_root_key(source), _favorite_relative_key(relative)
def _name_suffix(name):
"""``Path(name).suffix`` as Python 3.10-3.13 compute it.
@@ -1056,8 +1127,9 @@ def _workspace_filters(query: Mapping) -> dict:
the linked positive's text, title and tag names. ``category`` is exact,
``tags`` is a comma list that must all be present, ``sort=rating`` ranks by
the linked positive, and ``favorites=1`` keeps only favorites (the
Favorites page, #249 D51). Raises ValueError for any unknown or oversized
value.
Favorites page, #249 D51). ``collection=<id>`` (1 to 9 digits) keeps only
that collection's members (#331); it wins over ``favorites`` (a page never
sets both). Raises ValueError for any unknown or oversized value.
"""
text = query.get("q", "").strip()
folded = text.casefold()
@@ -1071,8 +1143,10 @@ def _workspace_filters(query: Mapping) -> dict:
{tag.strip() for tag in query.get("tags", "").split(",") if tag.strip()}
)
favorites = query.get("favorites", "0")
collection = query.get("collection", "")
if (
favorites not in ("0", "1")
or (collection != "" and _COLLECTION_ID.fullmatch(collection) is None)
or len(text) > 500
or len(folded) > 500
or media not in ("all", "image", "video")
@@ -1095,7 +1169,8 @@ def _workspace_filters(query: Mapping) -> dict:
"path": exact.casefold(),
"category": category,
"tags": tags,
"favorites": favorites == "1",
"favorites": favorites == "1" and not collection,
"collection": int(collection) if collection else None,
}
@@ -1113,6 +1188,8 @@ LISTING_FILTER_KEYS = frozenset(
"media",
"sort",
"favorites",
# A collection page's filter (#331): bulk actions act on its members.
"collection",
}
)
MAX_LISTING_FILTER_VALUE = 2048
@@ -1169,22 +1246,29 @@ def _admitted_entries(
entries: list,
filters: dict,
links: dict | None,
favorites: frozenset | None = None,
members: frozenset | None = None,
image_tags: dict | None = None,
own_categories: dict | None = None,
):
"""Yield ``(position, entry, link match)`` for each entry the filters admit.
*links* maps normalized paths to (hit, rating, positive id) for every
unambiguously linked image (see PromptDatabase.get_image_link_matches); a
category or tag filter admits only those, and ``q`` admits a filename match
or a linked-positive hit. *favorites*, when given, admits only those listed
paths (``favorites=1``). *image_tags* maps normalized paths to their
_ImageTagMatch: a tag filter also admits a file whose current image tags
(with its linked positive's tags) carry every tag, and ``q`` also admits an
image tag name hit.
*links* maps normalized paths to (hit, rating, positive id, inherited
category) for every unambiguously linked image (see
PromptDatabase.get_image_link_matches); a tag filter admits only those,
and ``q`` admits a filename match or a linked-positive hit. A category
admits a file whose effective category equals it (#332 MC5): its own
category (*own_categories*, ``{normalized path: category}``) when it has one,
else the inherited one, so an unlinked file with its own category is
admitted. *members*, when given, admits only those listed paths (the
favorites of ``favorites=1`` or a ``collection``'s members). *image_tags*
maps normalized paths to their _ImageTagMatch: a tag filter also admits a
file whose current image tags (with its linked positive's tags) carry every
tag, and ``q`` also admits an image tag name hit.
"""
query, media, folder = filters["query"], filters["media"], filters["folder"]
members_only = filters["category"] is not None or bool(filters["tags"])
category = filters["category"]
members_only = bool(filters["tags"])
own = own_categories or {}
native_folder = folder.replace("/", os.sep)
nested = native_folder + os.sep
exact = filters.get("path", "")
@@ -1197,12 +1281,18 @@ def _admitted_entries(
continue
if media != "all" and (media == "video") != entry.video:
continue
if favorites is not None and entry.path not in favorites:
if members is not None and entry.path not in members:
continue
match = None if links is None else links.get(entry.normalized)
image = None if image_tags is None else image_tags.get(entry.normalized)
if members_only and match is None and not (image and image.tags):
continue
if category is not None:
effective = own.get(entry.normalized) if own else None
if effective is None and match is not None:
effective = match[3]
if effective != category:
continue
if (
query
and query not in entry.folded
@@ -1219,8 +1309,9 @@ def _filtered_page(
size: int,
filters: dict,
links: dict | None,
favorites: frozenset | None = None,
members: frozenset | None = None,
image_tags: dict | None = None,
own_categories: dict | None = None,
) -> tuple[int, list, int | None]:
"""Count the *entries* the filters admit and take *size* from *start* on.
@@ -1230,7 +1321,7 @@ def _filtered_page(
"""
total, page, positives = 0, [], set()
for position, entry, match in _admitted_entries(
entries, filters, links, favorites, image_tags
entries, filters, links, members, image_tags, own_categories
):
total += 1
if match is not None:
@@ -1934,6 +2025,7 @@ class ImageRoutesMixin:
filters["category"],
filters["tags"],
filters["favorites"],
filters["collection"],
]
).encode()
).hexdigest()
@@ -1944,17 +2036,19 @@ class ImageRoutesMixin:
except ValueError as exc:
return bad_request(str(exc))
listings = [
(index, root, await self._get_gallery_files(root))
for index, root in enumerate(output_dirs)
if root_index is None or index == root_index
]
try:
if filters["collection"] is not None and not await self._run_in_executor(
self.db.collection_exists, filters["collection"]
):
raise CollectionMissing("No such collection")
listings = await self._workspace_listings(output_dirs, root_index)
page, total, has_more, next_cursor, positive_count = (
await self._run_in_executor(
self._select_workspace_page, listings, filters, after, limit, scope
)
)
except CollectionMissing:
return collection_not_found()
except Exception:
# A filter that needs the prompt links or the image tags cannot be
# served without them. Never echo the exception: it can carry SQL
@@ -2170,9 +2264,10 @@ class ImageRoutesMixin:
stat'ed (never every listed file), and a candidate's tags count only
while a byte record is verified at its present ``(size, mtime_ns)``
(spec amendment S1). Under a Tags filter the candidates' linked
positives are read too, so their tags and category combine.
positives are read too, so their tags combine; the category is checked
once, on the effective category, by _admitted_entries (#332).
"""
tags, query, category = filters["tags"], filters["query"], filters["category"]
tags, query = filters["tags"], filters["query"]
wanted = set(tags)
located = self._workspace_locators(listings)
# The tag index answers which files carry a relevant name (#379); no
@@ -2202,11 +2297,9 @@ class ImageRoutesMixin:
link = facts.get(entry.normalized)
carried = set(names).union(link[4] if link else ())
matches[entry.normalized] = _ImageTagMatch(
bool(wanted)
and wanted <= carried
and (category is None or (link is not None and link[3] == category)),
bool(wanted) and wanted <= carried,
bool(query) and any(query in name.casefold() for name in names),
link[:3] if link else None,
tuple(link[:4]) if link else None,
)
return matches
@@ -2215,10 +2308,16 @@ class ImageRoutesMixin:
) -> tuple:
"""The ordered entries and what admits them (blocking; run in executor).
Returns ``(entries, key, links, image_tags, favorites)``: the listing in
the filters' order, its keyset key, and the link, image-tag and favorite
maps _admitted_entries reads. One SELECT (get_image_link_matches) reads
every linked positive when *read_links*. A failed read raises when the
Returns ``(entries, key, links, image_tags, members, own)``: the
listing in the filters' order, its keyset key, and the link, image-tag,
member and own-category maps _admitted_entries reads (members: a
collection's, else the favorites, else None). A collection deleted since
the route checked it raises CollectionMissing. One SELECT
(get_image_link_matches) reads every linked positive when *read_links*;
under a category it keeps only links whose raw recipe or positive
category is that value (the inherited one is decided in Python), and
the links of files whose own category is that value are read by path,
so their rating and positive still count (#332). A failed read raises when the
filters need it (category, tags, the rating sort) or *strict*; else the
listing stays usable without links (a search then matches filenames
only). A failed image-tag read raises under a Tags filter; a search then
@@ -2230,10 +2329,10 @@ class ImageRoutesMixin:
try:
links = self.db.get_image_link_matches(
filters["text"],
category,
tags,
every_link=True,
wanted=self._workspace_wanted(listings),
category=category,
)
except Exception:
if strict or category is not None or tags or order == "rating":
@@ -2247,6 +2346,13 @@ class ImageRoutesMixin:
),
exc_info=True,
)
own = None
if category is not None:
# A failed read raises: an unfiltered page must never pass for the
# category (DC6).
own = self._own_category_map(self._get_all_output_dirs(), listings)
if links is not None:
links = self._own_category_links(links, own, filters)
image_tags = None
if filters["query"] or tags:
try:
@@ -2290,10 +2396,37 @@ class ImageRoutesMixin:
def key(entry):
return _workspace_key(entry, order)
favorites = (
self._workspace_favorite_paths(listings) if filters["favorites"] else None
return (
entries,
key,
links,
image_tags,
self._workspace_members(listings, filters),
own,
)
return entries, key, links, image_tags, favorites
def _own_category_links(self, links: dict, own: dict, filters: dict) -> dict:
"""*links* plus the links of files whose own category is the filter's.
Their inherited category differs, so the category's link read left
them out; read by path (get_image_link_facts), kept only when their
positive has every filter tag, as the link read would.
"""
category, wanted = filters["category"], set(filters["tags"])
missing = [
path
for path, value in own.items()
if value == category and path not in links
]
if not missing:
return links
facts = self.db.get_image_link_facts(missing, filters["text"])
extra = {
path: tuple(fact[:4])
for path, fact in facts.items()
if wanted <= set(fact[4])
}
return {**links, **extra} if extra else links
def _select_workspace_page(
self,
@@ -2322,7 +2455,7 @@ class ImageRoutesMixin:
or bool(filters["tags"])
or order == "rating"
)
entries, key, links, image_tags, favorites = self._workspace_admission(
entries, key, links, image_tags, members, own = self._workspace_admission(
listings, filters, read_links=needed or after is None
)
start = 0
@@ -2333,7 +2466,7 @@ class ImageRoutesMixin:
else:
start = bisect_right(entries, after, key=key)
total, page, positive_count = _filtered_page(
entries, start, limit + 1, filters, links, favorites, image_tags
entries, start, limit + 1, filters, links, members, image_tags, own
)
if not needed and after is not None:
positive_count = _COUNT_OMITTED
@@ -2355,12 +2488,12 @@ class ImageRoutesMixin:
The links are read strictly: the positives must be the listing's own.
"""
entries, _key, links, image_tags, favorites = self._workspace_admission(
entries, _key, links, image_tags, members, own = self._workspace_admission(
listings, filters, read_links=True, strict=True
)
targets, positives, unlinked, scopes = [], {}, 0, {}
for _position, entry, match in _admitted_entries(
entries, filters, links, favorites, image_tags
entries, filters, links, members, image_tags, own
):
if len(targets) == cap:
return ListingSelection(targets, list(positives), unlinked, True)
@@ -2382,9 +2515,11 @@ class ImageRoutesMixin:
) -> ListingSelection:
"""The images a parse_listing_filter result admits right now (#392).
Raises ValueError for a root outside the current roots, and
Raises ValueError for a root outside the current roots,
ListingSourceChanged when the filter's ``root_scope`` no longer names
its root (the sources were renumbered since the listing).
its root (the sources were renumbered since the listing), and
CollectionMissing when its ``collection`` no longer exists (checked
before any scan and again after the members are read).
"""
output_dirs = list(self._get_all_output_dirs())
root_index = _workspace_root(params, output_dirs)
@@ -2394,11 +2529,11 @@ class ImageRoutesMixin:
raise ListingSourceChanged()
raise ValueError("Invalid gallery root scope")
filters = _workspace_filters(params)
listings = [
(index, root, await self._get_gallery_files(root))
for index, root in enumerate(output_dirs)
if root_index is None or index == root_index
]
if filters["collection"] is not None and not await self._run_in_executor(
self.db.collection_exists, filters["collection"]
):
raise CollectionMissing("No such collection")
listings = await self._workspace_listings(output_dirs, root_index)
return await self._run_in_executor(
self._select_workspace_matches, listings, filters, cap
)
@@ -3583,25 +3718,58 @@ class ImageRoutesMixin:
self._favorite_listing_cache[cache_key] = (files, listing)
return listing
def _workspace_favorite_paths(self, listings: list[tuple]) -> frozenset:
"""The listed paths of every favorite in the selected roots (blocking).
def _workspace_members(self, listings: list[tuple], filters: dict):
"""The member set the filters ask for, or None (blocking)."""
collection = filters.get("collection")
if collection is not None:
members = self._workspace_member_paths(
listings,
lambda root_key: self.db.collection_members(collection, root_key),
)
# Checked after the reads: a collection deleted before or while
# they ran is a 404, never an empty page passing for it.
if not self.db.collection_exists(collection):
raise CollectionMissing("No such collection")
return members
if filters["favorites"]:
return self._workspace_favorite_paths(listings)
return None
One primary-key read per root (get_root_favorites), matched against the
root's cached ``{relative key: listing spelling}``, so the work follows
the favorites rather than the files, and a key folds case exactly as
the toggle stored it (Windows). A favorite whose file is gone matches
nothing. A failed read raises: an unfiltered page must never pass for
the favorites.
def _workspace_favorite_paths(self, listings: list[tuple]) -> frozenset:
"""The listed paths of every favorite in the selected roots (blocking)."""
return self._workspace_member_paths(listings, self.db.get_root_favorites)
def _workspace_member_paths(
self, listings: list[tuple], keys_for_source: Callable[[str], Iterable[str]]
) -> frozenset:
"""The listed paths of the stored keys of each selected root (blocking).
One read per root (``keys_for_source(root key)``: the favorites' or a
collection's primary-key range), matched against the root's cached
``{relative key: listing spelling}``, so the work follows the members
rather than the files, and a key folds case exactly as it was stored
(Windows). A member whose file is gone matches nothing. A failed read
raises: an unfiltered page must never pass for the members.
"""
paths = set()
for _index, root, files in listings:
listing = self._favorite_listing(root, files)
for key in self.db.get_root_favorites(_favorite_root_key(root)):
for key in keys_for_source(_favorite_root_key(root)):
listed = listing.get(key)
if listed is not None:
paths.add(root / listed)
return frozenset(paths)
async def _workspace_listings(
self, output_dirs: list[Path], root_index: int | None
) -> list[tuple]:
"""``(index, root, scan)`` of every selected root (cached scans)."""
return [
(index, root, await self._get_gallery_files(root))
for index, root in enumerate(output_dirs)
if root_index is None or index == root_index
]
def _workspace_favorites(self, page: list[tuple]) -> set[tuple[int, str]]:
"""``(root index, listing path)`` of the favorites among *page* entries."""
by_root: dict[tuple[int, Path], dict[str, str]] = {}
@@ -3679,17 +3847,25 @@ class ImageRoutesMixin:
except Exception:
# The link stays known; only its title is withheld.
self.logger.warning("Inspection link title could not be read")
# An output file's own category (#332); LoRA roots have none.
own = (
self._read_own_category(outputs, path)
if index < len(outputs)
else None
)
try:
observed = inspect_current_file(
path,
source_label=_public_path(roots[index]),
saved=saved,
library=self._inspection_library(
saved, unavailable=unavailable
library=apply_own_category(
self._inspection_library(saved, unavailable=unavailable),
own,
),
image_tags=self._inspection_image_tags(
roots[index], Path(relative).parts
),
recorded=recorded_output_lookup(self.db),
)
except CurrentFileChanged:
return None
@@ -3752,6 +3928,19 @@ class ImageRoutesMixin:
return read
def _read_own_category(self, sources, path):
"""One file's own category, None, or OWN_CATEGORY_UNREAD (blocking).
Keyed by ``image_category_key`` over the output sources, as Images
writes it. A failed read is logged and hides the category (DC5).
"""
try:
key = image_category_key(sources, path)
return None if key is None else self.db.image_category(*key)
except Exception:
self.logger.warning("Image category could not be read", exc_info=True)
return OWN_CATEGORY_UNREAD
def _inspection_library(self, saved: dict | None, *, unavailable: bool) -> dict:
"""The inspector's Library facts; the route sets the title afterwards."""
result = {
+14 -3
View File
@@ -3,7 +3,7 @@
import os
import stat
from pathlib import Path
from typing import Callable
from typing import Any, Callable
try:
from ...utils.prompt_inspector import (
@@ -100,18 +100,23 @@ def inspect_current_file(
saved: dict | None,
library: dict,
image_tags: Callable[[int, int], dict] | None = None,
recorded: Callable[[Path], Any] | None = None,
) -> dict:
"""Current-file evidence read between two equal stats of *path*.
*image_tags*, when given, is called once with that stable
``(st_size, st_mtime_ns)`` and its facet is returned as ``image_tags``
(spec amendment S1: tags of a byte record verified at exactly that stat).
*recorded*, when given, is asked within the same stat window for what
recorded run history knows about *path* (``RecordedOutput`` or None: the
save node that wrote it and the texts PromptManager recorded, #386).
"""
before = _current_file_facts(path)
try:
inspection = read_inspection_file(path)
except MetadataReadError:
inspection = None
run = recorded(path) if recorded is not None and inspection is not None else None
after = _current_file_facts(path)
fields = ("st_dev", "st_ino", "st_mode", "st_size", "st_mtime_ns", "st_ctime_ns")
if any(getattr(before, key) != getattr(after, key) for key in fields):
@@ -129,9 +134,15 @@ def inspect_current_file(
profile = (
PREVIEW_INSPECTION_PROFILE if _in_host_temp(path) else INSPECTION_PROFILE
)
# The file name anchors on the saver that wrote it (#381).
# The file name anchors on the saver that wrote it (#381); recorded run
# history names that saver when the file alone cannot (#386).
evidence = inspect_prompt_evidence(
inspection.metadata, role_saved, profile=profile, file_name=path.name
inspection.metadata,
role_saved,
profile=profile,
file_name=path.name,
recorded_saver_id=run.saver_id if run is not None else None,
recorded_texts=run.texts if run is not None else None,
)
return {
"file": {
+16 -1
View File
@@ -47,6 +47,11 @@ import time
import uuid
from pathlib import PurePosixPath, PureWindowsPath
try:
from ...utils import exit_signal
except ImportError: # tests import the package top-level
from utils import exit_signal
logger = logging.getLogger(__name__)
# Start and finish order: wall-clock times can tie (coarse Windows clocks).
_sequence = itertools.count()
@@ -195,11 +200,21 @@ class Job:
self.counters = dict.fromkeys(groups, 0)
self.outcomes = {group: [] for group in groups}
self.log = []
self.cancelled = False
self._cancel_requested = False
self.last_event = None
self.task = None
self._subscribers = set()
@property
def cancelled(self):
"""Cancel was requested, or the interpreter is exiting (#376): every
check a job already makes between files also stops it at exit."""
return self._cancel_requested or exit_signal.requested()
@cancelled.setter
def cancelled(self, value):
self._cancel_requested = bool(value)
@property
def record_changing(self):
return self.admission == "records"
+54
View File
@@ -0,0 +1,54 @@
"""The listing identity of one media file, as a request names it.
Shared by the image-annotation routes and the collection routes (#331): a
target is ``{root_index, root_scope, path}``, checked here without touching
the filesystem or the database. The scope is compared with the current roots
by the caller.
"""
import os
import re
from .images import _forbidden, _has_traversal
from .prompts import bad_request
try:
from ...utils.media_capture import CAPTURE_EXTENSIONS
except ImportError:
from utils.media_capture import CAPTURE_EXTENSIONS
MAX_TARGET_PATH = 1024
TARGET_FIELDS = frozenset({"root_index", "root_scope", "path"})
_HEX64 = re.compile(r"[0-9a-f]{64}")
TAGGING_MEDIA_MESSAGE = "Only media files can be tagged"
def media_target(data: object, *, media_message: str = TAGGING_MEDIA_MESSAGE):
"""``((index, scope, relative), None)`` for a target, else ``(None, error)``.
Only the listing identity is accepted: an int root index, the root's
64-hex scope and a bounded root-relative POSIX media path, never an
absolute, drive-qualified, backslash or ``..`` path (a 400 for a malformed
shape, a 403 for an unsafe or non-media path, carrying *media_message*).
"""
if not isinstance(data, dict) or set(data) != TARGET_FIELDS:
return None, bad_request("target must name root_index, root_scope and path")
index, scope, relative = data["root_index"], data["root_scope"], data["path"]
if type(index) is not int or index < 0:
return None, bad_request("Invalid gallery root")
if not isinstance(scope, str) or _HEX64.fullmatch(scope) is None:
return None, bad_request("Invalid gallery root scope")
if not isinstance(relative, str) or not relative or len(relative) > MAX_TARGET_PATH:
return None, bad_request("path must be a relative media path")
parts = relative.split("/")
if (
"\0" in relative
or "\\" in relative
or relative.startswith("/")
or _has_traversal(relative)
or any(part in ("", ".", "..") or ":" in part for part in parts)
):
return None, _forbidden()
if os.path.splitext(parts[-1])[1].lower() not in CAPTURE_EXTENSIONS:
return None, _forbidden(media_message)
return (index, scope, relative), None
+4
View File
@@ -1227,7 +1227,9 @@ class PromptRoutesMixin:
"""
from .images import (
MAX_LISTING_SELECTION,
CollectionMissing,
ListingSourceChanged,
collection_not_found,
parse_listing_filter,
)
@@ -1252,6 +1254,8 @@ class PromptRoutesMixin:
},
status=409,
)
except CollectionMissing:
return collection_not_found()
except ValueError as exc:
return refused(str(exc))
if selection.truncated:
+7
View File
@@ -92,6 +92,11 @@ try:
except ImportError: # tests import the package top-level
from utils.parallel import map_parallel, note_refusal, refused_at_shutdown
try:
from ...utils.exit_signal import ExitRequested
except ImportError: # tests import the package top-level
from utils.exit_signal import ExitRequested
# The placeholder frame a video without a preview is answered with (JPEG
# bytes, never written to disk); written previews follow the settings.
PLACEHOLDER_SIZE = (thumbnail_settings.DEFAULT_EDGE, thumbnail_settings.DEFAULT_EDGE)
@@ -460,6 +465,8 @@ class AutoThumbnailer:
def _run(self):
try:
self._drain()
except ExitRequested:
pass # #376: the interpreter is exiting; queued files are dropped
except Exception:
self._logger.warning("The automatic preview worker stopped", exc_info=True)
finally:
+60
View File
@@ -7,6 +7,7 @@ Refactored from standalone_tagger.py for integration with PromptManager API.
import errno
import gc
import inspect
import os
import sys
from pathlib import Path
@@ -32,6 +33,54 @@ try:
except ImportError:
from py.tagging_defaults import DEFAULT_TAG_TEMPLATE
try:
from ..utils import exit_signal
except ImportError:
from utils import exit_signal
def abort_hub_downloads() -> None:
"""Abort huggingface_hub's Xet transfers (run by the exit hook, #376).
Xet downloads report progress from native code, which ignores an
exception raised in ``tqdm.update``; huggingface_hub aborts them itself
on Ctrl+C with ``abort_xet_session``. Nothing happens when
huggingface_hub was never imported or has no Xet support.
"""
if "huggingface_hub" not in sys.modules:
return
try:
from huggingface_hub.utils import _xet
except ImportError:
return
abort = getattr(_xet, "abort_xet_session", None)
if callable(abort):
abort()
def _stop_at_exit(download) -> Dict[str, Any]:
"""Keyword arguments that stop a huggingface_hub download at its next
chunk once the interpreter is exiting (#376), or none when this
huggingface_hub does not take ``tqdm_class``.
huggingface_hub 1.x removes its partial temp file on any failure; older
versions keep their resumable ``.incomplete`` file. The model's own path
is only written by the final rename either way.
"""
try:
if "tqdm_class" not in inspect.signature(download).parameters:
return {}
from huggingface_hub.utils import tqdm as hub_tqdm
except (ImportError, TypeError, ValueError):
return {}
class StopAtExit(hub_tqdm):
def update(self, n=1):
exit_signal.checkpoint()
return super().update(n)
return {"tqdm_class": StopAtExit}
def _autotag_config():
"""The saved defaults; imported late, as config needs ComfyUI's server."""
@@ -397,6 +446,7 @@ class AutoTagService:
filename=filename,
local_dir=str(model_dir),
local_dir_use_symlinks=False,
**_stop_at_exit(hf_hub_download),
)
if progress_callback:
@@ -541,6 +591,13 @@ class AutoTagService:
else:
return self._download_hf_model(progress_callback)
except Exception as e:
if exit_signal.requested():
# ExitRequested between chunks, or the Xet session aborted
# by the exit hook (abort_hub_downloads).
self.logger.info(
f"Model download stopped: ComfyUI is exiting ({model_type})"
)
return False
self.logger.error(f"Download failed: {e}")
if progress_callback:
progress_callback(f"Error: {str(e)}", 0)
@@ -570,6 +627,7 @@ class AutoTagService:
filename=config["filename"],
local_dir=str(gguf_dir),
local_dir_use_symlinks=False,
**_stop_at_exit(hf_hub_download),
)
self.logger.info("GGUF model downloaded")
@@ -587,6 +645,7 @@ class AutoTagService:
filename=config["mmproj_filename"],
local_dir=str(gguf_dir),
local_dir_use_symlinks=False,
**_stop_at_exit(hf_hub_download),
)
self.logger.info("mmproj downloaded")
@@ -613,6 +672,7 @@ class AutoTagService:
repo_id=config["repo"],
local_dir=str(model_dir),
local_dir_use_symlinks=False,
**_stop_at_exit(snapshot_download),
)
self.logger.info("HF model downloaded")
+5 -2
View File
@@ -44,8 +44,10 @@ from .lora_manager_layout import (
)
try:
from ..utils import exit_signal
from ..utils.hashing import generate_prompt_hash
except ImportError: # pragma: no cover - flat test imports
from utils import exit_signal
from utils.hashing import generate_prompt_hash
LORA_CATEGORY = "lora-manager"
@@ -189,14 +191,15 @@ def stage(
downloads its CivitAI example images. A sidecar that cannot be read, or
whose record fails to build, is skipped by its name, as are the
centralized sidecars discovery could not parse. ``cancelled()`` is asked
before each sidecar (raises ImportCancelled); ``on_record(index,
before each sidecar, and so is the exit signal (#376: the legacy scan's
``cancelled`` does not read it); either raises ImportCancelled. ``on_record(index,
record_or_None, staged)`` follows each one.
"""
sidecars = discovery.sidecars
staged = StagedImport(total=len(sidecars))
staged.skipped.extend(discovery.unreadable)
for index, sidecar in enumerate(sidecars):
if cancelled():
if cancelled() or exit_signal.requested():
raise ImportCancelled()
record = None
path = sidecar.metadata_path
+4 -1
View File
@@ -676,6 +676,9 @@ test('collection viewer uses real public transforms and resets on navigation',as
async function slideshowFixture(page,request,preferences={videoAutoplay:false,videoMute:true,videoLoop:false}) {
await expect(await request.post('/__test__/reset',{data:{full_inspector:true,include_video:true}})).toBeOK();
await page.goto('/prompt_manager/workspace#/images');
// Install the fake clock before the fixture wraps setTimeout: the wrapped timer (which the viewer session captures)
// must be the fake one, so page.clock.runFor alone decides when the 3500 ms slideshow deadline fires.
await page.clock.install();
await page.evaluate(async preferences=>{
if(preferences) localStorage.setItem('gallerySettings',JSON.stringify(preferences));
const [{createMediaViewer},{createInspectionController},{createScope},{createPromptSender}]=await Promise.all([
@@ -695,7 +698,7 @@ async function slideshowFixture(page,request,preferences={videoAutoplay:false,vi
},preferences);
await expect(mediaDialog(page).getByRole('button',{name:'Zoom in',exact:true})).toHaveAttribute('aria-disabled','false');
await expect(mediaDialog(page).locator('#viewer-details > div > p.small')).toHaveText('complete.png');
await page.clock.install();await page.clock.pauseAt(new Date(Date.now()+1000));
await page.clock.pauseAt(new Date(Date.now()+1000));
}
const slideshowButton=page=>mediaDialog(page).locator('#viewer-toolbar').getByRole('button',{name:/^(Play|Pause|Stop) slideshow$/});
const timerCount=page=>page.evaluate(()=>window.slideshowTimers.size);
+3 -79
View File
@@ -1,11 +1,9 @@
// #328: the legacy admin and gallery auto-tag tools keep working beside System › Auto-tagging. The model download is
// the `model_download` Activity job; the pages' stream attaches to it. No model is ever downloaded: the fixture's
// fake download waits until the test releases it (`/__test__/tagging/download`).
// #328: the legacy admin and gallery auto-tag batch and review keep working beside System › Auto-tagging. Since #330
// the pages download and unload models only through System (legacy-system-links.spec.mjs); a batch never downloads
// one. The fixture's fake download (`/__test__/tagging/download`) is held so nothing is ever fetched.
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
import { test, expect } from './fixtures.mjs';
const TITLE = 'Tagging model download · WD14 ViT';
async function post(request, url, data) {
await expect(await request.post(url, data === undefined ? {} : { data })).toBeOK();
}
@@ -33,61 +31,6 @@ test.afterEach(async ({ page, request, browserName }) => {
await post(request, '/__test__/jobs/restart');
});
async function startDownload(page, request) {
await page.locator('#autoTagBtn').click();
const status = page.locator('#wd14VitModelStatus');
await status.getByRole('button', { name: 'Download', exact: true }).click();
const modal = page.locator('#autoTagDownloadModal');
await expect(modal).toBeVisible();
await expect(page.locator('#downloadStatus')).toHaveText('Downloading…');
await expect(page.locator('#downloadProgressBar')).toHaveClass(/animate-pulse/);
await expect(modal).toContainText('The download continues in Activity.');
await expect(page.locator('#cancelDownloadBtn')).toHaveText('Close');
const running = (await downloads(request)).filter(job => job.state === 'running');
expect(running.map(job => [job.title, job.cancellable])).toEqual([[TITLE, false]]);
return { modal, status };
}
test('admin download modal closes on complete and Activity lists the job without Cancel', async ({ page, request }) => {
await page.goto('/prompt_manager/admin');
const { modal, status } = await startDownload(page, request);
await page.goto('/prompt_manager/workspace#/activity');
const item = page.locator('#workspace-main section.surface > .activity-item')
.filter({ has: page.getByRole('link', { name: new RegExp(`${TITLE}$`) }) });
await expect(item.locator('.chip')).toHaveText('Running');
await expect(item.locator('progress')).not.toHaveAttribute('value');
await expect(item.getByRole('button', { name: 'Cancel job', exact: true })).toHaveCount(0);
// Back on the admin page a second download of the same model attaches to the running job.
await page.goto('/prompt_manager/admin');
await startDownload(page, request);
await post(request, '/__test__/tagging/download', { mode: 'succeed' });
await expect(modal).toBeHidden();
await expect(status).toContainText('Downloaded');
expect((await downloads(request)).map(job => [job.title, job.state])).toEqual([[TITLE, 'finished']]);
});
test('gallery download modal closes on complete', async ({ page, request }) => {
await page.goto('/prompt_manager/gallery');
const { modal, status } = await startDownload(page, request);
await post(request, '/__test__/tagging/download', { mode: 'succeed' });
await expect(modal).toBeHidden();
await expect(page.locator('#autoTagModal')).toBeVisible();
await expect(status).toContainText('Downloaded');
expect((await downloads(request)).map(job => job.state)).toEqual(['finished']);
});
test('a failed download closes the modal and says so', async ({ page, request }) => {
await page.goto('/prompt_manager/gallery');
const { modal, status } = await startDownload(page, request);
await post(request, '/__test__/tagging/download', { mode: 'fail' });
await expect(modal).toBeHidden();
await expect(page.getByText('Download failed: Download failed')).toBeVisible();
await expect(status.getByRole('button', { name: 'Download', exact: true })).toBeVisible();
expect((await downloads(request)).map(job => job.state)).toEqual(['failed']);
});
// (wd14-vit is not downloaded in these tests: beforeEach marks it so the download tests can offer it.)
// The gallery's batch reader keyed on `data.status` values the server never sends, so its modal never moved or
// closed (#328 Problem 5). It now reads `type` frames as the admin page does.
@@ -171,25 +114,6 @@ test('admin review opens the first image', async ({ page }) => {
await expect(page.getByText('Failed to start review mode')).toHaveCount(0);
});
test('gallery unload works and is refused while the gate is held', async ({ page, request }) => {
await page.goto('/prompt_manager/gallery');
await runGalleryBatch(page); // keeps the model in memory
await page.locator('#autoTagBtn').click();
const unload = page.locator('#unloadModelBtn');
await expect(unload).toBeVisible();
// A held download owns the model gate: Unload is refused with the server's error.
const held = await request.post('/prompt_manager/jobs/model-download', { data: { model_type: 'gguf' } });
expect(held.status()).toBe(202);
await unload.click();
await expect(page.getByText(/^Failed to unload model: A tagging model is busy/)).toBeVisible();
await post(request, '/__test__/tagging/download', { mode: 'succeed' });
await expect.poll(async () => (await downloads(request)).map(job => job.state)).toEqual(['finished']);
await unload.click();
await expect(page.getByText('WD14-SWINV2 model unloaded successfully')).toBeVisible();
});
/** The admin review lists every row, videos too: Skip any image that is not one until the confirmation shows. */
async function untilAutoTagged(page, confirm) {
// Each image's skip check answers before its confirmation or its error shows.
+141 -18
View File
@@ -15,9 +15,40 @@ const LABELS = Object.freeze({
diagnostics: 'Diagnostics & logs',
});
async function openPageOptions(page) {
await page.locator('#pageOptionsBtn').click();
await expect(page.locator('#settingsModal')).toBeVisible();
}
/** The auto-tag modal with its model rows read from /autotag/models. */
async function openAutoTag(page) {
const read = page.waitForResponse(item => item.url().endsWith('/prompt_manager/autotag/models'));
await page.locator('#autoTagBtn').click();
await read;
await expect(page.locator('#autoTagModal')).toBeVisible();
}
const notDownloaded = (page, request) => post(request, '/__test__/tagging/models', { unavailable: ['gguf'] });
const modelLoaded = (page, request) => post(request, '/__test__/tagging/loaded', { model: 'wd14-vit' });
// One row per System link of the design's §2 map: the page, the control's label, how to find it, the section it opens
// and what a test must arrange before the control shows.
const LINKS = Object.freeze([
{ page: 'admin', label: 'Settings', locate: page => page.locator('#settingsBtn'), section: 'overview' },
{ page: 'admin', label: 'Maintenance', locate: page => page.locator('#maintenanceBtn'), section: 'maintenance' },
{ page: 'admin', label: 'Backup DB', locate: page => page.locator('#backupBtn'), section: 'backups' },
{ page: 'admin', label: 'Restore DB', locate: page => page.locator('#restoreBtn'), section: 'backups' },
{ page: 'gallery', label: 'Settings', locate: page => page.locator('#settingsBtn'), section: 'appearance' },
{ page: 'admin', label: 'Manage in System (model row)', arrange: notDownloaded, reveal: openAutoTag,
locate: page => page.locator('#ggufModelStatus a'), section: 'tagging' },
{ page: 'gallery', label: 'Manage in System (model row)', arrange: notDownloaded, reveal: openAutoTag,
locate: page => page.locator('#ggufModelStatus a'), section: 'tagging' },
{ page: 'gallery', label: 'Manage in System (Unload Model)', arrange: modelLoaded, reveal: openAutoTag,
locate: page => page.locator('#unloadModelBtn'), section: 'tagging' },
{ page: 'gallery', label: 'Page options Maintenance', reveal: openPageOptions,
locate: page => page.locator('#pageOptionsMaintenanceLink'), section: 'maintenance' },
{ page: 'gallery', label: 'Page options Sources & scans', reveal: openPageOptions,
locate: page => page.locator('#pageOptionsSourcesLink'), section: 'sources' },
{ page: 'admin', label: 'Diagnostics', locate: page => page.locator('#diagnosticsBtn'), section: 'diagnostics' },
{ page: 'admin', label: 'View Logs', locate: page => page.locator('#logsBtn'), section: 'diagnostics' },
]);
@@ -54,17 +85,31 @@ async function openLink(page, request, link) {
return control;
}
// Each row opens its section in a new tab and sends no request from the legacy page (the design's `a System link sends
// no request from the legacy page`, checked per row so one slow page never times out the whole table).
for (const link of LINKS) {
test(`legacy ${link.page} ${link.label} opens System ${link.section}`, async ({ page, context, request }) => {
// The gallery's thumbnail startup check (K9) reads its status on a timer of its own; it is off here.
await page.addInitScript(() => {
localStorage.setItem('gallerySettings', JSON.stringify({ checkThumbnailsAtStartup: false }));
});
const control = await openLink(page, request, link);
const legacyUrl = page.url();
await expect(control).toHaveAttribute('href', `workspace#/system/${link.section}`);
await expect(control).toHaveAttribute('target', '_blank');
await expect(control).toHaveAttribute('rel', 'noopener');
// The page's own start-up reads (list, film strips, settings) finish first; then the click must add none.
// Lazy thumbnails may still load whenever the layout changes; they are not the click's.
await page.waitForLoadState('networkidle');
const sent = [];
const record = item => { if (item.resourceType() !== 'image') sent.push(`${item.method()} ${item.url()}`); };
page.on('request', record);
const opened = context.waitForEvent('page');
await control.click();
const system = await opened;
await expect(system).toHaveURL(new RegExp(`/prompt_manager/workspace#/system/${link.section}$`));
page.off('request', record);
expect(sent, 'a System link sends no request from the legacy page').toEqual([]);
// System's page heading is "System"; the open section is the pressed button of its section switcher.
const main = system.locator('#workspace-main');
await expect(main.getByRole('heading', { level: 1, name: 'System', exact: true })).toBeVisible();
@@ -75,24 +120,6 @@ for (const link of LINKS) {
});
}
test('a System link sends no request from the legacy page', async ({ page, context, request }) => {
for (const link of LINKS) {
const control = await openLink(page, request, link);
// The page's own start-up reads (list, film strips, settings) finish first; then the click must add none.
await page.waitForLoadState('networkidle');
const sent = [];
const record = item => sent.push(`${item.method()} ${item.url()}`);
page.on('request', record);
const opened = context.waitForEvent('page');
await control.click();
const system = await opened;
await system.waitForLoadState('domcontentloaded');
page.off('request', record);
expect(sent, `${link.page} ${link.label}`).toEqual([]);
await system.close();
}
});
test('kept admin flows still work', async ({ page, context }) => {
await page.goto(PAGES.admin);
const texts = page.locator('#resultsList .prompt-text');
@@ -161,3 +188,99 @@ test('kept admin flows still work', async ({ page, context }) => {
await tab.close();
}
});
test('Page options keeps the page-only settings admin', async ({ page, request }) => {
const before = (await (await request.get('/prompt_manager/settings')).json()).settings;
await page.goto(PAGES.admin);
await expect(page.locator('#resultsList .prompt-text').first()).toBeVisible();
await page.locator('#pageOptionsBtn').click();
const dialog = page.locator('#settingsModal');
await expect(dialog).toBeVisible();
await expect(dialog.locator('h3')).toHaveText('Page options');
// Exactly the page-only setting (G5): the list's Infinite Scrolling default.
expect(await dialog.locator('input, select, textarea').evaluateAll(fields => fields.map(field => field.id)))
.toEqual(['infiniteScrollDefault']);
await expect(page.locator('#infiniteScrollDefault')).not.toBeChecked();
await page.locator('#infiniteScrollDefault').check();
const saved = page.waitForRequest(item => item.method() === 'POST' && item.url().endsWith('/prompt_manager/settings'));
await page.locator('#saveSettings').click();
expect((await saved).postDataJSON()).toEqual({ infinite_scroll: true });
await expect(dialog).toBeHidden();
// The list header follows the saved default at once; nothing else changed.
await expect(page.locator('#infiniteScrollSession')).toBeChecked();
const after = (await (await request.get('/prompt_manager/settings')).json()).settings;
expect(after).toEqual({ ...before, infinite_scroll: true });
await page.reload();
// The dialog shows the saved default once the page has read it (the header follows it first).
await expect(page.locator('#infiniteScrollSession')).toBeChecked();
await page.locator('#pageOptionsBtn').click();
await expect(page.locator('#infiniteScrollDefault')).toBeChecked();
});
test('Page options keeps the page-only settings gallery', async ({ page }) => {
// Native Appearance keys and the startup prompt's own keys are in the shared object before the page opens.
const stored = { imageQuality: 'high', showImageInfo: false, videoMute: false,
autoLoadMetadata: false, workspaceLayout: 'masonry', thumbnailPromptDismissed: 'sig', thumbnailsGenerated: true,
checkThumbnailsAtStartup: false, defaultViewMode: 'list' };
await page.addInitScript(value => {
if (!sessionStorage.getItem('seeded')) {
sessionStorage.setItem('seeded', '1');
localStorage.setItem('gallerySettings', value);
}
}, JSON.stringify(stored));
const listed = page.waitForResponse(item => item.url().includes('/prompt_manager/images/output?'));
await page.goto(PAGES.gallery);
await listed;
// With no Page options control, a stored default view still opens the page in it.
await expect(page.locator('#listViewBtn')).toHaveClass(/bg-pm-accent/);
await openPageOptions(page);
const dialog = page.locator('#settingsModal');
await expect(dialog.locator('h3')).toHaveText('Page options');
// Exactly the page-only settings (G5); the default view is System › Appearance's Default layout.
expect(await dialog.locator('input, select, textarea').evaluateAll(fields => fields.map(field => field.id))).toEqual([
'checkThumbnailsAtStartup', 'lazyLoadingToggle', 'defaultLimitSelect', 'gridColumnsSelect',
'infiniteScrollToggle', 'debugModeToggle', 'apiTimeoutInput']);
await expect(dialog.getByText('Default View Mode')).toHaveCount(0);
await page.locator('#defaultLimitSelect').selectOption('50');
await page.locator('#debugModeToggle').check();
await page.locator('#saveSettingsBtn').click();
await expect(dialog).toBeHidden();
const read = () => page.evaluate(() => JSON.parse(localStorage.getItem('gallerySettings')));
let saved = await read();
expect(saved).toMatchObject({ ...stored, defaultLimit: 50, debugMode: true });
// Reset puts back only the page options' defaults; the Appearance, startup-prompt and default-view keys stay.
await openPageOptions(page);
page.once('dialog', dialog => dialog.accept());
await page.locator('#resetSettingsBtn').click();
await expect(page.getByText('Page options reset to defaults')).toBeVisible();
await expect(page.locator('#defaultLimitSelect')).toHaveValue('100');
saved = await read();
expect(saved).toMatchObject({ defaultViewMode: 'list', imageQuality: 'high', showImageInfo: false, videoMute: false,
autoLoadMetadata: false, workspaceLayout: 'masonry', thumbnailPromptDismissed: 'sig', thumbnailsGenerated: true,
defaultLimit: 100, debugMode: false, checkThumbnailsAtStartup: true, lazyLoading: true, gridColumns: 8,
infiniteScroll: false, apiTimeout: 30 });
});
test('the legacy auto-tag modal stays open after Manage in System', async ({ page, context, request }) => {
for (const legacy of ['admin', 'gallery']) {
await post(request, '/__test__/tagging/models', { unavailable: ['gguf'] });
await page.goto(PAGES[legacy]);
await openAutoTag(page);
const link = page.locator('#ggufModelStatus').getByRole('link', { name: 'Manage in System', exact: true });
await expect(link).toHaveAttribute('title', 'Opens System › Auto-tagging in a new tab');
// A downloaded model keeps its status text.
await expect(page.locator('#wd14VitModelStatus')).toHaveText('\u2713 Downloaded');
const opened = context.waitForEvent('page');
await link.click();
const system = await opened;
await expect(system).toHaveURL(/#\/system\/tagging$/);
await expect(page.locator('#autoTagModal')).toBeVisible();
await system.close();
// Downloaded in System meanwhile: reopening the modal reads the models again (CUJ 4).
await post(request, '/__test__/tagging/models', { unavailable: [] });
await page.goto(PAGES[legacy]);
await openAutoTag(page);
await expect(page.locator('#ggufModelStatus')).toHaveText('\u2713 Downloaded');
}
});
+3
View File
@@ -262,12 +262,15 @@ def create_app(root: Path) -> web.Application:
with target.db.model.get_connection() as conn:
for table in (
"saved_views",
"image_collection_items",
"image_collections",
"image_favorites",
"media_review_decisions",
"media_review_items",
"media_image_tags",
"media_annotation_bytes",
"media_annotation_files",
"image_categories",
"image_prompt_pairs",
"prompt_pair_tags",
"prompt_pairs",
+1 -236
View File
@@ -104,33 +104,6 @@ for (const leave of ['beforeunload', 'pagehide']) {
});
}
// #222: the gallery's Clear cache clears the server's thumbnails by its own request and keeps this browser's
// preferences and other tools' storage; nothing of Prompt Manager's is cached in a browser, and the toast says so.
test('the gallery cache clear keeps preferences and other tools caches', async ({ page }) => {
await page.addInitScript(() => {
if (sessionStorage.getItem('seeded')) return;
sessionStorage.setItem('seeded', '1');
localStorage.setItem('pm_workspace_theme', 'light');
localStorage.setItem('other-tool-setting', 'keep');
localStorage.setItem('pm_other_tool', 'keep');
});
await page.goto('/prompt_manager/gallery');
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(2);
await page.evaluate(async () => { await (await caches.open('other-tool-gallery-thumbnails')).put('/x', new Response('x')); });
const clears = [];
page.on('request', req => { if (req.url().includes('/clear-thumbnails')) clears.push(req.method()); });
page.once('dialog', dialog => dialog.accept());
await page.locator('#settingsBtn').click();
await page.locator('#clearCacheBtn').click();
await expect(page.getByText("Thumbnails cleared; original images untouched. Nothing of Prompt Manager's is cached in this browser."))
.toBeVisible();
expect(clears).toEqual(['POST']);
expect(await page.evaluate(async () => ({
theme: localStorage.getItem('pm_workspace_theme'), other: localStorage.getItem('other-tool-setting'),
prefixed: localStorage.getItem('pm_other_tool'), notice: localStorage.getItem('pm_update_notice_dismissed'),
caches: await caches.keys(),
}))).toEqual({ theme: 'light', other: 'keep', prefixed: 'keep', notice: '3.2.0', caches: ['other-tool-gallery-thumbnails'] });
});
// #379: the reviewed tags are that image's own tags; no prompt changes.
test('actual gallery review generates and persists tags for the selected root\'s image', async ({ page, request }) => {
@@ -226,7 +199,7 @@ test('gallery infinite scroll appends a real backend page while retaining existi
await expect(await request.post('/__test__/reset', { data: { extra_images: 108 } })).toBeOK();
await page.goto('/prompt_manager/gallery');
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(100);
await page.locator('#settingsBtn').click();
await page.locator('#pageOptionsBtn').click();
await page.locator('#infiniteScrollToggle').check();
await page.locator('#saveSettingsBtn').click();
await expect(page.locator('#settingsModal')).toBeHidden();
@@ -246,37 +219,6 @@ async function promptSnapshot(request) {
})).sort((a, b) => a.id - b.id);
}
async function downloadBackup(page) {
const downloading = page.waitForEvent('download');
await page.locator('#backupBtn').click();
const download = await downloading;
expect(download.suggestedFilename()).toMatch(/^prompts_backup_.*\.db$/);
expect(await download.failure()).toBeNull();
const chunks = [];
for await (const chunk of await download.createReadStream()) chunks.push(chunk);
const buffer = Buffer.concat(chunks);
expect(buffer.subarray(0, 16).toString()).toBe('SQLite format 3\0');
return { name: download.suggestedFilename(), mimeType: 'application/octet-stream', buffer };
}
async function chooseRestore(page, file) {
await page.locator('#restoreBtn').click();
await expect(page.locator('#restoreModal')).toBeVisible();
await page.locator('#restoreFileInput').setInputFiles(file);
}
// Confirms the chosen restore and waits for the server's answer, not for a timer: a restore is disk-bound and can take
// longer than an expect timeout on a slow runner (#359). The modal closes only after that answer.
async function confirmRestoreAndWait(page) {
const restored = page.waitForResponse(response => response.url().endsWith('/prompt_manager/restore'));
await page.locator('#confirmRestore').click();
const response = await restored;
const body = await response.text();
expect(response.status(), body).toBe(200);
expect(JSON.parse(body)).toMatchObject({ success: true, prompt_count: 2, backup_created: expect.any(String) });
await expect(page.locator('#restoreModal')).toBeHidden();
}
function confirmNextDialog(page, accept, message) {
page.once('dialog', async dialog => {
expect(dialog.type()).toBe('confirm');
@@ -292,88 +234,6 @@ async function deletePromptButton(page, text) {
return page.locator(`[data-action="delete"][data-prompt-id="${id}"]`);
}
test('downloaded backup restores deleted prompts, tags, and image links through the actual upload form', async ({ page, request }) => {
await page.goto('/prompt_manager/admin');
await expect(page.locator('.prompt-text')).toHaveCount(2);
const before = await promptSnapshot(request);
expect(before.map(prompt => prompt.tags)).toEqual([['cats'], ['dogs']]);
expect(before.every(prompt => prompt.images.length === 1)).toBe(true);
const backup = await downloadBackup(page);
const button = await deletePromptButton(page, 'Cat portrait');
confirmNextDialog(page, true, 'delete this prompt');
await button.click();
await expect(page.locator('.prompt-text')).toHaveCount(1);
await expect(await request.post('/prompt_manager/save', { data: { text: 'Added after backup', tags: ['temporary'] } })).toBeOK();
expect(await promptSnapshot(request)).not.toEqual(before);
await chooseRestore(page, backup);
confirmNextDialog(page, true, backup.name);
await confirmRestoreAndWait(page);
await expect(page.locator('.prompt-text')).toHaveCount(2);
await page.reload();
await expect(page.locator('.prompt-text').filter({ hasText: 'Cat portrait' })).toBeVisible();
await expect(page.locator('.prompt-text').filter({ hasText: 'Added after backup' })).toHaveCount(0);
expect(await promptSnapshot(request)).toEqual(before);
await page.goto('/prompt_manager/admin#/tags');
await page.locator('[data-tag-name="cats"]').click();
await expect(page.locator('#tagPromptsGrid')).toContainText('Cat portrait');
await page.locator('#tagPromptsGrid img').first().click();
await expect(page.locator('.viewer-canvas img')).toBeVisible();
await expect.poll(() => page.locator('.viewer-canvas img').evaluate(image => image.complete && image.naturalWidth > 0)).toBe(true);
});
test('restore validation and cancellation preserve live data and allow a valid retry', async ({ page, request }) => {
await page.goto('/prompt_manager/admin');
const before = await promptSnapshot(request);
const backup = await downloadBackup(page);
let uploads = 0;
page.on('request', request => { if (request.url().endsWith('/prompt_manager/restore')) uploads++; });
await page.locator('#restoreBtn').click();
await page.locator('#confirmRestore').click();
await expect(page.getByText('Please select a database file to restore', { exact: true })).toBeVisible();
await page.locator('#restoreFileInput').setInputFiles({ ...backup, name: 'backup.txt' });
await page.locator('#confirmRestore').click();
await expect(page.getByText('Please select a valid .db file', { exact: true })).toBeVisible();
expect(uploads).toBe(0);
await page.locator('#restoreFileInput').setInputFiles(backup);
confirmNextDialog(page, false, 'replace your current database');
await page.locator('#confirmRestore').click();
await expect(page.locator('#restoreModal')).toBeVisible();
expect(uploads).toBe(0);
await page.locator('#restoreFileInput').setInputFiles({ name: 'corrupt.db', mimeType: 'application/octet-stream', buffer: Buffer.from('not a SQLite database') });
confirmNextDialog(page, true, 'corrupt.db');
const rejected = page.waitForResponse(response => response.url().endsWith('/prompt_manager/restore'));
await page.locator('#confirmRestore').click();
expect((await rejected).status()).toBe(400);
await expect(page.getByText(/Restore failed:/)).toBeVisible();
await expect(page.locator('#restoreModal')).toBeVisible();
await expect(page.locator('#confirmRestore')).toBeEnabled();
expect(await promptSnapshot(request)).toEqual(before);
await page.locator('#restoreFileInput').setInputFiles(backup);
confirmNextDialog(page, true, backup.name);
await confirmRestoreAndWait(page);
expect(await promptSnapshot(request)).toEqual(before);
});
test('backup and restore recover from interrupted network requests without losing data', async ({ page, request }) => {
await page.goto('/prompt_manager/admin');
const before = await promptSnapshot(request);
await page.route('**/prompt_manager/backup', route => route.abort('failed'), { times: 1 });
await page.locator('#backupBtn').click();
await expect(page.getByText(/Backup failed:/)).toBeVisible();
const backup = await downloadBackup(page);
await chooseRestore(page, backup);
await page.route('**/prompt_manager/restore', route => route.abort('failed'), { times: 1 });
confirmNextDialog(page, true, backup.name);
await page.locator('#confirmRestore').click();
await expect(page.getByText(/Restore failed:/)).toBeVisible();
await expect(page.locator('#confirmRestore')).toBeEnabled();
await expect(page.locator('#restoreModal')).toBeVisible();
expect(await promptSnapshot(request)).toEqual(before);
confirmNextDialog(page, true, backup.name);
await confirmRestoreAndWait(page);
expect(await promptSnapshot(request)).toEqual(before);
});
test('prompt deletion requires confirmation and can retry after a network failure', async ({ page, request }) => {
await page.goto('/prompt_manager/admin');
const before = await promptSnapshot(request);
@@ -401,98 +261,3 @@ test('prompt deletion requires confirmation and can retry after a network failur
await expect(await request.get(url)).toBeOK();
}
});
// The duplicate scan reads the output root only, but the fixture's input root holds a different
// shared/same.png. A bare /images/serve/<path> URL tries every gallery root in turn, so once the
// output copy is deleted it answers 200 with the input copy: check each file in its own root.
const serveIn = (root, path) => `/prompt_manager/images/serve/${path.split(/[\\/]/).map(encodeURIComponent).join('/')}?root=${root}`;
const rootOf = url => new URL(url, 'http://gallery.test').searchParams.get('root');
// Which copy the scan keeps depends on file times: the fixture stamps shared/same.png with the
// wall clock, and the kernel can stamp its bulk copies a fraction of a millisecond earlier.
// listOrder rewrites the finished scan's groups so a test can pick the order it needs.
async function openDuplicateScan(page, request, listOrder = null) {
await expect(await request.post('/__test__/reset', { data: { extra_images: 2 } })).toBeOK();
if (listOrder) {
await page.route('**/prompt_manager/scan_duplicates/stream', async route => {
const response = await route.fetch();
const body = (await response.text()).replace(/^data: (.*)$/gm, (line, json) => {
const event = JSON.parse(json);
if (event.type !== 'complete') return line;
for (const group of event.duplicates) group.images = listOrder(group.images);
return `data: ${JSON.stringify(event)}`;
});
await route.fulfill({ response, body });
});
}
await page.goto('/prompt_manager/gallery');
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(4);
const originals = await page.locator('#galleryGrid img[data-original]').evaluateAll(images => images.map(image => image.dataset.original));
const output = rootOf(originals.find(url => url.includes('/bulk/')));
const input = rootOf(originals.find(url => url.includes('/shared/same.png') && rootOf(url) !== output));
expect(output).not.toBeNull();
expect(input).not.toBeNull();
await page.locator('#settingsBtn').click();
await page.locator('#scanDuplicatesBtn').click();
await expect(page.locator('#duplicatesCount')).toHaveText('1');
const candidates = page.locator('.duplicate-checkbox');
await expect(candidates).toHaveCount(2);
return { originals, output, input, candidates };
}
async function expectOnlySelectedGone(page, request, { originals, output, input }, selected, retained) {
expect((await request.get(serveIn(output, selected))).status()).toBe(404);
await expect(await request.get(serveIn(output, retained))).toBeOK();
await expect(await request.get(serveIn(input, 'shared/same.png'))).toBeOK();
await page.reload();
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(3);
const remaining = await page.locator('#galleryGrid img[data-original]').evaluateAll(images => images.map(image => image.dataset.original));
expect(remaining).toHaveLength(originals.length - 1);
for (const url of remaining) {
expect(originals).toContain(url);
await expect(await request.get(url)).toBeOK();
}
}
test('duplicate cleanup cancels safely and deletes only the selected file after a retry', async ({ page, request }) => {
const scan = await openDuplicateScan(page, request);
const { output, candidates } = scan;
const selected = await candidates.first().getAttribute('data-image-path');
const retained = await candidates.last().getAttribute('data-image-path');
await candidates.first().check();
let deletes = 0;
page.on('request', request => { if (request.url().endsWith('/delete_duplicate_images')) deletes++; });
confirmNextDialog(page, false, 'permanently delete 1 duplicate image files');
await page.locator('#removeDuplicatesBtn').click();
expect(deletes).toBe(0);
await expect(await request.get(serveIn(output, selected))).toBeOK();
await page.route('**/prompt_manager/delete_duplicate_images', route => route.abort('failed'), { times: 1 });
confirmNextDialog(page, true, 'permanently delete 1 duplicate image files');
await page.locator('#removeDuplicatesBtn').click();
await expect(page.getByText('Failed to delete duplicate images', { exact: true })).toBeVisible();
await expect(page.locator('#removeDuplicatesBtn')).toBeEnabled();
await expect(await request.get(serveIn(output, selected))).toBeOK();
confirmNextDialog(page, true, 'permanently delete 1 duplicate image files');
const deleted = page.waitForResponse(response => response.url().endsWith('/delete_duplicate_images'));
await page.locator('#removeDuplicatesBtn').click();
expect((await (await deleted).json())).toMatchObject({ success: true, deleted_count: 1, failed_count: 0 });
await expect(page.locator('#duplicatesModal')).toBeHidden();
await expectOnlySelectedGone(page, request, scan, selected, retained);
});
test('deleting a duplicate whose path another gallery root also holds removes only the scanned copy', async ({ page, request }) => {
// The order CI saw: a bulk copy stamped first, so shared/same.png is the first Delete candidate.
const ciOrder = ['bulk/0000.png', 'shared/same.png', 'bulk/0001.png'];
const rank = image => ciOrder.indexOf(image.path.replace(/\\/g, '/'));
const scan = await openDuplicateScan(page, request, images => [...images].sort((a, b) => rank(a) - rank(b)));
const selected = await scan.candidates.first().getAttribute('data-image-path');
const retained = await scan.candidates.last().getAttribute('data-image-path');
expect(selected.replace(/\\/g, '/')).toBe('shared/same.png');
await scan.candidates.first().check();
confirmNextDialog(page, true, 'permanently delete 1 duplicate image files');
const deleted = page.waitForResponse(response => response.url().endsWith('/delete_duplicate_images'));
await page.locator('#removeDuplicatesBtn').click();
expect((await (await deleted).json())).toMatchObject({ success: true, deleted_count: 1, failed_count: 0 });
await expect(page.locator('#duplicatesModal')).toBeHidden();
await expectOnlySelectedGone(page, request, scan, selected, retained);
});
@@ -395,14 +395,15 @@ test('a failed settings read leaves the policy disabled with the reason', async
await expect(row.locator('input')).toBeDisabled();
});
test('the legacy admin Backup and Restore keep working', async ({ page, request }) => {
test('the legacy Backup and Restore routes keep working and the admin page links to System', async ({ page, request }) => {
const bytes = await backupBytes(request);
const restored = await request.post('/prompt_manager/restore', { multipart: { database_file: {
name: 'prompts_backup.db', mimeType: 'application/octet-stream', buffer: bytes } } });
await expect(restored).toBeOK();
expect((await restored.json()).success).toBe(true);
// #330: the admin page's Backup DB and Restore DB open System › Backup & export.
await page.goto('/prompt_manager/admin');
await expect(page.locator('#restoreFileInput')).toHaveCount(1);
for (const id of ['#backupBtn', '#restoreBtn']) await expect(page.locator(id)).toHaveAttribute('href', 'workspace#/system/backups');
});
/** Start a fake running job (a metadata scan another tab started), which an exclusive restore must wait for. */
@@ -181,6 +181,7 @@ _TABLES = (
"media_image_tags",
"media_annotation_bytes",
"media_annotation_files",
"image_categories",
"image_prompt_pairs",
"prompt_pair_tags",
"prompt_pairs",
@@ -33,6 +33,7 @@ _TABLES = (
"media_image_tags",
"media_annotation_bytes",
"media_annotation_files",
"image_categories",
"image_prompt_pairs",
"prompt_pair_tags",
"prompt_pairs",
@@ -85,6 +85,7 @@ _TABLES = (
"media_image_tags",
"media_annotation_bytes",
"media_annotation_files",
"image_categories",
"image_prompt_pairs",
"prompt_pair_tags",
"prompt_pairs",
@@ -51,6 +51,7 @@ _TABLES = (
"media_image_tags",
"media_annotation_bytes",
"media_annotation_files",
"image_categories",
"image_prompt_pairs",
"prompt_pair_tags",
"prompt_pairs",
@@ -247,6 +248,22 @@ def register_tagging_commands(
model.unavailable = set(data["unavailable"])
return web.json_response({"success": True})
@routes.post("/__test__/tagging/loaded")
async def loaded(request: web.Request) -> web.Response:
"""Mark one fake model loaded (``{"model": "<key>"}``), as a batch that
kept it in memory leaves it; ``/autotag/models`` then reports
``model_loaded``. Reset unloads it."""
data = await request.json() if request.can_read_body else None
if (
request.query
or not isinstance(data, dict)
or set(data) != {"model"}
or data["model"] not in MODELS
):
raise web.HTTPBadRequest()
model.loaded = data["model"]
return web.json_response({"success": True})
@routes.post("/__test__/tagging/download")
async def download_mode(request: web.Request) -> web.Response:
"""Switch the model's download to a fake (never a real download).
@@ -119,6 +119,7 @@ def seed_workspace_visual(api: PromptManagerAPI, root: Path) -> None:
"media_image_tags",
"media_annotation_bytes",
"media_annotation_files",
"image_categories",
"image_prompt_pairs",
"prompt_pair_tags",
"prompt_pairs",
@@ -86,6 +86,7 @@ _TABLES = (
"media_image_tags",
"media_annotation_bytes",
"media_annotation_files",
"image_categories",
"image_prompt_pairs",
"prompt_pair_tags",
"prompt_pairs",
+165 -3
View File
@@ -11,7 +11,7 @@ import { installHostHooks } from './host-fixture.mjs';
// The worker's owned host must be up before the shared host hooks open the page.
test.beforeEach(async ({ execution }) => { expect(execution.databasePath).toBeTruthy(); });
installHostHooks(test, ['PromptManager', 'PromptManagerText', 'PromptManagerFocused', 'PromptManagerPair', 'PMHistoryInputs']);
installHostHooks(test, ['PromptManager', 'PromptManagerText', 'PromptManagerFocused', 'PromptManagerPair', 'PMHistoryInputs', 'PMCycledText']);
// `primitive`: a core string node the queue-time resolver follows. `legacy`: the frontend's virtual Primitive node,
// whose value is queued as literal text. `preview`: PrimitiveString -> PreviewAny, the shape of the maintainer's
@@ -37,6 +37,7 @@ function build(page, { kind, positive, negative, prefix }) {
primitive.widgets[0].value = side.text;
return;
}
if (side.via === 'cycled') { connect(add('PMCycledText', 0, y), 0, prompt, name); return; }
const text = add('PrimitiveString', 0, y);
text.widgets.find(widget => widget.name === 'value').value = side.text;
let source = text;
@@ -75,9 +76,9 @@ function build(page, { kind, positive, negative, prefix }) {
}, { kind, positive, negative, prefix });
}
async function run(execution, page, testCase) {
async function run(execution, page, testCase, built) {
const prefix = `textlink/${randomUUID()}`;
const { output, workflow, promptNode } = await build(page, { ...testCase, prefix });
const { output, workflow, promptNode } = built ?? await build(page, { ...testCase, prefix });
const id = randomUUID();
const response = await execution.request.post('/prompt', { data: { prompt: output, prompt_id: id, extra_data: { extra_pnginfo: { workflow } } } });
expect(response.ok(), await response.text()).toBe(true);
@@ -193,3 +194,164 @@ for (const [label, linked, negativeVia] of [['positive', ['text'], 'primitive'],
await expectImagesLinked(execution, evidence, sides.text.via === 'preview');
});
}
// #347: a cached PromptManager node counts every job, as an executing one does. Jobs after the first change only
// the sampler seed, so ComfyUI serves the PromptManager node (and its PreviewAny source) from its output cache.
const COUNTS = `import json, sqlite3, sys, pathlib
db = sqlite3.connect(pathlib.Path(sys.argv[1]).resolve().as_uri() + "?mode=ro", uri=True)
one = lambda sql, key: (db.execute(sql, (int(key),)).fetchone() or [None])[0] if key != "-" else None
print(json.dumps({"positive": one("SELECT run_count FROM prompts WHERE id=?", sys.argv[2]),
"negative": one("SELECT run_count FROM prompts WHERE id=?", sys.argv[3]),
"pair": one("SELECT run_count FROM prompt_pairs WHERE id=?", sys.argv[4])}))`;
async function runCounts(execution, context) {
const key = value => value == null ? '-' : String(value);
const { stdout } = await promisify(execFile)(process.env.PM_TEST_PYTHON || 'python',
['-c', COUNTS, execution.databasePath, key(context.positive_prompt_id), key(context.negative_prompt_id), key(context.pair_id)], { timeout: 5000 });
return JSON.parse(stdout);
}
async function setCache(execution, mode) {
expect((await execution.request.post('/pm_execution_fixture/cache', { data: { mode } })).ok()).toBe(true);
}
// Queue an already built workflow again with another sampler seed; returns the job's history item and cached nodes.
async function queueAgain(execution, first, { seed, edit, targets, status = 'success' } = {}) {
const output = structuredClone(first.output);
const sampler = Object.values(output).find(node => node.class_type === 'KSampler');
sampler.inputs.seed = seed;
edit?.(output);
const id = randomUUID();
const data = { prompt: output, prompt_id: id, extra_data: { extra_pnginfo: { workflow: first.workflow } } };
if (targets) data.partial_execution_targets = targets;
const response = await execution.request.post('/prompt', { data });
expect(response.ok(), await response.text()).toBe(true);
let item;
await expect.poll(async () => {
item = (await (await execution.request.get(`/history/${id}`)).json())[id];
return item?.status?.status_str ?? null;
}, { timeout: 20000 }).toBe(status);
const cached = item.status.messages.find(([event]) => event === 'execution_cached')?.[1]?.nodes ?? [];
return { id, item, cached };
}
// The recorded context of one job, once its capture reached SQLite.
async function capturedContext(execution, id) {
let result;
await expect.poll(async () => {
result = await readExecutionRuns(execution.databasePath, id);
const row = result.rows[0];
return row ? result.captures[row.run_id].state : null;
}, { timeout: 10000 }).toMatch(/^(complete|partial)$/);
return result.captures[result.rows[0].run_id].contexts[0];
}
async function firstRun(execution, page, testCase) {
const prefix = `textlink/${randomUUID()}`;
const built = await build(page, { ...testCase, prefix });
const evidence = await run(execution, page, testCase, built);
return { ...evidence, workflow: built.workflow };
}
for (const kind of ['PromptManager', 'PromptManagerFocused', 'PromptManagerSplit', 'PromptManagerText', 'PromptManagerPair']) {
test(`${kind} counts each cached re-run of text linked from PreviewAny (#347)`, async ({ execution, page }) => {
test.setTimeout(90000);
await setCache(execution, 'RAM_PRESSURE');
const pair = kind === 'PromptManagerPair';
const testCase = { kind, positive: { typed: 'Typed under the link', text: unique(`Cached ${kind}`), via: 'preview' },
negative: pair ? { typed: unique('Typed negative') } : undefined };
const first = await firstRun(execution, page, testCase);
for (const seed of [2, 3]) {
const again = await queueAgain(execution, first, { seed });
expect(again.cached, 'ComfyUI served the PromptManager node from its cache').toContain(first.promptNode);
}
const counts = await runCounts(execution, first.context);
expect(counts.positive).toBe(3);
if (pair) expect(counts).toMatchObject({ negative: 0, pair: 3 });
});
}
test.describe('cached runs in one renderer (#347)', () => {
// Cache behavior does not depend on the renderer: one project runs these.
test.beforeEach(({}, testInfo) => { test.skip(testInfo.project.name !== 'classic-start', 'covered in classic-start'); });
const linked = () => ({ kind: 'PromptManager', positive: { typed: 'Typed under the link', text: unique('Cached'), via: 'preview' } });
for (const mode of ['CLASSIC', 'LRU', 'RAM_PRESSURE', 'NONE']) {
test(`counts 3 jobs under the ${mode} cache`, async ({ execution, page }) => {
test.setTimeout(90000);
await setCache(execution, mode);
try {
const first = await firstRun(execution, page, linked());
for (const seed of [2, 3]) {
const again = await queueAgain(execution, first, { seed });
if (mode === 'NONE') expect(again.cached).not.toContain(first.promptNode);
else expect(again.cached).toContain(first.promptNode);
}
expect((await runCounts(execution, first.context)).positive).toBe(3);
} finally { await setCache(execution, 'RAM_PRESSURE'); }
});
}
test('alternating upstream values A, B, A under LRU each count their own prompt', async ({ execution, page }) => {
test.setTimeout(90000);
await setCache(execution, 'LRU');
try {
const [a, b] = [unique('Cycled A'), unique('Cycled B')];
const cycle = async value => expect((await execution.request.post('/pm_execution_fixture/cycled', { data: { value } })).ok()).toBe(true);
await cycle(a);
const first = await firstRun(execution, page, { kind: 'PromptManager', positive: { typed: 'Typed under the link', text: a, via: 'cycled' } });
expect(first.context.positive.raw).toBe(a);
await cycle(b);
const second = await queueAgain(execution, first, { seed: 2 });
expect(second.cached).not.toContain(first.promptNode);
await cycle(a);
const third = await queueAgain(execution, first, { seed: 3 });
expect(third.cached, 'LRU kept A').toContain(first.promptNode);
expect((await runCounts(execution, first.context)).positive).toBe(2);
const bContext = await capturedContext(execution, second.id);
expect(bContext.positive.raw).toBe(b);
expect((await runCounts(execution, bContext)).positive).toBe(1);
} finally { await setCache(execution, 'RAM_PRESSURE'); }
});
test('typed text served from the cache still counts once per job', async ({ execution, page }) => {
test.setTimeout(90000);
await setCache(execution, 'RAM_PRESSURE');
const first = await firstRun(execution, page, { kind: 'PromptManager', positive: { typed: unique('Typed cached') } });
for (const seed of [2, 3]) expect((await queueAgain(execution, first, { seed })).cached).toContain(first.promptNode);
expect((await runCounts(execution, first.context)).positive).toBe(3);
});
test('a job failing downstream of a cached node counts it once; a job that does not need it, never', async ({ execution, page }) => {
test.setTimeout(90000);
await setCache(execution, 'RAM_PRESSURE');
const first = await firstRun(execution, page, linked());
const failing = await queueAgain(execution, first, { seed: 2, status: 'error', edit: output => {
output.pm_failure = { class_type: 'PMHistoryOutput', inputs: { value: [Object.keys(output).find(id => output[id].class_type === 'KSampler'), 0], mode: 'failure' } };
} });
expect(failing.cached).toContain(first.promptNode);
expect((await runCounts(execution, first.context)).positive).toBe(2);
// "Run selected output": only an unrelated output node; ComfyUI still reports the cached PromptManager node
const unrelated = await queueAgain(execution, first, { seed: 2, targets: ['pm_unrelated'], edit: output => {
output.pm_unrelated = { class_type: 'PMHistoryCached', inputs: { value: unique('unrelated') } };
} });
expect(unrelated.cached).toContain(first.promptNode);
expect((await runCounts(execution, first.context)).positive).toBe(2);
});
test('a cached node rewired to the negative input does not count the positive prompt', async ({ execution, page }) => {
test.setTimeout(90000);
await setCache(execution, 'RAM_PRESSURE');
const first = await firstRun(execution, page, linked());
const rewired = await queueAgain(execution, first, { seed: 2, edit: output => {
const sampler = Object.values(output).find(node => node.class_type === 'KSampler');
[sampler.inputs.positive, sampler.inputs.negative] = [sampler.inputs.negative, sampler.inputs.positive];
} });
expect(rewired.cached).toContain(first.promptNode);
expect((await runCounts(execution, first.context)).positive).toBe(1);
// The history capture of a cached job is unchanged: effective text from the cache, no stored-prompt hint
const cachedContext = await capturedContext(execution, (await queueAgain(execution, first, { seed: 3 })).id);
expect(cachedContext.positive.effective).toBe(first.context.positive.effective);
expect(cachedContext.positive_prompt_id).toBeNull();
expect((await runCounts(execution, first.context)).positive).toBe(2);
});
});
+33
View File
@@ -201,6 +201,39 @@ NODE_CLASS_MAPPINGS.update(
{"PMHistoryImages": PMHistoryImages, "PMHistoryImageOutput": PMHistoryImageOutput}
)
# #347: text whose value (and cache signature) follows a value the test sets, so
# one queued graph can serve A, then B, then A again (each cached under LRU).
_CYCLED = {"value": "cycled"}
class PMCycledText:
@classmethod
def INPUT_TYPES(cls) -> dict:
return {"required": {}}
RETURN_TYPES = ("STRING",)
FUNCTION = "execute"
CATEGORY = "testing/execution-history"
@classmethod
def IS_CHANGED(cls, **kwargs: object) -> str:
return _CYCLED["value"]
def execute(self) -> tuple:
return (_CYCLED["value"],)
NODE_CLASS_MAPPINGS["PMCycledText"] = PMCycledText
@PromptServer.instance.routes.post("/pm_execution_fixture/cycled")
async def cycled_text(request: web.Request) -> web.Response:
value = (await request.json())["value"]
if type(value) is not str:
raise web.HTTPBadRequest()
_CYCLED["value"] = value
return web.json_response({"value": value})
@PromptServer.instance.routes.post("/pm_execution_fixture/metadata")
async def metadata(request: web.Request) -> web.Response:
+23
View File
@@ -140,3 +140,26 @@ test('an output saved without metadata links once its recorded run names a store
expect((await request.post('/pm_execution_fixture/metadata', { data: { disabled: false } })).ok()).toBe(true);
}
});
// #332 on the real host: the migration ran at startup and the image category routes
// set, list and remove an own category on a generated output.
test('an output gets its own image category and loses it again', async ({ execution }, testInfo) => {
const request = execution.request;
const filename = await saved(request, graph({ text: `category host ${randomUUID()}`, prefix: 'pm332/category' }));
let image;
await expect.poll(async () => {
const listing = await (await request.get('/prompt_manager/images/output?mode=workspace&q=pm332&rescan=1')).json();
image = (listing.images || []).find(entry => entry.filename === filename || entry.path?.endsWith(`/${filename}`));
return Boolean(image);
}, { timeout: 20000, message: `${filename} is listed by Images` }).toBe(true);
const items = [{ root_index: image.root_index, root_scope: image.root_scope, path: image.path }];
const value = `PM332 host ${testInfo.project.name}`;
const set = await request.put('/prompt_manager/images/category', { data: { items, category: value } });
expect(set.status()).toBe(200);
expect(await set.json()).toEqual({ success: true, category: value, updated: 1, unchanged: 0, not_found: 0 });
expect((await (await request.get('/prompt_manager/images/categories')).json()).categories).toContain(value);
const cleared = await request.delete('/prompt_manager/images/category', { data: { items } });
expect(cleared.status()).toBe(200);
expect(await cleared.json()).toEqual({ success: true, cleared: 1 });
expect((await (await request.get('/prompt_manager/images/categories')).json()).categories).not.toContain(value);
});
+9
View File
@@ -1,7 +1,10 @@
{
"routes": [
"DELETE /prompt_manager/collections/items",
"DELETE /prompt_manager/collections/{collection_id}",
"DELETE /prompt_manager/delete/{prompt_id}",
"DELETE /prompt_manager/image-annotations/capture/{token}",
"DELETE /prompt_manager/images/category",
"DELETE /prompt_manager/images/favorite",
"DELETE /prompt_manager/images/{image_id}",
"DELETE /prompt_manager/pairs/{pair_id}",
@@ -14,6 +17,8 @@
"GET /prompt_manager/backup",
"GET /prompt_manager/backup/status",
"GET /prompt_manager/categories",
"GET /prompt_manager/collections",
"GET /prompt_manager/collections/membership",
"GET /prompt_manager/diagnostics",
"GET /prompt_manager/export",
"GET /prompt_manager/gallery",
@@ -26,6 +31,7 @@
"GET /prompt_manager/history/runs/{run_id}/outputs",
"GET /prompt_manager/image-annotations/capture/{token}",
"GET /prompt_manager/images/all",
"GET /prompt_manager/images/categories",
"GET /prompt_manager/images/inspect/{image_path}",
"GET /prompt_manager/images/output",
"GET /prompt_manager/images/prompt/{image_path}",
@@ -90,6 +96,7 @@
"GET /prompt_manager/vocabulary/tags/{tag_id}/impact",
"GET /prompt_manager/web",
"GET /prompt_manager/workspace",
"PATCH /prompt_manager/collections/{collection_id}",
"PATCH /prompt_manager/library/{positive_id}",
"PATCH /prompt_manager/pairs/{pair_id}",
"PATCH /prompt_manager/saved-views/{view_id}",
@@ -104,6 +111,7 @@
"POST /prompt_manager/bulk/delete",
"POST /prompt_manager/bulk/tags",
"POST /prompt_manager/cleanup",
"POST /prompt_manager/collections/items",
"POST /prompt_manager/delete_duplicate_images",
"POST /prompt_manager/diagnostics/test-link",
"POST /prompt_manager/gallery/rescan",
@@ -156,6 +164,7 @@
"POST /prompt_manager/vocabulary/tags",
"POST /prompt_manager/vocabulary/tags/{tag_id}/delete",
"PUT /prompt_manager/image-annotations/tags",
"PUT /prompt_manager/images/category",
"PUT /prompt_manager/images/favorite",
"PUT /prompt_manager/prompts/{prompt_id}",
"PUT /prompt_manager/prompts/{prompt_id}/rating",
+85 -19
View File
@@ -76,35 +76,101 @@ test("the workspace keys have their own defaults and validation", () => {
assert.equal(normalizeGallerySettings({ defaultViewMode: "masonry" }).defaultViewMode, "grid");
});
test("a legacy gallery Save keeps the workspace keys and every other stored key", () => {
// #330 S13: the legacy gallery's Page options edit only the page-only keys (PAGE_OPTION_KEYS); System › Appearance owns the rest, so a Save
// or a Reset from the legacy page keeps every other stored key.
function legacyGallery(stored, controls, extra = {}) {
const fs = require("node:fs");
const path = require("node:path");
const vm = require("node:vm");
const settings = require("../../web/js/gallery-settings.js");
const storage = new Map([["gallerySettings", JSON.stringify({ workspaceLayout: "masonry", workspaceDensity: "compact", someOtherTool: 1 })]]);
const controls = {
lazyLoadingToggle: { checked: true }, imageQualitySelect: { value: "high" }, defaultViewModeSelect: { value: "list" },
defaultLimitSelect: { value: "200" }, gridColumnsSelect: { value: "6" }, showImageInfoToggle: { checked: false },
autoLoadMetadataToggle: { checked: true }, cacheMetadataToggle: { checked: true }, showFilePathsToggle: { checked: false },
sidebarCollapsedByDefaultToggle: { checked: true }, infiniteScrollToggle: { checked: false }, debugModeToggle: { checked: false },
apiTimeoutInput: { value: "30" }, checkThumbnailsAtStartup: { checked: true }, videoAutoplayToggle: { checked: true },
videoMuteToggle: { checked: false }, videoLoopToggle: { checked: true },
};
const storage = new Map([["gallerySettings", JSON.stringify(stored)]]);
const context = vm.createContext({
document: { addEventListener() {}, getElementById: id => controls[id] },
window: {}, console: { debug() {}, warn() {}, error() {} }, GallerySettings: settings,
localStorage: { getItem: key => storage.get(key) ?? null, setItem: (key, value) => storage.set(key, value) },
localStorage: { getItem: key => storage.get(key) ?? null, setItem: (key, value) => storage.set(key, value),
removeItem: key => storage.delete(key) },
...extra,
});
const file = path.join(__dirname, "../../web/js/gallery.js");
const Manager = vm.runInContext(`${fs.readFileSync(file, "utf8")}\nGalleryManager;`, context, { filename: file });
const manager = Object.create(Manager.prototype);
Object.assign(manager, { showNotification() {}, hideSettings() {}, applySettings() {} });
const calls = [];
Object.assign(manager, { showNotification: (message, type) => calls.push([type, message]), hideSettings() {},
applySettings: applied => calls.push(["apply", applied.lazyLoading]) });
return { manager, storage, calls, saved: () => JSON.parse(storage.get("gallerySettings")) };
}
const PAGE_OPTION_CONTROLS = () => ({
checkThumbnailsAtStartup: { checked: false }, lazyLoadingToggle: { checked: false },
defaultLimitSelect: { value: "200" },
gridColumnsSelect: { value: "6" }, infiniteScrollToggle: { checked: true }, debugModeToggle: { checked: true },
apiTimeoutInput: { value: "45" },
});
test("a legacy gallery Save keeps the workspace keys and every other stored key", () => {
const { manager, saved } = legacyGallery({ workspaceLayout: "masonry", workspaceDensity: "compact", someOtherTool: 1,
imageQuality: "high", showImageInfo: false, videoMute: false, thumbnailPromptDismissed: "sig", defaultViewMode: "list" },
PAGE_OPTION_CONTROLS());
manager.saveSettings();
const saved = JSON.parse(storage.get("gallerySettings"));
assert.equal(saved.workspaceLayout, "masonry");
assert.equal(saved.workspaceDensity, "compact");
assert.equal(saved.someOtherTool, 1);
assert.equal(saved.defaultViewMode, "list");
assert.equal(saved.showImageInfo, false);
assert.equal(saved.videoMute, false);
const result = saved();
assert.equal(result.workspaceLayout, "masonry");
assert.equal(result.workspaceDensity, "compact");
assert.equal(result.someOtherTool, 1);
assert.equal(result.imageQuality, "high");
assert.equal(result.showImageInfo, false);
assert.equal(result.videoMute, false);
assert.equal(result.thumbnailPromptDismissed, "sig");
// Page options no longer edit the default view (S13 G5), so a stored one stays.
assert.equal(result.defaultViewMode, "list");
// The page-only keys come from the dialog.
assert.equal(result.lazyLoading, false);
assert.equal(result.debugMode, true);
assert.equal(result.checkThumbnailsAtStartup, false);
assert.equal(result.defaultLimit, 200);
assert.equal(result.gridColumns, 6);
assert.equal(result.infiniteScroll, true);
assert.equal(result.apiTimeout, 45);
});
test("a legacy gallery Reset restores only the page options and keeps every other stored key", () => {
const controls = PAGE_OPTION_CONTROLS();
const { manager, saved, calls } = legacyGallery({ workspaceLayout: "masonry", imageQuality: "high", videoMute: false,
thumbnailPromptDismissed: "sig", thumbnailsGenerated: true, lazyLoading: false, debugMode: true, defaultLimit: 200,
gridColumns: 6, infiniteScroll: true, apiTimeout: 45, checkThumbnailsAtStartup: false, defaultViewMode: "list" }, controls,
{ confirm: message => (calls.push(["confirm", message]), true) });
manager.resetSettings();
const result = saved();
const { DEFAULT_SETTINGS } = require("../../web/js/gallery-settings.js");
for (const key of ["checkThumbnailsAtStartup", "lazyLoading", "defaultLimit", "gridColumns",
"infiniteScroll", "debugMode", "apiTimeout"]) {
assert.equal(result[key], DEFAULT_SETTINGS[key], key);
}
// The default view is not a page option (S13 G5): Reset keeps a stored one.
assert.equal(result.defaultViewMode, "list");
assert.equal(result.workspaceLayout, "masonry");
assert.equal(result.imageQuality, "high");
assert.equal(result.videoMute, false);
assert.equal(result.thumbnailPromptDismissed, "sig");
assert.equal(result.thumbnailsGenerated, true);
// The dialog shows the defaults again and the page applies them.
assert.equal(controls.lazyLoadingToggle.checked, DEFAULT_SETTINGS.lazyLoading);
assert.equal(controls.apiTimeoutInput.value, DEFAULT_SETTINGS.apiTimeout);
assert.deepEqual(calls, [["confirm", "Reset these page options to their defaults?"], ["apply", DEFAULT_SETTINGS.lazyLoading],
["success", "Page options reset to defaults"]]);
});
test("Page options no longer hold the default view; System › Appearance's Default layout is the one control", () => {
const { manager } = legacyGallery({}, {});
assert.deepEqual([...manager.constructor.PAGE_OPTION_KEYS], ["checkThumbnailsAtStartup", "lazyLoading", "defaultLimit",
"gridColumns", "infiniteScroll", "debugMode", "apiTimeout"]);
});
test("a viewer mute or loop click remembers only that key and never reads the Page options dialog", () => {
const { manager, saved, calls } = legacyGallery({ defaultLimit: 200, lazyLoading: false, videoMute: true }, {});
manager.storeSetting("videoMute", false);
const result = saved();
assert.equal(result.videoMute, false);
assert.equal(result.defaultLimit, 200);
assert.equal(result.lazyLoading, false);
assert.deepEqual(calls, []);
});
+15 -12
View File
@@ -112,11 +112,12 @@ test('saving gallery settings retains the remembered thumbnail dismissal', () =>
});
// #249 System S5: the stream's work is now an Activity job that outlives the page's connection, so the gallery's
// Cancel asks the server to stop the job. The stream's first frame, `job {id, started}`, names the job; the gallery
// Cancel asks the server to stop the job. Since #330 the startup dialog is the gallery's only thumbnail stream (System
// › Sources & scans generates the rest). The stream's first frame, `job {id, started}`, names the job; the gallery
// cancels by that id, only a job this stream started, and a Cancel before the id waits for it (review I1).
const JOB = 'b'.repeat(32);
function streamFixture({ modal = false, postFails = false } = {}) {
function streamFixture({ postFails = false } = {}) {
const posts = [];
const warnings = [];
const gets = [];
@@ -146,18 +147,20 @@ function streamFixture({ modal = false, postFails = false } = {}) {
return { success: true };
} } });
const emit = (type, data) => (listeners[type] || []).forEach(handler => handler({ data: JSON.stringify(data) }));
// The dialog's Cancel runs the handler and from then on reads as cancelled, as the startup dialog does.
let cancelled = false;
const ui = { addStatusMessage() {}, progressText: element(), progressPercent: element(), progressBar: element(),
progressDetails: element(), progressETA: element(), isCancelled: () => false,
onCancel: handler => { ui.cancel = handler; } };
const result = modal ? manager.generateThumbnailsWithProgressInModal(ui)
: manager.generateThumbnailsWithProgress({ quality: 'medium' });
progressDetails: element(), progressETA: element(), isCancelled: () => cancelled,
onCancel: handler => { ui.cancel = () => { cancelled = true; handler(); }; } };
const result = manager.generateThumbnailsWithProgressInModal(ui);
const outcome = result.then(value => ({ value }), error => ({ error: error.message }));
const cancel = () => (modal ? ui.cancel() : elements.cancelThumbnailsBtn.click());
const cancel = () => ui.cancel();
const settle = () => new Promise(resolve => setImmediate(resolve));
return { emit, source, posts, gets, warnings, outcome, cancel, settle };
}
const cancelUrl = id => `/prompt_manager/jobs/${id}/cancel`;
const CANCELLED = 'Thumbnail generation cancelled; previews already written are kept';
test('Cancel before the job id arrives cancels that job once it does', async () => {
const f = streamFixture();
@@ -165,8 +168,8 @@ test('Cancel before the job id arrives cancels that job once it does', async ()
await f.settle();
assert.deepEqual(f.posts, []);
assert.equal(f.source.closed, 0); // kept open until the stream names its job
assert.deepEqual(await f.outcome, { error: 'Thumbnail generation cancelled by user' });
f.emit('status', { phase: 'scanning', message: 'Scanning…' });
assert.deepEqual(await f.outcome, { error: CANCELLED });
// The job frame is the stream's first (the dialog closes the stream on any later frame once cancelled).
f.emit('job', { id: JOB, started: true });
await f.settle();
assert.equal(JSON.stringify(f.posts), JSON.stringify([[cancelUrl(JOB), {}]]));
@@ -185,7 +188,7 @@ test('Cancel after the id cancels it at once, and a double Cancel posts once', a
await f.settle();
assert.equal(JSON.stringify(f.posts), JSON.stringify([[cancelUrl(JOB), {}]]));
assert.equal(f.source.closed, 1);
assert.deepEqual(await f.outcome, { error: 'Thumbnail generation cancelled by user' });
assert.deepEqual(await f.outcome, { error: CANCELLED });
assert.deepEqual(f.gets, []);
});
@@ -216,11 +219,11 @@ test('a cancel post that fails is reported, not thrown', async () => {
await f.settle();
assert.equal(f.posts.length, 1);
assert.equal(f.warnings.length, 1);
assert.deepEqual(await f.outcome, { error: 'Thumbnail generation cancelled by user' });
assert.deepEqual(await f.outcome, { error: CANCELLED });
});
test('the startup dialog cancels its own job by id too', async () => {
const f = streamFixture({ modal: true });
const f = streamFixture();
f.cancel();
await f.settle();
assert.deepEqual(f.posts, []);
+19 -113
View File
@@ -65,109 +65,12 @@ function load(file, className, fetchUrl, extra = {}, dom = elements()) {
const BUSY = 'Scan for metadata already running';
test('gallery cache clear reports a refused server clear and never claims success', async () => {
const calls = [];
const { page, dom, notifications, storage } = load('gallery.js', 'GalleryManager', async url => {
calls.push(url);
return jsonResponse(409, { success: false, error: BUSY });
});
page.getSettings = () => ({ thumbnailsGenerated: true });
dom.getElementById('thumbnailStatus').textContent = 'Generated';
await page.clearCache();
assert.deepEqual(calls, ['/prompt_manager/images/clear-thumbnails']);
assert.deepEqual(notifications, [['error', `Thumbnails were not cleared: ${BUSY}`]]);
assert.equal(dom.nodes.thumbnailStatus.textContent, 'Generated');
assert.equal(storage.get('gallerySettings'), undefined); // thumbnailsGenerated is not reset
});
test('gallery cache clear says what it cleared and keeps preferences and other data (#222)', async () => {
const calls = [];
const { page, dom, notifications, storage } = load('gallery.js', 'GalleryManager', async url => {
calls.push(url);
return jsonResponse(200, { success: true, cleared_count: 3 });
});
page.getSettings = () => ({ thumbnailsGenerated: true });
await page.clearCache();
// The server's thumbnails are cleared by their own explicit call; nothing Prompt Manager owns is cached in this
// browser, so no browser entry is removed and the toast says exactly that.
assert.deepEqual(calls, ['/prompt_manager/images/clear-thumbnails']);
assert.deepEqual(notifications, [['success',
"Thumbnails cleared; original images untouched. Nothing of Prompt Manager's is cached in this browser."]]);
assert.equal(dom.nodes.thumbnailStatus.textContent, 'Not generated');
assert.equal(JSON.parse(storage.get('gallerySettings')).thumbnailsGenerated, false);
assert.equal(storage.get('userData'), 'keep');
assert.equal(storage.get('pm_cached'), '1');
assert.equal(storage.get('pm_workspace_theme'), 'light');
assert.equal(storage.get('pm_update_notice_dismissed'), '3.3.0');
});
test("gallery cache clear never deletes another tool's gallery or thumbnail caches (#222)", async () => {
const names = new Set(['comfy-gallery', 'thumbnail-cache']);
const caches = { keys: async () => [...names], delete: async name => names.delete(name) };
const { page, notifications } = load('gallery.js', 'GalleryManager',
async () => jsonResponse(200, { success: true }), { caches, window: { caches } });
page.getSettings = () => ({});
await page.clearCache();
assert.deepEqual([...names], ['comfy-gallery', 'thumbnail-cache']);
assert.equal(notifications.at(-1)[0], 'success');
});
test('gallery cache clear reports a lost server clear instead of success', async () => {
const { page, dom, notifications, storage } = load('gallery.js', 'GalleryManager', async () => {
throw new TypeError('Failed to fetch');
});
page.getSettings = () => ({ thumbnailsGenerated: true });
dom.getElementById('thumbnailStatus').textContent = 'Generated';
await page.clearCache();
assert.deepEqual(notifications, [['error', 'Thumbnails were not cleared: the server could not be reached.']]);
assert.equal(dom.nodes.thumbnailStatus.textContent, 'Generated');
assert.equal(storage.get('gallerySettings'), undefined);
});
test('gallery cache clear says so when browser storage is denied', async () => {
const denied = { get length() { throw new Error('SecurityError'); }, key() { throw new Error('SecurityError'); },
getItem() { throw new Error('SecurityError'); }, setItem() { throw new Error('SecurityError'); },
removeItem() { throw new Error('SecurityError'); } };
const { page, notifications } = load('gallery.js', 'GalleryManager',
async () => jsonResponse(200, { success: true }), { localStorage: denied });
page.getSettings = () => ({});
await page.clearCache();
assert.deepEqual(notifications, [['warning',
'Thumbnails cleared; original images untouched. This browser refused access to its storage, so nothing there was cleared.']]);
});
test('admin LoRA import shows a refusal frame and closes its modal', async () => {
const { page, dom, timers, notifications } = load('admin.js', 'PromptAdmin',
async () => sseResponse([{ type: 'error', status: BUSY, progress: 0, message: BUSY }]));
page.saveLoraSettings = async () => true;
page.loadStatistics = () => assert.fail('a refused import must not reload statistics');
await page.runLoraImport();
assert.equal(dom.nodes.loraImportStatus.textContent, `Error: ${BUSY}`);
assert.deepEqual(notifications.filter(([type]) => type === 'error'), [['error', `LoRA import failed: ${BUSY}`]]);
assert.equal(timers.length, 1);
timers[0]();
assert.deepEqual(notifications.at(-1), ['hide', 'loraImportModal']);
});
test('admin settings save shows why config.json refused it', async () => {
const message = 'config.json cannot be read, so nothing was saved. Fix or remove the file, then save again.';
const { page, notifications } = load('admin.js', 'PromptAdmin',
async () => jsonResponse(409, { success: false, persisted: false, code: 'config-unreadable', error: message }));
page.settings = {};
page._collectScanPaths = () => [];
page.saveLoraSettings = async () => assert.fail('nothing else is saved after a refusal');
page.setInfiniteScroll = () => assert.fail('nothing changes after a refusal');
await page.saveSettings();
@@ -178,27 +81,30 @@ test('admin settings save keeps its plain message when the reply has no reason',
const { page, notifications } = load('admin.js', 'PromptAdmin', async () => ({ ok: false, status: 500,
json: async () => { throw new SyntaxError('not json'); } }));
page.settings = {};
page._collectScanPaths = () => [];
await page.saveSettings();
assert.deepEqual(notifications, [['error', 'Failed to save settings']]);
});
test('admin maintenance names the job that refused it', async () => {
const { page, dom, notifications } = load('admin.js', 'PromptAdmin',
async () => jsonResponse(409, { success: false, error: BUSY }), {
escapeHtml: text => String(text),
document: { readyState: 'loading', addEventListener() {}, getElementById: id => dom.getElementById(id),
querySelectorAll: selector => (selector === '.maintenance-option:checked' ? [{ id: 'vacuum' }] : []) },
});
// #330 S13-9: the admin page's reduced Page options save sends only the list's scroll default, so a stale legacy form
// can no longer overwrite scan folders or other settings edited in System.
test('admin Page options save posts only infinite_scroll', async () => {
const sent = [];
const { page, dom, notifications } = load('admin.js', 'PromptAdmin', async (url, init) => {
sent.push([url, init.method, JSON.parse(init.body)]);
return jsonResponse(200, { success: true, restart_required: true });
});
page.settings = { infiniteScroll: false };
const scrolled = [];
page.setInfiniteScroll = (on, options) => scrolled.push([on, { ...options }]);
dom.getElementById('infiniteScrollDefault').checked = true;
await page.runMaintenance();
await page.saveSettings();
const results = dom.nodes.maintenanceResults.innerHTML;
assert.match(results, /Maintenance Failed/);
assert.match(results, new RegExp(BUSY));
assert.doesNotMatch(results, /HTTP 409/);
assert.deepEqual(notifications, [['error', '❌ Maintenance failed']]);
assert.equal(dom.nodes.runMaintenanceBtn.disabled, false);
assert.deepEqual(sent, [['/prompt_manager/settings', 'POST', { infinite_scroll: true }]]);
assert.deepEqual(scrolled, [[true, { session: false }]]);
assert.equal(page.settings.infiniteScroll, true);
assert.deepEqual(notifications, [['success', 'Page options saved'], ['hide', 'settingsModal']]);
});
+49
View File
@@ -185,6 +185,26 @@ test('legacy fill applies to local records', async () => {
assert.notEqual(zeroed.prompts.positive.text, '');
});
test('a branch named by recorded run history shows the Recorded run badge', async () => {
// #386: the server labels sides whose save node recorded history named; Execution recorded stays stronger.
const base = await localRecord(new File([fs.readFileSync(new URL('f09.png', PARITY))], 'f09.png'));
const recorded = {...base, prompts:{positive:{...base.prompts.positive, source:'recorded-run'}, negative:{...base.prompts.negative, source:'recorded-run'}}};
let view = renderFull(recorded);
try {
assert.deepEqual(view.texts(n => n.className === 'chip blue'), ['Recorded run']);
assert.ok(view.texts(n => n.tagName === 'p').includes('Prompt evidence: recorded run'));
} finally { view.done(); }
view = renderFull({...recorded, prompts:{...recorded.prompts, negative:{...recorded.prompts.negative, source:'execution'}}});
try {
assert.deepEqual(view.texts(n => n.className === 'chip blue'), ['Execution recorded']);
} finally { view.done(); }
// A recorded branch whose negative left out a run-time part is still inferred from the workflow.
view = renderFull({...recorded, prompts:{...recorded.prompts, negative:{...recorded.prompts.negative, source:'runtime-parts-omitted'}}});
try {
assert.deepEqual(view.texts(n => n.className === 'chip blue'), ['Workflow inferred']);
} finally { view.done(); }
});
// #249 PR C Task C3 (D88): one chip row, the file's current image tags first, then D21's linked tags not already shown.
function tagged(imageTags, library = {status:'linked', tags:['Minimal', 'warm_light', 'editorial'], tags_scope:'positive'}) {
return {...record(), imageTags, library};
@@ -363,3 +383,32 @@ test('the D117 wrap rule targets only the settings summary button', async () =>
assert.deepEqual(rules.map(line => line.trim()), ['.ca-tools #viewer-details .settings-summary { white-space: normal; /* D117 */ }']);
assert.doesNotMatch(css, /#viewer-details > div > button/);
});
test('category chip names its origin and shows an own category for any link status (#332)', async () => {
const base = await localRecord(localPng('a.png'));
const chipTexts = record => {
const {f, done} = renderFull(record);
try {
return f.walk(f.container, n => /(^| )chip( |$)/.test(n.className ?? '')).map(n => n.textContent);
} finally { done(); }
};
const cases = [
[{status:'none', category:'Scans', category_scope:'image'}, 'Image category: Scans'],
[{status:'ambiguous', category:'Scans', category_scope:'image'}, 'Image category: Scans'],
[{status:'linked', category:'Portraits', category_scope:'image'}, 'Image category: Portraits'],
[{status:'linked', category:'Landscapes', category_scope:'positive'}, 'Positive prompt category: Landscapes'],
[{status:'linked', category:'Recipes', category_scope:'pair'}, 'Recipe category: Recipes'],
];
for (const [library, text] of cases) {
const chips = chipTexts({...base, library:{...base.library, tags:[], tags_scope:null, ...library}});
assert.ok(chips.includes(text), `${JSON.stringify(library)} -> ${chips}`);
}
for (const library of [
{status:'none', category:'Inherited', category_scope:'positive'},
{status:'linked', category:null, category_scope:null},
{status:'none', category:null, category_scope:null},
]) {
const chips = chipTexts({...base, library:{...base.library, tags:[], tags_scope:null, ...library}});
assert.equal(chips.some(chip => /category: /.test(chip)), false, JSON.stringify(library));
}
});
@@ -278,6 +278,7 @@ test('record rejects root path media and invalid versioned fields with retryable
data => { data.resolution.sampler_id = {}; },
data => { data.library.tags_scope = 'recipe'; },
data => { data.library.category_scope = 'recipe'; },
data => { data.library.tags_scope = 'image'; },
data => { delete data.library.category_scope; },
data => { delete data.library.tags_scope; },
]) {
@@ -942,3 +943,14 @@ test("names the server quotes (parentheses, '!*') still match their original and
assert.equal(record.download.url, entry.url);
assert.deepEqual(record.preview, { url: entry.thumbnail_url });
});
test('record accepts an image\'s own category on any link status (#332)', () => {
const toRecord = contract('toInspectionRecord');
for (const status of ['linked', 'none', 'ambiguous', 'unavailable']) {
const data = envelope();
Object.assign(data.library, { status, category: 'Scans', category_scope: 'image' });
const record = toRecord(data, image(), extractGenerationFacts);
assert.equal(record.library.category, 'Scans');
assert.equal(record.library.category_scope, 'image');
}
});
+2 -2
View File
@@ -63,9 +63,9 @@ test('status line shows only server-reported backups (D103)', async () => {
const PREVIEW = Object.freeze({
success: true, token: 'r'.repeat(43), expires_at: '2026-10-25T12:15:00+00:00',
summary: { prompts: 1234, positives: 900, pairs: 3, archived_pairs: 0, image_links: 56, runs: 1, saved_views: 2 },
summary: { prompts: 1234, positives: 900, pairs: 3, archived_pairs: 0, image_links: 56, runs: 1, saved_views: 2, collections: 1 },
schema: { migrated: true, library_state: true }, integrity: 'ok', documents: { repaired: 0, cleared: 0, images: 0 },
live: { prompts: 1, positives: 1, pairs: 0, archived_pairs: 0, image_links: 0, runs: 0, saved_views: 0 }, size: 4096,
live: { prompts: 1, positives: 1, pairs: 0, archived_pairs: 0, image_links: 0, runs: 0, saved_views: 0, collections: 0 }, size: 4096,
safety_backup: { destination_label: 'prompts.db.backup_<date and time>' },
});
+35
View File
@@ -264,6 +264,41 @@ class TestRestoreRoute(BackupRestoreTestCase):
texts = sorted(p["text"] for p in self.api.db.search_prompts())
self.assertEqual(texts, ["restored a", "restored b"])
async def test_restore_round_trip_keeps_tags_and_image_links(self):
"""A downloaded backup brings back every prompt with its tags and images.
The legacy admin page's Restore form left for System (#330); this keeps
its browser test's data check at the route.
"""
db = self.api.db
originals = {}
for text, tags, image in (
("round trip cat", ["cats", "pets"], "cat.png"),
("round trip dog", ["dogs"], "dog.png"),
):
prompt_id = db.save_prompt(
text=text, tags=tags, prompt_hash=generate_prompt_hash(text)
)
path = self._write(image, b"\x89PNG\r\n\x1a\n")
self.assertTrue(db.link_image_to_prompt(str(prompt_id), path))
originals[text] = (prompt_id, sorted(tags), [path])
resp = await self.client.get("/prompt_manager/backup")
self.assertEqual(resp.status, 200)
backup = self._write("downloaded.db", await resp.read())
self.assertTrue(db.delete_prompt(originals["round trip dog"][0]))
resp = await self._restore(backup)
self.assertEqual(resp.status, 200, await resp.text())
restored = {p["text"]: p for p in self.api.db.search_prompts()}
self.assertEqual(sorted(restored), sorted(originals))
for text, (prompt_id, tags, paths) in originals.items():
self.assertEqual(restored[text]["id"], prompt_id)
self.assertEqual(sorted(restored[text]["tags"]), tags)
images = self.api.db.get_prompt_images(str(prompt_id))
self.assertEqual([image["image_path"] for image in images], paths)
async def test_restore_leaves_no_temp_files_behind(self):
upload = self._make_db_file(["restored"])
created = []
+601
View File
@@ -0,0 +1,601 @@
"""Collections (#331): the routes under /prompt_manager/collections."""
import json
import ntpath
import sqlite3
from unittest.mock import patch
try:
from test_api_images import ImageAPITestCase, make_png
from open_handles import assert_closed
except ImportError:
from tests.test_api_images import ImageAPITestCase, make_png
from tests.open_handles import assert_closed
BASE = "/prompt_manager/collections"
ITEMS = BASE + "/items"
MEMBERSHIP = BASE + "/membership"
SOURCES_CHANGED = "The sources changed. Refresh the images and try again."
class CollectionRoutesTestCase(ImageAPITestCase):
async def tearDownAsync(self):
self.api.db.close_all()
assert_closed(self, self._temp_db.name, self._output_tmp.name)
await super().tearDownAsync()
def make(self, *names, root=None):
for name in names:
make_png((root or self.output_dir) / name)
def scope(self, index=0):
roots = list(self.api._get_all_output_dirs())
return self.api._gallery_root_scope(roots[index], index)
def item(self, path, index=0, scope=None):
return {
"root_index": index,
"root_scope": scope or self.scope(index),
"path": path,
}
def two_sources(self):
second = self.comfy_dir / "second"
second.mkdir()
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
return second
async def call(self, method, url, body=None, status=200, **params):
response = await self.client.request(method, url, json=body, params=params)
text = await response.text()
self.assertEqual(response.status, status, text)
data = json.loads(text)
self.assertNotIn(self.comfy_dir.as_posix(), text)
self.assertNotIn(str(self.comfy_dir).replace("\\", "\\\\"), text)
return data
async def add(self, status=200, **body):
return await self.call("POST", ITEMS, body, status)
async def listed(self):
return (await self.call("GET", BASE))["collections"]
class TestCollectionRoutes(CollectionRoutesTestCase):
async def test_each_route_succeeds(self):
self.make("a.png", "b.png", "c.png")
self.assertEqual(
await self.call("GET", BASE), {"success": True, "collections": []}
)
made = await self.add(
name=" LoRa testing ", items=[self.item("a.png"), self.item("b.png")]
)
grids = made["collection"]
self.assertEqual(grids["name"], "LoRa testing")
self.assertEqual(
made,
{
"success": True,
"collection": grids,
"created": True,
"added": 2,
"already": 0,
"not_found": 0,
},
)
again = await self.add(
collection_id=grids["id"], items=[self.item("b.png"), self.item("c.png")]
)
self.assertEqual(
(again["created"], again["added"], again["already"]), (False, 1, 1)
)
self.assertEqual(
await self.listed(),
[{"id": grids["id"], "name": "LoRa testing", "count": 3, "missing": 0}],
)
membership = await self.call(
"GET", MEMBERSHIP, root_index="0", root_scope=self.scope(), path="a.png"
)
self.assertEqual(membership, {"success": True, "collection_ids": [grids["id"]]})
removed = await self.call(
"DELETE",
ITEMS,
{"collection_id": grids["id"], "items": [self.item("a.png")]},
)
self.assertEqual(removed, {"success": True, "collection": grids, "removed": 1})
renamed = await self.call(
"PATCH", f"{BASE}/{grids['id']}", {"name": "LoRA grids"}
)
self.assertEqual(
renamed,
{"success": True, "collection": {"id": grids["id"], "name": "LoRA grids"}},
)
deleted = await self.call("DELETE", f"{BASE}/{grids['id']}")
self.assertEqual(deleted, {"success": True, "deleted": True})
self.assertEqual(await self.listed(), [])
# Files are never touched.
self.assertTrue((self.output_dir / "a.png").is_file())
async def test_names_are_unique_ignoring_case(self):
self.make("a.png", "b.png")
first = (await self.add(name="LoRa testing", items=[self.item("a.png")]))[
"collection"
]
second = await self.add(name="lora TESTING", items=[self.item("b.png")])
self.assertEqual(second["collection"], first)
self.assertFalse(second["created"])
self.assertEqual(len(await self.listed()), 1)
async def test_bad_requests_are_400(self):
self.make("a.png")
made = await self.add(name="Set", items=[self.item("a.png")])
cid = made["collection"]["id"]
good = [self.item("a.png")]
for body in (
{"items": good},
{"name": "x", "collection_id": cid, "items": good},
{"name": "x"},
{"name": "x", "items": good, "extra": 1},
{"collection_id": "1", "items": good},
{"collection_id": True, "items": good},
{"collection_id": 0, "items": good},
{"collection_id": 10**10, "items": good},
{"name": "", "items": good},
{"name": "x" * 101, "items": good},
{"name": "bell\x07", "items": good},
{"name": 5, "items": good},
{"name": "x", "items": []},
{"name": "x", "items": "a.png"},
{"name": "x", "items": [self.item("a.png")] * 501},
{"name": "x", "items": [{"root_index": 0, "path": "a.png"}]},
{"name": "x", "items": [{**self.item("a.png"), "root_index": "0"}]},
{"name": "x", "items": [{**self.item("a.png"), "root_scope": "xyz"}]},
):
with self.subTest(body=body):
data = await self.add(400, **body)
self.assertFalse(data["success"])
for body in ({"items": good}, {"collection_id": cid}):
with self.subTest(remove=body):
await self.call("DELETE", ITEMS, body, 400)
for raw in ("abc", "0", "1234567890", "-1", "1.0"):
with self.subTest(path_id=raw):
await self.call("PATCH", f"{BASE}/{raw}", {"name": "x"}, 400)
await self.call("DELETE", f"{BASE}/{raw}", None, 400)
await self.call("PATCH", f"{BASE}/{cid}", {"name": ""}, 400)
await self.call("PATCH", f"{BASE}/{cid}", {"label": "x"}, 400)
for params in (
{"root_index": "x", "root_scope": self.scope(), "path": "a.png"},
{"root_index": "0", "path": "a.png"},
{"root_index": "0", "root_scope": "nope", "path": "a.png"},
):
with self.subTest(membership=params):
await self.call("GET", MEMBERSHIP, status=400, **params)
self.assertEqual((await self.listed())[0]["count"], 1)
async def test_unsafe_paths_are_403(self):
self.make("a.png")
for path in (
"../a.png",
"sub/../../a.png",
"C:/a.png",
"a\\b.png",
"/a.png",
"notes.txt",
):
with self.subTest(path=path):
await self.add(403, name="Set", items=[self.item(path)])
await self.call(
"GET",
MEMBERSHIP,
status=403,
root_index="0",
root_scope=self.scope(),
path=path,
)
self.assertEqual(await self.listed(), [])
async def test_a_body_over_the_cap_is_413(self):
response = await self.client.post(
ITEMS, data=b'{"name":"x","items":[' + b" " * (1024 * 1024 + 1) + b"]}"
)
self.assertEqual(response.status, 413)
async def test_missing_collection_and_files_are_404(self):
self.make("a.png")
data = await self.add(404, collection_id=999, items=[self.item("a.png")])
self.assertEqual(data["code"], "collection-not-found")
data = await self.call(
"DELETE", ITEMS, {"collection_id": 999, "items": [self.item("a.png")]}, 404
)
self.assertEqual(data["code"], "collection-not-found")
data = await self.call("PATCH", f"{BASE}/999", {"name": "x"}, 404)
self.assertEqual(data["code"], "collection-not-found")
data = await self.call("DELETE", f"{BASE}/999", None, 404)
self.assertEqual(data["code"], "collection-not-found")
data = await self.add(404, name="New", items=[self.item("never.png")])
self.assertEqual(data["code"], "files-not-found")
self.assertEqual(await self.listed(), [])
async def test_unlisted_files_count_as_not_found(self):
self.make("a.png")
made = await self.add(
name="Set", items=[self.item("a.png"), self.item("later.png")]
)
self.assertEqual((made["added"], made["not_found"]), (1, 1))
again = await self.add(
collection_id=made["collection"]["id"],
items=[self.item("a.png"), self.item("later.png")],
)
self.assertEqual(
(again["added"], again["already"], again["not_found"]), (0, 1, 1)
)
async def test_name_taken_is_409(self):
self.make("a.png")
first = (await self.add(name="First", items=[self.item("a.png")]))["collection"]
second = (await self.add(name="Second", items=[self.item("a.png")]))[
"collection"
]
data = await self.call(
"PATCH", f"{BASE}/{second['id']}", {"name": "FIRST"}, 409
)
self.assertEqual(
(data["code"], data["collection_id"]), ("name-taken", first["id"])
)
renamed = await self.call("PATCH", f"{BASE}/{second['id']}", {"name": "SECOND"})
self.assertEqual(renamed["collection"]["name"], "SECOND")
async def test_stale_scope_is_409_after_reordering_sources(self):
second = self.two_sources()
self.make("a.png")
self.make("b.png", root=second)
stale = self.item("b.png", 1)
made = await self.add(name="Set", items=[self.item("a.png")])
cid = made["collection"]["id"]
self.api._get_all_output_dirs = lambda: [second, self.output_dir]
for method, body in (
("POST", {"collection_id": cid, "items": [self.item("b.png", 0), stale]}),
("DELETE", {"collection_id": cid, "items": [stale]}),
):
with self.subTest(method=method):
data = await self.call(method, ITEMS, body, 409)
self.assertEqual(data["code"], "source-changed")
self.assertEqual(data["error"], SOURCES_CHANGED)
data = await self.call(
"GET",
MEMBERSHIP,
status=409,
root_index="1",
root_scope=stale["root_scope"],
path="b.png",
)
self.assertEqual(data["code"], "source-changed")
# Nothing was written by the refused add; the member follows its folder.
listed = await self.listed()
self.assertEqual((listed[0]["count"], listed[0]["missing"]), (1, 0))
membership = await self.call(
"GET", MEMBERSHIP, root_index="1", root_scope=self.scope(1), path="a.png"
)
self.assertEqual(membership["collection_ids"], [cid])
async def test_counts_exclude_a_missing_member(self):
second = self.two_sources()
self.make("a.png", "b.png")
self.make("c.png", root=second)
made = await self.add(
name="Grids",
items=[self.item("a.png"), self.item("b.png"), self.item("c.png", 1)],
)
other = await self.add(name="Other", items=[self.item("c.png", 1)])
(self.output_dir / "b.png").unlink()
self.api.invalidate_gallery_cache()
listed = await self.listed()
self.assertEqual(
listed,
[
{
"id": made["collection"]["id"],
"name": "Grids",
"count": 2,
"missing": 1,
},
{
"id": other["collection"]["id"],
"name": "Other",
"count": 1,
"missing": 0,
},
],
)
# The member is removable although it is no longer listed.
removed = await self.call(
"DELETE",
ITEMS,
{"collection_id": made["collection"]["id"], "items": [self.item("b.png")]},
)
self.assertEqual(removed["removed"], 1)
async def test_counts_are_null_when_a_source_scan_fails(self):
self.make("a.png")
await self.add(name="Set", items=[self.item("a.png")])
async def broken(_root):
raise OSError("scan failed at /secret/place")
with patch.object(self.api, "_get_gallery_files", side_effect=broken):
data = await self.call("GET", BASE)
self.assertEqual(
data["collections"][0] | {"id": 0},
{"id": 0, "name": "Set", "count": None, "missing": None},
)
self.assertNotIn("secret", json.dumps(data))
async def test_windows_case_spellings_are_one_member(self):
self.make("Sub/A.png")
with patch("py.api.images._path_flavor", ntpath):
made = await self.add(
name="Set", items=[self.item("Sub/A.png"), self.item("SUB/a.PNG")]
)
self.assertEqual((made["added"], made["already"]), (1, 0))
(self.output_dir / "Sub" / "A.png").rename(
self.output_dir / "Sub" / "a.png"
)
(self.output_dir / "Sub").rename(self.output_dir / "sub")
self.api.invalidate_gallery_cache()
listed = await self.listed()
self.assertEqual((listed[0]["count"], listed[0]["missing"]), (1, 0))
membership = await self.call(
"GET",
MEMBERSHIP,
root_index="0",
root_scope=self.scope(),
path="sub/a.png",
)
self.assertEqual(membership["collection_ids"], [made["collection"]["id"]])
response = await self.client.get(
"/prompt_manager/images/output",
params={
"mode": "workspace",
"collection": str(made["collection"]["id"]),
},
)
page = await response.json()
self.assertEqual([i["path"] for i in page["images"]], ["sub/a.png"])
async def test_posix_case_distinct_files_stay_two_members(self):
self.make("A.png", "a.png")
if len(list(self.output_dir.iterdir())) != 2:
self.skipTest("case-insensitive filesystem")
made = await self.add(
name="Set", items=[self.item("A.png"), self.item("a.png")]
)
self.assertEqual(made["added"], 2)
async def test_a_500_never_carries_the_exception_text(self):
self.make("a.png")
made = await self.add(name="Set", items=[self.item("a.png")])
cid = made["collection"]["id"]
def broken(*_args, **_kwargs):
raise sqlite3.OperationalError("disk I/O error at /secret/prompts.db")
cases = (
(
"list_collections",
"GET",
BASE,
None,
{},
"Collections could not be read.",
),
(
"file_collections",
"GET",
MEMBERSHIP,
None,
{"root_index": "0", "root_scope": self.scope(), "path": "a.png"},
"Collections could not be read.",
),
(
"add_to_collection",
"POST",
ITEMS,
{"collection_id": cid, "items": [self.item("a.png")]},
{},
"The collection could not be changed.",
),
(
"remove_from_collection",
"DELETE",
ITEMS,
{"collection_id": cid, "items": [self.item("a.png")]},
{},
"The collection could not be changed.",
),
(
"rename_collection",
"PATCH",
f"{BASE}/{cid}",
{"name": "x"},
{},
"The collection could not be changed.",
),
(
"delete_collection",
"DELETE",
f"{BASE}/{cid}",
None,
{},
"The collection could not be changed.",
),
)
for method_name, method, url, body, params, message in cases:
with self.subTest(method=method_name):
with patch.object(self.api.db, method_name, side_effect=broken):
data = await self.call(method, url, body, 500, **params)
self.assertEqual(data, {"success": False, "error": message})
class TestCollectionListingFilter(CollectionRoutesTestCase):
"""#392's "all N matching" filter carries ``collection`` like the listing."""
async def test_a_filter_with_a_collection_resolves_only_members(self):
second = self.two_sources()
self.make("a.png", "b.png")
self.make("c.png", root=second)
made = await self.add(
name="Set", items=[self.item("a.png"), self.item("c.png", 1)]
)
cid = str(made["collection"]["id"])
selection = await self.api.resolve_listing_filter({"collection": cid})
self.assertEqual(
sorted((index, path) for index, _scope, path in selection.targets),
[(0, "a.png"), (1, "c.png")],
)
selection = await self.api.resolve_listing_filter(
{"collection": cid, "root": "1", "root_scope": self.scope(1)}
)
self.assertEqual([t[2] for t in selection.targets], ["c.png"])
for bad in ("0", "x", "1234567890"):
with self.subTest(bad=bad):
with self.assertRaises(ValueError):
from py.api.images import parse_listing_filter
parse_listing_filter({"collection": bad})
async def test_bulk_category_acts_on_collection_members_only(self):
self.make("a.png", "b.png")
inside = self._save_prompt("Inside", category="Old")
outside = self._save_prompt("Outside", category="Old")
self.api.db.link_image_to_prompt(inside, str(self.output_dir / "a.png"))
self.api.db.link_image_to_prompt(outside, str(self.output_dir / "b.png"))
made = await self.add(name="Set", items=[self.item("a.png")])
cid = str(made["collection"]["id"])
data = await self.call(
"POST",
"/prompt_manager/bulk/category",
{"images_filter": {"collection": cid}, "category": "New"},
)
self.assertEqual((data["updated_count"], data["prompt_count"]), (1, 1))
self.assertEqual(self.api.db.get_prompt_by_id(inside)["category"], "New")
self.assertEqual(self.api.db.get_prompt_by_id(outside)["category"], "Old")
async def test_a_deleted_collection_in_a_bulk_filter_is_a_404(self):
self.make("a.png")
data = await self.call(
"POST",
"/prompt_manager/bulk/category",
{"images_filter": {"collection": "77"}, "category": "New"},
status=404,
)
self.assertEqual(data["code"], "collection-not-found")
with self.assertRaises(Exception) as caught:
await self.api.resolve_listing_filter({"collection": "77"})
self.assertEqual(type(caught.exception).__name__, "CollectionMissing")
class TestReviewFindings(CollectionRoutesTestCase):
"""PR #400 review: loop work, counts snapshot, id and query rules."""
async def test_no_realpath_runs_on_the_event_loop(self):
import os
import threading
names = [f"img{n}.png" for n in range(20)]
self.make(*names)
await self.call("GET", BASE) # warm the listing cache
loop_thread = threading.current_thread()
on_loop = []
real = os.path.realpath
def watched(*args, **kwargs):
if threading.current_thread() is loop_thread:
on_loop.append(args[0])
return real(*args, **kwargs)
items = [self.item(name) for name in names]
cid = None
with patch("os.path.realpath", side_effect=watched):
made = await self.add(name="Set", items=items)
cid = made["collection"]["id"]
await self.call("DELETE", ITEMS, {"collection_id": cid, "items": items[:5]})
await self.call(
"GET",
MEMBERSHIP,
root_index="0",
root_scope=items[0]["root_scope"],
path="img9.png",
)
await self.call("GET", BASE)
self.assertEqual(on_loop, [])
async def test_missing_is_never_negative(self):
self.make("a.png", "b.png")
made = await self.add(
name="Set", items=[self.item("a.png"), self.item("b.png")]
)
cid = made["collection"]["id"]
real = self.api.db.collection_count_snapshot
def raced(root_keys):
members, stored = real(root_keys)
return members, {cid: 1} # a remove landed between two reads
with patch.object(self.api.db, "collection_count_snapshot", side_effect=raced):
listed = await self.listed()
self.assertEqual((listed[0]["count"], listed[0]["missing"]), (2, 0))
async def test_path_ids_follow_the_listing_rule(self):
self.make("a.png")
made = await self.add(name="Set", items=[self.item("a.png")])
cid = made["collection"]["id"]
for raw in ("0", "007", f"0{cid}"):
with self.subTest(raw=raw):
await self.call("PATCH", f"{BASE}/{raw}", {"name": "x"}, 400)
await self.call("DELETE", f"{BASE}/{raw}", None, 400)
self.assertEqual(len(await self.listed()), 1)
async def test_membership_refuses_repeated_query_keys(self):
self.make("a.png", "b.png")
response = await self.client.get(
MEMBERSHIP
+ f"?root_index=0&root_scope={self.scope()}&path=a.png&path=b.png"
)
self.assertEqual(response.status, 400)
response = await self.client.get(
MEMBERSHIP
+ f"?root_index=0&root_index=0&root_scope={self.scope()}&path=a.png"
)
self.assertEqual(response.status, 400)
async def test_deleting_a_collection_leaves_everything_else(self):
from py.api.images import _favorite_root_key
try:
from tests.test_annotation_schema import insert_group_rows, table_rows
except ImportError:
from test_annotation_schema import insert_group_rows, table_rows
self.make("a.png")
db = self.api.db
db.set_image_favorite(_favorite_root_key(self.output_dir), "a.png")
prompt = self._save_prompt("Linked")
db.link_image_to_prompt(prompt, str(self.output_dir / "a.png"))
with db.model.get_connection() as conn:
insert_group_rows(conn, locator=str(self.output_dir / "a.png"))
made = await self.add(name="Set", items=[self.item("a.png")])
tables = (
"image_favorites",
"generated_images",
"prompts",
"tags",
"media_annotation_files",
"media_annotation_bytes",
"media_image_tags",
)
before = table_rows(db.model.get_connection(), tables)
file_bytes = (self.output_dir / "a.png").read_bytes()
await self.call("DELETE", f"{BASE}/{made['collection']['id']}")
self.assertEqual(table_rows(db.model.get_connection(), tables), before)
self.assertEqual((self.output_dir / "a.png").read_bytes(), file_bytes)
self.assertTrue(all(before[t] for t in tables))
+36
View File
@@ -693,6 +693,39 @@ class TestCleanupJob(DuplicateFilesTestCase):
"Moved 2 database links to the kept copies", "\n".join(cleanup.log)
)
async def test_duplicate_removal_moves_collection_memberships(self):
self.trio()
db = self.api.db
root_key = duplicate_files.favorite_root_key(self.output_dir)
grids = db.add_to_collection(name="LoRA grids", items=[(root_key, "b.png")])
grids = grids["collection"]["id"]
other = db.add_to_collection(
name="Other", items=[(root_key, "sub/c.png"), (root_key, "d.png")]
)["collection"]["id"]
db.add_to_collection(collection_id=grids, items=[(root_key, "sub/c.png")])
_plan, cleanup = await self.remove_all_but("a.png")
# The union: a.png joins both; c.png's second grids row was a repeat.
self.assertEqual(db.collection_members(grids, root_key), {"a.png"})
self.assertEqual(db.collection_members(other, root_key), {"a.png", "d.png"})
self.assertEqual(sorted(db.file_collections(root_key, "a.png")), [grids, other])
self.assertEqual(db.file_collections(root_key, "b.png"), [])
self.assertEqual(db.file_collections(root_key, "sub/c.png"), [])
self.assertIn(
"Moved 2 collection memberships to the kept copies", "\n".join(cleanup.log)
)
def test_one_collection_membership_is_singular(self):
job = type(
"Job", (), {"log": [], "log_line": lambda s, line: s.log.append(line)}
)()
totals = dict.fromkeys(
("moved", "merged", "favorites", "tags", "review", "collections"), 0
)
duplicate_files.DuplicateFilesMixin._log_moved_records(
job, {**totals, "collections": 1}
)
self.assertEqual(job.log, ["Moved 1 collection membership to the kept copies"])
async def test_a_failed_removal_leaves_its_links_on_the_kept_copy(self):
self.trio()
prompt = self.save("Stays on c")
@@ -849,8 +882,11 @@ class TestLinkQueries(DuplicateFilesTestCase):
with self.subTest(sql=sql):
self.assertNotIn("SCAN generated_images", plan)
self.assertNotIn("SCAN image_favorites", plan)
self.assertNotIn("SCAN image_collection_items", plan)
finally:
conn.close()
statements = duplicate_files.query_plan_sql()
self.assertTrue(any("image_collection_items" in sql for sql in statements))
class TestReviewHelpers(DuplicateFilesTestCase):
+87
View File
@@ -21,6 +21,8 @@ from tests.test_api_images import ImageAPITestCase, make_png
from tests.test_api_prompt_inspector import graph, png
from tests.test_execution_context_database import bundle as context_bundle
from tests.test_execution_history_database import JOB, SESSION
from tests.saver_anchor_graphs import BRANCH_TEXTS, shared_prefix
from tests.test_recorded_saver_anchor import LLM_TEXT, text_generate_run
from utils.execution_contexts import ContextBundle
from utils.execution_history_events import make_finish, make_start
from utils.execution_outputs import OutputBundle, OutputReference, snapshot_output_roots
@@ -1505,3 +1507,88 @@ class HistoryOutputAPITests(ImageAPITestCase):
body = await response.json()
self.assertEqual(response.status, 200, body)
self.assertEqual(body["prompts"]["positive"]["text"], "Readable portrait")
async def recorded_shared_prefix_run(self):
"""#386: the dataset shape, branch 3's saver wrote the file in a run."""
path = png(self.output_dir / "IMG_GEN_00371_.png", prompt=shared_prefix())
# make_start/make_finish observe the run at 0 ms: the file was written then.
os.utime(path, ns=(1_000_000, 1_000_000))
start = self.capture(
[self.ref(node="33", filename=path.name), self.ref("43", 1, "other.png")]
)
return start
async def test_both_inspect_routes_name_the_branch_recorded_history_saved(self):
start = await self.recorded_shared_prefix_run()
_, body = await self.outputs(start)
descriptor = next(
item["current_candidate"]["media"]
for item in body["items"]
if item["occurrence"]["output_node_id"] == "33"
)
response = await self.current(descriptor["inspection_url"])
recorded = await response.json()
self.assertEqual(response.status, 200, recorded)
response = await self.client.get(
"/prompt_manager/images/inspect/IMG_GEN_00371_.png", params={"root": "0"}
)
images = await response.json()
self.assertEqual(response.status, 200, images)
for result in (recorded, images):
self.assertEqual(result["resolution"]["state"], "resolved")
self.assertEqual(
result["prompts"]["positive"],
{
"status": "resolved",
"text": BRANCH_TEXTS[2],
"source": "recorded-run",
},
)
async def test_a_file_rewritten_after_its_run_keeps_several_prompts(self):
await self.recorded_shared_prefix_run()
os.utime(self.output_dir / "IMG_GEN_00371_.png") # written now: not that run
response = await self.client.get(
"/prompt_manager/images/inspect/IMG_GEN_00371_.png", params={"root": "0"}
)
body = await response.json()
self.assertEqual(response.status, 200, body)
self.assertEqual(body["resolution"]["state"], "ambiguous")
self.assertEqual(body["prompts"]["positive"]["status"], "ambiguous")
async def test_a_failing_history_lookup_judges_the_file_alone(self):
await self.recorded_shared_prefix_run()
with (
patch(
"database.run_attribution.recorded_output",
side_effect=sqlite3.OperationalError("locked"),
),
self.assertLogs("prompt_manager.run_attribution", "WARNING"),
):
response = await self.client.get(
"/prompt_manager/images/inspect/IMG_GEN_00371_.png",
params={"root": "0"},
)
body = await response.json()
self.assertEqual(response.status, 200, body)
self.assertEqual(body["resolution"]["state"], "ambiguous")
async def test_a_language_model_text_shows_what_promptmanager_recorded(self):
path = png(self.output_dir / "moss_00026_.png", prompt=text_generate_run())
os.utime(path, ns=(1_000_000, 1_000_000))
start = self.capture([self.ref(node="9", filename=path.name)])
self.api.db.write_execution_contexts(
start.run_id,
start.storage_epoch,
context_bundle("", LLM_TEXT, anchor="7"),
)
response = await self.client.get(
"/prompt_manager/images/inspect/moss_00026_.png", params={"root": "0"}
)
body = await response.json()
self.assertEqual(response.status, 200, body)
self.assertEqual(
body["prompts"]["positive"],
{"status": "resolved", "text": LLM_TEXT, "source": "execution"},
)
self.assertEqual(body["prompts"]["negative"]["text"], "")
+375
View File
@@ -0,0 +1,375 @@
"""Image category routes (#332): PUT/DELETE /images/category, GET /images/categories."""
import json
import ntpath
import sqlite3
from unittest.mock import patch
try:
from test_api_images import ImageAPITestCase, make_png
from open_handles import assert_closed
from test_image_categories import category_checksum
except ImportError:
from tests.test_api_images import ImageAPITestCase, make_png
from tests.open_handles import assert_closed
from tests.test_image_categories import category_checksum
from database.image_categories import MAX_KEY_LENGTH, TABLE
from py.api.images import _favorite_root_key, image_category_key
CATEGORY = "/prompt_manager/images/category"
CATEGORIES = "/prompt_manager/images/categories"
class ImageCategoryTestCase(ImageAPITestCase):
async def tearDownAsync(self):
self.api.db.close_all()
assert_closed(self, self._temp_db.name, self._output_tmp.name)
await super().tearDownAsync()
def target(self, path, index=0):
roots = self.api._get_all_output_dirs()
return {
"root_index": index,
"root_scope": self.api._gallery_root_scope(roots[index], index),
"path": path,
}
async def send(self, method, body, status=200):
response = await self.client.request(method, CATEGORY, json=body)
text = await response.text()
self.assertEqual(response.status, status, text)
return json.loads(text)
async def put(self, paths, category, status=200):
items = [p if isinstance(p, dict) else self.target(p) for p in paths]
return await self.send("PUT", {"items": items, "category": category}, status)
async def delete(self, paths, status=200):
items = [p if isinstance(p, dict) else self.target(p) for p in paths]
return await self.send("DELETE", {"items": items}, status)
async def listed(self, status=200):
response = await self.client.get(CATEGORIES)
text = await response.text()
self.assertEqual(response.status, status, text)
return json.loads(text)
def rows(self):
return [
tuple(row)
for row in self.api.db.model.get_connection().execute(
f"SELECT root_key, relative_path, category FROM {TABLE} "
"ORDER BY root_key, relative_path"
)
]
def key(self, relative, root=None):
roots = self.api._get_all_output_dirs()
return image_category_key(roots, (root or self.output_dir) / relative)
class TestSetAndClear(ImageCategoryTestCase):
async def test_put_sets_listed_and_counts_missing(self):
make_png(self.output_dir / "a.png")
make_png(self.output_dir / "sub" / "b.png")
data = await self.put(["a.png", "sub/b.png", "missing.png"], " Portraits ")
self.assertEqual(
data,
{
"success": True,
"category": "Portraits",
"updated": 2,
"unchanged": 0,
"not_found": 1,
},
)
root_key = _favorite_root_key(self.output_dir)
self.assertEqual(
self.rows(),
[(root_key, "a.png", "Portraits"), (root_key, "sub/b.png", "Portraits")],
)
data = await self.put(["a.png", "sub/b.png"], "Portraits")
self.assertEqual((data["updated"], data["unchanged"]), (0, 2))
data = await self.put(["a.png"], "Scans")
self.assertEqual((data["updated"], data["unchanged"]), (1, 0))
self.assertEqual(self.api.db.image_category(*self.key("a.png")), "Scans")
async def test_put_with_no_listed_file_is_404_and_writes_nothing(self):
data = await self.put(["missing.png", "gone/x.png"], "Scans", status=404)
self.assertEqual(data["code"], "files-not-found")
self.assertEqual(self.rows(), [])
async def test_duplicate_targets_are_written_and_counted_once(self):
make_png(self.output_dir / "a.png")
item = self.target("a.png")
data = await self.put([item, dict(item)], "Scans")
self.assertEqual((data["updated"], data["not_found"]), (1, 0))
async def test_delete_clears_listed_and_missing_files(self):
make_png(self.output_dir / "a.png")
await self.put(["a.png"], "Scans")
self.api.db.set_image_categories([self.key("gone.png")], "Scans")
data = await self.delete(["a.png", "gone.png", "never.png"])
self.assertEqual(data, {"success": True, "cleared": 2})
self.assertEqual(self.rows(), [])
self.assertEqual((await self.delete(["a.png"]))["cleared"], 0)
async def test_put_never_writes_prompts(self):
make_png(self.output_dir / "a.png")
make_png(self.output_dir / "b.png")
positive = self._save_prompt("castle at dusk", category="Landscapes")
self.api.db.link_image_to_prompt(positive, str(self.output_dir / "a.png"))
negative = self._save_prompt("blurry", category="Negatives")
self.api.db.save_prompt_pair("castle at dusk", "blurry", category="Recipes")
conn = self.api.db.model.get_connection()
before = category_checksum(conn)
links = conn.execute("SELECT * FROM generated_images").fetchall()
await self.put(["a.png", "b.png"], "Portraits")
await self.listed()
await self.delete(["a.png"])
self.assertEqual(category_checksum(conn), before)
self.assertEqual(
conn.execute("SELECT * FROM generated_images").fetchall(), links
)
self.assertIsNotNone(negative)
async def test_two_sources_in_one_request(self):
second = self.comfy_dir / "renders"
make_png(self.output_dir / "a.png")
make_png(second / "a.png")
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
data = await self.put([self.target("a.png", 0), self.target("a.png", 1)], "X")
self.assertEqual(data["updated"], 2)
self.assertEqual(
{row[0] for row in self.rows()},
{_favorite_root_key(self.output_dir), _favorite_root_key(second)},
)
self.assertEqual((await self.delete([self.target("a.png", 1)]))["cleared"], 1)
async def test_nested_sources_key_by_the_deepest(self):
inner = self.output_dir / "nested"
make_png(inner / "a.png")
self.api._get_all_output_dirs = lambda: [self.output_dir, inner]
# Named through the outer source, stored under the inner one.
data = await self.put([self.target("nested/a.png", 0)], "Scans")
self.assertEqual(data["updated"], 1)
self.assertEqual(self.rows(), [(_favorite_root_key(inner), "a.png", "Scans")])
# The same file through the inner source is the same key.
data = await self.put([self.target("a.png", 1)], "Scans")
self.assertEqual(data["unchanged"], 1)
self.assertEqual((await self.listed())["categories"], ["Scans"])
self.assertEqual((await self.delete([self.target("a.png", 1)]))["cleared"], 1)
async def test_video_files_can_have_a_category(self):
(self.output_dir / "clip.mp4").write_bytes(b"\x00\x00\x00\x18ftypmp42")
data = await self.put(["clip.mp4"], "Clips")
self.assertEqual(data["updated"], 1)
async def test_windows_case_folds_the_key(self):
make_png(self.output_dir / "Sub" / "A.png")
with patch("py.api.images._path_flavor", ntpath):
data = await self.put(["Sub/A.png"], "Scans")
self.assertEqual(data["updated"], 1)
data = await self.put(["sub/a.PNG"], "Scans")
self.assertEqual(data["unchanged"], 1)
self.assertEqual(self.rows()[0][1], "sub/a.png")
self.assertEqual((await self.delete(["SUB/A.png"]))["cleared"], 1)
async def test_a_key_beyond_the_stored_bound_is_403(self):
make_png(self.output_dir / "a.png")
too_long = "/" + "r" * MAX_KEY_LENGTH
with patch("py.api.images._favorite_root_key", return_value=too_long):
data = await self.put(["a.png"], "Scans", status=403)
self.assertNotIn(too_long, json.dumps(data))
await self.delete(["a.png"], status=403)
self.assertEqual(self.rows(), [])
async def test_an_unexpected_failure_is_a_logged_500(self):
make_png(self.output_dir / "a.png")
def broken(*_args, **_kwargs):
raise sqlite3.OperationalError("secret detail")
self.api.db.set_image_categories = broken
self.api.db.clear_image_categories = broken
with self.assertLogs("prompt_manager", level="ERROR"):
data = await self.put(["a.png"], "Scans", status=500)
self.assertEqual(data["error"], "The category could not be changed.")
self.assertNotIn("secret", json.dumps(data))
with self.assertLogs("prompt_manager", level="ERROR"):
data = await self.delete(["a.png"], status=500)
self.assertNotIn("secret", json.dumps(data))
async def test_a_symlinked_output_folder_keys_like_its_favorites(self):
"""A detected output folder can be a symlink: keys use its real path."""
real = self.comfy_dir / "real-output"
make_png(real / "a.png")
link = self.comfy_dir / "linked-output"
self._make_symlink(real, link, target_is_directory=True)
self.api._get_all_output_dirs = lambda: [link]
data = await self.put(["a.png"], "Scans")
self.assertEqual(data["updated"], 1)
self.assertEqual(self.rows(), [(_favorite_root_key(link), "a.png", "Scans")])
self.assertEqual((await self.listed())["categories"], ["Scans"])
self.assertEqual((await self.delete(["a.png"]))["cleared"], 1)
async def test_lone_surrogates_are_400_not_500(self):
make_png(self.output_dir / "a.png")
await self.send(
"PUT", {"items": [self.target("a.png")], "category": "x\ud800"}, 400
)
await self.send("DELETE", {"items": [self.target("a\ud800.png")]}, 400)
self.assertEqual(self.rows(), [])
async def test_delete_category_never_reaches_delete_image(self):
make_png(self.output_dir / "a.png")
calls = []
async def recorded(request):
calls.append(request.match_info)
raise AssertionError("delete_image reached")
self.api.delete_image = recorded
await self.put(["a.png"], "Scans")
data = await self.delete(["a.png"])
self.assertEqual(data["cleared"], 1)
self.assertEqual(calls, [])
class TestRefusals(ImageCategoryTestCase):
async def test_malformed_bodies_are_400(self):
make_png(self.output_dir / "a.png")
item = self.target("a.png")
cases = [
{"items": [item]},
{"category": "x"},
{"items": [item], "category": "x", "extra": 1},
{"items": item, "category": "x"},
{"items": [], "category": "x"},
{"items": [item] * 501, "category": "x"},
{"items": [item], "category": 3},
{"items": [item], "category": None},
{"items": [item], "category": " \t "},
{"items": [item], "category": "x" * 101},
{"items": [item], "category": "a\x07b"},
{"items": [{**item, "root_index": "0"}], "category": "x"},
{"items": [{**item, "root_scope": "nothex"}], "category": "x"},
{"items": [{"root_index": 0, "path": "a.png"}], "category": "x"},
{"items": [{**item, "root_index": 1}], "category": "x"},
]
for body in cases:
with self.subTest(body=str(body)[:80]):
data = await self.send("PUT", body, status=400)
self.assertFalse(data["success"])
await self.send("DELETE", {"items": [item], "category": "x"}, status=400)
await self.send("DELETE", {"items": []}, status=400)
response = await self.client.request(
"PUT", CATEGORY, data="not json", headers={"Content-Type": "text/plain"}
)
self.assertEqual(response.status, 400)
self.assertEqual(self.rows(), [])
async def test_unsafe_paths_are_403(self):
(self.output_dir / "notes.txt").write_text("not media")
for path in ("../outside.png", "/a.png", "C:/a.png", "a\\b.png", "a/./b.png"):
with self.subTest(path=path):
await self.put([self.target(path)], "x", status=403)
data = await self.put([self.target("notes.txt")], "x", status=403)
self.assertEqual(data["error"], "Only media files can have a category")
await self.delete([self.target("notes.txt")], status=403)
async def test_a_stale_scope_is_409(self):
make_png(self.output_dir / "a.png")
item = {**self.target("a.png"), "root_scope": "0" * 64}
data = await self.put([item], "x", status=409)
self.assertEqual(data["code"], "source-changed")
data = await self.delete([item], status=409)
self.assertEqual(data["code"], "source-changed")
async def test_an_oversized_body_is_413(self):
response = await self.client.request(
"PUT",
CATEGORY,
data=b'{"items": [], "category": "' + b"x" * (2 * 1024 * 1024) + b'"}',
headers={"Content-Type": "application/json"},
)
self.assertEqual(response.status, 413)
class TestCategoryList(ImageCategoryTestCase):
async def test_lists_only_values_whose_file_is_listed(self):
make_png(self.output_dir / "a.png")
make_png(self.output_dir / "c.png")
await self.put(["a.png"], "B")
await self.put(["c.png"], "A")
self.api.db.set_image_categories([self.key("gone.png")], "Z")
self.assertEqual(
await self.listed(), {"success": True, "categories": ["A", "B"]}
)
async def test_empty_library_lists_nothing(self):
self.assertEqual((await self.listed())["categories"], [])
async def test_a_failed_read_is_500(self):
make_png(self.output_dir / "a.png")
def broken(*_args, **_kwargs):
raise sqlite3.OperationalError("secret detail")
self.api.db.own_categories_of = broken
with self.assertLogs("prompt_manager", level="ERROR"):
data = await self.listed(status=500)
self.assertEqual(data["error"], "Image categories could not be read.")
self.assertNotIn("secret", json.dumps(data))
class TestAllMatchingFilter(ImageCategoryTestCase):
"""#392 "Select all N matching": the listing filter, resolved now."""
async def test_filter_sets_and_clears_every_match(self):
for name in ("bulk/a.png", "bulk/b.png", "other/c.png"):
make_png(self.output_dir / name)
data = await self.send(
"PUT", {"filter": {"subfolder": "bulk"}, "category": "Scans"}
)
self.assertEqual((data["updated"], data["not_found"]), (2, 0))
self.assertEqual(
sorted(row[1] for row in self.rows()), ["bulk/a.png", "bulk/b.png"]
)
data = await self.send("DELETE", {"filter": {"subfolder": "bulk"}})
self.assertEqual(data["cleared"], 2)
data = await self.send(
"PUT", {"filter": {"subfolder": "none"}, "category": "Scans"}, 404
)
self.assertEqual(data["code"], "files-not-found")
self.assertEqual(
(await self.send("DELETE", {"filter": {"q": "zz"}}))["cleared"], 0
)
async def test_filter_refusals(self):
make_png(self.output_dir / "a.png")
item = self.target("a.png")
for body in (
{"filter": {}, "items": [item], "category": "x"},
{"filter": {"unknown": "1"}, "category": "x"},
{"filter": "q=a", "category": "x"},
{"filter": {"root": "0", "root_scope": "nothex"}, "category": "x"},
):
with self.subTest(body=body):
await self.send("PUT", body, 400)
data = await self.send(
"PUT",
{"filter": {"root": "0", "root_scope": "0" * 64}, "category": "x"},
409,
)
self.assertEqual(data["code"], "source-changed")
self.assertEqual(self.rows(), [])
async def test_more_than_500_matching_is_refused(self):
for number in range(501):
make_png(self.output_dir / f"many/{number:03}.png", size=(1, 1))
data = await self.send("PUT", {"filter": {}, "category": "x"}, 400)
self.assertEqual(data["code"], "too-many-images")
self.assertEqual(data["error"], "Select at most 500 images.")
self.assertEqual(self.rows(), [])
+74 -2
View File
@@ -62,12 +62,16 @@ class TestVocabularyAPI(ImageAPITestCase):
},
)
self.assertEqual(
categories["categories"], [{"value": "Animal", "prompt_count": 1}]
categories["categories"],
[{"value": "Animal", "prompt_count": 1, "image_count": 0}],
)
_, rest = await self.get(
"categories", limit=1, cursor=categories["next_cursor"]
)
self.assertEqual(rest["categories"], [{"value": "Pet", "prompt_count": 1}])
self.assertEqual(
rest["categories"],
[{"value": "Pet", "prompt_count": 1, "image_count": 0}],
)
self.assertFalse(rest["has_more"])
async def test_vocabulary_cursor_is_signed_epoch_scoped(self):
@@ -488,6 +492,7 @@ class TestVocabularyAPI(ImageAPITestCase):
"other_prompts": 0,
"negative_variants": 1,
"archived_negative_variants": 0,
"images": 0,
}
self.assertEqual(
preview,
@@ -528,6 +533,73 @@ class TestVocabularyAPI(ImageAPITestCase):
self.assertEqual(response.status, 409)
self.assertEqual(replay["code"], "vocabulary-conflict")
async def test_category_rows_report_images_and_image_only_values(self):
self._save_prompt("Owner", category="Prompt value")
_, options_before = await self.get_options()
self.api.db.set_image_categories(
[("root", "a.png"), ("root", "b.png")], "Prompt value"
)
self.api.db.set_image_categories([("root", "c.png")], "Image value")
response, first = await self.get("categories", limit=1)
self.assertEqual(response.status, 200)
self.assertEqual(
first["categories"],
[{"value": "Prompt value", "prompt_count": 1, "image_count": 2}],
)
self.assertEqual(first["total"], 2)
_, rest = await self.get("categories", limit=1, cursor=first["next_cursor"])
self.assertEqual(
rest["categories"],
[{"value": "Image value", "prompt_count": 0, "image_count": 1}],
)
self.assertFalse(rest["has_more"])
# The Prompt library's category list does not gain image values.
_, options_after = await self.get_options()
self.assertEqual(options_after, options_before)
self.assertEqual(options_after["categories"], ["Prompt value"])
async def get_options(self):
response = await self.client.get("/prompt_manager/library/options")
return response, await response.json()
async def test_category_routes_rename_image_only_value(self):
epoch = self.api.db.get_execution_epoch()
self._save_prompt("Bystander", category="Kept")
self.api.db.set_image_categories(
[("root", "a.png"), ("second", "b.png")], "Image value"
)
prompts_before = self.category_rows()
body = {"expected_epoch": epoch, "source": "Image value", "target": "Moved"}
response, preview = await self.send("POST", "categories/preview", body)
self.assertEqual(response.status, 200, preview)
impact = {
"positive_prompts": 0,
"other_prompts": 0,
"negative_variants": 0,
"archived_negative_variants": 0,
"images": 2,
}
self.assertEqual(preview["impact"], impact)
stale = preview["impact_token"]
self.api.db.set_image_categories([("root", "late.png")], "Image value")
response, refused = await self.send(
"POST", "categories/apply", {**body, "impact_token": stale}
)
self.assertEqual(response.status, 409)
self.assertEqual(refused["code"], "vocabulary-conflict")
self.assertEqual(self.api.db.image_category("root", "a.png"), "Image value")
_, preview = await self.send("POST", "categories/preview", body)
response, applied = await self.send(
"POST",
"categories/apply",
{**body, "impact_token": preview["impact_token"]},
)
self.assertEqual(response.status, 200, applied)
self.assertEqual(applied["impact"], {**impact, "images": 3})
for key in (("root", "a.png"), ("second", "b.png"), ("root", "late.png")):
self.assertEqual(self.api.db.image_category(*key), "Moved")
self.assertEqual(self.category_rows(), prompts_before)
async def test_preview_validation(self):
epoch = self.api.db.get_execution_epoch()
self._save_prompt("Owner", category="Scenery")
@@ -0,0 +1,262 @@
"""Images filters by the effective category (#332 MC4, MC5, DC6).
An image's effective category is its own category when it has one, else the
category its link supplies (the recipe's when set, else the positive's), else
none. The category filter, "select all N matching" and saved views match it.
"""
import os
import sqlite3
from unittest.mock import patch
try:
import test_api_workspace_images as workspace
from test_image_categories import category_checksum
except ImportError:
from tests import test_api_workspace_images as workspace
from tests.test_image_categories import category_checksum
from py.api.images import image_category_key, parse_listing_filter
class WorkspaceCategoryCase(workspace.ImageAPITestCase):
"""Listing helpers plus own-category writes keyed as the routes key them."""
media = workspace.TestWorkspaceImages.media
get = workspace.TestWorkspaceImages.get
tag = workspace.TestWorkspaceImageTags.tag
names = workspace.TestWorkspaceImageTags.names
async def tearDownAsync(self):
self.api.db.close_all()
workspace.assert_closed(self, self._temp_db.name, self._output_tmp.name)
await super().tearDownAsync()
def own(self, path, value, sources=None):
sources = sources or self.api._get_all_output_dirs()
self.api.db.set_image_categories([image_category_key(sources, path)], value)
async def listed(self, **params):
response, body = await self.get(sort="name", **params)
self.assertEqual(response.status, 200, body)
return self.names(body), body
class TestEffectiveCategoryFilter(WorkspaceCategoryCase):
def batch(self):
"""Four images of "castle at dusk" (Landscapes); one is overridden."""
castle = self._save_prompt("castle at dusk", category="Landscapes", rating=4)
files = [self.media(f"castle_{n}.png", mtime=100 + n) for n in range(4)]
for path in files:
self.api.db.link_image_to_prompt(castle, str(path))
self.own(files[0], "Portraits")
return castle, files
async def test_category_filter_prefers_own_category(self):
castle, _files = self.batch()
conn = self.api.db.model.get_connection()
before = category_checksum(conn)
names, body = await self.listed(category="Portraits")
self.assertEqual(names, ["castle_0.png"])
# The caption still counts the prompt behind the image shown.
self.assertEqual(body["positive_count"], 1)
names, body = await self.listed(category="Landscapes")
self.assertEqual(names, ["castle_1.png", "castle_2.png", "castle_3.png"])
self.assertEqual(body["total"], 3)
names, _ = await self.listed()
self.assertEqual(len(names), 4)
self.assertEqual(category_checksum(conn), before)
self.assertEqual(self.api.db.get_prompt_by_id(castle)["category"], "Landscapes")
async def test_category_filter_lists_unlinked_with_own(self):
scans = [self.media(f"scan_{n}.png", mtime=100 + n) for n in range(2)]
self.media("plain.png", mtime=50)
for path in scans:
self.own(path, "Scans")
names, body = await self.listed(category="Scans")
self.assertEqual(names, ["scan_0.png", "scan_1.png"])
self.assertEqual(body["positive_count"], 0)
async def test_recipe_category_is_inherited(self):
"""MC4: a recipe's own category, not the positive's, is inherited."""
db = self.api.db
self._save_prompt("Pair positive", category="Night")
pair = db.save_prompt_pair("Pair positive", "Noise", category="Recipes")
blank = db.save_prompt_pair("Pair positive", "Other noise", category=" ")
recipe = self.media("recipe.png", mtime=103)
blank_recipe = self.media("blank_recipe.png", mtime=102)
positive = db.get_prompt_pair(pair)["positive_prompt_id"]
for path, pair_id in ((recipe, pair), (blank_recipe, blank)):
image = db.link_image_to_prompt(positive, str(path))
db.link_image_to_pair(pair_id, image, {"source": "fixture"})
self.assertEqual((await self.listed(category="Recipes"))[0], ["recipe.png"])
# A blank recipe category falls back to the positive's.
self.assertEqual((await self.listed(category="Night"))[0], ["blank_recipe.png"])
# An own category still wins over the recipe's.
self.own(recipe, "Mine")
self.assertEqual((await self.listed(category="Recipes"))[0], [])
self.assertEqual((await self.listed(category="Mine"))[0], ["recipe.png"])
async def test_ambiguous_files_inherit_nothing_but_keep_their_own(self):
first = self._save_prompt("Owner one", category="Night")
second = self._save_prompt("Owner two", category="Night")
ambiguous = self.media("ambiguous.png")
for owner in (first, second):
self.api.db.link_image_to_prompt(owner, str(ambiguous))
self.assertEqual((await self.listed(category="Night"))[0], [])
self.own(ambiguous, "Night")
self.assertEqual((await self.listed(category="Night"))[0], ["ambiguous.png"])
async def test_combines_with_tags_search_and_rating(self):
stone = self._save_prompt(
"castle at dusk", category="Landscapes", tags=["stone"], rating=4
)
files = [self.media(f"castle_{n}.png", mtime=100 + n) for n in range(2)]
for path in files:
self.api.db.link_image_to_prompt(stone, str(path))
self.own(files[0], "Portraits")
other = self._save_prompt("harbor at noon", category="Portraits", rating=2)
harbor = self.media("harbor.png", mtime=90)
self.api.db.link_image_to_prompt(other, str(harbor))
names, _ = await self.listed(category="Portraits", tags="stone")
self.assertEqual(names, ["castle_0.png"])
names, _ = await self.listed(category="Portraits", q="harbor")
self.assertEqual(names, ["harbor.png"])
names, _ = await self.listed(category="Portraits", q="castle_0")
self.assertEqual(names, ["castle_0.png"])
response, body = await self.get(category="Portraits", sort="rating")
self.assertEqual(response.status, 200, body)
self.assertEqual(self.names(body), ["castle_0.png", "harbor.png"])
self.assertEqual(body["positive_count"], 2)
async def test_image_tag_admitted_file_keeps_its_inherited_category(self):
night = self._save_prompt("Night owl", category="Night", tags=["a"])
tagged = self.media("tagged.png", mtime=103)
unlinked = self.media("unlinked.png", mtime=102)
self.api.db.link_image_to_prompt(night, str(tagged))
self.tag(tagged, ["b"])
self.tag(unlinked, ["a", "b"])
names, _ = await self.listed(tags="a,b", category="Night")
self.assertEqual(names, ["tagged.png"])
self.own(unlinked, "Night")
names, _ = await self.listed(tags="a,b", category="Night")
self.assertEqual(names, ["tagged.png", "unlinked.png"])
async def test_cursor_pages_walk_only_the_category(self):
files = [self.media(f"img_{n:02}.png", mtime=100 + n) for n in range(7)]
for path in files[::2]:
self.own(path, "Even")
seen, cursor = [], None
while True:
params = {"category": "Even", "limit": 2, "sort": "name"}
if cursor:
params["cursor"] = cursor
response, body = await self.get(**params)
self.assertEqual(response.status, 200, body)
seen += self.names(body)
self.assertEqual(body["total"], 4)
cursor = body["next_cursor"]
if not cursor:
break
self.assertEqual(seen, ["img_00.png", "img_02.png", "img_04.png", "img_06.png"])
async def test_an_outer_source_lists_a_nested_sources_own_category(self):
inner = self.output_dir / "nested"
nested = self.media("nested/a.png")
self.media("b.png")
self.api._get_all_output_dirs = lambda: [self.output_dir, inner]
self.own(nested, "Scans")
names, _ = await self.listed(category="Scans", root="0")
self.assertEqual(names, ["nested/a.png"])
names, _ = await self.listed(category="Scans", root="1")
self.assertEqual(names, ["a.png"])
# The outer source's own key for that path is not the file's key.
from py.api.images import _favorite_root_key
self.api.db.clear_image_categories([image_category_key([inner], nested)])
self.api.db.set_image_categories(
[(_favorite_root_key(self.output_dir), "nested/a.png")], "Stale"
)
self.assertEqual((await self.listed(category="Stale"))[0], [])
async def test_a_gone_files_category_is_not_listed(self):
gone = self.media("gone.png")
self.own(gone, "Scans")
os.unlink(gone)
self.api.invalidate_gallery_cache()
self.assertEqual((await self.listed(category="Scans"))[0], [])
async def test_a_failed_own_category_read_is_listing_failed(self):
self.media("a.png")
def broken(*_args, **_kwargs):
raise sqlite3.OperationalError("secret detail")
with patch.object(self.api.db, "own_categories_of", side_effect=broken):
with self.assertLogs("prompt_manager", level="ERROR"):
response, body = await self.get(category="Scans")
self.assertEqual(response.status, 500)
self.assertNotIn("secret", str(body))
# Without a category nothing new is read.
response, body = await self.get()
self.assertEqual(response.status, 200, body)
async def test_select_all_matching_uses_the_effective_category(self):
_castle, files = self.batch()
resolved = await self.api.resolve_listing_filter(
parse_listing_filter({"category": "Portraits"}), 100
)
self.assertEqual([t[2] for t in resolved.targets], ["castle_0.png"])
resolved = await self.api.resolve_listing_filter(
parse_listing_filter({"category": "Landscapes"}), 100
)
self.assertEqual(len(resolved.targets), 3)
self.assertEqual(len(files), 4)
async def test_saved_view_category_means_the_effective_category(self):
"""CUJ 8: a saved view's category is the listing's query value."""
_castle, files = self.batch()
self.own(files[1], "Portraits")
names, _ = await self.listed(category="Portraits")
self.assertEqual(names, ["castle_0.png", "castle_1.png"])
class TestFilterAndChipAgreeOnBlanks(WorkspaceCategoryCase):
"""The filter and the inspector use one inherited-category rule."""
async def chip(self, path):
response = await self.client.get(
f"/prompt_manager/images/inspect/{path}", params={"root": "0"}
)
body = await response.json()
self.assertEqual(response.status, 200, body)
return body["library"]["category"]
async def test_blank_recipe_values_agree(self):
db = self.api.db
positive = self._save_prompt("Pair positive", category="Night")
blanks = {"spaces": " ", "tab": "\t", "nbsp": " ", "set": "Recipes"}
for index, (label, value) in enumerate(blanks.items()):
pair = db.save_prompt_pair("Pair positive", f"noise {label}")
path = self.media(f"{label}.png", mtime=100 + index)
image = db.link_image_to_prompt(positive, str(path))
db.link_image_to_pair(pair, image, {"source": "fixture"})
with db.model.get_connection() as conn:
conn.execute(
"UPDATE prompt_pairs SET category=? WHERE id=?", (value, pair)
)
blob = db.save_prompt_pair("Pair positive", "noise blob")
path = self.media("blob.png", mtime=90)
image = db.link_image_to_prompt(positive, str(path))
db.link_image_to_pair(blob, image, {"source": "fixture"})
with db.model.get_connection() as conn:
conn.execute(
"UPDATE prompt_pairs SET category=CAST(' ' AS BLOB) WHERE id=?",
(blob,),
)
for name in ("spaces", "tab", "nbsp", "set", "blob"):
with self.subTest(name=name):
shown = await self.chip(f"{name}.png")
for value in ("Night", "Recipes"):
listed = f"{name}.png" in (await self.listed(category=value))[0]
self.assertEqual(listed, shown == value, (name, value, shown))
+273 -13
View File
@@ -1347,7 +1347,7 @@ class TestWorkspaceImages(ImageAPITestCase):
_, data = await self.get(q="harbor")
self.assertEqual(self._names(data), ["night_a.png"])
self.assertEqual(
self.api.db.get_image_paths_matching("zebra", None, []),
self.api.db.get_image_paths_matching("zebra", []),
{
normalize_image_path(str(files[name]))
for name in ("night_ab", "day_ab", "titled")
@@ -1407,8 +1407,10 @@ class TestWorkspaceImages(ImageAPITestCase):
for query in statements
if query.lstrip().upper().startswith(("SELECT", "WITH"))
]
# Plus the one stream of tagged image records q and tags read (B3).
self.assertEqual(len(selects), 4, selects)
# Plus the one stream of tagged image records q and tags read (B3),
# and the category's own-category read of the one source (#332).
self.assertEqual(len(selects), 5, selects)
self.assertEqual(sum("image_categories" in q for q in selects), 1, selects)
self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects)
self.assertEqual(sum("media_image_tags" in q for q in selects), 1, selects)
@@ -1528,15 +1530,16 @@ class TestWorkspaceImages(ImageAPITestCase):
elsewhere = self.api.db.get_prompt_id_for_image("/elsewhere/7.png")
self.assertEqual(
self.api.db.get_image_link_matches(
"", "Night", [], every_link=True, wanted=wanted
"", [], every_link=True, wanted=wanted, category="Night"
),
{
normalize_image_path(str(files["night_ab"])): (
True,
5,
self.api.db.get_prompt_id_for_image(str(files["night_ab"])),
"Night",
),
"/elsewhere/7.png": (True, 5, elsewhere),
"/elsewhere/7.png": (True, 5, elsewhere, "Night"),
},
)
@@ -1588,6 +1591,24 @@ class TestWorkspaceImages(ImageAPITestCase):
for i in range(links)
],
)
# 2,000 images (one in 25 at 50,000 files) carry their own category
# (#332), so a category filter also reads and matches those.
from py.api.images import _favorite_root_key
root_key = _favorite_root_key(self.output_dir)
with self.api.db.model.get_connection() as conn:
conn.executemany(
"INSERT INTO image_categories VALUES(?,?,?,?)",
[
(
root_key,
listed[i][0].relative_to(self.output_dir).as_posix(),
f"cat-{i % 10}",
"2026-10-27T00:00:00+00:00",
)
for i in range(1, files, 25)[:2_000]
],
)
return listed
_BUDGET_REQUESTS = (
@@ -1615,7 +1636,7 @@ class TestWorkspaceImages(ImageAPITestCase):
import math
from py.api import images as images_module
from database.operations import link_filter_query
from database.operations import link_filter_query, link_matches_json_query
listed = await asyncio.get_running_loop().run_in_executor(
None, self._budget_fixture, 5_000, 2_000
@@ -1678,7 +1699,13 @@ class TestWorkspaceImages(ImageAPITestCase):
# q and tags also stream the tagged image records once
# (#249 B3); none are tagged here, so nothing is stat'ed.
streamed = "q" in params or "tags" in params
self.assertEqual(len(found), 2 if streamed else 1, found)
# A category also reads the own categories, one primary
# key range per source (#332).
own = "category" in params
self.assertEqual(len(found), 1 + streamed + own, found)
self.assertEqual(
sum("image_categories" in query for query in found), int(own)
)
self.assertEqual(
sum("media_image_tags" in query for query in found),
int(streamed),
@@ -1738,13 +1765,42 @@ class TestWorkspaceImages(ImageAPITestCase):
self.assertEqual(self.api._workspace_listed_cache, {})
self.assertEqual(self.api._workspace_order_cache, {})
# The per-row ambiguity probe is an index search, never a scan.
sql, params = link_filter_query("lantern", "cat-3", ["tag-1"], every_link=True)
with self.api.db.model.get_connection() as conn:
plan = [row[3] for row in conn.execute("EXPLAIN QUERY PLAN " + sql, params)]
self.assertIn(
"SEARCH other USING COVERING INDEX idx_image_file_path (file_path=?)", plan
# The read runs it wrapped in one JSON aggregate (#354): same plan.
sql, params = link_filter_query("lantern", ["tag-1"], every_link=True)
for query in (sql, link_matches_json_query(sql)):
with self.subTest(wrapped=query is not sql):
with self.api.db.model.get_connection() as conn:
plan = [
row[3]
for row in conn.execute("EXPLAIN QUERY PLAN " + query, params)
]
self.assertIn(
"SEARCH other USING COVERING INDEX idx_image_file_path "
"(file_path=?)",
plan,
)
self.assertFalse(
[step for step in plan if step.startswith("SCAN other")]
)
# The category's narrowing clause (#332) keeps every join an index
# search, plain and wrapped with the raw categories.
sql, params = link_filter_query(
"", ["tag-1"], every_link=True, category="cat-3"
)
self.assertFalse([step for step in plan if step.startswith("SCAN other")])
for query in (sql, link_matches_json_query(sql, categories=True)):
with self.subTest(category_wrapped=query is not sql):
with self.api.db.model.get_connection() as conn:
plan = [
row[3]
for row in conn.execute("EXPLAIN QUERY PLAN " + query, params)
]
self.assertIn(
"SEARCH other USING COVERING INDEX idx_image_file_path "
"(file_path=?)",
plan,
)
scans = [step for step in plan if step.startswith("SCAN")]
self.assertTrue(all(step.startswith("SCAN gi") for step in scans), plan)
@unittest.skipUnless(
PERF_BUDGET and not TRACED,
@@ -2494,3 +2550,207 @@ class TestListingFilterSelection(ImageAPITestCase):
self.assertEqual(response.status, 409)
self.assertEqual((await response.json())["code"], "source-changed")
class TestWorkspaceCollectionFilter(ImageAPITestCase):
"""#331: ``collection=<id>`` keeps only that collection's listed members."""
media = TestWorkspaceImages.media
get = TestWorkspaceImages.get
_names = TestWorkspaceImages._names
async def tearDownAsync(self):
self.api.db.close_all()
assert_closed(self, self._temp_db.name, self._output_tmp.name)
await super().tearDownAsync()
def keys(self, root, relative):
from py.api.images import _favorite_relative_key, _favorite_root_key
return (_favorite_root_key(root), _favorite_relative_key(relative))
def collect(self, name, *members):
"""A collection holding ``(root, relative)`` members; its id."""
result = self.api.db.add_to_collection(
name=name, items=[self.keys(root, rel) for root, rel in members]
)
return result["collection"]["id"]
def two_sources(self):
second = self.comfy_dir / "second"
second.mkdir()
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
return second
async def test_listing_keeps_only_collection_members(self):
second = self.two_sources()
self.media("a.png", 300)
self.media("b.png", 200)
self.media("sub/c.png", 100, root=second)
self.media("d.png", 50, root=second)
grids = self.collect(
"Grids",
(self.output_dir, "a.png"),
(second, "sub/c.png"),
(self.output_dir, "gone.png"), # stored, never listed
)
self.collect("Other", (self.output_dir, "b.png"))
response, data = await self.get(collection=str(grids))
self.assertEqual(response.status, 200, data)
self.assertEqual(self._names(data), ["a.png", "c.png"])
self.assertEqual(data["total"], 2)
_, data = await self.get(collection=str(grids), root="1")
self.assertEqual(self._names(data), ["c.png"])
# Without the filter the listing is unchanged.
_, data = await self.get()
self.assertEqual(data["total"], 4)
async def test_collection_filter_combines_with_the_other_filters(self):
night = self._save_prompt("Neon zebra", category="Night", tags=["a"])
day = self._save_prompt("Morning field", category="Day", tags=["a"])
files = {
name: self.media(name, 100 + index)
for index, name in enumerate(
("night.png", "day.png", "zebra_plain.png", "clip.mp4", "out.png")
)
}
self.api.db.link_image_to_prompt(night, str(files["night.png"]))
self.api.db.link_image_to_prompt(day, str(files["day.png"]))
members = [(self.output_dir, name) for name in files if name != "out.png"]
grids = str(self.collect("Grids", *members))
cases = (
({}, ["clip.mp4", "day.png", "night.png", "zebra_plain.png"]),
({"q": "zebra"}, ["night.png", "zebra_plain.png"]),
({"category": "Night"}, ["night.png"]),
({"tags": "a"}, ["day.png", "night.png"]),
({"media": "video"}, ["clip.mp4"]),
({"media": "image", "q": "day"}, ["day.png"]),
)
for params, expected in cases:
with self.subTest(**params):
response, data = await self.get(collection=grids, sort="name", **params)
self.assertEqual(response.status, 200, data)
self.assertEqual(self._names(data), expected)
self.assertEqual(data["total"], len(expected))
async def test_collection_filter_combines_with_an_images_own_category(self):
"""#331 with #332: a member is admitted by its effective category."""
from py.api.images import image_category_key
night = self._save_prompt("Neon zebra", category="Night")
files = {
name: self.media(name, 100 + index)
for index, name in enumerate(
("linked.png", "overridden.png", "own.png", "outside.png")
)
}
for name in ("linked.png", "overridden.png"):
self.api.db.link_image_to_prompt(night, str(files[name]))
sources = self.api._get_all_output_dirs()
def own(name, value):
key = image_category_key(sources, files[name])
self.api.db.set_image_categories([key], value)
own("overridden.png", "Portraits")
own("own.png", "Night") # unlinked: only its own category
own("outside.png", "Night") # not a member
members = [(self.output_dir, name) for name in files if name != "outside.png"]
grids = str(self.collect("Grids", *members))
for category, expected in (
("Night", ["linked.png", "own.png"]),
("Portraits", ["overridden.png"]),
):
with self.subTest(category=category):
response, data = await self.get(
collection=grids, category=category, sort="name"
)
self.assertEqual(response.status, 200, data)
self.assertEqual(self._names(data), expected)
self.assertEqual(data["total"], len(expected))
async def test_invalid_collection_is_a_400(self):
for value in ("x", "0", "-1", "1.5", "1234567890", " 1"):
with self.subTest(value=value):
response, data = await self.get(collection=value)
self.assertEqual(response.status, 400, data)
async def test_unknown_collection_is_a_404(self):
self.media("a.png")
response, data = await self.get(collection="77")
self.assertEqual(response.status, 404)
self.assertEqual(data["code"], "collection-not-found")
self.assertNotIn("images", data)
self.api._get_all_output_dirs = lambda: []
response, data = await self.get(collection="77")
self.assertEqual((response.status, data["code"]), (404, "collection-not-found"))
async def test_a_collection_deleted_during_the_read_is_a_404(self):
self.media("a.png")
grids = self.collect("Grids", (self.output_dir, "a.png"))
real = self.api.db.collection_members
def deleted_meanwhile(collection_id, root_key):
self.api.db.delete_collection(collection_id)
return real(collection_id, root_key)
with patch.object(
self.api.db, "collection_members", side_effect=deleted_meanwhile
):
response, data = await self.get(collection=str(grids))
self.assertEqual((response.status, data["code"]), (404, "collection-not-found"))
async def test_a_failed_members_read_is_a_500_never_an_unfiltered_page(self):
self.media("a.png")
self.media("b.png")
grids = self.collect("Grids", (self.output_dir, "a.png"))
def broken(*_args, **_kwargs):
raise sqlite3.OperationalError("database is locked")
with patch.object(self.api.db, "collection_members", side_effect=broken):
response, data = await self.get(collection=str(grids))
self.assertEqual(response.status, 500)
self.assertNotIn("images", data)
self.assertNotIn("locked", json.dumps(data))
async def test_reordering_sources_keeps_members(self):
second = self.two_sources()
self.media("a.png", root=second)
self.media("b.png")
grids = str(self.collect("Grids", (second, "a.png")))
_, data = await self.get(collection=grids)
self.assertEqual(
[(i["root_index"], i["path"]) for i in data["images"]], [(1, "a.png")]
)
self.api._get_all_output_dirs = lambda: [second, self.output_dir]
_, data = await self.get(collection=grids)
self.assertEqual(
[(i["root_index"], i["path"]) for i in data["images"]], [(0, "a.png")]
)
async def test_cursor_from_another_collection_is_refused(self):
for index, name in enumerate(("a.png", "b.png", "c.png")):
self.media(name, 100 + index)
everything = [(self.output_dir, name) for name in ("a.png", "b.png", "c.png")]
first = str(self.collect("First", *everything))
second = str(self.collect("Second", *everything))
_, page = await self.get(collection=first, limit=1)
token = page["next_cursor"]
self.assertIsInstance(token, str)
response, _ = await self.get(collection=first, limit=1, cursor=token)
self.assertEqual(response.status, 200)
for params in ({"collection": second}, {}, {"favorites": "1"}):
with self.subTest(**params):
response, _ = await self.get(cursor=token, limit=1, **params)
self.assertEqual(response.status, 400)
async def test_collection_wins_over_favorites(self):
self.media("a.png")
self.media("b.png")
grids = str(self.collect("Grids", (self.output_dir, "a.png")))
from py.api.images import _favorite_root_key
self.api.db.set_image_favorite(_favorite_root_key(self.output_dir), "b.png")
_, data = await self.get(collection=grids, favorites="1")
self.assertEqual(self._names(data), ["a.png"])
+823
View File
@@ -0,0 +1,823 @@
"""Cached re-runs of runtime-linked prompt text count once per job (#347).
Jobs are simulated end to end: the queue hook (submit + dequeue), the
``execution_cached`` replay, the real node executions and ComfyUI's output
cache (fakes of ``execution`` and ``comfy_execution.caching``, as in
tests/test_execution_context_host.py).
"""
import logging
import os
import sys
import tempfile
import threading
import types
import unittest
from collections import namedtuple
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 utils import cached_runs as module
from utils.cached_runs import (
CachedRuns,
Facts,
Role,
Seen,
graph_role,
graph_seen,
own_counts,
register_cached_runs,
stored_facts,
)
from utils.execution_context_host import cached_string_outputs
from utils.hashing import generate_prompt_hash
from utils.run_keys import prompt_node_key, upstream_closure, upstream_nodes
from utils.usage_tracking import (
PendingFirstUse,
QueuedRuns,
find_pair_id,
find_pair_id_by_hashes,
handle_queued_prompt,
on_dequeue,
)
Entry = namedtuple("CacheEntry", "ui outputs")
LLM = "a lighthouse the LLM wrote"
class _Cache:
def __init__(self):
self.entries = {}
def get_local(self, node_id):
return self.entries.get(node_id)
class HierarchicalCache(_Cache):
pass
class LRUCache(_Cache):
pass
class RAMPressureCache(LRUCache):
pass
class NullCache(_Cache):
pass
def host_modules():
"""Fake ``execution`` and ``comfy_execution.caching`` for the cache reader."""
execution = types.ModuleType("execution")
execution.CacheEntry = Entry
class PromptExecutor:
def add_message(self, event, data, broadcast):
return None
execution.PromptExecutor = PromptExecutor
caching = types.ModuleType("comfy_execution.caching")
for cls in (HierarchicalCache, LRUCache, RAMPressureCache, NullCache):
setattr(caching, cls.__name__, cls)
package = types.ModuleType("comfy_execution")
package.caching = caching
return {
"execution": execution,
"comfy_execution": package,
"comfy_execution.caching": caching,
}
def graph(
kind="PromptManager",
text=None,
role="positive",
sampler=True,
negative_text="",
seed=1,
):
"""PrimitiveString -> PreviewAny -> ``kind`` -> KSampler (runtime text).
``text`` set: typed text instead of the runtime link. ``role`` "negative":
the node feeds the sampler's negative input. Pair: its negative editor holds
``negative_text`` and it feeds both sampler inputs.
"""
g = {
"str": {"class_type": "PrimitiveString", "inputs": {"value": "seed text"}},
"llm": {"class_type": "PreviewAny", "inputs": {"source": ["str", 0]}},
"ckpt": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt": "m"}},
"pm": {
"class_type": kind,
"inputs": {"text": ["llm", 0] if text is None else text},
},
"other": {
"class_type": "CLIPTextEncode",
"inputs": {"text": "", "clip": ["ckpt", 1]},
},
"save": {"class_type": "SaveImage", "inputs": {"images": ["ks", 0]}},
}
if kind != "PromptManagerText":
g["pm"]["inputs"]["clip"] = ["ckpt", 1]
if kind == "PromptManagerPair":
g["pm"]["inputs"]["negative_text"] = negative_text
source = ["pm", 0]
if kind == "PromptManagerText":
g["enc"] = {
"class_type": "CLIPTextEncode",
"inputs": {"text": ["pm", 0], "clip": ["ckpt", 1]},
}
source = ["enc", 0]
positive, negative = (source, ["other", 0])
if role == "negative":
positive, negative = negative, positive
if kind == "PromptManagerPair":
positive, negative = ["pm", 0], ["pm", 1]
if sampler:
g["ks"] = {
"class_type": "KSampler",
"inputs": {
"model": ["ckpt", 0],
"positive": positive,
"negative": negative,
"seed": seed,
},
}
else:
g["save"]["inputs"]["images"] = positive
return g
class JobCase(unittest.TestCase):
"""A database, the queue hook stashes and one ``CachedRuns`` per test."""
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, "runs.db"))
self.addCleanup(self.db.close_all)
patcher = mock.patch.dict(sys.modules, host_modules())
patcher.start()
self.addCleanup(patcher.stop)
self.queued, self.pair_queued = QueuedRuns(), QueuedRuns()
self.pending, self.pair_pending = PendingFirstUse(), PendingFirstUse()
self.context = (None, None)
self.runs = CachedRuns(context=lambda: self.context)
self.queue = object()
self.executor = types.SimpleNamespace(
server=types.SimpleNamespace(prompt_queue=self.queue),
caches=types.SimpleNamespace(outputs=HierarchicalCache()),
)
self.item_id = 0
self.trackers = []
# -- nodes ----------------------------------------------------------------
def node(self, kind):
if kind == "PromptManagerPair":
from prompt_manager_pair import PromptManagerPair as cls
elif kind == "PromptManagerText":
from prompt_manager_text import PromptManagerText as cls
else:
import prompt_manager
cls = getattr(prompt_manager, kind)
node = cls.__new__(cls)
node.db = self.db
node.logger = logging.getLogger("test.cached_runs")
node.prompt_tracker = mock.Mock()
self.trackers.append(node.prompt_tracker)
node.comfyui_integration = mock.Mock()
node._inject_lora_trigger_words = lambda text: text
node.pending_first_use = self.pending
node.pending_pair_first_use = self.pair_pending
node.cached_runs = self.runs
return node
def execute(self, g, node_id="pm", text=LLM, negative="", index=None):
"""Run the node as ComfyUI would; returns its STRING outputs."""
kind = g[node_id]["class_type"]
node = self.node(kind)
self.context = (self.context[0], index)
clip = mock.Mock()
if kind == "PromptManagerText":
return node.process_text(text, prompt=g, unique_id=node_id)
if kind == "PromptManagerPair":
return node.encode_prompt(
clip, text, negative_text=negative, prompt=g, unique_id=node_id
)
return node.encode_prompt(clip, text, prompt=g, unique_id=node_id)
def cache(self, node_id, *results):
"""Store executed results (one per list index) as ComfyUI's cache entry."""
width = len(results[0])
outputs = [[result[slot] for result in results] for slot in range(width)]
self.executor.caches.outputs.entries[node_id] = Entry({}, outputs)
# -- jobs -----------------------------------------------------------------
def job(self, g, cached=(), run=None, outputs=("save",), prompt_id=None):
"""One queued job: hook, dequeue, ``execution_cached``, executions, done."""
self.item_id += 1
prompt_id = prompt_id or f"job-{self.item_id}"
request = handle_queued_prompt(
{"prompt": g, "prompt_id": prompt_id},
self.queued,
pair_queued=self.pair_queued,
)
on_dequeue(
request["prompt_id"],
self.queued,
lambda: self.db,
self.pending,
pair_queued=self.pair_queued,
pair_pending=self.pair_pending,
)
self.runs.begin_attempt(
((self.item_id, prompt_id, g, {}, list(outputs), {}), self.item_id)
)
self.context = (prompt_id, None)
self.runs.replay(self.executor, prompt_id, list(cached), lambda: self.db)
result = run() if run else None
self.runs.end_attempt(self.item_id)
self.context = (None, None)
return result
def first(self, g, text=LLM, negative="", node_id="pm"):
"""A job that executes the node and fills the cache."""
result = self.job(
g, run=lambda: self.execute(g, node_id, text=text, negative=negative)
)
self.cache(node_id, result)
return result
def count(self, text):
row = self.db.get_prompt_by_hash(generate_prompt_hash(text))
return row["run_count"] if row else None
def pair_count(self, positive, negative):
pair_id = find_pair_id(self.db, positive, negative)
return self.db.get_prompt_pair(pair_id)["run_count"] if pair_id else None
class TestCachedJobsCount(JobCase):
def test_runtime_text_counts_every_cached_job_for_each_node_kind(self):
for kind in (
"PromptManager",
"PromptManagerFocused",
"PromptManagerSplit",
"PromptManagerText",
):
with self.subTest(kind=kind):
text = f"{LLM} for {kind}"
g = graph(kind)
self.first(g, text=text)
self.job(graph(kind, seed=2), cached=["pm"])
self.job(graph(kind, seed=3), cached=["pm"])
self.assertEqual(self.count(text), 3)
def test_typed_text_still_counts_once_per_job(self):
g = graph(text="typed castle")
self.first(g, text="typed castle")
self.job(g, cached=["pm"])
self.job(g, cached=["pm"])
self.assertEqual(self.count("typed castle"), 3)
def test_rewired_to_negative_while_cached_counts_nothing(self):
self.first(graph())
self.job(graph(role="negative"), cached=["pm"])
self.assertEqual(self.count(LLM), 1)
# Rewired back to positive: the receipt still proves the text
self.job(graph(), cached=["pm"])
self.assertEqual(self.count(LLM), 2)
def test_no_sampler_counts_every_node_on_replay(self):
g = graph(sampler=False)
self.first(g)
self.job(g, cached=["pm"])
self.assertEqual(self.count(LLM), 2)
def test_empty_editor_replays_nothing(self):
g = graph()
self.first(g, text="")
self.job(g, cached=["pm"])
self.assertEqual(self.db.search_prompts(), [])
def test_alternating_upstream_values_each_count_their_own_prompt(self):
g = graph()
self.first(g, text="value A")
self.first(g, text="value B")
self.cache("pm", self.execute_outputs("value A"))
self.job(g, cached=["pm"])
self.assertEqual((self.count("value A"), self.count("value B")), (2, 1))
def execute_outputs(self, text):
return (mock.ANY, text)
def test_list_items_replay_each_item_in_order(self):
g = graph()
results = []
def run():
for index, text in enumerate(("one", "two", "three")):
results.append(self.execute(g, text=text, index=index))
self.job(g, run=run)
self.cache("pm", *results)
self.job(g, cached=["pm"])
self.assertEqual([self.count(t) for t in ("one", "two", "three")], [2, 2, 2])
def test_a_repeated_id_or_message_never_counts_a_node_twice(self):
g = graph()
self.first(g)
self.job(g, cached=["pm", "pm"])
self.assertEqual(self.count(LLM), 2)
self.item_id += 1
self.runs.begin_attempt(((1, "again", g, {}, ["save"], {}), self.item_id))
for _ in range(2):
self.runs.replay(self.executor, "again", ["pm"], lambda: self.db)
self.assertEqual(self.count(LLM), 3)
def test_a_shorter_list_keeps_its_receipt_after_a_longer_one(self):
g = graph()
self.first(g, text="alpha")
results = []
def run():
for index, text in enumerate(("alpha", "beta")):
results.append(self.execute(g, text=text, index=index))
self.job(g, run=run)
self.cache("pm", (mock.ANY, "alpha"))
self.job(g, cached=["pm"])
self.assertEqual((self.count("alpha"), self.count("beta")), (3, 1))
def test_a_job_reusing_a_prompt_id_counts_and_executes_normally(self):
g = graph()
self.first(g)
self.job(g, cached=["pm"], prompt_id="same")
self.job(g, run=lambda: self.execute(g), prompt_id="same")
self.job(g, cached=["pm"], prompt_id="same")
self.assertEqual(self.count(LLM), 4)
class TestFailClosed(JobCase):
def setUp(self):
super().setUp()
self.g = graph()
self.first(self.g)
def assertNotCounted(self):
self.assertEqual(self.count(LLM), 1)
def test_changed_cached_text_is_not_counted(self):
self.cache("pm", (mock.ANY, LLM + " edited"))
self.job(self.g, cached=["pm"])
self.assertNotCounted()
def test_no_receipt_in_this_process_is_not_counted(self):
self.runs = CachedRuns(context=lambda: self.context)
self.job(self.g, cached=["pm"])
self.assertNotCounted()
def test_missing_cache_entry_is_not_counted(self):
self.executor.caches.outputs.entries.clear()
self.job(self.g, cached=["pm"])
self.assertNotCounted()
def test_unsupported_cache_is_not_counted(self):
entries = self.executor.caches.outputs.entries
self.executor.caches.outputs = NullCache()
self.executor.caches.outputs.entries = entries
self.job(self.g, cached=["pm"])
self.assertNotCounted()
def test_more_list_values_than_the_bound_are_not_counted(self):
self.cache("pm", *[(mock.ANY, LLM)] * 65)
self.job(self.g, cached=["pm"])
self.assertNotCounted()
def test_missing_list_index_is_not_counted(self):
self.cache("pm", (mock.ANY, LLM), (mock.ANY, LLM))
self.job(self.g, cached=["pm"])
self.assertNotCounted()
def test_node_the_job_does_not_run_is_not_counted(self):
self.job(self.g, cached=["pm"], outputs=("other",))
self.assertNotCounted()
def test_message_for_another_job_is_not_counted(self):
self.item_id += 1
self.runs.begin_attempt(((1, "this", self.g, {}, ["save"], {}), self.item_id))
self.runs.replay(self.executor, "another", ["pm"], lambda: self.db)
self.assertNotCounted()
def test_unkeyable_node_is_not_counted(self):
g = graph()
g["pm"]["inputs"]["extra"] = ["missing", 0]
self.job(g, cached=["pm"])
self.assertIsNone(prompt_node_key(g, "pm"))
self.assertNotCounted()
def test_deleted_prompt_is_skipped(self):
self.db.delete_prompt(
self.db.get_prompt_by_hash(generate_prompt_hash(LLM))["id"]
)
self.job(self.g, cached=["pm"])
self.assertIsNone(self.count(LLM))
def test_database_errors_never_raise(self):
def broken():
raise RuntimeError("database unavailable")
self.item_id += 1
self.runs.begin_attempt(((1, "j", self.g, {}, ["save"], {}), self.item_id))
self.runs.replay(self.executor, "j", ["pm"], broken)
self.assertNotCounted()
def test_malformed_messages_and_results_are_ignored(self):
self.runs.begin_attempt(None)
self.runs.begin_attempt((("x",), "not an int"))
self.runs.replay(self.executor, None, ["pm"], lambda: self.db)
self.item_id += 1
self.runs.begin_attempt(((1, "j", self.g, {}, ["save"], {}), self.item_id))
self.runs.replay(self.executor, "j", "pm", lambda: self.db)
self.runs.replay(self.executor, "j", ["pm"] * 501, lambda: self.db)
self.runs.replay(self.executor, "j", [7, None], lambda: self.db)
self.assertNotCounted()
class TestReplayThenExecution(JobCase):
def test_execution_after_replay_in_the_same_job_is_not_counted_again(self):
# RAM-pressure frees the entry after execution_cached: the node executes
g = graph()
self.first(g)
self.job(g, cached=["pm"], run=lambda: self.execute(g))
self.assertEqual(self.count(LLM), 2)
# The suppressed execution still leaves a receipt: the next job counts
self.job(g, cached=["pm"])
self.assertEqual(self.count(LLM), 3)
def test_suppression_keeps_the_pending_catch_up_and_the_tracker_queue(self):
g = graph()
self.first(g)
prompt_id = self.db.get_prompt_by_hash(generate_prompt_hash(LLM))["id"]
def run():
# The replay counted, then the prompt was deleted and re-saved
self.db.delete_prompt(prompt_id)
self.pending.add(generate_prompt_hash(LLM))
return self.execute(g)
self.job(g, cached=["pm"], run=run)
self.assertEqual(self.count(LLM), 1) # the catch-up, not the own count
push = self.trackers[-1].set_current_prompt.call_args.kwargs["push_to_queue"]
self.assertTrue(push)
def test_without_host_context_the_node_counts_as_before(self):
g = graph()
self.runs = CachedRuns(context=lambda: (None, None))
self.first(g)
self.job(g, cached=["pm"])
self.job(g, run=lambda: self.execute(g))
self.assertEqual(self.count(LLM), 2)
class TestPairReplay(JobCase):
def test_runtime_positive_counts_positive_and_pair_every_job(self):
g = graph("PromptManagerPair", negative_text="blurry")
self.first(g, negative="blurry")
self.job(graph("PromptManagerPair", negative_text="blurry", seed=2), ["pm"])
self.job(graph("PromptManagerPair", negative_text="blurry", seed=3), ["pm"])
self.assertEqual(self.count(LLM), 3)
self.assertEqual(self.pair_count(LLM, "blurry"), 3)
self.assertEqual(self.count("blurry"), 0)
def test_empty_negative_pair_replays_its_pair_count(self):
g = graph("PromptManagerPair")
self.first(g)
self.job(g, cached=["pm"])
self.assertEqual(self.pair_count(LLM, ""), 2)
self.assertEqual(self.count(LLM), 2)
def test_pair_execution_after_replay_is_not_counted_again(self):
g = graph("PromptManagerPair", negative_text="blurry")
self.first(g, negative="blurry")
self.job(g, ["pm"], run=lambda: self.execute(g, negative="blurry"))
self.assertEqual(self.count(LLM), 2)
self.assertEqual(self.pair_count(LLM, "blurry"), 2)
def test_pair_role_follows_the_current_graph(self):
g = graph("PromptManagerPair", negative_text="blurry")
self.first(g, negative="blurry")
rewired = graph("PromptManagerPair", negative_text="blurry")
rewired["ks"]["inputs"]["positive"] = ["other", 0]
self.job(rewired, cached=["pm"])
self.assertEqual(self.count(LLM), 1)
self.assertEqual(self.pair_count(LLM, "blurry"), 1)
class TestOwnCounts(unittest.TestCase):
"""The node path and replay decide through the same function."""
def test_single_node_table(self):
facts = Facts("p")
for used, seen, expected in (
({"text"}, False, [("prompt", "p")]),
({"text"}, True, []),
(set(), False, []),
):
with self.subTest(used=used, seen=seen):
counts = own_counts(
"PromptManager", facts, Role(frozenset(used)), Seen(text=seen)
)
self.assertEqual(counts, expected)
self.assertEqual(
own_counts("PromptManager", Facts(None), Role(frozenset({"text"})), Seen()),
[],
)
def test_pair_table(self):
both = Role(frozenset({"text", "negative_text"}), True)
self.assertEqual(
own_counts("PromptManagerPair", Facts("p", "n"), both, Seen()),
[("prompt", "p"), ("prompt", "n"), ("pair", "p", "n")],
)
self.assertEqual(
own_counts("PromptManagerPair", Facts("p", "p"), both, Seen()),
[("prompt", "p"), ("pair", "p", "p")],
)
self.assertEqual(
own_counts("PromptManagerPair", Facts("p"), both, Seen(text=True)),
[("pair", "p", None)],
)
self.assertEqual(
own_counts(
"PromptManagerPair", Facts("p", "n"), both, Seen(True, True, True)
),
[],
)
def test_node_and_replay_agree_on_role_and_seen(self):
for g, node_id in (
(graph(), "pm"),
(graph(text="typed"), "pm"),
(graph(role="negative"), "pm"),
(graph(sampler=False), "pm"),
(graph("PromptManagerPair", negative_text="n"), "pm"),
(graph("PromptManagerPair", text="typed", negative_text="n"), "pm"),
):
kind = g[node_id]["class_type"]
text = "typed" if isinstance(g[node_id]["inputs"]["text"], str) else LLM
facts = stored_facts(text, g[node_id]["inputs"].get("negative_text"))
with self.subTest(kind=kind, inputs=g[node_id]["inputs"]):
replay = own_counts(
kind,
facts,
graph_role(g, node_id, kind),
graph_seen(g, node_id, kind, facts),
)
self.assertEqual(replay, self.node_counts(g, node_id, kind, text))
@staticmethod
def node_counts(g, node_id, kind, text):
"""What the executing node counted itself, from a recording database."""
db = mock.Mock()
db.get_prompt_by_hash.side_effect = lambda h: {"id": h}
db.save_prompt_pair.return_value = 1
negative = g[node_id]["inputs"].get("negative_text", "")
db.get_prompt_pair.return_value = {
"positive_prompt_id": generate_prompt_hash(text),
"negative_prompt_id": (
generate_prompt_hash(negative) if negative.strip() else None
),
}
case = JobCase()
case.db = db
case.pending, case.pair_pending = PendingFirstUse(), PendingFirstUse()
case.runs = CachedRuns(context=lambda: (None, None))
case.trackers = []
case.context = (None, None)
if kind == "PromptManagerPair":
node = case.node(kind)
node.encode_prompt(
mock.Mock(), text, negative_text=negative, prompt=g, unique_id=node_id
)
else:
node = case.node(kind)
node._track_prompt_execution(text, text, None, None, {}, g, node_id)
counts = [("prompt", c.args[0]) for c in db.record_prompt_use.call_args_list]
if kind == "PromptManagerPair":
counts += [
("pair", generate_prompt_hash(text), db.get_prompt_pair()[k])
for k in ["negative_prompt_id"]
if db.record_prompt_pair_use.called
]
return counts
class TestReceipts(unittest.TestCase):
def setUp(self):
self.context = ("j1", None)
self.runs = CachedRuns(
context=lambda: self.context, max_keys=2, max_receipts=2, max_items=3
)
self.g = {
name: {"class_type": "PromptManager", "inputs": {"text": name}}
for name in ("a", "b", "c")
}
def begin(self, item_id, prompt_id="j1"):
self.context = (prompt_id, None)
self.runs.begin_attempt(
((item_id, prompt_id, self.g, {}, ["a", "b", "c"], {}), item_id)
)
def keys(self):
return list(self.runs._receipts)
def test_least_recently_used_run_keys_are_dropped(self):
self.begin(1)
for name in ("a", "b", "c"):
self.runs.record_item(name, (name,), Facts("h"))
self.assertEqual(self.keys(), [prompt_node_key(self.g, n) for n in ("b", "c")])
def test_receipts_per_key_are_bounded_and_equal_outputs_replace(self):
key = prompt_node_key(self.g, "a")
for item_id, text in ((1, "x"), (2, "y"), (3, "z"), (4, "y")):
self.begin(item_id)
self.runs.record_item("a", (text,), Facts(text))
self.runs.end_attempt(item_id)
receipts = self.runs._receipts[key]
self.assertEqual([r.item_id for r in receipts], [4, 3])
def test_list_indices_past_the_bound_are_not_recorded(self):
self.begin(1)
self.context = ("j1", 3)
self.runs.record_item("a", ("x",), Facts("x"))
self.assertEqual(self.keys(), [])
def test_non_string_outputs_and_unknown_nodes_are_not_recorded(self):
self.begin(1)
self.runs.record_item("a", (None,), Facts("x"))
self.runs.record_item("zzz", ("x",), Facts("x"))
self.assertEqual(self.keys(), [])
def test_task_done_clears_only_its_own_attempt(self):
self.begin(1)
self.runs.end_attempt(2)
self.runs.record_item("a", ("x",), Facts("x"))
self.runs.end_attempt(1)
self.runs.record_item("b", ("x",), Facts("x"))
self.assertEqual(self.keys(), [prompt_node_key(self.g, "a")])
def test_job_nodes_follow_the_outputs_to_execute(self):
self.runs.begin_attempt(((1, "j1", self.g, {}, ["b"], {}), 1))
self.assertEqual(list(self.runs._attempt.nodes), ["b"])
class TestCacheReader(unittest.TestCase):
def setUp(self):
patcher = mock.patch.dict(sys.modules, host_modules())
patcher.start()
self.addCleanup(patcher.stop)
self.cache = LRUCache()
self.executor = types.SimpleNamespace(
caches=types.SimpleNamespace(outputs=self.cache)
)
def read(self, outputs, kind="PromptManager"):
self.cache.entries["n"] = Entry({}, outputs)
return cached_string_outputs(self.executor, "n", kind)
def test_reads_each_string_slot_in_order(self):
self.assertEqual(self.read([["c"], ["a", "b"]]), (("a", "b"),))
self.assertEqual(self.read([["t"]], "PromptManagerText"), (("t",),))
self.assertEqual(
self.read([[1], [2], ["p"], ["n"]], "PromptManagerPair"), (("p",), ("n",))
)
def test_fails_closed(self):
for outputs, kind in (
([["c"], ["a", 1]], "PromptManager"),
([["c"]], "PromptManager"),
([["c"], ["a"] * 65], "PromptManager"),
([[1], [2], ["p", "q"], ["n"]], "PromptManagerPair"),
([["c"], "text"], "PromptManager"),
([["c"], ["a"]], "CLIPTextEncode"),
([[]] * 5, "PromptManager"),
):
with self.subTest(outputs=outputs, kind=kind):
self.assertIsNone(self.read(outputs, kind))
self.cache.entries["m"] = ({}, [["c"], ["a"]]) # not a CacheEntry
self.assertIsNone(cached_string_outputs(self.executor, "m", "PromptManager"))
self.executor.caches.outputs = NullCache()
self.assertIsNone(self.read([["c"], ["a"]]))
class TestGraphHelpers(unittest.TestCase):
def test_upstream_nodes_and_closure(self):
g = graph()
self.assertEqual(
set(upstream_nodes(g, ["save"])),
{"save", "ks", "pm", "llm", "str", "ckpt", "other"},
)
self.assertEqual(upstream_nodes(g, ["missing"]), [])
self.assertEqual(upstream_nodes(None, ["save"]), [])
self.assertIsNone(upstream_closure(g, ["save"], limit=2))
self.assertIsNone(upstream_closure({"a": {"inputs": {"x": ["b", 0]}}}, ["a"]))
def test_find_pair_by_hashes_matches_the_text_lookup(self):
with tempfile.TemporaryDirectory() as tmp:
db = PromptDatabase(os.path.join(tmp, "p.db"))
try:
pair_id = db.save_prompt_pair("cat", "")
self.assertEqual(
find_pair_id_by_hashes(db, generate_prompt_hash("cat"), None),
pair_id,
)
self.assertIsNone(
find_pair_id_by_hashes(
db, generate_prompt_hash("cat"), generate_prompt_hash("x")
)
)
finally:
db.close_all()
assert_closed(self, tmp)
class TestRegistration(unittest.TestCase):
def setUp(self):
self.modules = host_modules()
patcher = mock.patch.dict(sys.modules, self.modules)
patcher.start()
self.addCleanup(patcher.stop)
flag = mock.patch.object(module, "_registered", threading.Event())
flag.start()
self.addCleanup(flag.stop)
class Queue:
def __init__(self):
self.done = []
def get(self, timeout=None):
return self.result
def task_done(self, item_id, *args, **kwargs):
self.done.append(item_id)
def test_registers_once_and_follows_the_job(self):
queue = self.Queue()
server = types.SimpleNamespace(prompt_queue=queue)
runs = mock.Mock()
self.assertTrue(register_cached_runs(server, lambda: "db", runs))
self.assertFalse(register_cached_runs(server, lambda: "db", runs))
queue.result = (("item",), 5)
queue.get()
runs.begin_attempt.assert_called_once_with((("item",), 5))
executor = self.modules["execution"].PromptExecutor()
executor.server = server
executor.add_message("execution_start", {"prompt_id": "j"}, broadcast=False)
executor.add_message(
"execution_cached", {"prompt_id": "j", "nodes": ["n"]}, broadcast=False
)
runs.replay.assert_called_once()
self.assertEqual(runs.replay.call_args.args[1:3], ("j", ["n"]))
other = self.modules["execution"].PromptExecutor()
other.server = types.SimpleNamespace(prompt_queue=object())
other.add_message("execution_cached", {"prompt_id": "j", "nodes": []}, False)
self.assertEqual(runs.replay.call_count, 1)
queue.task_done(5, None, status=None)
runs.end_attempt.assert_called_once_with(5)
def test_hosts_without_the_interfaces_register_nothing(self):
class NoTaskDone:
def get(self):
return None
self.assertFalse(
register_cached_runs(
types.SimpleNamespace(prompt_queue=NoTaskDone()), lambda: None
)
)
self.assertFalse(register_cached_runs(types.SimpleNamespace(), lambda: None))
del self.modules["execution"].PromptExecutor.add_message
self.assertFalse(
register_cached_runs(
types.SimpleNamespace(prompt_queue=self.Queue()), lambda: None
)
)
if __name__ == "__main__":
unittest.main()
+44
View File
@@ -614,6 +614,50 @@ class TestRestoreFromFile(BackupTestCase):
)
assert_closed(self, str(backups[0]))
def test_backup_and_restore_keep_collections(self):
with self.db.model.get_connection() as conn:
conn.execute(
"INSERT INTO image_collections (id, name, name_key, created_at) "
"VALUES (4, 'LoRa testing', 'lora testing', '2026-10-27T00:00:00Z')"
)
conn.execute(
"INSERT INTO image_collection_items VALUES "
"(4, 'root', 'grid.png', '2026-10-27T00:00:00Z')"
)
with tempfile.TemporaryDirectory() as directory:
# The rows still sit in the WAL; a backup must carry them.
backup = os.path.join(directory, "backup.db")
self.assertTrue(self.db.model.backup_database(backup))
self.assertEqual(
PromptModel.verify_database_file(backup, require_library_state=True),
(True, "ok"),
)
with self.db.model.get_connection() as conn:
conn.execute("DELETE FROM image_collections")
self.db.model.restore_from_file(backup)
conn = self.db.model.get_connection()
self.assertEqual(
[
tuple(r)
for r in conn.execute("SELECT id, name FROM image_collections")
],
[(4, "LoRa testing")],
)
self.assertEqual(
[
tuple(r)
for r in conn.execute(
"SELECT collection_id, root_key, relative_path "
"FROM image_collection_items"
)
],
[(4, "root", "grid.png")],
)
PromptModel.close_all_instances()
assert_closed(self, directory)
def test_restore_without_a_live_database_returns_empty_backup_path(self):
src = self._make_source_db(["fresh"])
self.db.close_all()
+456
View File
@@ -0,0 +1,456 @@
"""The prompt-link read of the workspace Images page (#354).
``PromptDatabase.get_image_link_matches`` returns every row of its filter
SELECT in one SQLite step (one JSON value) instead of one step per row, so the
read does not have to win the GIL back once per linked image while ComfyUI
generates. Its result must equal the streamed read it replaces, which stays as
the fallback, on every filter and on odd values.
"""
import itertools
import os
import sqlite3
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database import operations # noqa: E402
from database.operations import PromptDatabase, link_filter_query # noqa: E402
try:
from tests.open_handles import assert_closed # noqa: E402
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed # noqa: E402
LOGGER = "prompt_manager.database"
MISSING_FUNCTION_SQL = (
"SELECT pm_missing_function(json_group_array(json_array("
"file_path, rating, matched, positive))) FROM ({select})"
)
QUERIES = ("", "lantern", "LANTERN title", "no-such-text")
CATEGORIES = (None, "cat-a", "cat-b", "cat-none")
TAG_SETS = ([], ["t1"], ["t1", "t2"], ["t-none"])
PROMPTS = (
# id, text, category, rating, tags, library title
(1, "red lantern at dusk", "cat-a", 5, ("t1", "t2"), None),
(2, "blue harbor", "cat-a", None, ("t1",), None),
(3, "green field", "cat-b", 3, (), "Lantern title"),
(4, "pair positive", "cat-b", 4, ("t2",), None),
(5, "pair negative lantern", "cat-a", 1, ("t1",), None),
)
QUOTED = '/out/it\'s "quoted".png'
BACKSLASHED = "C:\\out\\back\\slash.png"
UNICODE = "/out/ünïcødé 画像.png"
NEWLINE = "/out/new\nline.png"
IMAGES = (
# prompt id, file path
(1, "/out/a.png"),
(2, QUOTED),
(2, BACKSLASHED),
(3, UNICODE),
(1, NEWLINE),
(1, "/out/ambiguous.png"),
(2, "/out/ambiguous.png"),
(5, "/out/pair.png"), # linked through pair 1: its positive is prompt 4
(3, "/out/two-pairs.png"), # two pairs: no single pair, so prompt 3
(1, None), # no file path: never listed
)
class LinkMatchesTestCase(unittest.TestCase):
def setUp(self):
self.temp = tempfile.TemporaryDirectory()
self.root = Path(os.path.realpath(self.temp.name))
self.db = PromptDatabase(str(self.root / "prompts.db"))
for name in ("_one_step_unsupported", "_one_step_warned"):
patcher = patch.object(operations, name, False)
patcher.start()
self.addCleanup(patcher.stop)
self.statements = []
original = self.db.model.get_connection
def traced_connection():
connection = original()
connection.set_trace_callback(self.statements.append)
return connection
patcher = patch.object(
self.db.model, "get_connection", side_effect=traced_connection
)
patcher.start()
self.addCleanup(patcher.stop)
def tearDown(self):
self.db.close_all()
assert_closed(self, self.root)
self.temp.cleanup()
def seed(self):
with self.db.model.get_connection() as conn:
for prompt_id, text, category, rating, tags, title in PROMPTS:
conn.execute(
"INSERT INTO prompts(id,text,category,rating,hash) "
"VALUES(?,?,?,?,?)",
(prompt_id, text, category, rating, f"hash-{prompt_id}"),
)
for tag in tags:
conn.execute("INSERT OR IGNORE INTO tags(name) VALUES(?)", (tag,))
conn.execute(
"INSERT INTO prompt_tags(prompt_id,tag_id) "
"SELECT ?, id FROM tags WHERE name=?",
(prompt_id, tag),
)
if title:
conn.execute(
"INSERT INTO prompt_library_state(prompt_id,title) "
"VALUES(?,?)",
(prompt_id, title),
)
ids = {}
for prompt_id, path in IMAGES:
cursor = conn.execute(
"INSERT INTO generated_images"
"(prompt_id,image_path,filename,file_path) VALUES(?,?,?,?)",
(prompt_id, path or "/elsewhere/x.png", "x.png", path),
)
ids.setdefault(path, cursor.lastrowid)
conn.execute(
"INSERT INTO prompt_pairs(id,positive_prompt_id,negative_prompt_id) "
"VALUES(1,4,5),(2,3,5),(3,3,NULL)"
)
conn.executemany(
"INSERT INTO image_prompt_pairs(image_id,pair_id,provenance) "
"VALUES(?,?,'{}')",
[
(ids["/out/pair.png"], 1),
(ids["/out/two-pairs.png"], 2),
(ids["/out/two-pairs.png"], 3),
],
)
def add_image(self, prompt_id, path):
with self.db.model.get_connection() as conn:
conn.execute(
"INSERT INTO generated_images(prompt_id,image_path,filename,file_path)"
" VALUES(?,?,?,?)",
(prompt_id, "/blob/x.png", "x.png", path),
)
def streamed(self, q, category, tags, every_link=False, wanted=None):
"""Today's read, called directly (the fallback)."""
sql, params = link_filter_query(
q, tags, every_link=every_link, category=category
)
with self.db.model.get_connection() as conn:
return operations._streamed_link_matches(
conn, sql, params, wanted, category is not None
)
def selects(self):
found = [s for s in self.statements if s.lstrip().upper().startswith("SELECT")]
self.statements.clear()
return found
class TestOneStepLinkRead(LinkMatchesTestCase):
def test_plain_read_returns_every_unambiguous_link(self):
self.seed()
matches = self.db.get_image_link_matches("", [], every_link=True)
self.assertEqual(
matches,
{
"/out/a.png": (True, 5, 1, None),
QUOTED: (True, None, 2, None),
BACKSLASHED: (True, None, 2, None),
UNICODE: (True, 3, 3, None),
NEWLINE: (True, 5, 1, None),
"/out/pair.png": (True, 4, 4, None),
"/out/two-pairs.png": (True, 3, 3, None),
},
)
self.assertIsNone(matches[QUOTED][1])
self.assertIsInstance(matches["/out/a.png"][1], int)
self.assertIsInstance(matches["/out/a.png"][2], int)
def test_search_hits_follow_text_title_and_tags(self):
self.seed()
self.assertEqual(
self.db.get_image_link_matches("lantern", []),
{
"/out/a.png": (True, 5, 1, None),
NEWLINE: (True, 5, 1, None),
UNICODE: (True, 3, 3, None),
"/out/two-pairs.png": (True, 3, 3, None),
},
)
self.assertEqual(
self.db.get_image_link_matches("lantern", [], category="cat-a"),
{
"/out/a.png": (True, 5, 1, "cat-a"),
NEWLINE: (True, 5, 1, "cat-a"),
QUOTED: (False, None, 2, "cat-a"),
BACKSLASHED: (False, None, 2, "cat-a"),
},
)
def test_every_filter_equals_the_streamed_read(self):
self.seed()
wanted_sets = (None, {"/out/a.png", UNICODE, "/out/pair.png", "/absent.png"})
for q, category, tags, every_link, wanted in itertools.product(
QUERIES, CATEGORIES, TAG_SETS, (False, True), wanted_sets
):
with self.subTest(
q=q, category=category, tags=tags, every=every_link, wanted=wanted
):
expected = self.streamed(q, category, tags, every_link, wanted)
self.statements.clear()
actual = self.db.get_image_link_matches(
q,
list(tags),
every_link=every_link,
wanted=wanted,
category=category,
)
self.assertEqual(actual, expected)
self.assertEqual(
{path: tuple(map(type, row)) for path, row in actual.items()},
{path: tuple(map(type, row)) for path, row in expected.items()},
)
found = self.selects()
self.assertEqual(len(found), 1, found)
self.assertIn("json_group_array", found[0])
def test_empty_database_returns_nothing_in_one_statement(self):
self.assertEqual(self.db.get_image_link_matches("", [], every_link=True), {})
self.assertEqual(self.db.get_image_link_matches("x", ["t"], category="c"), {})
found = self.selects()
self.assertEqual(len(found), 2, found)
self.assertTrue(all("json_group_array" in query for query in found))
def test_wanted_limits_the_result_to_the_listing(self):
self.seed()
self.assertEqual(
self.db.get_image_link_matches(
"", [], every_link=True, wanted={"/out/pair.png", "/x.png"}
),
{"/out/pair.png": (True, 4, 4, None)},
)
self.assertEqual(
self.db.get_image_link_matches("", [], every_link=True, wanted=set()),
{},
)
def test_category_read_carries_the_inherited_category_in_one_step(self):
# A recipe's own category wins over its positive's (#332 MC4); the
# read only narrows to links holding the value, the caller compares.
self.seed()
with self.db.model.get_connection() as conn:
conn.execute("UPDATE prompt_pairs SET category='cat-a' WHERE id=1")
self.statements.clear()
for category, expected in (
("cat-a", {"/out/pair.png": (True, 4, 4, "cat-a")}),
("cat-b", {"/out/pair.png": (True, 4, 4, "cat-a")}),
):
with self.subTest(category=category):
actual = self.db.get_image_link_matches(
"", [], every_link=True, category=category, wanted={"/out/pair.png"}
)
self.assertEqual(actual, expected)
self.assertEqual(
actual, self.streamed("", category, [], True, {"/out/pair.png"})
)
found = [query for query in self.selects() if "json_group_array" in query]
self.assertEqual(len(found), 2, found)
self.assertTrue(all("pair_category" in query for query in found))
def test_plain_read_leaves_the_categories_out_of_the_json(self):
self.seed()
self.statements.clear()
self.db.get_image_link_matches("", [], every_link=True)
(query,) = self.selects()
aggregate = query[query.index("json_array(") :].split(")", 1)[0]
self.assertEqual(aggregate, "json_array(file_path, rating, matched, positive")
def test_read_releases_its_snapshot(self):
# A statement left open after the read would pin its WAL snapshot, and
# a TRUNCATE checkpoint from another connection would report busy.
self.seed()
self.assertTrue(self.db.get_image_link_matches("", [], every_link=True))
other = sqlite3.connect(str(self.root / "prompts.db"))
try:
other.execute("UPDATE prompts SET rating=2 WHERE id=1")
other.commit()
busy, _log, _done = other.execute(
"PRAGMA wal_checkpoint(TRUNCATE)"
).fetchone()
finally:
other.close()
self.assertEqual(busy, 0)
self.assertEqual(os.path.getsize(self.root / "prompts.db-wal"), 0)
def test_wrapped_query_keeps_the_ambiguity_probe_an_index_search(self):
self.seed()
sql, params = link_filter_query(
"lantern", ["t1"], every_link=True, category="cat-a"
)
with self.db.model.get_connection() as conn:
plan = [
row[3]
for row in conn.execute(
"EXPLAIN QUERY PLAN "
+ operations.link_matches_json_query(sql, categories=True),
params,
)
]
self.assertIn(
"SEARCH other USING COVERING INDEX idx_image_file_path (file_path=?)", plan
)
self.assertFalse([step for step in plan if step.startswith("SCAN other")])
class TestStreamedFallback(LinkMatchesTestCase):
def test_missing_json_functions_stream_and_stop_trying(self):
self.seed()
expected = self.streamed("", None, [], True)
self.statements.clear()
wrapper = operations.link_matches_json_query
with (
patch.object(operations, "_LINK_MATCHES_JSON_SQL", MISSING_FUNCTION_SQL),
patch.object(
operations, "link_matches_json_query", side_effect=wrapper
) as attempts,
self.assertLogs(LOGGER, "WARNING") as logged,
):
first = self.db.get_image_link_matches("", [], every_link=True)
found = self.selects()
self.assertEqual(attempts.call_count, 1)
second = self.db.get_image_link_matches("", [], every_link=True)
again = self.selects()
self.assertEqual(first, expected)
self.assertEqual(second, expected)
self.assertEqual(len(logged.records), 1, logged.output)
self.assertIn("no such function: pm_missing_function", logged.output[0])
self.assertTrue(operations._one_step_unsupported)
# SQLite refuses the wrapped statement before it runs (not traced).
self.assertEqual(len(found), 1, found)
self.assertNotIn("json_group_array", found[0])
# The next read skips the attempt for the rest of the process.
self.assertEqual(attempts.call_count, 1)
self.assertEqual(len(again), 1, again)
self.assertNotIn("json_group_array", again[0])
def test_blob_path_streams_that_read_and_keeps_trying(self):
self.seed()
self.add_image(2, b"/out/blob.png")
expected = self.streamed("", None, [], True)
self.assertEqual(expected[b"/out/blob.png"], (True, None, 2, None))
self.statements.clear()
with self.assertLogs(LOGGER, "WARNING") as logged:
first = self.db.get_image_link_matches("", [], every_link=True)
second = self.db.get_image_link_matches("", [], every_link=True)
self.assertEqual(first, expected)
self.assertEqual(second, expected)
self.assertEqual(len(logged.records), 1, logged.output)
self.assertFalse(operations._one_step_unsupported)
found = self.selects()
self.assertEqual(sum("json_group_array" in query for query in found), 2)
self.assertEqual(len(found), 4, found)
def test_blob_that_reads_as_json_still_streams(self):
# SQLite 3.45+ reads a BLOB that is valid JSONB as JSON (b"\x00" is
# null) instead of refusing it; the streamed read keeps the bytes.
self.seed()
self.add_image(2, b"\x00")
expected = self.streamed("", None, [], True)
self.assertIn(b"\x00", expected)
with self.assertLogs(LOGGER, "WARNING"):
self.assertEqual(
self.db.get_image_link_matches("", [], every_link=True), expected
)
self.assertFalse(operations._one_step_unsupported)
def test_blob_rating_streams(self):
self.seed()
with self.db.model.get_connection() as conn:
# Only a database written around the CHECK can hold one.
conn.execute("PRAGMA ignore_check_constraints=ON")
conn.execute("UPDATE prompts SET rating=? WHERE id=3", (b"\x00",))
conn.execute("PRAGMA ignore_check_constraints=OFF")
expected = self.streamed("", None, [], True)
self.assertEqual(expected[UNICODE], (True, b"\x00", 3, None))
with self.assertLogs(LOGGER, "WARNING"):
self.assertEqual(
self.db.get_image_link_matches("", [], every_link=True), expected
)
def test_blob_category_streams_a_category_read(self):
self.seed()
with self.db.model.get_connection() as conn:
# Admitted by its recipe's text category; the positive's is a BLOB.
conn.execute("UPDATE prompts SET category=? WHERE id=4", (b"cat-b",))
conn.execute("UPDATE prompt_pairs SET category='cat-b' WHERE id=1")
expected = self.streamed("", "cat-b", [], True)
self.assertEqual(expected["/out/pair.png"], (True, 4, 4, "cat-b"))
with self.assertLogs(LOGGER, "WARNING") as logged:
actual = self.db.get_image_link_matches(
"", [], every_link=True, category="cat-b"
)
self.assertEqual(actual, expected)
self.assertIn("BLOB", logged.output[0])
self.assertFalse(operations._one_step_unsupported)
def test_json_text_over_the_length_limit_streams(self):
self.seed()
expected = self.streamed("", None, [], True)
conn = self.db.model.get_connection()
limit = conn.setlimit(sqlite3.SQLITE_LIMIT_LENGTH, 200)
try:
with self.assertLogs(LOGGER, "WARNING") as logged:
actual = self.db.get_image_link_matches("", [], every_link=True)
finally:
conn.setlimit(sqlite3.SQLITE_LIMIT_LENGTH, limit)
self.assertEqual(actual, expected)
self.assertIn("too big", logged.output[0])
self.assertFalse(operations._one_step_unsupported)
def test_an_interrupted_read_raises_without_a_fallback(self):
# Move library interrupts a retired connection: the read fails as
# before, and the one warning is kept for a real fallback.
self.seed()
conn = self.db.model.get_connection()
conn.set_progress_handler(lambda: 1, 1)
self.statements.clear()
try:
with (
self.assertNoLogs(LOGGER, "WARNING"),
self.assertRaisesRegex(sqlite3.OperationalError, "^interrupted$"),
):
self.db.get_image_link_matches("", [], every_link=True)
finally:
conn.set_progress_handler(None, 1)
self.assertFalse(operations._one_step_warned)
self.assertFalse(operations._one_step_unsupported)
self.assertFalse([s for s in self.statements if "json_group_array" not in s])
def test_the_maintenance_link_rows_reader_is_not_shadowed(self):
from database.maintenance import MaintenanceOperations
self.assertIs(PromptDatabase._link_rows, MaintenanceOperations._link_rows)
def test_a_failed_streamed_read_still_raises(self):
with (
patch.object(operations, "_LINK_MATCHES_JSON_SQL", MISSING_FUNCTION_SQL),
patch.object(operations, "_LINK_FILTER_SQL", "SELECT broken FROM nowhere"),
self.assertLogs(LOGGER, "WARNING"),
):
with self.assertRaises(sqlite3.OperationalError):
self.db.get_image_link_matches("", [], every_link=True)
if __name__ == "__main__":
unittest.main()
+184
View File
@@ -0,0 +1,184 @@
"""Duplicate cleanup keeps images' own categories (#332 MC11).
The kept copy's own category wins; a kept copy without one receives the
removed copy's. The removed copy's row is never deleted, so a give-back
(the kept copy changed) and a failed delete lose nothing.
"""
import asyncio
import os
from unittest.mock import patch
try:
from tests.test_api_duplicate_files import WAIT, DuplicateFilesTestCase
except ImportError: # discovered with tests/ as the top-level directory
from test_api_duplicate_files import WAIT, DuplicateFilesTestCase
from py.api import duplicate_files # noqa: E402
from py.api.images import image_category_key # noqa: E402
from py.config import GalleryConfig # noqa: E402
class TestDuplicateCleanupKeepsCategories(DuplicateFilesTestCase):
def key(self, path):
return image_category_key(self.api._get_all_output_dirs(), path)
def category(self, path):
return self.api.db.image_category(*self.key(path))
def set_category(self, path, value):
self.api.db.set_image_categories([self.key(path)], value)
async def run_cleanup(self, keep="a.png"):
job, group = await self.found_group()
plan = await self.review(job.id, [self.keep_of(group, keep)])
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
return cleanup
async def test_kept_copy_without_a_category_receives_the_removed_one(self):
self.trio()
b = self.output_dir / "b.png"
self.set_category(b, "Portraits")
cleanup = await self.run_cleanup()
self.assertFalse(b.exists())
self.assertEqual(self.category(self.output_dir / "a.png"), "Portraits")
# The removed copy's row stays (hidden while its file is gone).
self.assertEqual(self.category(b), "Portraits")
self.assertIn(
"Kept 1 image category on the kept copies", "\n".join(cleanup.log)
)
async def test_kept_copy_keeps_its_own_category(self):
self.trio()
self.set_category(self.output_dir / "a.png", "Landscapes")
self.set_category(self.output_dir / "b.png", "Portraits")
self.set_category(self.output_dir / "sub" / "c.png", "Scans")
cleanup = await self.run_cleanup()
self.assertEqual(self.category(self.output_dir / "a.png"), "Landscapes")
self.assertNotIn("image categor", "\n".join(cleanup.log))
async def test_two_removed_copies_give_the_first_category(self):
self.trio()
self.set_category(self.output_dir / "b.png", "Portraits")
self.set_category(self.output_dir / "sub" / "c.png", "Scans")
cleanup = await self.run_cleanup()
self.assertIn(self.category(self.output_dir / "a.png"), {"Portraits", "Scans"})
self.assertIn(
"Kept 1 image category on the kept copies", "\n".join(cleanup.log)
)
async def test_give_back_loses_no_category(self):
self.trio()
self.set_category(self.output_dir / "b.png", "Portraits")
with patch.object(duplicate_files, "still_kept", return_value=False):
cleanup = await self.run_cleanup()
self.assertEqual(cleanup.counters["removed"], 0)
rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]}
self.assertEqual(
rows["b.png"], duplicate_files.ROW_DETAILS[duplicate_files.KEEPER_GONE]
)
self.assertTrue((self.output_dir / "b.png").is_file())
self.assertEqual(self.category(self.output_dir / "b.png"), "Portraits")
async def test_a_failed_delete_keeps_both_categories(self):
self.trio()
self.set_category(self.output_dir / "b.png", "Portraits")
real = duplicate_files._unlink
def unlink(path):
if os.path.basename(path) == "b.png":
raise PermissionError(13, "denied")
real(path)
with patch.object(duplicate_files, "_unlink", side_effect=unlink):
await self.run_cleanup()
self.assertTrue((self.output_dir / "b.png").is_file())
self.assertEqual(self.category(self.output_dir / "b.png"), "Portraits")
self.assertEqual(self.category(self.output_dir / "a.png"), "Portraits")
async def test_nested_sources_use_the_deepest_source_keys(self):
self.write("sub/a.png", b"same", mtime=1)
self.write("sub/b.png", b"same", mtime=2)
# The subfolder is listed first, so the cleanup's listed root is the
# file key's source; the outer source still holds the files.
GalleryConfig.MONITORING_DIRECTORIES = [
self._root(self.output_dir / "sub"),
self._root(self.output_dir),
]
sub_key = duplicate_files.favorite_root_key(self.output_dir / "sub")
self.api.db.set_image_categories([(sub_key, "b.png")], "Portraits")
await self.run_cleanup()
self.assertFalse((self.output_dir / "sub" / "b.png").exists())
self.assertEqual(self.api.db.image_category(sub_key, "a.png"), "Portraits")
outer_key = duplicate_files.favorite_root_key(self.output_dir)
self.assertIsNone(self.api.db.image_category(outer_key, "sub/a.png"))
async def test_outer_source_listed_first_still_keys_by_the_deepest(self):
self.write("sub/a.png", b"same", mtime=1)
self.write("sub/b.png", b"same", mtime=2)
GalleryConfig.MONITORING_DIRECTORIES = [
self._root(self.output_dir),
self._root(self.output_dir / "sub"),
]
sub_key = duplicate_files.favorite_root_key(self.output_dir / "sub")
self.api.db.set_image_categories([(sub_key, "b.png")], "Portraits")
job, group = await self.found_group()
[kept] = [f for f in group["files"] if f["path"].endswith("a.png")][:1]
plan = await self.review(
job.id, [self.keep_of(group, kept["path"], kept["root_index"])]
)
body = await self.apply(plan["token"])
cleanup = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
self.assertFalse((self.output_dir / "sub" / "b.png").exists())
self.assertEqual(self.api.db.image_category(sub_key, "a.png"), "Portraits")
async def test_a_symlinked_output_folder_still_moves_the_category(self):
link = self.output_dir.parent / "linked-output"
try:
os.symlink(self.output_dir, link, target_is_directory=True)
except (OSError, NotImplementedError):
self.skipTest("symlinks are not available")
self.trio()
with patch.object(self.api, "_get_all_output_dirs", return_value=[link]):
key = image_category_key([link], link / "b.png")
self.api.db.set_image_categories([key], "Portraits")
await self.run_cleanup()
kept = image_category_key([link], link / "a.png")
self.assertFalse((self.output_dir / "b.png").exists())
self.assertEqual(self.api.db.image_category(*kept), "Portraits")
async def test_a_copy_with_a_collection_and_a_category_moves_both(self):
# #331 + #332: one removed copy carries both records; one cleanup
# moves its membership and copies its category onto the kept copy.
self.trio()
db = self.api.db
root_key = duplicate_files.favorite_root_key(self.output_dir)
b = self.output_dir / "b.png"
grids = db.add_to_collection(name="LoRA grids", items=[(root_key, "b.png")])
grids = grids["collection"]["id"]
self.set_category(b, "Portraits")
seen = {}
real_log = duplicate_files.DuplicateFilesMixin._log_moved_records
def log_moved_records(job, totals):
seen.update(totals)
real_log(job, totals)
with patch.object(
duplicate_files.DuplicateFilesMixin,
"_log_moved_records",
staticmethod(log_moved_records),
):
cleanup = await self.run_cleanup()
self.assertFalse(b.exists())
self.assertEqual(db.collection_members(grids, root_key), {"a.png"})
self.assertEqual(db.file_collections(root_key, "b.png"), [])
self.assertEqual(self.category(self.output_dir / "a.png"), "Portraits")
self.assertEqual(seen["collections"], 1)
self.assertEqual(seen["categories"], 1)
log = "\n".join(cleanup.log)
self.assertIn("Moved 1 collection membership to the kept copies", log)
self.assertIn("Kept 1 image category on the kept copies", log)
+463
View File
@@ -0,0 +1,463 @@
"""Work running at Ctrl+C stops at its next safe point (#376).
ComfyUI catches Ctrl+C in ``main.py`` and returns without closing its event
loop. Python then joins every executor thread (``concurrent.futures``'s
``_python_exit``, a ``threading._register_atexit`` callback) before any
``atexit`` callback runs, so a call still running on the default executor
used to keep the process alive until it returned. The package registers its
own ``threading._register_atexit`` hook, which runs before that join: it sets
the exit signal that ``Job.cancelled`` and ``map_parallel`` read, and kills
the encoder processes.
The child-process test starts three long calls on the default executor (a
registry job checking ``job.cancelled`` between files, a ``map_parallel``
batch and an encoder process that ignores the exit, as ffmpeg does on
Windows), returns from its main module and must exit far sooner than any of
them would end.
"""
import json
import os
import subprocess
import sys
import tempfile
import textwrap
import threading
import unittest
from unittest import mock
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
try:
from utils import exit_signal # noqa: E402
except ImportError: # before #376: only the child-process test can run
exit_signal = None
# The parent never waits longer than this for the child (the hang watchdog
# allows 120 s per test). Each fake job alone would run 300 s or more.
CHILD_CEILING_SECONDS = 60
# From the child's return to its last atexit callback: the stopped calls end
# at their next file, item or chunk (the encoder is killed).
EXIT_BOUND_SECONDS = 15
JOB_SECONDS = 600
# The fake encoder outlives the ceiling, yet a failing run leaks it only
# briefly (the parent cannot reach the child's own child to kill it).
ENCODER_SECONDS = 90
REPORT_PREFIX = "PM_EXIT_REPORT "
CHILD = textwrap.dedent("""
import asyncio
import atexit
import contextlib
import importlib
import io
import json
import os
import sys
import threading
import time
root, tmp = sys.argv[1], sys.argv[2]
job_seconds, encoder_seconds = float(sys.argv[3]), int(sys.argv[4])
sys.path.insert(0, root)
state = {"stopped": {}}
def report():
returned = state.get("returned")
print("PM_EXIT_REPORT " + json.dumps({
"exit_seconds": None if returned is None else time.monotonic() - returned,
"stopped": state["stopped"],
}))
sys.stdout.flush()
# Registered first, so it runs after every other atexit callback, and
# every atexit callback runs after the executor threads were joined.
atexit.register(report)
from tests.test_init_package import _import_package, fake_server_module
output = os.path.join(tmp, "output")
with contextlib.redirect_stdout(io.StringIO()):
package, _ = _import_package(
"pm_exit_child",
os.path.join(tmp, "prompts.db"),
output,
fake_server_module(queue_task_done=True),
)
parallel = importlib.import_module("pm_exit_child.utils.parallel")
encode = importlib.import_module("pm_exit_child.py.api.thumbnail_encode")
registry = package.api._jobs
entered = {name: threading.Event() for name in ("job", "batch", "encoder")}
def job_work(job):
# A job's thread work: one "file" at a time, checking its cancel
# between files the way the thumbnail and import jobs do.
entered["job"].set()
deadline = time.monotonic() + job_seconds
while time.monotonic() < deadline:
if job.cancelled:
state["stopped"]["job"] = True
return
time.sleep(0.05)
def batch_item(index):
entered["batch"].set()
time.sleep(0.5)
return index
def batch_work():
try:
parallel.map_parallel(batch_item, range(int(job_seconds)), 2)
except Exception as error:
state["stopped"]["batch"] = type(error).__name__
def encoder_work():
# Stands for ffmpeg on Windows, which runs without a console and
# never sees Ctrl+C: only a kill ends it before its timeout.
cmd = [sys.executable, "-c", "import time; time.sleep(%d)" % encoder_seconds]
try:
encode.run_encoder(
cmd,
timeout=encoder_seconds,
track=lambda _process: entered["encoder"].set(),
)
except Exception as error:
state["stopped"]["encoder"] = type(error).__name__
async def begin():
loop = asyncio.get_running_loop()
async def runner(job):
await loop.run_in_executor(None, job_work, job)
registry.start("exit-test", "Exit test", admission="read", runner=runner)
loop.run_in_executor(None, batch_work)
loop.run_in_executor(None, encoder_work)
for name, event in entered.items():
if not await loop.run_in_executor(None, event.wait, 30):
sys.exit("not started: " + name)
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
loop.run_until_complete(begin())
state["returned"] = time.monotonic()
# Return with the three calls running and the loop open, as ComfyUI does
# after Ctrl+C.
""")
class TestExitStopsRunningWork(unittest.TestCase):
def test_running_executor_work_stops_at_exit(self):
with tempfile.TemporaryDirectory() as tmp:
os.makedirs(os.path.join(tmp, "output"))
script = os.path.join(tmp, "child.py")
with open(script, "w", encoding="utf-8") as handle:
handle.write(CHILD)
env = dict(os.environ)
env.pop("PROMPT_MANAGER_CONFIG_PATH", None)
env["PROMPT_MANAGER_DATA_DIR"] = tmp
env["PYTHONIOENCODING"] = "utf-8"
try:
completed = subprocess.run(
[
sys.executable,
script,
ROOT,
tmp,
str(JOB_SECONDS),
str(ENCODER_SECONDS),
],
cwd=ROOT,
env=env,
capture_output=True,
timeout=CHILD_CEILING_SECONDS,
encoding="utf-8",
errors="replace",
)
except subprocess.TimeoutExpired as hung:
self.fail(
f"the child did not exit within {CHILD_CEILING_SECONDS} s\n"
f"stdout:\n{hung.stdout}\nstderr:\n{hung.stderr}"
)
detail = f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}"
self.assertEqual(completed.returncode, 0, detail)
reports = [
line[len(REPORT_PREFIX) :]
for line in completed.stdout.splitlines()
if line.startswith(REPORT_PREFIX)
]
self.assertEqual(len(reports), 1, detail)
report = json.loads(reports[0])
self.assertLess(report["exit_seconds"], EXIT_BOUND_SECONDS, detail)
self.assertEqual(
report["stopped"],
{"job": True, "batch": "ExitRequested", "encoder": "EncoderClosed"},
detail,
)
for text in ("Traceback", "Background work not stopped in time"):
self.assertNotIn(text, completed.stderr, detail)
class _Isolated(unittest.TestCase):
"""Each test starts with the signal clear and no hook installed."""
def setUp(self):
exit_signal.reset()
self.addCleanup(exit_signal.reset)
class TestInstall(_Isolated):
def test_the_hook_is_registered_once_before_the_executor_join(self):
calls = []
with mock.patch.object(
exit_signal.threading,
"_register_atexit",
side_effect=lambda func: calls.append(func),
create=True,
):
self.assertTrue(exit_signal.install())
self.assertTrue(exit_signal.install())
self.assertEqual(len(calls), 1)
self.assertFalse(exit_signal.requested())
calls[0]()
self.assertTrue(exit_signal.requested())
def test_a_missing_internal_hook_degrades_to_the_old_exit(self):
with mock.patch.object(exit_signal, "threading") as fake:
del fake._register_atexit
self.assertFalse(exit_signal.install())
self.assertFalse(exit_signal.requested())
def test_a_refused_registration_degrades_to_the_old_exit(self):
refuse = RuntimeError("can't register atexit after shutdown")
with mock.patch.object(
exit_signal.threading, "_register_atexit", side_effect=refuse, create=True
):
self.assertFalse(exit_signal.install())
def test_a_failing_exit_action_never_escapes_the_hook(self):
# An exception out of a threading atexit callback would abort
# threading._shutdown and skip the executor join itself.
calls = []
def boom():
raise KeyboardInterrupt # a second Ctrl+C
with mock.patch.object(
exit_signal.threading,
"_register_atexit",
side_effect=lambda func: calls.append(func),
create=True,
):
exit_signal.install(boom)
with self.assertLogs(exit_signal.logger, "ERROR"):
calls[0]()
self.assertTrue(exit_signal.requested())
class TestCheckpoints(_Isolated):
def test_checked_runs_until_the_signal_then_raises(self):
seen = []
record = exit_signal.checked(seen.append)
record(1)
exit_signal.request()
with self.assertRaises(exit_signal.ExitRequested):
record(2)
self.assertEqual(seen, [1])
def test_map_parallel_starts_no_item_after_the_signal(self):
from utils.parallel import map_parallel
started = []
gate = threading.Event()
def item(index):
started.append(index)
if index == 0:
exit_signal.request()
gate.set()
return index
with self.assertRaises(exit_signal.ExitRequested):
map_parallel(item, range(50), 1)
self.assertTrue(gate.is_set())
self.assertEqual(started, [0])
def test_job_cancelled_follows_the_signal_and_its_own_cancel(self):
from py.api.jobs import Job
job = Job("scan", "Scan", admission="records")
self.assertFalse(job.cancelled)
job.cancelled = True
self.assertTrue(job.cancelled)
other = Job("scan", "Scan", admission="records", cancellable=False)
exit_signal.request()
self.assertTrue(other.cancelled)
def test_the_scan_walk_stops_between_directories(self):
from py.api.admin import _collect_media_files
with tempfile.TemporaryDirectory() as tmp:
for name in ("a", "b"):
os.makedirs(os.path.join(tmp, name))
with open(os.path.join(tmp, name, "x.png"), "wb") as image:
image.write(b"x")
self.assertEqual(len(_collect_media_files([tmp], [".png"])), 2)
exit_signal.request()
with self.assertRaises(exit_signal.ExitRequested):
_collect_media_files([tmp], [".png"])
class TestLoraImportStopsAtExit(_Isolated):
def test_staging_stops_before_the_next_sidecar_whatever_the_caller_checks(self):
import types
from py import lora_import
discovery = types.SimpleNamespace(
sidecars=[object(), object()], unreadable=[], layout=None
)
exit_signal.request()
with tempfile.TemporaryDirectory() as tmp:
with self.assertRaises(lora_import.ImportCancelled):
lora_import.stage(
discovery,
tmp,
media_extensions=frozenset(),
cancelled=lambda: False, # the legacy scan's own check
)
class _FakeTqdm:
"""huggingface_hub.utils.tqdm as far as a download uses it."""
def __init__(self, *args, **kwargs):
self.n = 0
def update(self, n=1):
self.n += n
def __enter__(self):
return self
def __exit__(self, *exc):
return False
def _fake_hub(chunks_before_exit):
"""A ``huggingface_hub`` whose download writes chunks to a temp file and,
like huggingface_hub 1.x, removes it on any failure before the rename."""
import types
def hf_hub_download(
repo_id, filename, local_dir=None, local_dir_use_symlinks=None, tqdm_class=None
):
final = os.path.join(local_dir, filename)
partial = final + ".abc.incomplete"
bar = (tqdm_class or _FakeTqdm)(total=10)
try:
with open(partial, "wb") as handle:
for index in range(10):
if index == chunks_before_exit:
exit_signal.request()
bar.update(1)
handle.write(b"x")
os.replace(partial, final)
finally:
if os.path.exists(partial):
os.unlink(partial)
return final
hub = types.ModuleType("huggingface_hub")
hub.hf_hub_download = hf_hub_download
hub.snapshot_download = lambda **kwargs: None
utils = types.ModuleType("huggingface_hub.utils")
utils.tqdm = _FakeTqdm
hub.utils = utils
return {"huggingface_hub": hub, "huggingface_hub.utils": utils}
class TestModelDownloadStopsAtExit(_Isolated):
def test_a_download_stops_between_chunks_and_leaves_no_file(self):
from py.autotag import AutoTagService
with tempfile.TemporaryDirectory() as tmp:
service = AutoTagService(models_dir=tmp)
with mock.patch.dict(sys.modules, _fake_hub(chunks_before_exit=3)):
with self.assertLogs(service.logger, "INFO") as logs:
self.assertFalse(service.download_model("wd14-vit"))
folder = os.path.join(tmp, "WD14-ViT")
self.assertEqual(os.listdir(folder), [])
self.assertTrue(any("exiting" in line for line in logs.output), logs.output)
self.assertFalse(any("Download failed" in line for line in logs.output))
def test_a_download_without_the_signal_finishes(self):
from py.autotag import AutoTagService
with tempfile.TemporaryDirectory() as tmp:
service = AutoTagService(models_dir=tmp)
with mock.patch.dict(sys.modules, _fake_hub(chunks_before_exit=None)):
self.assertTrue(service.download_model("wd14-vit"))
folder = os.path.join(tmp, "WD14-ViT")
self.assertEqual(
sorted(os.listdir(folder)), ["model.onnx", "selected_tags.csv"]
)
def test_xet_transfers_are_aborted_the_way_hub_aborts_on_ctrl_c(self):
import types
from py.autotag import abort_hub_downloads
aborted = []
xet = types.ModuleType("huggingface_hub.utils._xet")
xet.abort_xet_session = lambda: aborted.append(True)
hub = _fake_hub(chunks_before_exit=None)
hub["huggingface_hub.utils._xet"] = xet
hub["huggingface_hub.utils"]._xet = xet
with mock.patch.dict(sys.modules, hub):
abort_hub_downloads()
self.assertEqual(aborted, [True])
with mock.patch.dict(sys.modules):
sys.modules.pop("huggingface_hub", None)
abort_hub_downloads() # never imported: nothing to abort
self.assertEqual(aborted, [True])
def test_an_aborted_xet_download_is_logged_as_stopped(self):
from py.autotag import AutoTagService
hub = _fake_hub(chunks_before_exit=None)
def aborted(*args, **kwargs):
exit_signal.request() # the hook ran, then hub raised
raise RuntimeError("xet session aborted")
hub["huggingface_hub"].hf_hub_download = aborted
with tempfile.TemporaryDirectory() as tmp:
service = AutoTagService(models_dir=tmp)
with mock.patch.dict(sys.modules, hub):
with self.assertLogs(service.logger, "INFO") as logs:
self.assertFalse(service.download_model("wd14-vit"))
self.assertTrue(any("exiting" in line for line in logs.output), logs.output)
self.assertFalse(any("Download failed" in line for line in logs.output))
def test_an_old_hub_without_tqdm_class_downloads_as_before(self):
from py.autotag import AutoTagService
hub = _fake_hub(chunks_before_exit=None)
full = hub["huggingface_hub"].hf_hub_download
def old_download(
repo_id, filename, local_dir=None, local_dir_use_symlinks=None
):
return full(repo_id, filename, local_dir, local_dir_use_symlinks)
hub["huggingface_hub"].hf_hub_download = old_download
with tempfile.TemporaryDirectory() as tmp:
service = AutoTagService(models_dir=tmp)
with mock.patch.dict(sys.modules, hub):
self.assertTrue(service.download_model("wd14-vit"))
if __name__ == "__main__":
unittest.main()
+523
View File
@@ -0,0 +1,523 @@
"""Image categories (#332): schema, migration, restore check, store and file key."""
import hashlib
import ntpath
import os
import sqlite3
import tempfile
import unittest
from contextlib import closing
from pathlib import Path
from unittest.mock import patch
try:
from open_handles import assert_closed
except ImportError:
from tests.open_handles import assert_closed
from database.image_categories import (
COPY_SQL,
INDEX,
STATEMENTS,
TABLE,
copy_image_category,
image_categories_schema_error,
migrate_image_categories,
)
from database.models import PromptModel
from database.operations import PromptDatabase
from py.api.images import image_category_key
_DDL = (
f"CREATE TABLE {TABLE} (root_key TEXT NOT NULL, relative_path TEXT NOT NULL, "
"category TEXT NOT NULL, updated_at TEXT NOT NULL, "
"PRIMARY KEY (root_key, relative_path))"
)
_INDEX_DDL = f"CREATE INDEX {INDEX} ON {TABLE}(category)"
def category_checksum(conn) -> tuple:
"""sha256 of prompts and prompt_pairs ``(id, category, updated_at)``."""
sums = []
for table in ("prompts", "prompt_pairs"):
digest = hashlib.sha256()
for row in conn.execute( # nosec B608 - fixed table names
f"SELECT id, category, updated_at FROM {table} ORDER BY id"
):
digest.update(repr(tuple(row)).encode())
sums.append(digest.hexdigest())
return tuple(sums)
def table_checksums(conn) -> dict:
"""``{name: (schema sql, sha256 of rows)}`` of every table but ours."""
found = {}
for name, sql in conn.execute(
"SELECT name, sql FROM sqlite_master WHERE type='table' AND name != ?",
(TABLE,),
).fetchall():
digest = hashlib.sha256()
rows = conn.execute(f'SELECT * FROM "{name}"') # nosec B608
for text in sorted(repr(tuple(row)) for row in rows):
digest.update(text.encode())
found[name] = (sql, digest.hexdigest())
return found
class _TempDir(unittest.TestCase):
def setUp(self):
self.tmpdir = tempfile.mkdtemp()
self.path = os.path.join(self.tmpdir, "categories.db")
def tearDown(self):
PromptModel.close_all_instances()
assert_closed(self, self.tmpdir)
for name in os.listdir(self.tmpdir):
os.unlink(os.path.join(self.tmpdir, name))
os.rmdir(self.tmpdir)
class TestImageCategoriesSchema(_TempDir):
def test_migration_creates_table_and_index(self):
db = PromptDatabase(self.path)
db.close_all()
PromptModel.close_all_instances()
with closing(sqlite3.connect(self.path)) as conn:
names = {
(kind, name)
for kind, name in conn.execute(
"SELECT type, name FROM sqlite_master WHERE tbl_name=?", (TABLE,)
)
}
self.assertEqual(
names,
{
("table", TABLE),
("index", INDEX),
("index", f"sqlite_autoindex_{TABLE}_1"),
},
)
self.assertEqual(
conn.execute(f"SELECT COUNT(*) FROM {TABLE}").fetchone()[0], 0
)
self.assertIsNone(image_categories_schema_error(conn, allow_missing=False))
def test_migration_leaves_prompt_categories_unchanged(self):
"""23,000 links; the migration reads and writes only its own objects."""
db = PromptDatabase(self.path)
db.close_all()
PromptModel.close_all_instances()
with closing(sqlite3.connect(self.path)) as conn:
conn.execute(f"DROP TABLE {TABLE}") # a library from before #332
categories = ["Landscapes", "Portraits", None, " ", "Scans"]
conn.executemany(
"INSERT INTO prompts (text, category, hash, updated_at) "
"VALUES (?, ?, ?, ?)",
[
(f"p{n}", categories[n % 5], f"h{n}", f"2026-01-01 00:00:{n % 60}")
for n in range(200)
],
)
conn.executemany(
"INSERT INTO prompt_pairs (positive_prompt_id, negative_prompt_id, "
"category) VALUES (?, ?, ?)",
[(n, n + 100, categories[(n + 1) % 5]) for n in range(1, 50)],
)
conn.executemany(
"INSERT INTO generated_images (prompt_id, image_path, filename, "
"file_path) VALUES (?, ?, ?, ?)",
[
(n % 200 + 1, f"/out/{n}.png", f"{n}.png", f"/out/{n}.png")
for n in range(23000)
],
)
conn.commit()
before = table_checksums(conn)
prompts_before = category_checksum(conn)
seen = []
def authorize(action, arg1, arg2, _db, _trigger):
seen.append((action, arg1, arg2))
return sqlite3.SQLITE_OK
conn.set_authorizer(authorize)
migrate_image_categories(conn)
conn.commit()
conn.set_authorizer(None)
own = {TABLE, "sqlite_master", "sqlite_schema", INDEX}
touched = set()
for action, arg1, arg2 in seen:
if action in (sqlite3.SQLITE_READ, sqlite3.SQLITE_INSERT):
touched.add(arg1)
if action == sqlite3.SQLITE_READ:
self.assertIn(arg1, own, (action, arg1, arg2))
elif action in (
sqlite3.SQLITE_INSERT,
sqlite3.SQLITE_UPDATE,
sqlite3.SQLITE_DELETE,
):
self.assertIn(arg1, own, (action, arg1, arg2))
elif action == sqlite3.SQLITE_CREATE_TABLE:
self.assertEqual(arg1, TABLE)
elif action == sqlite3.SQLITE_CREATE_INDEX:
self.assertIn(arg1, (INDEX, f"sqlite_autoindex_{TABLE}_1"))
self.assertEqual(arg2, TABLE)
elif action == sqlite3.SQLITE_PRAGMA:
self.assertEqual((arg1, arg2), ("table_info", TABLE))
self.assertIn(TABLE, touched) # the authorizer saw the work
self.assertEqual(table_checksums(conn), before)
self.assertEqual(category_checksum(conn), prompts_before)
schema = conn.execute(
"SELECT sql FROM sqlite_master ORDER BY name"
).fetchall()
migrate_image_categories(conn) # a second run changes nothing
conn.commit()
self.assertEqual(
conn.execute("SELECT sql FROM sqlite_master ORDER BY name").fetchall(),
schema,
)
# The whole startup sequence keeps them too.
PromptDatabase(self.path).close_all()
PromptModel.close_all_instances()
with closing(sqlite3.connect(self.path)) as conn:
self.assertEqual(category_checksum(conn), prompts_before)
self.assertEqual(
conn.execute(f"SELECT COUNT(*) FROM {TABLE}").fetchone()[0], 0
)
def test_structure_accepted_and_refused(self):
variants = {
"missing": None,
"extra column": (
f"CREATE TABLE {TABLE} (root_key TEXT NOT NULL, relative_path TEXT "
"NOT NULL, category TEXT NOT NULL, updated_at TEXT NOT NULL, x TEXT, "
"PRIMARY KEY (root_key, relative_path))",
_INDEX_DDL,
),
"wrong index": (_DDL, f"CREATE INDEX {INDEX} ON {TABLE}(updated_at)"),
"extra index": (
_DDL,
_INDEX_DDL,
f"CREATE INDEX extra ON {TABLE}(root_key)",
),
"trigger": (
_DDL,
_INDEX_DDL,
f"CREATE TRIGGER t AFTER INSERT ON {TABLE} BEGIN SELECT 1; END",
),
"wrong ddl": (
f"CREATE TABLE {TABLE} (root_key TEXT, relative_path TEXT, "
"category TEXT, updated_at TEXT, "
"PRIMARY KEY (root_key, relative_path))",
_INDEX_DDL,
),
"view": (f"CREATE VIEW {TABLE} AS SELECT 1",),
}
for label, statements in variants.items():
with self.subTest(label), closing(sqlite3.connect(":memory:")) as conn:
for sql in statements or ():
conn.execute(sql)
error = image_categories_schema_error(conn)
if statements is None:
self.assertIsNone(error)
self.assertIsNotNone(
image_categories_schema_error(conn, allow_missing=False)
)
continue
self.assertIsNotNone(error)
self.assertIsNotNone(
image_categories_schema_error(conn, check_values=False)
)
with self.assertRaises(sqlite3.DatabaseError):
migrate_image_categories(conn)
# A table without its category index (never written by this code) is
# refused by the check but repaired by the install.
with closing(sqlite3.connect(":memory:")) as conn:
conn.execute(_DDL)
conn.execute(f"INSERT INTO {TABLE} VALUES ('r', 'a.png', 'Scans', 'now')")
self.assertIsNotNone(image_categories_schema_error(conn))
migrate_image_categories(conn)
self.assertIsNone(image_categories_schema_error(conn, allow_missing=False))
self.assertEqual(
conn.execute(f"SELECT COUNT(*) FROM {TABLE}").fetchone()[0], 1
)
with closing(sqlite3.connect(":memory:")) as conn:
migrate_image_categories(conn)
migrate_image_categories(conn) # idempotent
self.assertIsNone(image_categories_schema_error(conn, allow_missing=False))
def test_values_refused_but_startup_checks_structure_only(self):
bad = {
"blank": "",
"101 characters": "x" * 101,
"control character": "a\tb",
"nul": "a\x00b",
"leading space": " Landscapes",
"trailing space": "Landscapes ",
}
for label, value in bad.items():
with self.subTest(label), closing(sqlite3.connect(":memory:")) as conn:
migrate_image_categories(conn)
conn.execute(
f"INSERT INTO {TABLE} VALUES ('r', 'a.png', ?, 'now')", (value,)
)
self.assertIsNotNone(image_categories_schema_error(conn))
self.assertIsNone(
image_categories_schema_error(conn, check_values=False)
)
migrate_image_categories(conn) # startup still opens it
with closing(sqlite3.connect(":memory:")) as conn:
migrate_image_categories(conn)
conn.execute(
f"INSERT INTO {TABLE} VALUES ('r', ?, 'Scans', 'now')", ("x" * 5000,)
)
self.assertIsNotNone(image_categories_schema_error(conn))
conn.execute(f"DELETE FROM {TABLE}")
conn.execute(f"INSERT INTO {TABLE} VALUES ('r', x'00ff', 'Scans', 'now')")
self.assertIsNotNone(image_categories_schema_error(conn))
conn.execute(f"DELETE FROM {TABLE}")
for value in ("x" * 100, "Scans", "Café · 風景", "a b"):
conn.execute(
f"INSERT OR REPLACE INTO {TABLE} VALUES ('r', 'a.png', ?, 'now')",
(value,),
)
self.assertIsNone(image_categories_schema_error(conn), value)
def test_restore_check(self):
PromptDatabase(self.path).close_all()
PromptModel.close_all_instances()
ok, reason = PromptModel.verify_database_file(
self.path, require_library_state=True
)
self.assertTrue(ok, reason)
with closing(sqlite3.connect(self.path)) as conn:
conn.execute(f"DROP TABLE {TABLE}")
conn.commit()
# A backup from before #332 is accepted ...
ok, reason = PromptModel.verify_database_file(self.path)
self.assertTrue(ok, reason)
# ... but a current database must have the table.
ok, reason = PromptModel.verify_database_file(
self.path, require_library_state=True
)
self.assertFalse(ok)
self.assertIn("image categories", reason.lower())
with closing(sqlite3.connect(self.path)) as conn:
conn.execute(_DDL)
conn.execute(_INDEX_DDL)
conn.execute(f"INSERT INTO {TABLE} VALUES ('r', 'a.png', ' x', 'now')")
conn.commit()
ok, reason = PromptModel.verify_database_file(self.path)
self.assertFalse(ok)
self.assertIn("image categories", reason.lower())
class TestImageCategoryStore(_TempDir):
def setUp(self):
super().setUp()
self.db = PromptDatabase(self.path)
def tearDown(self):
self.db.close_all()
super().tearDown()
def rows(self):
return (
self.db.model.get_connection()
.execute(
f"SELECT root_key, relative_path, category, updated_at FROM {TABLE} "
"ORDER BY root_key, relative_path"
)
.fetchall()
)
def test_set_counts_updated_and_unchanged(self):
keys = [("/out", "a.png"), ("/other", "a.png"), ("/out", "a.png")]
self.assertEqual(
self.db.set_image_categories(keys, "Portraits"),
{"updated": 2, "unchanged": 0},
)
stamps = {(r[0], r[1]): r[3] for r in self.rows()}
self.assertEqual(
self.db.set_image_categories(
[("/out", "a.png"), ("/out", "b.png")], "Portraits"
),
{"updated": 1, "unchanged": 1},
)
# An unchanged file keeps its updated_at.
self.assertEqual(
{(r[0], r[1]): r[3] for r in self.rows()}[("/out", "a.png")],
stamps[("/out", "a.png")],
)
self.assertEqual(
self.db.set_image_categories([("/out", "a.png")], "Scans"),
{"updated": 1, "unchanged": 0},
)
self.assertEqual(self.db.image_category("/out", "a.png"), "Scans")
self.assertEqual(self.db.image_category("/other", "a.png"), "Portraits")
self.assertIsNone(self.db.image_category("/out", "missing.png"))
self.assertEqual(
self.db.own_categories_of(["/out", "/out", "/none"]),
{("/out", "a.png"): "Scans", ("/out", "b.png"): "Portraits"},
)
def test_clear_counts_removed_rows(self):
self.db.set_image_categories([("/out", "a.png"), ("/b", "c.png")], "Scans")
self.assertEqual(
self.db.clear_image_categories(
[("/out", "a.png"), ("/b", "c.png"), ("/out", "gone.png")]
),
2,
)
self.assertEqual(self.rows(), [])
self.assertEqual(self.db.clear_image_categories([("/out", "a.png")]), 0)
def test_a_failed_write_writes_nothing(self):
self.db.set_image_categories([("/out", "a.png")], "Scans")
keys = [("/out", "b.png"), ("/out", None)] # NOT NULL fails the second
with self.assertRaises(sqlite3.IntegrityError):
self.db.set_image_categories(keys, "Portraits")
self.assertEqual([r[:3] for r in self.rows()], [("/out", "a.png", "Scans")])
def test_store_never_touches_prompt_categories(self):
conn = self.db.model.get_connection()
with conn:
conn.execute(
"INSERT INTO prompts (text, category, hash) VALUES ('p', 'L', 'h')"
)
before = category_checksum(conn)
seen = set()
conn.set_authorizer(
lambda action, arg1, *_: seen.add(arg1) or sqlite3.SQLITE_OK
)
try:
self.db.set_image_categories([("/out", "a.png")], "Scans")
self.db.own_categories_of(["/out"])
self.db.image_category("/out", "a.png")
self.db.clear_image_categories([("/out", "a.png")])
finally:
conn.set_authorizer(None)
self.assertFalse(seen & {"prompts", "prompt_pairs", "generated_images"}, seen)
self.assertEqual(category_checksum(conn), before)
def test_copy_keeps_kept_copy_category(self):
self.db.set_image_categories([("/out", "removed.png")], "Scans")
self.db.set_image_categories([("/out", "kept.png")], "Portraits")
with self.db._pair_transaction() as conn:
inserted = copy_image_category(
conn, ("/out", "removed.png"), ("/out", "kept.png")
)
self.assertEqual(inserted, 0)
self.assertEqual(self.db.image_category("/out", "kept.png"), "Portraits")
with self.db._pair_transaction() as conn:
inserted = copy_image_category(
conn, ("/out", "removed.png"), ("/out", "new.png")
)
self.assertEqual(inserted, 1)
self.assertEqual(self.db.image_category("/out", "new.png"), "Scans")
# The removed copy's row is never deleted.
self.assertEqual(self.db.image_category("/out", "removed.png"), "Scans")
with self.db._pair_transaction() as conn:
self.assertEqual(
copy_image_category(conn, ("/out", "none.png"), ("/out", "x.png")), 0
)
self.assertEqual(
copy_image_category(
conn, ("/out", "removed.png"), ("/out", "removed.png")
),
0,
)
def test_statements_never_scan_the_table(self):
conn = self.db.model.get_connection()
for sql in STATEMENTS:
with self.subTest(sql):
params = (None,) * sql.count("?")
plan = " ".join(
row[-1] for row in conn.execute(f"EXPLAIN QUERY PLAN {sql}", params)
)
self.assertNotRegex(plan, rf"SCAN {TABLE}\b", plan)
plan = " ".join(
row[-1]
for row in conn.execute(
f"EXPLAIN QUERY PLAN SELECT rowid FROM {TABLE} WHERE category=? "
"ORDER BY rowid",
("x",),
)
)
self.assertIn(INDEX, plan)
self.assertIn(COPY_SQL, STATEMENTS)
class TestImageCategoryKey(unittest.TestCase):
def setUp(self):
self._tmp = tempfile.TemporaryDirectory()
self.base = Path(os.path.realpath(self._tmp.name))
self.outer = self.base / "output"
self.inner = self.outer / "nested"
(self.inner / "sub").mkdir(parents=True)
def tearDown(self):
self._tmp.cleanup()
def test_equals_the_favorite_key_without_nesting(self):
from py.api.images import _favorite_relative_key, _favorite_root_key
key = image_category_key([self.outer], self.outer / "sub" / "a.png")
self.assertEqual(
key, (_favorite_root_key(self.outer), _favorite_relative_key("sub/a.png"))
)
def test_deepest_source_wins_in_any_order(self):
from py.api.images import _favorite_root_key
path = self.inner / "sub" / "a.png"
for sources in ([self.outer, self.inner], [self.inner, self.outer]):
with self.subTest(sources=sources):
self.assertEqual(
image_category_key(sources, path),
(_favorite_root_key(self.inner), "sub/a.png"),
)
# A file of the outer source only is keyed by it.
self.assertEqual(
image_category_key([self.outer, self.inner], self.outer / "b.png"),
(_favorite_root_key(self.outer), "b.png"),
)
def test_outside_every_source_has_no_key(self):
self.assertIsNone(image_category_key([self.inner], self.outer / "a.png"))
self.assertIsNone(image_category_key([], self.outer / "a.png"))
# A sibling sharing a name prefix is not inside.
self.assertIsNone(
image_category_key([self.outer], self.base / "output2" / "a.png")
)
self.assertIsNone(image_category_key([self.outer], self.outer))
def test_symlinked_subfolder_is_not_resolved(self):
target = self.base / "elsewhere"
target.mkdir()
link = self.outer / "linked"
try:
os.symlink(target, link, target_is_directory=True)
except (OSError, NotImplementedError):
self.skipTest("symlinks are not available")
key = image_category_key([self.outer], link / "a.png")
self.assertEqual(key[1], "linked/a.png")
def test_windows_folds_case_of_the_relative_key(self):
with patch("py.api.images._path_flavor", ntpath):
upper = image_category_key([self.outer], self.outer / "Sub" / "A.PNG")
lower = image_category_key([self.outer], self.outer / "sub" / "a.png")
self.assertEqual(upper[1], "sub/a.png")
self.assertEqual(upper, lower)
def test_posix_keeps_case_distinct(self):
if os.path.normcase("A") == "a":
self.skipTest("this OS folds case")
upper = image_category_key([self.outer], self.outer / "A.png")
lower = image_category_key([self.outer], self.outer / "a.png")
self.assertNotEqual(upper, lower)
if __name__ == "__main__":
unittest.main()
+268
View File
@@ -0,0 +1,268 @@
"""Collections (#331): the database operations."""
import sqlite3
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
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.models import PromptModel
from database.operations import PromptDatabase
try: # guarded so a missing module fails assertions, not test loading
from database.image_collections import (
CollectionFilesMissing,
CollectionInvalid,
CollectionMissing,
CollectionNameTaken,
READ_SQL,
)
except ImportError:
CollectionFilesMissing = CollectionMissing = CollectionNameTaken = None
CollectionInvalid = ValueError
READ_SQL = None
ROOT = "/real/output"
OTHER = "/real/other"
class ImageCollectionOperationsTests(unittest.TestCase):
def setUp(self):
self.temp = tempfile.TemporaryDirectory()
self.addCleanup(self.temp.cleanup)
self.addCleanup(assert_closed, self, self.temp.name)
self.addCleanup(PromptModel.reset_schema_cache)
self.db = PromptDatabase(str(Path(self.temp.name) / "prompts.db"))
self.addCleanup(self.db.close_all)
def require(self):
self.assertTrue(
hasattr(self.db, "add_to_collection"),
"ImageCollectionOperations is not implemented",
)
def add(self, items, **target):
return self.db.add_to_collection(items=items, **target)
def test_add_creates_a_collection_by_name(self):
self.require()
result = self.add(
[(ROOT, "a.png"), (ROOT, "b.png"), (OTHER, "a.png")], name="LoRa testing"
)
collection = result["collection"]
self.assertEqual(collection["name"], "LoRa testing")
self.assertEqual(
result,
{"collection": collection, "created": True, "added": 3, "already": 0},
)
self.assertEqual(self.db.list_collections(), [collection])
self.assertTrue(self.db.collection_exists(collection["id"]))
self.assertEqual(self.db.get_collection(collection["id"]), collection)
self.assertEqual(
self.db.collection_members(collection["id"], ROOT), {"a.png", "b.png"}
)
self.assertEqual(
self.db.root_collection_members(OTHER), {collection["id"]: {"a.png"}}
)
self.assertEqual(self.db.collection_counts(), {collection["id"]: 3})
self.assertEqual(self.db.file_collections(ROOT, "a.png"), [collection["id"]])
self.assertEqual(self.db.file_collections(ROOT, "zzz.png"), [])
def test_add_by_name_ignores_case_and_unicode_form(self):
self.require()
first = self.add([(ROOT, "a.png")], name="Café grids")["collection"]
# NFD spelling (macOS) in another case joins the same collection.
second = self.add([(ROOT, "b.png")], name="CAFÉ GRIDS")
self.assertEqual(second["collection"], first)
self.assertFalse(second["created"])
self.assertEqual(len(self.db.list_collections()), 1)
self.assertEqual(
self.db.collection_members(first["id"], ROOT), {"a.png", "b.png"}
)
def test_add_counts_already_present(self):
self.require()
made = self.add([(ROOT, "a.png"), (ROOT, "b.png")], name="Set")
result = self.add(
[(ROOT, "b.png"), (ROOT, "c.png"), (ROOT, "c.png")],
collection_id=made["collection"]["id"],
)
self.assertEqual(
(result["created"], result["added"], result["already"]), (False, 1, 1)
)
self.assertEqual(self.db.collection_counts(), {made["collection"]["id"]: 3})
def test_add_with_a_new_name_and_nothing_inserted_creates_nothing(self):
self.require()
with self.assertRaises(CollectionFilesMissing):
self.add([], name="Empty")
self.assertEqual(self.db.list_collections(), [])
# An existing name with nothing to add is not an error.
made = self.add([(ROOT, "a.png")], name="Kept")["collection"]
self.assertEqual(
self.add([], name="kept"),
{"collection": made, "created": False, "added": 0, "already": 0},
)
def test_add_to_a_deleted_collection_raises_missing(self):
self.require()
made = self.add([(ROOT, "a.png")], name="Gone")["collection"]
self.db.delete_collection(made["id"])
with self.assertRaises(CollectionMissing):
self.add([(ROOT, "b.png")], collection_id=made["id"])
with self.assertRaises(CollectionMissing):
self.db.remove_from_collection(made["id"], [(ROOT, "a.png")])
with self.assertRaises(CollectionMissing):
self.db.rename_collection(made["id"], "Back")
self.assertFalse(self.db.collection_exists(made["id"]))
self.assertIsNone(self.db.get_collection(made["id"]))
def test_remove_returns_the_count(self):
self.require()
made = self.add([(ROOT, "a.png"), (ROOT, "b.png")], name="Set")["collection"]
result = self.db.remove_from_collection(
made["id"], [(ROOT, "a.png"), (ROOT, "a.png"), (ROOT, "never.png")]
)
self.assertEqual(result, {"collection": made, "removed": 1})
self.assertEqual(self.db.collection_members(made["id"], ROOT), {"b.png"})
# Removing the last member keeps the (now empty) collection.
self.db.remove_from_collection(made["id"], [(ROOT, "b.png")])
self.assertEqual(self.db.list_collections(), [made])
self.assertEqual(self.db.collection_counts(), {made["id"]: 0})
def test_rename_refuses_a_taken_name(self):
self.require()
first = self.add([(ROOT, "a.png")], name="First")["collection"]
second = self.add([(ROOT, "a.png")], name="Second")["collection"]
with self.assertRaises(CollectionNameTaken) as caught:
self.db.rename_collection(second["id"], " first ")
self.assertEqual(caught.exception.collection_id, first["id"])
self.assertEqual(self.db.get_collection(second["id"]), second)
def test_rename_accepts_a_case_variant_of_its_own_name(self):
self.require()
made = self.add([(ROOT, "a.png")], name="lora grids")["collection"]
renamed = self.db.rename_collection(made["id"], "LoRA grids ")
self.assertEqual(renamed, {"id": made["id"], "name": "LoRA grids"})
self.assertEqual(self.db.list_collections(), [renamed])
other = self.db.rename_collection(made["id"], "Portraits")
self.assertEqual(other["name"], "Portraits")
def test_delete_unknown_raises_missing(self):
self.require()
made = self.add([(ROOT, "a.png")], name="Doomed")["collection"]
self.assertTrue(self.db.delete_collection(made["id"]))
self.assertEqual(self.db.root_collection_members(ROOT), {})
with self.assertRaises(CollectionMissing):
self.db.delete_collection(made["id"])
with self.assertRaises(CollectionMissing):
self.db.delete_collection(999)
def test_list_is_in_name_key_order(self):
self.require()
for name in ("beta", "Alpha", "été", "Zeta", "alpha two"):
self.add([(ROOT, "a.png")], name=name)
self.assertEqual(
[row["name"] for row in self.db.list_collections()],
["Alpha", "alpha two", "beta", "Zeta", "été"],
)
def test_names_are_trimmed_and_collapsed(self):
self.require()
made = self.add([(ROOT, "a.png")], name=" LoRa \t testing\n ")["collection"]
self.assertEqual(made["name"], "LoRa testing")
for bad in ("", " ", "x" * 101, "bell\x07", "c1\x81", "del\x7f", 5, None):
with self.subTest(name=bad):
with self.assertRaises(CollectionInvalid):
self.add([(ROOT, "a.png")], name=bad)
self.assertEqual(len(self.db.list_collections()), 1)
def test_targets_and_items_are_validated(self):
self.require()
made = self.add([(ROOT, "a.png")], name="Set")["collection"]
for target in (
{},
{"name": "x", "collection_id": made["id"]},
{"collection_id": True},
{"collection_id": 0},
{"collection_id": "1"},
):
with self.subTest(target=target):
with self.assertRaises(CollectionInvalid):
self.add([(ROOT, "a.png")], **target)
for items in ([("root",)], [(ROOT, 5)], "a.png", [(ROOT, "x" * 4097)]):
with self.subTest(items=items):
with self.assertRaises(CollectionInvalid):
self.add(items, collection_id=made["id"])
def test_a_failure_after_the_membership_move_rolls_it_back(self):
self.require()
made = self.add([(ROOT, "b.png")], name="Grids")["collection"]
move = {
"removed": [],
"kept": [],
"targets": {},
"kept_path": "/real/output/a.png",
"favorites": [(ROOT, "b.png", ROOT, "a.png")],
"removed_locator": "/real/output/b.png",
"kept_locator": "/real/output/a.png",
"sha256": "0" * 64,
"kept_size": 1,
"kept_mtime_ns": 1,
}
with patch.object(
type(self.db), "_move_copy_tags", side_effect=RuntimeError("boom")
):
with self.assertRaises(RuntimeError):
self.db.move_copy_records(move)
self.assertEqual(self.db.file_collections(ROOT, "b.png"), [made["id"]])
self.assertEqual(self.db.file_collections(ROOT, "a.png"), [])
# Without the failure the same move succeeds.
self.assertEqual(self.db.move_copy_records(move)["collections"], 1)
self.assertEqual(self.db.file_collections(ROOT, "a.png"), [made["id"]])
def test_count_snapshot_reads_in_one_transaction(self):
self.require()
made = self.add([(ROOT, "a.png"), (OTHER, "b.png")], name="Set")["collection"]
statements = []
conn = self.db.model.get_connection()
conn.set_trace_callback(statements.append)
try:
members, stored = self.db.collection_count_snapshot([ROOT, OTHER, ROOT])
finally:
conn.set_trace_callback(None)
self.assertEqual(
members, {ROOT: {made["id"]: {"a.png"}}, OTHER: {made["id"]: {"b.png"}}}
)
self.assertEqual(stored, {made["id"]: 2})
self.assertEqual(statements[0], "BEGIN")
self.assertEqual(statements[-1], "ROLLBACK")
self.assertFalse(conn.in_transaction)
def test_reads_use_an_index(self):
self.require()
self.add([(ROOT, "a.png")], name="Set")
conn = sqlite3.connect(self.db.model.db_path)
try:
for sql in READ_SQL:
plan = " ".join(
row[3]
for row in conn.execute(
"EXPLAIN QUERY PLAN " + sql, [None] * sql.count("?")
)
)
with self.subTest(sql=sql):
self.assertNotIn("SCAN image_collection_items", plan)
if "?" in sql: # a lookup, not the whole list
self.assertNotIn("SCAN image_collections", plan)
finally:
conn.close()
if __name__ == "__main__":
unittest.main()
+521
View File
@@ -0,0 +1,521 @@
"""Collections schema (#331): install, validation and restore safety."""
import re
import sqlite3
import tempfile
import unittest
import unicodedata
from contextlib import closing
from pathlib import Path
from unittest.mock import patch
try:
from tests.open_handles import assert_closed
from tests.test_annotation_schema import LEGACY_TABLES, table_rows
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed
from test_annotation_schema import LEGACY_TABLES, table_rows
from database.models import PromptModel
try: # guarded so a missing module fails assertions, not test loading
from database import image_collections
from database.image_collections import (
image_collections_schema_error,
migrate_image_collections,
)
except ImportError:
image_collections = None
image_collections_schema_error = None
migrate_image_collections = None
COLLECTIONS = "image_collections"
ITEMS = "image_collection_items"
INDEX = "idx_image_collection_items_file"
TABLES = (COLLECTIONS, ITEMS)
DDL = {
COLLECTIONS: (
f"CREATE TABLE {COLLECTIONS} (id INTEGER PRIMARY KEY AUTOINCREMENT, "
"name TEXT NOT NULL, name_key TEXT NOT NULL UNIQUE, "
"created_at TEXT NOT NULL)"
),
ITEMS: (
f"CREATE TABLE {ITEMS} (collection_id INTEGER NOT NULL REFERENCES "
f"{COLLECTIONS}(id) ON DELETE CASCADE, root_key TEXT NOT NULL, "
"relative_path TEXT NOT NULL, added_at TEXT NOT NULL, "
"PRIMARY KEY (collection_id, root_key, relative_path))"
),
}
INDEX_DDL = f"CREATE INDEX {INDEX} ON {ITEMS}(root_key, relative_path)"
NOW = "2026-10-27T00:00:00.000000+00:00"
IMAGES = 23_000
def _compact(sql):
return re.sub(r"\s+", "", sql or "").lower()
def expected_catalog():
return {
("table", COLLECTIONS, _compact(DDL[COLLECTIONS])),
("table", ITEMS, _compact(DDL[ITEMS])),
("index", INDEX, _compact(INDEX_DDL)),
("index", f"sqlite_autoindex_{COLLECTIONS}_1", ""),
("index", f"sqlite_autoindex_{ITEMS}_1", ""),
}
def catalog(conn):
return {
(kind, name, _compact(sql))
for kind, name, sql in conn.execute(
"SELECT type,name,sql FROM sqlite_master "
"WHERE name IN (?,?,?) OR tbl_name IN (?,?)",
(*TABLES, INDEX, *TABLES),
)
}
def key(name):
return unicodedata.normalize("NFC", unicodedata.normalize("NFC", name).casefold())
def schema(conn, exclude=()):
return sorted(
tuple(row)
for row in conn.execute(
"SELECT type,name,tbl_name,sql FROM sqlite_master ORDER BY name"
)
if row[1] not in exclude and row[2] not in exclude
)
def row_counts(conn):
tables = [
name
for (name,) in conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' ORDER BY name"
)
]
return {
table: conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] # nosec
for table in tables
}
class ImageCollectionsSchemaTests(unittest.TestCase):
def setUp(self):
self.temp = tempfile.TemporaryDirectory()
self.addCleanup(self.temp.cleanup)
self.addCleanup(assert_closed, self, self.temp.name)
self.addCleanup(PromptModel.reset_schema_cache)
self.models = []
self.addCleanup(self.close_models)
self.path = str(Path(self.temp.name) / "live.db")
def close_models(self):
for model in self.models:
model.close_all()
def model(self, path=None):
model = PromptModel(path or self.path)
self.models.append(model)
return model
def require_module(self):
self.assertIsNotNone(
image_collections_schema_error,
"database.image_collections is not implemented",
)
def closed_current(self, name="current.db", *statements):
"""A closed current database, altered afterwards by *statements*."""
path = str(Path(self.temp.name) / name)
model = self.model(path)
model.close_all()
PromptModel.reset_schema_cache()
with closing(sqlite3.connect(path)) as raw, raw:
raw.execute("PRAGMA foreign_keys=OFF")
for statement in statements:
if isinstance(statement, tuple):
raw.execute(*statement)
else:
raw.execute(statement)
return path
@staticmethod
def drop_group(path):
with closing(sqlite3.connect(path)) as raw, raw:
raw.execute(f"DROP TABLE {ITEMS}")
raw.execute(f"DROP TABLE {COLLECTIONS}")
def test_migration_creates_tables_and_index(self):
self.require_module()
model = self.model()
with model.get_connection() as conn:
self.assertEqual(catalog(conn), expected_catalog())
self.assertIsNone(image_collections_schema_error(conn, allow_missing=False))
self.assertEqual(
conn.execute(f"SELECT COUNT(*) FROM {COLLECTIONS}").fetchone()[0], 0
)
self.assertEqual(
conn.execute(f"SELECT COUNT(*) FROM {ITEMS}").fetchone()[0], 0
)
self.assertEqual(
PromptModel.verify_database_file(self.path, require_library_state=True),
(True, "ok"),
)
def test_migration_is_idempotent(self):
self.require_module()
model = self.model()
with model.get_connection() as conn:
conn.execute(
f"INSERT INTO {COLLECTIONS} (name, name_key, created_at) "
"VALUES ('Kept', 'kept', ?)",
(NOW,),
)
conn.execute(f"INSERT INTO {ITEMS} VALUES (1, 'root', 'a.png', ?)", (NOW,))
before = schema(conn)
model.close_all()
for _ in range(2):
PromptModel.reset_schema_cache()
model = self.model()
with model.get_connection() as conn:
migrate_image_collections(conn)
self.assertEqual(schema(conn), before)
self.assertEqual(
[tuple(r) for r in conn.execute(f"SELECT * FROM {ITEMS}")],
[(1, "root", "a.png", NOW)],
)
model.close_all()
def test_migration_reads_only_its_own_tables(self):
self.require_module()
model = self.model()
with model.get_connection() as conn:
conn.execute("INSERT INTO prompts(id,text,hash) VALUES(1,'positive','p')")
conn.executemany(
"INSERT INTO generated_images(prompt_id,image_path,filename,"
"file_path) VALUES(1,?,?,?)",
((f"/out/{n}.png", f"{n}.png", f"/out/{n}.png") for n in range(IMAGES)),
)
conn.execute("INSERT INTO image_favorites VALUES('root','a.png',?)", (NOW,))
model.close_all()
self.drop_group(self.path)
with closing(sqlite3.connect(self.path)) as raw:
raw.execute("PRAGMA foreign_keys=ON")
counts_before = row_counts(raw)
schema_before = schema(raw)
legacy_before = table_rows(raw, LEGACY_TABLES)
self.assertEqual(counts_before["generated_images"], IMAGES)
touched = []
def record(action, arg1, arg2, _db, _trigger):
touched.append((action, arg1, arg2))
return sqlite3.SQLITE_OK
raw.set_authorizer(record)
try:
migrate_image_collections(raw)
finally:
raw.set_authorizer(None)
raw.commit()
reads = {
arg1 for action, arg1, _arg2 in touched if action == sqlite3.SQLITE_READ
}
self.assertTrue(reads)
self.assertLessEqual(reads, {"sqlite_master", COLLECTIONS, ITEMS})
writes = {
arg1
for action, arg1, _arg2 in touched
if action
in (sqlite3.SQLITE_INSERT, sqlite3.SQLITE_UPDATE, sqlite3.SQLITE_DELETE)
}
self.assertLessEqual(writes, {"sqlite_master"})
self.assertEqual(catalog(raw), expected_catalog())
counts_after = row_counts(raw)
self.assertEqual(counts_after.pop(COLLECTIONS), 0)
self.assertEqual(counts_after.pop(ITEMS), 0)
self.assertEqual(counts_after, counts_before)
self.assertEqual(schema(raw, exclude=TABLES), schema_before)
self.assertEqual(table_rows(raw, LEGACY_TABLES), legacy_before)
# A second run is a no-op.
after_first = schema(raw)
changes = raw.total_changes
migrate_image_collections(raw)
raw.commit()
self.assertEqual(schema(raw), after_first)
self.assertEqual(raw.total_changes, changes)
def test_migration_refuses_a_foreign_object_of_the_same_name(self):
self.require_module()
for statements in (
(f"CREATE TABLE {COLLECTIONS} (x)",),
(f"CREATE VIEW {COLLECTIONS} AS SELECT 1",),
(f"CREATE INDEX {INDEX} ON prompts(text)",),
):
with self.subTest(statements=statements):
path = str(Path(self.temp.name) / f"foreign{len(self.models)}.db")
self.model(path).close_all()
PromptModel.reset_schema_cache()
self.drop_group(path)
with closing(sqlite3.connect(path)) as raw, raw:
for statement in statements:
raw.execute(statement)
with closing(sqlite3.connect(path)) as raw:
before = schema(raw)
self.assertIsNotNone(image_collections_schema_error(raw))
with self.assertRaises(sqlite3.DatabaseError):
self.model(path).get_connection()
PromptModel.close_all_instances()
PromptModel.reset_schema_cache()
with closing(sqlite3.connect(path)) as raw:
self.assertEqual(schema(raw), before)
self.assertIsNone(
raw.execute(
"SELECT 1 FROM sqlite_master WHERE name=?", (ITEMS,)
).fetchone()
)
self.assertFalse(PromptModel.verify_database_file(path)[0])
def test_migration_rolls_back_on_error(self):
self.require_module()
self.model().close_all()
PromptModel.reset_schema_cache()
self.drop_group(self.path)
with closing(sqlite3.connect(self.path)) as raw:
with patch.object(
image_collections,
"image_collections_schema_error",
side_effect=[None, "Image collections went wrong"],
):
with self.assertRaisesRegex(sqlite3.DatabaseError, "went wrong"):
migrate_image_collections(raw)
raw.commit()
self.assertEqual(catalog(raw), set())
def test_startup_accepts_a_value_saved_under_older_rules(self):
self.require_module()
path = self.closed_current(
"older.db",
(
f"INSERT INTO {COLLECTIONS} (name, name_key, created_at) "
"VALUES (?, 'not the key', ?)",
(" " + "x" * 150 + " ", NOW),
),
)
model = self.model(path)
with model.get_connection() as conn:
self.assertIsNone(image_collections_schema_error(conn, check_values=False))
self.assertIsNotNone(image_collections_schema_error(conn))
self.assertEqual(
conn.execute(f"SELECT COUNT(*) FROM {COLLECTIONS}").fetchone()[0], 1
)
model.close_all()
self.assertFalse(
PromptModel.verify_database_file(path, require_library_state=True)[0]
)
def test_restore_accepts_an_old_backup_without_collections(self):
self.require_module()
path = self.closed_current("old.db")
self.drop_group(path)
with closing(sqlite3.connect(path)) as raw:
self.assertIsNone(image_collections_schema_error(raw))
self.assertEqual(
image_collections_schema_error(raw, allow_missing=False),
"Missing image collections schema",
)
self.assertEqual(PromptModel.verify_database_file(path), (True, "ok"))
model = self.model(path)
with model.get_connection() as conn:
self.assertEqual(catalog(conn), expected_catalog())
model.close_all()
self.assertEqual(
PromptModel.verify_database_file(path, require_library_state=True),
(True, "ok"),
)
def test_restore_requires_collections_on_a_current_stage(self):
self.require_module()
path = self.closed_current("stage.db")
self.drop_group(path)
self.assertEqual(
PromptModel.verify_database_file(path, require_library_state=True),
(False, "Missing image collections schema"),
)
# One table without the other is never acceptable.
partial = self.closed_current("partial.db", f"DROP TABLE {ITEMS}")
self.assertFalse(PromptModel.verify_database_file(partial)[0])
def assert_refused(self, name, *statements):
self.require_module()
path = self.closed_current(
f"{name}.db",
(
f"INSERT INTO {COLLECTIONS} (id, name, name_key, created_at) "
"VALUES (1, 'Good', 'good', ?)",
(NOW,),
),
(f"INSERT INTO {ITEMS} VALUES (1, 'root', 'a.png', ?)", (NOW,)),
*statements,
)
with closing(sqlite3.connect(path)) as raw:
self.assertIsNone(image_collections_schema_error(raw, check_values=False))
error = image_collections_schema_error(raw)
self.assertIsNotNone(error)
self.assertIn("collections", error.lower())
ok, reason = PromptModel.verify_database_file(path)
self.assertFalse(ok)
self.assertEqual(reason, error)
self.assertFalse(
PromptModel.verify_database_file(path, require_library_state=True)[0]
)
def test_a_valid_current_database_is_accepted(self):
self.require_module()
path = self.closed_current(
"valid.db",
(
f"INSERT INTO {COLLECTIONS} (id, name, name_key, created_at) "
"VALUES (1, ?, ?, ?)",
("LoRa testing", key("LoRa testing"), NOW),
),
(f"INSERT INTO {ITEMS} VALUES (1, 'root', ?, ?)", ("é" * 4096, NOW)),
)
self.assertEqual(
PromptModel.verify_database_file(path, require_library_state=True),
(True, "ok"),
)
def test_restore_refuses_a_non_text_name(self):
# Column affinity turns a number into text; a blob stays a blob.
self.assert_refused("blob-name", f"UPDATE {COLLECTIONS} SET name = x'61'")
def test_restore_refuses_an_empty_name(self):
self.assert_refused("empty", f"UPDATE {COLLECTIONS} SET name = '', name_key=''")
def test_restore_refuses_a_name_over_100_characters(self):
name = "é" * 101
self.assert_refused(
"long",
(f"UPDATE {COLLECTIONS} SET name = ?, name_key = ?", (name, key(name))),
)
def test_restore_refuses_a_control_character(self):
name = "a\x07b"
self.assert_refused(
"control",
(f"UPDATE {COLLECTIONS} SET name = ?, name_key = ?", (name, key(name))),
)
def test_restore_refuses_a_leading_or_trailing_space(self):
for name in (" Good", "Good "):
with self.subTest(name=name):
self.assert_refused(
f"padded{len(name)}{name[0] == ' '}",
(
f"UPDATE {COLLECTIONS} SET name = ?, name_key = ?",
(name, key(name)),
),
)
def test_restore_accepts_a_name_key_from_other_case_tables(self):
# Python versions fold some characters differently; a key written
# by another version must never make a backup unrestorable.
self.require_module()
path = self.closed_current(
"other-key.db",
(
f"INSERT INTO {COLLECTIONS} (id, name, name_key, created_at) "
"VALUES (1, 'Good', 'Good', ?)",
(NOW,),
),
)
self.assertEqual(
PromptModel.verify_database_file(path, require_library_state=True),
(True, "ok"),
)
def test_restore_refuses_an_empty_or_non_text_name_key(self):
self.assert_refused("empty-key", f"UPDATE {COLLECTIONS} SET name_key = ''")
self.assert_refused(
"blob-name-key", f"UPDATE {COLLECTIONS} SET name_key = x'61'"
)
self.assert_refused(
"long-name-key",
(f"UPDATE {COLLECTIONS} SET name_key = ?", ("k" * 1201,)),
)
def test_restore_refuses_a_c1_control_character(self):
for name in ("a\x81b", "a\x7fb"):
with self.subTest(name=repr(name)):
self.assert_refused(
f"c1-{ord(name[1])}",
(
f"UPDATE {COLLECTIONS} SET name = ?, name_key = ?",
(name, key(name)),
),
)
def test_restore_refuses_a_non_text_key(self):
# Column affinity turns a number into text; a blob stays a blob.
self.assert_refused("blob-key", f"UPDATE {ITEMS} SET relative_path = x'61'")
def test_restore_refuses_a_key_over_4096_characters(self):
self.assert_refused(
"long-key", (f"UPDATE {ITEMS} SET root_key = ?", ("é" * 4097,))
)
def test_restore_refuses_a_timestamp_over_64_bytes(self):
self.assert_refused(
"stamp", (f"UPDATE {COLLECTIONS} SET created_at = ?", ("t" * 65,))
)
self.assert_refused(
"item-stamp", (f"UPDATE {ITEMS} SET added_at = ?", ("t" * 65,))
)
def test_restore_refuses_an_item_of_a_missing_collection(self):
self.assert_refused("orphan", f"INSERT INTO {ITEMS} VALUES (9, 'r', 'b', 'n')")
def test_restore_refuses_a_non_integer_collection_id(self):
self.assert_refused("text-id", f"UPDATE {ITEMS} SET collection_id = 'one'")
def test_restore_refuses_a_trigger_or_extra_index(self):
self.require_module()
for name, statement in (
(
"trigger",
f"CREATE TRIGGER t AFTER INSERT ON {ITEMS} BEGIN SELECT 1; END",
),
("extra", f"CREATE INDEX extra ON {COLLECTIONS}(name)"),
):
with self.subTest(name=name):
path = self.closed_current(f"{name}.db", statement)
with closing(sqlite3.connect(path)) as raw:
self.assertIsNotNone(image_collections_schema_error(raw))
self.assertFalse(PromptModel.verify_database_file(path)[0])
def test_deleting_a_collection_deletes_its_members(self):
self.require_module()
model = self.model()
with model.get_connection() as conn:
conn.execute(
f"INSERT INTO {COLLECTIONS} (id, name, name_key, created_at) "
"VALUES (1, 'A', 'a', ?), (2, 'B', 'b', ?)",
(NOW, NOW),
)
conn.executemany(
f"INSERT INTO {ITEMS} VALUES (?, 'root', ?, ?)",
((1, "a.png", NOW), (1, "b.png", NOW), (2, "a.png", NOW)),
)
conn.execute(f"DELETE FROM {COLLECTIONS} WHERE id = 1")
self.assertEqual(
[tuple(r) for r in conn.execute(f"SELECT collection_id FROM {ITEMS}")],
[(2,)],
)
if __name__ == "__main__":
unittest.main()
+10
View File
@@ -884,6 +884,16 @@ class TestFilterStart(ImageTagJobTestCase):
self.assertEqual((status, data["code"]), (400, "no-images"))
self.assertEqual(self.api._jobs.list(), [])
async def test_filter_with_a_deleted_collection_is_a_404_without_a_job(self):
self.png("a.png")
status, data = await self.start(
{"mode": "add", "filter": {"collection": "77"}, "tags": ["x"]}
)
self.assertEqual((status, data["code"]), (404, "collection-not-found"))
self.assertEqual(self.api._jobs.list(), [])
async def test_filter_over_the_cap_is_refused_without_a_job(self):
from py.api import image_tag_jobs
+277
View File
@@ -0,0 +1,277 @@
"""The inspector and History show an image's own category (#332 §5, DC4, DC5)."""
import os
import sqlite3
import unittest
from unittest.mock import patch
from tests import test_api_history_outputs as history
from tests.test_api_images import make_png
from py.api.images import _favorite_root_key, image_category_key
from utils.inspection_response import (
OWN_CATEGORY_UNREAD,
_library_projection,
apply_own_category,
)
class InspectCase(history.ImageAPITestCase):
async def tearDownAsync(self):
self.api.db.close_all()
await super().tearDownAsync()
def own(self, path, value, sources=None):
sources = sources or self.api._get_all_output_dirs()
self.api.db.set_image_categories([image_category_key(sources, path)], value)
async def library(self, path, root="0"):
response = await self.client.get(
f"/prompt_manager/images/inspect/{path}", params={"root": root}
)
body = await response.json()
self.assertEqual(response.status, 200, body)
return body["library"]
class TestInspectorOwnCategory(InspectCase):
async def test_inspect_reports_image_scope(self):
path = make_png(self.output_dir / "image.png")
self.own(path, "Portraits")
library = await self.library("image.png")
self.assertEqual(library["status"], "none")
self.assertEqual(
(library["category"], library["category_scope"]), ("Portraits", "image")
)
async def test_own_overrides_linked_recipe_and_ambiguous(self):
db = self.api.db
linked = make_png(self.output_dir / "linked.png")
self._link_image(linked, "castle at dusk")
prompt = db.get_prompt_id_for_image(str(linked))
with db.model.get_connection() as conn:
conn.execute(
"UPDATE prompts SET category='Landscapes' WHERE id=?", (prompt,)
)
library = await self.library("linked.png")
self.assertEqual(
(library["category"], library["category_scope"]),
("Landscapes", "positive"),
)
self.own(linked, "Portraits")
library = await self.library("linked.png")
self.assertEqual(library["status"], "linked")
self.assertEqual(
(library["category"], library["category_scope"]), ("Portraits", "image")
)
recipe = make_png(self.output_dir / "recipe.png")
self._save_prompt("Pair positive", category="Night")
pair = db.save_prompt_pair("Pair positive", "Noise", category="Recipes")
positive = db.get_prompt_pair(pair)["positive_prompt_id"]
image = db.link_image_to_prompt(positive, str(recipe))
db.link_image_to_pair(pair, image, {"source": "fixture"})
library = await self.library("recipe.png")
self.assertEqual(
(library["category"], library["category_scope"]), ("Recipes", "pair")
)
self.own(recipe, "Mine")
library = await self.library("recipe.png")
self.assertEqual(
(library["category"], library["category_scope"]), ("Mine", "image")
)
ambiguous = make_png(self.output_dir / "ambiguous.png")
self._link_image(ambiguous, "owner one")
self._link_image(ambiguous, "owner two")
self.own(ambiguous, "Scans")
library = await self.library("ambiguous.png")
self.assertEqual(library["status"], "ambiguous")
self.assertEqual(
(library["category"], library["category_scope"]), ("Scans", "image")
)
async def test_a_failed_read_hides_the_category(self):
linked = make_png(self.output_dir / "linked.png")
self._link_image(linked, "castle at dusk")
def broken(*_args, **_kwargs):
raise sqlite3.OperationalError("secret")
with patch.object(self.api.db, "image_category", side_effect=broken):
with self.assertLogs("prompt_manager", level="WARNING"):
library = await self.library("linked.png")
self.assertEqual(library["status"], "linked")
self.assertIsNone(library["category"])
self.assertIsNone(library["category_scope"])
async def test_a_lora_root_reads_no_own_category(self):
lora = self.comfy_dir / "loras"
make_png(lora / "preview.png")
roots = [self.output_dir, lora]
reads = []
real = self.api.db.image_category
def recorded(*args):
reads.append(args)
return real(*args)
with (
patch.object(self.api, "_media_roots", lambda outputs: roots),
patch.object(self.api, "_linked_only_error", self._no_error),
patch.object(self.api.db, "image_category", side_effect=recorded),
):
response = await self.client.get(
"/prompt_manager/images/inspect/preview.png", params={"root": "1"}
)
body = await response.json()
self.assertEqual(response.status, 200, body)
self.assertEqual(reads, [])
self.assertIsNone(body["library"]["category"])
@staticmethod
async def _no_error(*_args):
return None
class TestProjection(unittest.TestCase):
def test_scope_image_is_kept_for_category_only(self):
base = {"status": "none", "scope": "saved-library", "tags": []}
projected = _library_projection(
{**base, "category": "Scans", "category_scope": "image"}
)
self.assertEqual(projected["category_scope"], "image")
projected = _library_projection({**base, "tags_scope": "image"})
self.assertNotIn("tags_scope", projected)
def test_apply_own_category_three_cases(self):
library = {"status": "linked", "category": "L", "category_scope": "positive"}
self.assertEqual(
apply_own_category(library, "Mine"),
{"status": "linked", "category": "Mine", "category_scope": "image"},
)
self.assertEqual(apply_own_category(library, None), library)
self.assertEqual(
apply_own_category(library, OWN_CATEGORY_UNREAD),
{"status": "linked", "category": None, "category_scope": None},
)
self.assertEqual(library["category"], "L") # never mutated
class TestHistoryOwnCategory(history.ImageAPITestCase):
"""Only the inspection request reads the own category, keyed as Images."""
# The History fixture's helpers, without inheriting its tests.
_base = history.HistoryOutputAPITests
ref = _base.ref
capture = _base.capture
outputs = _base.outputs
current = _base.current
media = _base.media
async def get_application(self):
app = await super().get_application()
self.temp_dir = self.comfy_dir / "temp"
self.temp_dir.mkdir()
self.host_output = str(self.output_dir)
self.host_temp = str(self.temp_dir)
self.patches = history.ExitStack()
self.addCleanup(self.patches.close)
for name, value in (
("folder_paths.get_output_directory", lambda: self.host_output),
("folder_paths.get_temp_directory", lambda: self.host_temp),
):
self.patches.enter_context(patch(name, value, create=True))
self.patches.enter_context(
patch.object(
history.GalleryConfig,
"allowed_gallery_parents",
lambda: [str(self.comfy_dir)],
)
)
return app
async def tearDownAsync(self):
self.api.close_history_reads()
self.patches.close()
await super().tearDownAsync()
async def inspect(self, start):
descriptor = await self.media(start)
response = await self.current(descriptor["inspection_url"])
body = await response.json()
self.assertEqual(response.status, 200, body)
return body["library"]
async def test_history_inspection_shows_own_category(self):
path = make_png(self.output_dir / "image.png")
self.api.db.set_image_categories(
[image_category_key([self.output_dir], path)], "Portraits"
)
start = self.capture([self.ref()])
library = await self.inspect(start)
self.assertEqual(
(library["category"], library["category_scope"]), ("Portraits", "image")
)
async def test_nested_sources_give_the_images_key(self):
inner = self.output_dir / "nested"
path = make_png(inner / "image.png")
self.api._get_all_output_dirs = lambda: [self.output_dir, inner]
self.api.db.set_image_categories(
[(_favorite_root_key(inner), "image.png")], "S"
)
start = self.capture([self.ref(subfolder="nested")])
library = await self.inspect(start)
self.assertEqual(library["category"], "S")
self.assertEqual(
image_category_key([self.output_dir, inner], path),
(_favorite_root_key(inner), "image.png"),
)
async def test_a_symlinked_output_folder_still_shows_it(self):
real = self.comfy_dir / "real-output"
real.mkdir()
link = self.comfy_dir / "linked-output"
self._make_symlink(real, link, target_is_directory=True)
make_png(real / "image.png")
self.host_output = str(link)
self.api._get_all_output_dirs = lambda: [real]
self.api.db.set_image_categories([(_favorite_root_key(real), "image.png")], "S")
start = self.capture([self.ref()])
library = await self.inspect(start)
self.assertEqual(library["category"], "S")
async def test_outputs_and_media_read_no_own_category(self):
make_png(self.output_dir / "image.png")
start = self.capture([self.ref()])
reads = []
real = self.api.db.image_category
def recorded(*args):
reads.append(args)
return real(*args)
with patch.object(self.api.db, "image_category", side_effect=recorded):
descriptor = await self.media(start)
for field in ("url", "thumbnail_url"):
response = await self.current(descriptor[field])
await response.read()
self.assertEqual(response.status, 200)
self.assertEqual(reads, [])
await self.current(descriptor["inspection_url"])
self.assertEqual(len(reads), 1)
async def test_a_temp_output_has_none(self):
reference = {"folder_type": "temp", "subfolder": "", "filename": "a.png"}
self.assertIsNone(self.api._history_own_category(reference, None))
async def test_a_failed_read_hides_it(self):
make_png(self.output_dir / "image.png")
start = self.capture([self.ref()])
def broken(*_args):
raise sqlite3.OperationalError("secret")
with patch.object(self.api.db, "image_category", side_effect=broken):
with self.assertLogs("prompt_manager", level="WARNING"):
library = await self.inspect(start)
self.assertIsNone(library["category"])
self.assertTrue(os.path.exists(self.output_dir / "image.png"))
+29
View File
@@ -1426,3 +1426,32 @@ class TestPerRecordBudgets(unittest.TestCase):
self.assertEqual(metadata["record_formats"], {"workflow": "limited"})
self.assertEqual(metadata["workflow_graph"]["status"], "limited")
self.assertEqual(metadata["prompt_graph"]["status"], "absent")
class TestInheritedCategory(unittest.TestCase):
"""#332 MC4: one inherited-category rule for the chip, filter and History."""
def test_recipe_set_blank_and_absent(self):
from utils.inspection_response import inherited_category, library_metadata
cases = [
(("Night", "Recipes", True), ("Recipes", "pair")),
(("Night", " ", True), ("Night", "positive")),
(("Night", " \t", True), ("Night", "positive")),
(("Night", b" ", True), ("Night", "positive")),
(("Night", None, True), ("Night", "positive")),
(("Night", "Recipes", False), ("Night", "positive")),
((" ", None, False), (" ", "positive")),
((None, None, False), (None, "positive")),
]
for args, expected in cases:
with self.subTest(args=args):
self.assertEqual(inherited_category(*args), expected)
positive, recipe, has_recipe = args
metadata = library_metadata(
positive, ["t"], (recipe, []) if has_recipe else None
)
self.assertEqual(
(metadata["category"], metadata["category_scope"]), expected
)
self.assertEqual(metadata["tags_scope"], "positive")
+17
View File
@@ -85,6 +85,23 @@ class TestPairUsageTracking(unittest.TestCase):
self.db.get_prompt_by_id(pair["negative_prompt_id"])["run_count"], 0
)
def test_execution_records_a_cached_run_receipt_for_both_editors(self):
from utils.cached_runs import Facts
node = self.node()
node.cached_runs = Mock()
node.cached_runs.replayed.return_value = False
node.encode_prompt(
Mock(),
"cat",
negative_text="",
prompt=pair_graph(negative=""),
unique_id="1",
)
node.cached_runs.record_item.assert_called_once_with(
"1", ("cat", ""), Facts(usage_tracking.generate_prompt_hash("cat"), None)
)
def test_validation_rejection_never_counts_or_saves(self):
self.submit(pair_graph())
self.assertEqual(self.db.list_prompt_pairs()["total"], 0)
+44
View File
@@ -278,5 +278,49 @@ class TestNodeRecordsUse(UsageTestCase):
self.assertEqual(self.db.get_prompt_by_id(pid)["run_count"], 2)
class TestNodeRecordsCachedRunReceipt(UsageTestCase):
"""Each execution leaves a receipt for cached jobs (#347)."""
GRAPH = {
"1": {"class_type": "PromptManagerText", "inputs": {"text": "typed text"}},
}
def _node(self):
from unittest import mock
from prompt_manager_text import PromptManagerText
node = PromptManagerText.__new__(PromptManagerText)
node.db = self.db
node.logger = __import__("logging").getLogger("test.prompt_usage")
node.prompt_tracker = mock.Mock()
node.comfyui_integration = mock.Mock()
node._inject_lora_trigger_words = lambda text: text
node.cached_runs = mock.Mock()
node.cached_runs.replayed.return_value = False
return node
def test_receipt_holds_the_returned_text_and_the_stored_hash(self):
from utils.cached_runs import Facts
node = self._node()
node.process_text(
"typed text", prepend_text="pre", prompt=self.GRAPH, unique_id="1"
)
node.cached_runs.record_item.assert_called_once_with(
"1", ("pre typed text",), Facts(generate_prompt_hash("typed text"))
)
# Hook-visible text: the node counts nothing itself
row = self.db.get_prompt_by_hash(generate_prompt_hash("typed text"))
self.assertEqual(row["run_count"], 0)
def test_empty_editor_records_empty_facts(self):
from utils.cached_runs import Facts
node = self._node()
node.process_text(" ", prompt=self.GRAPH, unique_id="1")
node.cached_runs.record_item.assert_called_once_with("1", (" ",), Facts(None))
if __name__ == "__main__":
unittest.main()
+379
View File
@@ -0,0 +1,379 @@
"""Recorded run history names the saver that wrote a file (#386, #240).
When a file's graph alone cannot name its prompt (sampler branches that differ
only in their prompt and share one saver prefix), the inspector asks recorded
history which save node wrote the file and walks from that node, as the #381
file-name anchor does. A result the file alone resolves is never changed; a
recorded saver that names nothing certain leaves the file "several possible
prompts". The resolved sides are labelled as coming from the recorded run.
"""
import json
import unittest
from utils import prompt_pair_resolver as resolver
from utils.prompt_inspector import inspect_prompt_evidence
try:
from tests import saver_anchor_graphs as shapes
except ImportError: # discovered with ``-s tests``
import saver_anchor_graphs as shapes
def _resolve(graph, **options):
return resolver.resolve_prompt_pair(
{"prompt": graph}, profile=resolver.INSPECTION_PROFILE, **options
)
def _positive(result):
side = result["positive"]
return side["text"] if isinstance(side, dict) else None
def _evidence(graph, **options):
return inspect_prompt_evidence({"prompt": json.dumps(graph)}, **options)
# ``branch_bench``'s savers are node ``N3`` of branch N (10s based).
def _saver(index):
return str(index * 10 + 3)
def guider_dataset():
"""The dataset run's shape: SamplerCustomAdvanced branches on BasicGuider
(no negative), every branch saved with one prefix."""
graph = {
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "m"}},
"2": {"class_type": "EmptyLatentImage", "inputs": {"width": 64}},
"4": {"class_type": "RandomNoise", "inputs": {"noise_seed": 1}},
"5": {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "euler"}},
"6": {
"class_type": "BasicScheduler",
"inputs": {"model": ["1", 0], "steps": 8, "scheduler": "simple"},
},
}
for index, text in enumerate(shapes.BRANCH_TEXTS, start=1):
base = index * 10
graph[str(base)] = {
"class_type": "CLIPTextEncode",
"inputs": {"text": text, "clip": ["1", 1]},
}
graph[str(base + 4)] = {
"class_type": "BasicGuider",
"inputs": {"model": ["1", 0], "conditioning": [str(base), 0]},
}
graph[str(base + 1)] = {
"class_type": "SamplerCustomAdvanced",
"inputs": {
"noise": ["4", 0],
"guider": [str(base + 4), 0],
"sampler": ["5", 0],
"sigmas": ["6", 0],
"latent_image": ["2", 0],
},
}
graph[str(base + 2)] = {
"class_type": "VAEDecode",
"inputs": {"samples": [str(base + 1), 0], "vae": ["1", 2]},
}
graph[str(base + 3)] = {
"class_type": "SaveImage",
"inputs": {
"filename_prefix": "dataset/IMG_GEN",
"images": [str(base + 2), 0],
},
}
return graph
class TestRecordedSaverInResolver(unittest.TestCase):
def test_the_recorded_saver_picks_its_branch_of_a_shared_prefix(self):
for index, text in enumerate(shapes.BRANCH_TEXTS, start=1):
with self.subTest(branch=index):
result = _resolve(
shapes.shared_prefix(),
file_name="IMG_GEN_00371_.png",
recorded_saver_id=_saver(index),
)
self.assertEqual(result["state"], "resolved")
self.assertEqual(_positive(result), text)
self.assertEqual(result["sampler_id"], str(index * 10 + 1))
def test_without_a_recorded_saver_the_shared_prefix_stays_several(self):
result = _resolve(shapes.shared_prefix(), file_name="IMG_GEN_00371_.png")
self.assertEqual(result["state"], "ambiguous")
def test_a_recorded_saver_the_graph_lacks_anchors_nothing(self):
for node in ("missing", "3", "2"):
with self.subTest(node=node):
result = _resolve(shapes.shared_prefix(), recorded_saver_id=node)
self.assertEqual(result["state"], "ambiguous")
def test_a_recorded_saver_of_an_unknown_class_anchors(self):
graph = shapes.branch_bench(
prefixes=["dataset/IMG_GEN"] * len(shapes.BRANCH_TEXTS),
saver="VendorSaveImage",
)
result = _resolve(graph, recorded_saver_id=_saver(2))
self.assertEqual(_positive(result), shapes.BRANCH_TEXTS[1])
def test_the_default_profile_ignores_a_recorded_saver(self):
result = resolver.resolve_prompt_pair(
{"prompt": shapes.shared_prefix()}, recorded_saver_id=_saver(1)
)
self.assertEqual(result["state"], "ambiguous")
class TestRecordedSaverInInspector(unittest.TestCase):
def test_the_inspector_shows_the_recorded_branch_labelled_as_recorded(self):
evidence = _evidence(
shapes.shared_prefix(),
file_name="IMG_GEN_00371_.png",
recorded_saver_id=_saver(3),
)
self.assertEqual(evidence["resolution"]["state"], "resolved")
positive = evidence["prompts"]["positive"]
self.assertEqual(positive["text"], shapes.BRANCH_TEXTS[2])
self.assertEqual(positive["source"], "recorded-run")
self.assertEqual(evidence["prompts"]["negative"]["source"], "recorded-run")
def test_a_guider_branch_without_a_negative_shows_its_positive(self):
self.assertEqual(
_evidence(guider_dataset(), file_name="IMG_GEN_00371_.png")["resolution"][
"state"
],
"ambiguous",
)
evidence = _evidence(
guider_dataset(), file_name="IMG_GEN_00371_.png", recorded_saver_id="23"
)
self.assertEqual(
evidence["prompts"]["positive"],
{
"status": "resolved",
"text": shapes.BRANCH_TEXTS[1],
"source": "recorded-run",
},
)
self.assertEqual(evidence["prompts"]["negative"]["status"], "absent")
def test_a_file_name_naming_another_saver_contradicts_history(self):
# The bench file name names branch 2's saver; history says branch 4's.
graph = shapes.branch_bench()
for index in (1, 3, 4):
graph[_saver(index)]["inputs"]["filename_prefix"] = "bench/run/shared"
result = _resolve(
graph, file_name="s45_branch_2_00001_.png", recorded_saver_id=_saver(4)
)
self.assertIs(result["recorded_anchor"], False)
evidence = _evidence(
shapes.shared_prefix(),
file_name="IMG_GEN_00001_.png",
recorded_saver_id=_saver(4),
)
self.assertEqual(
evidence["prompts"]["positive"]["text"], shapes.BRANCH_TEXTS[3]
)
def test_a_file_that_resolves_alone_keeps_its_own_result(self):
# The file name names branch 2; a recorded saver never overrides it.
evidence = _evidence(
shapes.branch_bench(),
file_name="s45_branch_2_00001_.png",
recorded_saver_id=_saver(4),
)
positive = evidence["prompts"]["positive"]
self.assertEqual(positive["text"], shapes.BRANCH_TEXTS[1])
self.assertEqual(positive["source"], "graph-inferred")
def test_a_recorded_saver_that_names_nothing_keeps_several_prompts(self):
evidence = _evidence(shapes.shared_prefix(), recorded_saver_id="missing")
self.assertEqual(evidence["resolution"]["state"], "ambiguous")
self.assertEqual(evidence["prompts"]["positive"]["status"], "ambiguous")
def test_no_recorded_saver_is_unchanged(self):
self.assertEqual(
_evidence(shapes.shared_prefix(), recorded_saver_id=None),
_evidence(shapes.shared_prefix()),
)
LLM_TEXT = "a red vintage car parked by the sea, golden hour"
def text_generate_run():
"""The maintainer's moss run: a language model writes PromptManager's text."""
return {
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "m"}},
"2": {"class_type": "EmptyLatentImage", "inputs": {"width": 64}},
"17": {"class_type": "PrimitiveStringMultiline", "inputs": {"value": "a car"}},
"14": {
"class_type": "PrimitiveStringMultiline",
"inputs": {"value": "You are a prompt engineer."},
},
"11": {
"class_type": "TextGenerate",
"inputs": {
"prompt": ["17", 0],
"system_prompt": ["14", 0],
"clip": ["1", 1],
"max_length": 4096,
"sampling_mode.seed": 0,
},
},
"36": {
"class_type": "PromptManagerPair",
"inputs": {
"text": ["11", 0],
"negative_text": "",
"prepend_text": "",
"append_text": "",
"clip": ["1", 1],
},
},
"7": {
"class_type": "KSampler",
"inputs": {
"seed": 1,
"steps": 20,
"cfg": 1.0,
"sampler_name": "euler",
"scheduler": "simple",
"denoise": 1.0,
"model": ["1", 0],
"positive": ["36", 0],
"negative": ["36", 1],
"latent_image": ["2", 0],
},
},
"8": {
"class_type": "VAEDecode",
"inputs": {"samples": ["7", 0], "vae": ["1", 2]},
},
"9": {
"class_type": "SaveImage",
"inputs": {"filename_prefix": "moss", "images": ["8", 0]},
},
}
def two_generated_branches():
"""``text_generate_run`` plus a second generated branch saved by an unknown
saver and also previewed."""
graph = text_generate_run()
graph["31"] = {**graph["11"], "inputs": {**graph["11"]["inputs"]}}
graph["37"] = {
"class_type": "PromptManagerPair",
"inputs": {**graph["36"]["inputs"], "text": ["31", 0]},
}
graph["20"] = {
"class_type": "KSampler",
"inputs": {
**graph["7"]["inputs"],
"positive": ["37", 0],
"negative": ["37", 1],
},
}
graph["21"] = {
"class_type": "VAEDecode",
"inputs": {"samples": ["20", 0], "vae": ["1", 2]},
}
graph["22"] = {
"class_type": "Image Save",
"inputs": {"filename_prefix": "other", "images": ["21", 0]},
}
graph["23"] = {"class_type": "PreviewImage", "inputs": {"images": ["21", 0]}}
return graph
class TestRecordedTexts(unittest.TestCase):
"""The texts PromptManager recorded at execution, per sampler of the run."""
def test_a_run_time_text_shows_what_the_run_recorded(self):
alone = _evidence(text_generate_run(), file_name="moss_00026_.png")
self.assertEqual(alone["resolution"]["reason_code"], "runtime-value")
evidence = _evidence(
text_generate_run(),
file_name="moss_00026_.png",
recorded_saver_id="9",
recorded_texts={"7": {"positive": LLM_TEXT, "negative": ""}},
)
self.assertEqual(
evidence["prompts"],
{
"positive": {
"status": "resolved",
"text": LLM_TEXT,
"source": "execution",
},
"negative": {"status": "resolved", "text": "", "source": "execution"},
},
)
self.assertEqual(evidence["resolution"]["state"], "resolved")
self.assertEqual(evidence["resolution"]["sampler_id"], "7")
def test_an_unrecorded_negative_stays_open(self):
evidence = _evidence(
text_generate_run(),
recorded_saver_id="9",
recorded_texts={"7": {"positive": LLM_TEXT, "negative": None}},
)
self.assertEqual(evidence["prompts"]["positive"]["text"], LLM_TEXT)
self.assertEqual(evidence["prompts"]["negative"]["status"], "unresolved")
def test_texts_that_do_not_name_the_open_sampler_change_nothing(self):
alone = _evidence(text_generate_run())
for texts in (
{"8": {"positive": LLM_TEXT, "negative": ""}},
{"7": {"positive": None, "negative": ""}},
{"7": {"positive": " ", "negative": ""}},
{"7": "not roles"},
{},
):
with self.subTest(texts=texts):
self.assertEqual(
_evidence(
text_generate_run(), recorded_saver_id="9", recorded_texts=texts
),
alone,
)
def test_texts_need_the_saver_history_recorded_to_lead_to_the_sampler(self):
texts = {"7": {"positive": LLM_TEXT, "negative": ""}}
alone = _evidence(text_generate_run())
for saver in (None, "missing", "2"):
with self.subTest(saver=saver):
self.assertEqual(
_evidence(
text_generate_run(),
recorded_saver_id=saver,
recorded_texts=texts,
),
alone,
)
def test_texts_follow_the_recorded_saver_not_another_sampler(self):
# Review: the texts are those of the sampler history's saver leads to.
graph = two_generated_branches()
texts = {
"7": {"positive": "STAGE X TEXT", "negative": ""},
"20": {"positive": "STAGE Y TEXT", "negative": ""},
}
for saver, text, sampler in (
("22", "STAGE Y TEXT", "20"),
("9", "STAGE X TEXT", "7"),
):
with self.subTest(saver=saver):
evidence = _evidence(
graph,
file_name="renamed_00001_.png",
recorded_saver_id=saver,
recorded_texts=texts,
)
self.assertEqual(evidence["prompts"]["positive"]["text"], text)
self.assertEqual(evidence["resolution"]["sampler_id"], sampler)
if __name__ == "__main__":
unittest.main()
+1
View File
@@ -49,6 +49,7 @@ SUMMARY_KEYS = {
"image_links",
"runs",
"saved_views",
"collections",
}
+43
View File
@@ -226,6 +226,7 @@ class TestStageAndPublish(StageTestCase):
for key in ("pairs", "archived_pairs", "image_links", "runs"):
self.assertEqual(staged.summary[key], 0, key)
self.assertEqual(staged.summary["saved_views"], 0)
self.assertEqual(staged.summary["collections"], 0)
self.assertEqual(staged.schema, {"migrated": False, "library_state": True})
self.assertEqual(
staged.digest,
@@ -239,6 +240,48 @@ class TestStageAndPublish(StageTestCase):
staged.discard()
class TestCollectionsSummary(StageTestCase):
"""#331: collections join the review's summary and the live fingerprint."""
def add_collection(self, db, name):
with db.model.get_connection() as conn:
conn.execute(
"INSERT INTO image_collections (name, name_key, created_at) "
"VALUES (?, ?, '2026-10-27T00:00:00+00:00')",
(name, name.casefold()),
)
def test_summary_counts_collections(self):
src_path = os.path.join(self.uploads, "with-collections.db")
src = PromptDatabase(src_path)
try:
self.save(src, "kept")
self.add_collection(src, "First")
self.add_collection(src, "Second")
finally:
src.close_all()
staged = self.db.model.stage_restore(src_path)
try:
self.assertEqual(staged.summary["collections"], 2)
finally:
staged.discard()
self.assertIn("collections", restore_stage.SUMMARY_KEYS)
def test_adding_a_collection_changes_the_fingerprint(self):
conn = self.db.model.get_connection()
summary, before = restore_stage.live_fingerprint(conn)
self.assertEqual(summary["collections"], 0)
self.add_collection(self.db, "After the review")
summary, after = restore_stage.live_fingerprint(conn)
self.assertEqual(summary["collections"], 1)
self.assertNotEqual(after, before)
# A delete and a re-add with the same count still changes it (rowid).
with self.db.model.get_connection() as write:
write.execute("DELETE FROM image_collections")
self.add_collection(self.db, "Replacement")
self.assertNotEqual(restore_stage.live_fingerprint(conn)[1], after)
class TestOlderBackups(StageTestCase):
"""Rollback-journal and older-schema backups restore exactly as before."""
+212 -17
View File
@@ -9,13 +9,22 @@ import unittest
from pathlib import Path
from unittest import mock
from database.run_attribution import stored_run_prompt
from database.run_attribution import (
RecordedOutput,
recorded_output,
stored_run_prompt,
)
from tests import test_execution_history_database as history_tests
from tests.background_stops import close_handler
from tests.test_api_admin import AdminAPITestCase
from tests.test_execution_context_database import snapshot
from tests.test_execution_context_database import known, snapshot
from tests.test_prompt_attribution import issue_graph, llm_text_graph, make_tracker
from utils.execution_contexts import ContextBundle, RecipeContext
from utils.execution_contexts import (
ContextBundle,
PromptSnapshot,
RecipeContext,
TextEvidence,
)
from utils.execution_outputs import (
OutputBundle,
OutputReference,
@@ -38,28 +47,34 @@ class RecordedRuns:
start = history_tests.HistoryDatabaseTests.start
finish = history_tests.HistoryDatabaseTests.finish
def record_run(self, db, output_root, files, positives, item=0, window=None):
def record_run(
self, db, output_root, files, positives, item=0, window=None, nodes=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)
if positives: # a run with no PromptManager node records no contexts
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",
),
"complete",
"captured",
),
)
)
root = snapshot_output_roots(STYLE, output_root, output_root).output_key
nodes = nodes or ["9"] * len(files)
references = tuple(
OutputReference("9", i, "output", sub, name, STYLE, root)
for i, (sub, name) in enumerate(files)
OutputReference(node, i, "output", sub, name, STYLE, root)
for i, (node, (sub, name)) in enumerate(zip(nodes, files))
)
db.write_execution_outputs_for_finish(
finish, OutputBundle(references, "complete", "reported")
@@ -171,6 +186,186 @@ class TestStoredRunPrompt(RecordedRuns, unittest.TestCase):
self.assertEqual(linked, [prompt_id])
class TestRecordedOutputNode(RecordedRuns, unittest.TestCase):
"""#386: the save node history says wrote a file, under the same rules."""
close_databases = history_tests.HistoryDatabaseTests.close_databases
database = history_tests.HistoryDatabaseTests.database
setUp = TestStoredRunPrompt.setUp
def node(self, path=None):
found = recorded_output(self.db, path or self.image, self.output)
return found.saver_id if found is not None else None
def test_one_run_one_saver_names_the_saver(self):
other = os.path.join(self.output, "sub", "img_00002_.png")
Path(other).write_bytes(b"png")
self.record_run(
self.db,
self.output,
[("sub", "img_00001_.png"), ("sub", "img_00002_.png")],
[],
nodes=["13", "23"],
)
self.assertEqual(self.node(), "13")
self.assertEqual(self.node(other), "23")
def test_a_run_without_prompt_manager_contexts_still_names_the_saver(self):
self.record_run(self.db, self.output, [("sub", "img_00001_.png")], [])
self.assertEqual(self.node(), "9")
def test_two_runs_reporting_the_file_name_nothing(self):
self.record_run(
self.db, self.output, [("sub", "img_00001_.png")], [], nodes=["13"]
)
self.record_run(
self.db, self.output, [("sub", "img_00001_.png")], [], item=1, nodes=["23"]
)
self.assertIsNone(self.node())
def test_one_run_reporting_the_file_from_two_savers_names_nothing(self):
self.record_run(
self.db,
self.output,
[("sub", "img_00001_.png"), ("sub", "img_00001_.png")],
[],
nodes=["13", "23"],
)
self.assertIsNone(self.node())
def test_a_file_written_outside_the_run_names_nothing(self):
self.record_run(
self.db,
self.output,
[("sub", "img_00001_.png")],
[],
window=(1_000_000, 2_000_000),
)
self.assertIsNone(self.node())
def test_unrecorded_outside_or_unknown_roots_name_nothing(self):
self.record_run(self.db, self.output, [("sub", "img_00001_.png")], [])
other = os.path.join(self.output, "sub", "other.png")
Path(other).write_bytes(b"png")
self.assertIsNone(self.node(other))
self.assertIsNone(recorded_output(self.db, self.image, None))
self.assertIsNone(recorded_output(self.db, "", self.output))
def test_the_run_texts_recorded_at_execution_come_with_the_saver(self):
self.record_run(
self.db, self.output, [("sub", "img_00001_.png")], ["an llm prompt"]
)
self.assertEqual(
recorded_output(self.db, self.image, self.output),
RecordedOutput("9", {"4": {"positive": "an llm prompt", "negative": ""}}),
)
def test_texts_only_inferred_from_the_graph_are_not_recorded(self):
start = self.start(self.db, 0, now_ms() - 5000)
self.db.write_execution_start(start)
finish = self.finish(start, time=now_ms() + 5000)
self.db.write_execution_finish(finish)
inferred = TextEvidence("graph text", "inferred", "graph_literal", None)
self.db.write_execution_contexts(
start.run_id,
start.storage_epoch,
ContextBundle(
(
RecipeContext(
"7",
PromptSnapshot(inferred, inferred, known(""), known("")),
snapshot("kept"),
"inferred",
),
),
"complete",
"captured",
),
)
root = snapshot_output_roots(STYLE, self.output, self.output).output_key
self.db.write_execution_outputs_for_finish(
finish,
OutputBundle(
(
OutputReference(
"9", 0, "output", "sub", "img_00001_.png", STYLE, root
),
),
"complete",
"reported",
),
)
self.assertEqual(
recorded_output(self.db, self.image, self.output).texts,
{"7": {"positive": None, "negative": "kept"}},
)
def test_several_savers_name_no_saver_but_keep_the_texts(self):
self.record_run(
self.db,
self.output,
[("sub", "img_00001_.png"), ("sub", "img_00001_.png")],
["text"],
nodes=["13", "23"],
)
found = recorded_output(self.db, self.image, self.output)
self.assertIsNone(found.saver_id)
self.assertEqual(found.texts, {"4": {"positive": "text", "negative": ""}})
def test_the_route_lookup_logs_a_failure_and_names_nothing(self):
from database.run_attribution import recorded_output_lookup
broken = mock.Mock()
broken.model.get_connection.return_value.execute.side_effect = (
sqlite3.OperationalError("locked")
)
broken.get_execution_contexts.side_effect = sqlite3.OperationalError("x")
with mock.patch(
"database.run_attribution.host_output_directory", return_value=self.output
):
self.assertIsNone(recorded_output_lookup(broken)(self.image))
with (
mock.patch(
"database.run_attribution._recorded_runs",
side_effect=sqlite3.OperationalError("locked"),
),
self.assertLogs("prompt_manager.run_attribution", "WARNING"),
):
self.assertIsNone(
recorded_output_lookup(broken, close=True)(self.image)
)
broken.model.close.assert_called_once_with()
def test_a_history_reference_spells_the_file_as_history_recorded_it(self):
# On Windows the history route's target is normcased (lower case);
# history matches names exactly, so the route passes its reference.
from database.run_attribution import recorded_output_lookup
upper = os.path.join(self.output, "sub", "IMG_00001_.png")
Path(upper).write_bytes(b"png")
self.record_run(self.db, self.output, [("sub", "IMG_00001_.png")], [])
target = os.path.join(self.output, "sub", "img_00001_.png")
reference = {
"folder_type": "output",
"subfolder": "sub",
"filename": "IMG_00001_.png",
}
with mock.patch(
"database.run_attribution.host_output_directory", return_value=self.output
):
found = recorded_output_lookup(self.db, reference=reference)(target)
self.assertEqual(found.saver_id, "9")
temp = {**reference, "folder_type": "temp"}
self.assertIsNone(recorded_output_lookup(self.db, reference=temp)(target))
def test_a_library_without_history_tables_names_nothing(self):
broken = mock.Mock()
broken.model.get_connection.return_value.execute.side_effect = (
sqlite3.OperationalError("no such table")
)
self.assertIsNone(recorded_output(broken, self.image, self.output))
class TestScanUsesHistory(RecordedRuns, AdminAPITestCase):
"""The output scan applies the same rules as the monitor."""
+291 -1
View File
@@ -13,6 +13,11 @@ from database.library_commands import LibraryCommandError
from database.operations import PromptDatabase
from utils.hashing import generate_prompt_hash
try:
from test_image_categories import category_checksum
except ImportError:
from tests.test_image_categories import category_checksum
class TestVocabularyReads(unittest.TestCase):
def setUp(self):
@@ -168,7 +173,8 @@ class TestVocabularyReads(unittest.TestCase):
self.assertEqual(page["total"], 4)
self.assertFalse(page["has_more"])
for row in page["categories"]:
self.assertEqual(set(row), {"value", "prompt_count"})
self.assertEqual(set(row), {"value", "prompt_count", "image_count"})
self.assertEqual(row["image_count"], 0)
self.assertEqual(
row["prompt_count"],
self.db.list_library(category=row["value"])["total"],
@@ -238,6 +244,111 @@ class TestVocabularyReads(unittest.TestCase):
with self.assertRaises(ValueError):
self.read("list_vocabulary_tags", **options)
def own(self, value, *paths, root="root"):
"""Give each relative path under ``root`` the own category ``value``."""
self.db.set_image_categories([(root, path) for path in paths], value)
def all_categories(self, limit=200):
rows, pages, after = [], [], None
while True:
page = self.read("list_vocabulary_categories", limit=limit, after=after)
pages.append(page)
rows.extend(page["categories"])
if not page["has_more"]:
return rows, pages
after = page["next_key"]
def test_category_list_includes_image_only_values_last(self):
self.save("Oldest", category="Zeta")
self.save("Newer", category="Alpha")
self.own("Zeta", "a.png", "b.png")
self.own("Only images B", "c.png")
self.own("Only images A", "d.png", "e.png", "gone.png")
rows, _pages = self.all_categories()
self.assertEqual(
rows,
[
{"value": "Zeta", "prompt_count": 1, "image_count": 2},
{"value": "Alpha", "prompt_count": 1, "image_count": 0},
{"value": "Only images A", "prompt_count": 0, "image_count": 3},
{"value": "Only images B", "prompt_count": 0, "image_count": 1},
],
)
page = self.read("list_vocabulary_categories")
self.assertEqual(page["total"], 4)
def test_category_cursor_crosses_into_image_only_values(self):
for index in range(3):
self.save(f"Owner {index}", category=f"P{index}")
self.own("P1", "shared.png")
self.own("I-b", "x.png")
self.own("I-a", "y.png")
self.own("I-c", "z.png", root="other")
for limit in (1, 2, 3, 4):
with self.subTest(limit=limit):
rows, pages = self.all_categories(limit=limit)
self.assertEqual(
[row["value"] for row in rows],
["P0", "P1", "P2", "I-a", "I-b", "I-c"],
)
self.assertEqual({page["total"] for page in pages}, {6})
self.assertEqual(
[row["image_count"] for row in rows], [0, 1, 0, 1, 1, 1]
)
first = self.read("list_vocabulary_categories", limit=3)
self.assertEqual(first["next_key"], [3, "P2"])
second = self.read("list_vocabulary_categories", limit=1, after=[3, "P2"])
self.assertEqual(second["next_key"], [2**63 - 1, "I-a"])
def test_image_rows_on_a_prompt_value_count_once(self):
self.save("Owner", category="Shared")
self.db.save_prompt_pair("Owner", "Noise", category="Shared")
self.own("Shared", *[f"{n}.png" for n in range(5)])
self.own("Shared", "elsewhere.png", root="second")
rows, _pages = self.all_categories()
self.assertEqual(
rows, [{"value": "Shared", "prompt_count": 1, "image_count": 6}]
)
def test_library_options_ignore_image_categories(self):
self.save("Owner", category="Prompt value")
before = self.db.get_library_options()
self.own("Image value", "a.png")
self.own("Prompt value", "b.png")
self.assertEqual(self.db.get_library_options(), before)
self.assertEqual(before["categories"], ["Prompt value"])
def test_category_reads_use_the_image_category_index(self):
from database import vocabulary
self.own("Value", "a.png")
statements = [
vocabulary._CATEGORY_ROWS + "SELECT COUNT(*) FROM counted",
vocabulary._CATEGORY_ROWS
+ "SELECT value,first_pid,prompt_count,image_count FROM counted "
"WHERE first_pid>? OR (first_pid=? AND value>?) "
"ORDER BY first_pid,value LIMIT ?",
]
plans = []
with self.db.model.get_connection() as conn:
for sql in statements:
params = [1] * sql.count("?")
plans.extend(
row[-1] for row in conn.execute("EXPLAIN QUERY PLAN " + sql, params)
)
for sql in (
vocabulary._IMAGE_CATEGORY_ROWIDS,
*vocabulary._CATEGORY_UPDATES,
):
params = [1] * sql.count("?")
plans.extend(
row[-1] for row in conn.execute("EXPLAIN QUERY PLAN " + sql, params)
)
image_steps = [step for step in plans if "image_categories" in step]
self.assertTrue(image_steps)
for step in image_steps:
self.assertIn("idx_image_categories_category", step, step)
def test_reads_share_one_snapshot_and_report_epoch(self):
self.save("Cat", category="Animal", tags=["cat"])
epoch = self.db.get_execution_epoch()
@@ -741,6 +852,7 @@ class TestVocabularyCategoryCommands(unittest.TestCase):
"other_prompts": 1,
"negative_variants": 2,
"archived_negative_variants": 1,
"images": 0,
},
)
self.assertFalse(result["target_exists"])
@@ -810,6 +922,7 @@ class TestVocabularyCategoryCommands(unittest.TestCase):
"other_prompts": 0,
"negative_variants": 0,
"archived_negative_variants": 0,
"images": 0,
},
)
self.assertEqual(self.rows("prompts")[variant]["category"], "Other")
@@ -902,6 +1015,171 @@ class TestVocabularyCategoryCommands(unittest.TestCase):
self.apply(bad)
self.assertEqual(self.snapshot(), before)
def own(self, value, *paths, root="root"):
self.db.set_image_categories([(root, path) for path in paths], value)
def image_rows(self):
with self.db.model.get_connection() as conn:
return {
(row[0], row[1]): (row[2], row[3], row[4])
for row in conn.execute(
"SELECT root_key,relative_path,category,updated_at,rowid "
"FROM image_categories ORDER BY rowid"
)
}
def checksum(self):
with self.db.model.get_connection() as conn:
return category_checksum(conn)
def test_preview_counts_image_rows(self):
self.owners()
self.own("Scenery", "a.png", "b.png", "missing.png")
self.own("Scenery", "c.png", root="second")
self.own("Other", "d.png")
before = self.image_rows()
result = self.preview()
self.assertEqual(result["impact"]["images"], 4)
self.assertEqual(self.image_rows(), before)
self.assertTrue(self.preview(target="Other")["target_exists"])
def test_apply_renames_image_categories(self):
self.owners()
bystander = self.save("Bystander", category="Other")
self.own("Scenery", "a.png", "b.png")
self.own("Scenery", "c.png", root="second")
self.own(" Scenery ", "padded.png")
self.own("Other", "d.png")
self.age_rows()
with self.db.model.get_connection() as conn:
conn.execute("UPDATE image_categories SET updated_at=?", (_OLD_STAMP,))
images_before = self.image_rows()
prompts_before, pairs_before = self.snapshot()
preview = self.preview()
receipt = self.apply(preview["impact_token"])
self.assertEqual(receipt["impact"]["images"], 3)
images_after = self.image_rows()
self.assertEqual(set(images_after), set(images_before))
renamed = {("root", "a.png"), ("root", "b.png"), ("second", "c.png")}
for key, (value, stamp, rowid) in images_after.items():
old_value, old_stamp, old_rowid = images_before[key]
self.assertEqual(rowid, old_rowid, key)
if key in renamed:
self.assertEqual(value, "Landscape")
self.assertNotEqual(stamp, _OLD_STAMP)
else:
self.assertEqual((value, stamp), (old_value, old_stamp))
prompts_after, pairs_after = self.snapshot()
self.assertEqual(prompts_after[bystander], prompts_before[bystander])
for before, after in (
(prompts_before, prompts_after),
(pairs_before, pairs_after),
):
for row_id, row in after.items():
if before[row_id]["category"] != "Scenery":
self.assertEqual(row, before[row_id])
# Merging into a value images already hold keeps every file key.
merge = self.preview(source="Landscape", target="Other")
self.assertTrue(merge["target_exists"])
self.assertEqual(merge["impact"]["images"], 3)
receipt = self.apply(merge["impact_token"], source="Landscape", target="Other")
self.assertTrue(receipt["consolidated"])
merged = self.image_rows()
self.assertEqual(set(merged), set(images_before))
self.assertEqual(
sorted(key for key, row in merged.items() if row[0] == "Other"),
sorted([*renamed, ("root", "d.png")]),
)
def test_image_only_value_can_be_renamed(self):
self.save("Bystander", category="Kept")
self.db.save_prompt_pair("Pair", "Noise", category="Kept")
self.own("Image only", "a.png", "b.png")
self.age_rows()
before = self.checksum()
preview = self.preview(source="Image only", target="Renamed")
self.assertEqual(
preview["impact"],
{
"positive_prompts": 0,
"other_prompts": 0,
"negative_variants": 0,
"archived_negative_variants": 0,
"images": 2,
},
)
self.assertFalse(preview["target_exists"])
receipt = self.apply(
preview["impact_token"], source="Image only", target="Renamed"
)
self.assertFalse(receipt["consolidated"])
self.assertEqual({row[0] for row in self.image_rows().values()}, {"Renamed"})
self.assertEqual(self.checksum(), before)
# An image-only target counts as existing.
self.own("Image target", "c.png")
self.assertTrue(
self.preview(source="Renamed", target="Image target")["target_exists"]
)
# Nothing holds the value any more: refused.
with self.assertRaises(LibraryCommandError) as caught:
self.preview(source="Image only", target="Again")
self.assertEqual(caught.exception.code, "vocabulary-conflict")
def test_token_refused_after_an_image_row_gains_or_loses_the_value(self):
self.owners()
self.own("Scenery", "a.png", "b.png")
def gain():
self.own("Scenery", "late.png")
def lose():
self.db.clear_image_categories([("root", "a.png")])
def move_away():
self.own("Elsewhere", "b.png")
def target_on_image():
self.own("Landscape", "target.png")
for mutate in (gain, lose, move_away, target_on_image):
with self.subTest(mutate.__name__):
token = self.preview()["impact_token"]
mutate()
before = self.snapshot(), self.image_rows()
with self.assertRaises(LibraryCommandError) as caught:
self.apply(token)
self.assertEqual(caught.exception.code, "vocabulary-conflict")
self.assertEqual((self.snapshot(), self.image_rows()), before)
def test_image_category_failure_rolls_back_prompts(self):
self.owners()
self.own("Scenery", "a.png")
token = self.preview()["impact_token"]
with self.db.model.get_connection() as conn:
conn.execute(
"CREATE TRIGGER abort_images BEFORE UPDATE ON image_categories "
"BEGIN SELECT RAISE(ABORT, 'boom'); END"
)
before = self.snapshot(), self.image_rows()
with self.assertRaises(sqlite3.Error):
self.apply(token)
self.assertEqual((self.snapshot(), self.image_rows()), before)
def test_category_exists_reads_the_image_index(self):
from database import vocabulary
with self.db.model.get_connection() as conn:
plan = [
row[-1]
for row in conn.execute(
"EXPLAIN QUERY PLAN " + vocabulary._CATEGORY_EXISTS, {"v": "x"}
)
]
steps = [step for step in plan if "image_categories" in step]
self.assertTrue(steps, plan)
for step in steps:
self.assertIn("idx_image_categories_category", step)
def test_category_apply_failure_rolls_back(self):
self.owners()
token = self.preview()["impact_token"]
@@ -942,6 +1220,18 @@ class TestSharedTokenHashing(unittest.TestCase):
),
"cefe7f6a5544b7f6a203f5527810899c6af568aa7abc1b75adffbfb5ae6251ae",
)
with_images = category_impact_token(
"e", "S", "T", True, [(1, True)], [], images=[9, 4]
)
self.assertRegex(with_images, r"\A[0-9a-f]{64}\Z")
self.assertEqual(
with_images,
category_impact_token("e", "S", "T", True, [(1, True)], [], images=[4, 9]),
)
self.assertNotEqual(
with_images,
category_impact_token("e", "S", "T", True, [(1, True)], [], images=[4]),
)
if __name__ == "__main__":
+6 -4
View File
@@ -33,15 +33,17 @@ test('execution script discovers only the new matrix and preserves all 78 origin
const executions = list('tools/vendor/playwright.execution.config.mjs');
assert.match(executions, /\[classic-start\]/);
assert.match(executions, /\[nodes2-start\]/);
// + the fresh-host library location case, + #342 image attribution, + #343 linked text
assert.match(executions, /Total: 90 tests in 5 files/);
// + the fresh-host library location case, + #342 image attribution, + #343 linked text,
// + #332 image category routes on a generated output (image-attribution.spec.mjs, 2 projects),
// + #347 cached runtime-linked runs: five node kinds and eight one-renderer cases (execution-text-link.spec.mjs, 2 projects)
assert.match(executions, /Total: 118 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.equal(cases.filter(line => line.includes('image-attribution.spec.mjs')).length, 3);
assert.equal(cases.filter(line => line.includes('execution-text-link.spec.mjs')).length, 32);
}
assert.doesNotMatch(executions, /prompt-manager.spec/);
});
+532
View File
@@ -0,0 +1,532 @@
"""Count a cached PromptManager node like the execution it reuses (#347).
The queue hook (``usage_tracking``) counts text it can read from the queued
graph. Text only known at run time (a PreviewAny or LLM link, a list node) is
counted by the node itself when it executes. When ComfyUI later serves the node
from its output cache the node does not execute, so without this module that
job is not counted.
When a PromptManager-family node executes it records a *receipt item*: the
role-independent facts its own count is made from (the stored text hashes)
and digests of the STRING outputs it returned. When ComfyUI reports the node
as cached for a later job (``execution_cached``), the receipt whose digests
equal the text in the cache entry is replayed once for that job. The node's
role (positive, negative, which Pair outputs are used) and what the hook
already saw are recomputed from that job's own graph, and both paths count
through the same ``own_counts`` so they cannot drift.
Fail closed: when the cached text cannot be matched to a receipt recorded in
this process, nothing is counted (the behavior before #347). In memory only and
bounded; starts no threads.
"""
import hashlib
import threading
from collections import OrderedDict
from typing import Any, Callable, Dict, FrozenSet, List, NamedTuple, Optional, Tuple
try:
from .hashing import generate_prompt_hash
from .logging_config import get_logger
from .prompt_graph import (
PROMPT_MANAGER_TYPES,
intrinsic_pair_nodes,
resolve_text,
run_prompt_components,
run_prompt_nodes,
)
from .prompt_pair_resolver import resolve_text_input
from .queue_observers import subscribe_dequeue, subscribe_task_done
from .run_keys import prompt_node_key, upstream_nodes
from .usage_tracking import find_pair_id_by_hashes, survive_switch
except ImportError:
from utils.hashing import generate_prompt_hash
from utils.logging_config import get_logger
from utils.prompt_graph import (
PROMPT_MANAGER_TYPES,
intrinsic_pair_nodes,
resolve_text,
run_prompt_components,
run_prompt_nodes,
)
from utils.prompt_pair_resolver import resolve_text_input
from utils.queue_observers import subscribe_dequeue, subscribe_task_done
from utils.run_keys import prompt_node_key, upstream_nodes
from utils.usage_tracking import find_pair_id_by_hashes, survive_switch
logger = get_logger("prompt_manager.cached_runs")
PAIR = "PromptManagerPair"
MAX_KEYS = 256 # run keys remembered (the same bound as MAX_KEYED_RUNS)
MAX_RECEIPTS = 4 # receipts per run key: a recurring upstream value (A, B, A)
MAX_ITEMS = 64 # list items per receipt (the cache reader's bound)
MAX_CACHED_NODES = 500 # the execution_cached guard the history capture uses
_EMPTY_HASH = generate_prompt_hash("")
Count = Tuple[Any, ...] # ("prompt", hash) or ("pair", positive, negative or None)
class Facts(NamedTuple):
"""What one execution stored: hashes of the stored text, None when nothing.
``positive_hash`` is None for an empty editor (or saving disabled);
``negative_hash`` (Pair only) is None for an empty negative, which the pair
stores as ``negative_prompt_id IS NULL``.
"""
positive_hash: Optional[str]
negative_hash: Optional[str] = None
class Role(NamedTuple):
"""How a job's graph uses the node: which editors reach a sampler's positive."""
used_fields: FrozenSet[str]
pair_used: bool = False
class Seen(NamedTuple):
"""Which of the node's texts the queue hook resolved, and so already counted."""
text: bool = False
negative_text: bool = False
pair: bool = False
class _Item(NamedTuple):
outputs: Tuple[str, ...]
facts: Facts
class _Receipt:
__slots__ = ("item_id", "items")
def __init__(self, item_id: int) -> None:
self.item_id = item_id
self.items: Dict[int, _Item] = {}
def outputs(self) -> Dict[int, Tuple[str, ...]]:
return {index: item.outputs for index, item in self.items.items()}
class _Attempt(NamedTuple):
item_id: int
prompt_id: str
graph: Dict[str, Any]
nodes: Dict[str, Tuple[str, str]] # node id -> (class, run key)
replayed: set
def text_digest(value: str) -> str:
"""Digest of one returned STRING value, exactly as returned."""
return hashlib.sha256(value.encode("utf-8", "surrogatepass")).hexdigest()
def stored_facts(text: Any, negative_text: Any = None) -> Facts:
"""The facts of an execution that stores ``text`` (and a Pair's negative)."""
positive = text if isinstance(text, str) and text.strip() else None
negative = (
negative_text
if isinstance(negative_text, str) and negative_text.strip()
else None
)
return Facts(
generate_prompt_hash(positive) if positive is not None else None,
generate_prompt_hash(negative) if negative is not None else None,
)
def pair_components(facts: Facts, role: Role) -> List[Tuple[str, str]]:
"""``(field, hash)`` of each used, stored Pair component; each hash once."""
found, counted = [], set()
for field, prompt_hash in (
("text", facts.positive_hash),
("negative_text", facts.negative_hash),
):
if field not in role.used_fields or prompt_hash is None:
continue
if prompt_hash in counted:
continue
counted.add(prompt_hash)
found.append((field, prompt_hash))
return found
def own_counts(node_type: str, facts: Facts, role: Role, seen: Seen) -> List[Count]:
"""What one execution of a node counts itself: text the hook could not see.
Shared by the executing node and by replay. Excludes the pending first-use
catch-up, which only an execution that saves the prompt can consume.
"""
if facts.positive_hash is None:
return []
if node_type != PAIR:
if "text" in role.used_fields and not seen.text:
return [("prompt", facts.positive_hash)]
return []
counts: List[Count] = [
("prompt", prompt_hash)
for field, prompt_hash in pair_components(facts, role)
if not getattr(seen, field)
]
if role.pair_used and not seen.pair:
counts.append(("pair", facts.positive_hash, facts.negative_hash))
return counts
def graph_role(graph: Any, node_id: str, node_type: str) -> Role:
"""The node's role in a queued graph, by the functions the nodes use."""
if node_type != PAIR:
used = {"text"} if node_id in run_prompt_nodes(graph) else set()
return Role(frozenset(used))
used = {
field for source, field in run_prompt_components(graph) if source == node_id
}
return Role(frozenset(used), node_id in intrinsic_pair_nodes(graph))
def _same(resolved: Optional[str], prompt_hash: Optional[str]) -> bool:
return resolved is not None and generate_prompt_hash(resolved) == prompt_hash
def graph_seen(graph: Any, node_id: str, node_type: str, facts: Facts) -> Seen:
"""Which texts the queue hook resolved from ``graph`` (as the node decides)."""
if node_type != PAIR:
return Seen(text=_same(resolve_text(graph, node_id), facts.positive_hash))
positive = _same(resolve_text_input(graph, node_id, "text"), facts.positive_hash)
negative = _same(
resolve_text_input(graph, node_id, "negative_text"),
facts.negative_hash or _EMPTY_HASH,
)
return Seen(positive, negative, positive and negative)
def _host_context() -> Tuple[Optional[str], Optional[int]]:
"""(prompt id, list index) of the node executing on this thread, if known."""
try:
from comfy_execution.utils import get_executing_context
except Exception:
return None, None
try:
context = get_executing_context()
except Exception:
return None, None
prompt_id = getattr(context, "prompt_id", None)
index = getattr(context, "list_index", None)
return (
prompt_id if isinstance(prompt_id, str) and prompt_id else None,
index if type(index) is int else None,
)
def write_count(db: Any, count: Count) -> bool:
"""Record one counted run; False when the prompt or pair no longer exists."""
def write() -> bool:
if count[0] == "prompt":
row = db.get_prompt_by_hash(count[1])
if row:
db.record_prompt_use(row["id"])
return bool(row)
pair_id = find_pair_id_by_hashes(db, count[1], count[2])
if pair_id:
db.record_prompt_pair_use(pair_id)
return bool(pair_id)
return bool(survive_switch(db, write))
class CachedRuns:
"""Receipts of executed nodes and the one job currently executing."""
def __init__(
self,
*,
max_keys: int = MAX_KEYS,
max_receipts: int = MAX_RECEIPTS,
max_items: int = MAX_ITEMS,
context: Callable[[], Tuple[Optional[str], Optional[int]]] = _host_context,
) -> None:
self._max_keys = max_keys
self._max_receipts = max_receipts
self._max_items = max_items
self._context = context
self._receipts: "OrderedDict[str, List[_Receipt]]" = OrderedDict()
self._attempt: Optional[_Attempt] = None
self._lock = threading.Lock()
# -- the job ------------------------------------------------------------
def begin_attempt(self, result: Any) -> None:
"""A job was dequeued: ``result`` is ``PromptQueue.get``'s return value."""
attempt = None
try:
item, item_id = result[0], result[1]
prompt_id, graph, outputs = item[1], item[2], item[4]
if type(item_id) is int and isinstance(prompt_id, str):
if isinstance(graph, dict):
attempt = _Attempt(
item_id, prompt_id, graph, _job_nodes(graph, outputs), set()
)
except Exception as error:
logger.debug(f"Dequeued job not tracked for cached runs: {error}")
with self._lock:
self._attempt = attempt
def end_attempt(self, item_id: Any) -> None:
"""The job finished: its receipts are complete and replace equal older ones.
Deduplicating only complete receipts keeps a list execution that is
still recording ([A] so far of [A, B]) from dropping an older [A].
"""
with self._lock:
if self._attempt is not None and self._attempt.item_id == item_id:
self._attempt = None
for run_key, receipts in self._receipts.items():
if receipts and receipts[0].item_id == item_id:
self._receipts[run_key] = _deduplicated(receipts)
def _current(self) -> Tuple[Optional[_Attempt], Optional[int]]:
"""The attempt the executing node belongs to, and its list index."""
prompt_id, index = self._context()
with self._lock:
attempt = self._attempt
if attempt is None or prompt_id is None or prompt_id != attempt.prompt_id:
return None, None
return attempt, index
# -- node execution -----------------------------------------------------
def replayed(self, node_id: Any) -> bool:
"""Whether this job already counted the executing node by replay."""
attempt, _ = self._current()
if attempt is None:
return False
with self._lock:
return str(node_id) in attempt.replayed
def record_item(self, node_id: Any, outputs: Tuple[Any, ...], facts: Facts) -> None:
"""Remember what this execution returned and was entitled to count."""
try:
attempt, index = self._current()
node = attempt.nodes.get(str(node_id)) if attempt else None
if node is None:
return
index = 0 if index is None else index
if not 0 <= index < self._max_items:
return
if any(not isinstance(value, str) for value in outputs):
return
item = _Item(tuple(text_digest(value) for value in outputs), facts)
self._store(node[1], attempt.item_id, index, item)
except Exception as error:
logger.debug(f"Could not record a cached-run receipt: {error}")
def _store(self, run_key: str, item_id: int, index: int, item: _Item) -> None:
with self._lock:
receipts = self._receipts.pop(run_key, [])
if receipts and receipts[0].item_id == item_id:
receipt = receipts[0]
else:
# The previous job's receipt is complete (its task_done may
# never have fired): it replaces older receipts it equals.
receipts = _deduplicated(receipts)
receipt = _Receipt(item_id)
receipts.insert(0, receipt)
receipt.items[index] = item
self._receipts[run_key] = receipts[: self._max_receipts]
while len(self._receipts) > self._max_keys:
self._receipts.popitem(last=False)
# -- cached nodes -------------------------------------------------------
def replay(
self,
executor: Any,
prompt_id: Any,
nodes: Any,
db_factory: Callable[[], Any],
reader: Optional[Callable[[Any, str, str], Any]] = None,
) -> None:
"""Count the job's cached nodes whose cached text a receipt proves."""
if reader is None:
reader = _host().cached_string_outputs
with self._lock:
attempt = self._attempt
if attempt is None or prompt_id != attempt.prompt_id:
return
if type(nodes) not in (list, tuple) or len(nodes) > MAX_CACHED_NODES:
return
db = None
for node_id in nodes:
node = attempt.nodes.get(node_id) if type(node_id) is str else None
if node is None:
continue
node_type, run_key = node
items = self._match(run_key, reader(executor, node_id, node_type))
if items is None:
logger.debug(
f"Cached node {node_id} not counted: no matching execution"
)
continue
with self._lock:
if node_id in attempt.replayed:
continue # a repeated message or id never counts twice
attempt.replayed.add(node_id)
counts: List[Count] = []
role = graph_role(attempt.graph, node_id, node_type)
for item in items:
seen = graph_seen(attempt.graph, node_id, node_type, item.facts)
counts.extend(own_counts(node_type, item.facts, role, seen))
if not counts:
continue
try:
db = db if db is not None else db_factory()
for count in counts:
if not write_count(db, count):
logger.info(
f"Cached node {node_id}: its {count[0]} no longer exists, "
"not counted"
)
except Exception as error:
logger.warning(f"Could not count cached node {node_id}: {error}")
def _match(self, run_key: str, values: Any) -> Optional[List[_Item]]:
"""The items of the one receipt whose outputs equal ``values``, in order."""
if not values:
return None
length = len(values[0])
if length == 0 or length > self._max_items:
return None
if any(len(slot) != length for slot in values):
return None
wanted = {
index: tuple(text_digest(slot[index]) for slot in values)
for index in range(length)
}
with self._lock:
receipts = self._receipts.get(run_key)
if not receipts:
return None
matches = [receipt for receipt in receipts if receipt.outputs() == wanted]
if not matches:
return None
self._receipts.move_to_end(run_key)
facts = {
tuple(receipt.items[index].facts for index in range(length))
for receipt in matches
}
if len(facts) != 1:
return None
return [matches[0].items[index] for index in range(length)]
def _deduplicated(receipts: List[_Receipt]) -> List[_Receipt]:
"""The newest receipt replaces older ones with the same outputs."""
if not receipts:
return receipts
outputs = receipts[0].outputs()
return [receipts[0]] + [r for r in receipts[1:] if r.outputs() != outputs]
def _job_nodes(graph: Dict[str, Any], outputs: Any) -> Dict[str, Tuple[str, str]]:
"""PromptManager-family nodes the job's outputs need, with their run keys."""
starts = (
[str(node) for node in outputs] if isinstance(outputs, (list, tuple)) else []
)
needed = upstream_nodes(graph, starts)
nodes = {}
for node_id in needed:
node = graph.get(node_id)
node_type = node.get("class_type") if isinstance(node, dict) else None
if node_type not in PROMPT_MANAGER_TYPES:
continue
run_key = prompt_node_key(graph, node_id)
if run_key is not None:
nodes[node_id] = (node_type, run_key)
return nodes
def _host() -> Any:
"""The host adapter module (imported late: it needs ComfyUI's modules)."""
try:
from . import execution_context_host
except ImportError:
from utils import execution_context_host
return execution_context_host
# Shared between the nodes and the host observers
CACHED_RUNS = CachedRuns()
_registered = threading.Event()
_register_lock = threading.Lock()
def register_cached_runs(
server: Any, db_factory: Callable[[], Any], runs: Optional[CachedRuns] = None
) -> bool:
"""Observe dequeue, ``execution_cached`` and ``task_done``; once per process.
Registers nothing when the host lacks any of them: counting is then as
before #347.
"""
runs = CACHED_RUNS if runs is None else runs
with _register_lock:
if _registered.is_set():
return False
queue = getattr(server, "prompt_queue", None)
if not callable(getattr(queue, "get", None)) or not callable(
getattr(queue, "task_done", None)
):
logger.info("Cached runs not counted: the prompt queue is not observable")
return False
try:
from execution import PromptExecutor
except Exception:
logger.info("Cached runs not counted: no ComfyUI executor")
return False
if not callable(getattr(PromptExecutor, "add_message", None)):
logger.info("Cached runs not counted: executor messages unavailable")
return False
subscribe_executor_messages = _host().subscribe_executor_messages
def on_message(executor: Any, args: tuple, kwargs: dict) -> None:
try:
event = args[0] if args else kwargs.get("event")
if event != "execution_cached":
return
server_of = getattr(executor, "server", None)
if getattr(server_of, "prompt_queue", None) is not queue:
return
data = args[1] if len(args) > 1 else kwargs.get("data")
if type(data) is not dict:
return
runs.replay(
executor, data.get("prompt_id"), data.get("nodes"), db_factory
)
except Exception as error:
logger.warning(f"Could not count cached nodes: {error}")
unsubscribers = []
try:
unsubscribers.append(
subscribe_dequeue(queue, "cached_runs", runs.begin_attempt)
)
unsubscribers.append(
subscribe_task_done(
queue,
"cached_runs",
lambda args, kwargs: runs.end_attempt(
args[0] if args else kwargs.get("item_id")
),
)
)
unsubscribers.append(subscribe_executor_messages("cached_runs", on_message))
except Exception:
for unsubscribe in unsubscribers:
unsubscribe()
raise
_registered.set()
logger.info("Counting cached runs of runtime-linked prompt text")
return True
+80 -18
View File
@@ -75,6 +75,82 @@ def subscribe_executor_messages(key: str, callback: Callable) -> Callable[[], No
return state.subscribe(key, callback)
# STRING output slots of each PromptManager-family node, in output order
STRING_SLOTS = {
"PromptManager": (1,),
"PromptManagerFocused": (1,),
"PromptManagerSplit": (1,),
"PromptManagerText": (0,),
"PromptManagerPair": (2, 3),
}
MAX_CACHED_VALUES = 64
def _local_cache_reader(executor: object) -> Callable | None:
"""``node_id -> outputs`` of the executor's local output cache, or None.
Only the supported cache classes, only ``get_local`` (never a provider or
an async lookup), only a real ``CacheEntry`` with at most four outputs.
"""
from comfy_execution import caching
from execution import CacheEntry
cache = executor.caches.outputs
supported = tuple(
cls
for name in ("HierarchicalCache", "LRUCache", "RAMPressureCache")
if isinstance(cls := getattr(caching, name, None), type)
)
if type(cache) not in supported:
return None
local = getattr(cache, "get_local", None)
if not callable(local):
return None
def read(node_id: str) -> list | tuple | None:
entry = local(node_id)
if type(entry) is not CacheEntry:
return None
outputs = entry.outputs
if type(outputs) not in (list, tuple) or not 0 < len(outputs) <= 4:
return None
return outputs
return read
def cached_string_outputs(
executor: object, node_id: str, node_type: str
) -> tuple[tuple[str, ...], ...] | None:
"""One tuple per STRING slot of a cached node's local entry, or None (#347).
None when the cache or entry is not supported, a slot is missing or holds
more than 64 values, any value is not a ``str``, or the slots differ in
length. Never runs a provider or an async lookup.
"""
try:
slots = STRING_SLOTS.get(node_type)
if slots is None or type(node_id) is not str:
return None
local = _local_cache_reader(executor)
outputs = local(node_id) if local is not None else None
if outputs is None or not max(slots) < len(outputs):
return None
found = []
for slot in slots:
values = outputs[slot]
if type(values) not in (list, tuple) or len(values) > MAX_CACHED_VALUES:
return None
if any(type(value) is not str for value in values):
return None
found.append(tuple(values))
if len({len(values) for values in found}) != 1:
return None
return tuple(found)
except Exception:
return None
def observe_cached_context(
executor: object, message_type: object, message_data: object
) -> None:
@@ -101,19 +177,8 @@ def observe_cached_context(
nodes = message_data.get("nodes")
if type(nodes) not in (list, tuple) or len(nodes) > 500:
return
from comfy_execution import caching
from execution import CacheEntry
cache = executor.caches.outputs
supported = tuple(
cls
for name in ("HierarchicalCache", "LRUCache", "RAMPressureCache")
if isinstance(cls := getattr(caching, name, None), type)
)
if type(cache) not in supported:
return
local = getattr(cache, "get_local", None)
if not callable(local):
local = _local_cache_reader(executor)
if local is None:
return
# Only source IDs retained by the bounded submitted graph are queried.
for source in cap.sources:
@@ -124,11 +189,8 @@ def observe_cached_context(
if source.node_type == "PromptManagerPair"
else 0 if source.node_type == "PromptManagerText" else 1
)
entry = local(source.node_id)
if type(entry) is not CacheEntry:
continue
outputs = entry.outputs
if type(outputs) not in (list, tuple) or not slot < len(outputs) <= 4:
outputs = local(source.node_id)
if outputs is None or not slot < len(outputs):
continue
values = outputs[slot]
if type(values) not in (list, tuple) or len(values) > 64:
+110
View File
@@ -0,0 +1,110 @@
"""The process exit signal: running work stops at its next safe point (#376).
ComfyUI catches Ctrl+C in ``main.py`` and returns without closing its event
loop. The interpreter then runs ``threading._shutdown``, which first calls the
``threading._register_atexit`` callbacks, newest first. One of them,
``concurrent.futures``'s ``_python_exit``, joins every executor thread, and
only after that do the ``atexit`` callbacks run (the package teardown). A call
still running on the default executor (a scan batch, a model download, a
thumbnail batch) used to keep the process alive until it returned.
``install()`` registers a callback in that same list after ``concurrent.futures``
did, so it runs before the join. It sets the signal that ``requested()`` reads
(``Job.cancelled`` and ``map_parallel`` read it) and runs ``on_exit``.
``threading._register_atexit`` is CPython-internal. It has the same contract
on 3.12, 3.13 and 3.14. Without it, or when it refuses the registration, exit
behaves as before: the call in flight runs to its end.
"""
import logging
import threading
logger = logging.getLogger("prompt_manager.exit")
_requested = threading.Event()
_install_lock = threading.Lock()
_installed = False
_on_exit = []
class ExitRequested(Exception):
"""The interpreter is exiting: work not yet started is not started."""
def requested():
"""True once the interpreter has started exiting."""
return _requested.is_set()
def request():
"""Set the signal (the exit hook; tests)."""
_requested.set()
def checkpoint():
"""Raise ExitRequested once the interpreter is exiting."""
if _requested.is_set():
raise ExitRequested("ComfyUI is exiting")
def checked(func):
"""``func`` behind a checkpoint: a call that starts after the signal
raises ExitRequested instead of running."""
def run(*args, **kwargs):
checkpoint()
return func(*args, **kwargs)
return run
def _hook():
"""The threading atexit callback. It never raises: an exception here
would end ``threading._shutdown`` before it shuts the executors."""
_requested.set()
for action in tuple(_on_exit):
try:
action()
except (Exception, KeyboardInterrupt):
logger.error("Stopping work at exit failed", exc_info=True)
def install(on_exit=None):
"""Register the exit hook once per process; True when it is registered.
``on_exit`` (optional) runs inside the hook, after the signal is set.
False (logged at DEBUG) when this Python has no ``_register_atexit`` or
refuses it; exit then waits for running work as before.
"""
global _installed
with _install_lock:
if on_exit is not None and on_exit not in _on_exit:
_on_exit.append(on_exit)
if _installed:
return True
# Imported first so its pool join is registered before our hook:
# the callbacks run newest first, so ours runs before that join.
import concurrent.futures.thread # noqa: F401
register = getattr(threading, "_register_atexit", None)
if not callable(register):
logger.debug("threading._register_atexit is missing; exit waits for jobs")
return False
try:
register(_hook)
except RuntimeError as error: # already shutting down
logger.debug(f"Exit hook not registered: {error}")
return False
_installed = True
return True
def reset():
"""Clear the signal and forget the registration (tests only; a hook
already registered stays in ``threading`` and sets the signal at exit)."""
global _installed
with _install_lock:
_requested.clear()
_installed = False
_on_exit.clear()
+56 -8
View File
@@ -300,6 +300,49 @@ def _blank(value: str | bytes | None) -> bool:
return value is None or not value.strip()
def inherited_category(
positive_category: str | bytes | None,
recipe_category: str | bytes | None,
has_recipe: bool,
) -> tuple:
"""``(value, scope)`` of the category an image's link supplies (#332 MC4).
The one rule for the inspector chip, the Images category filter and
History (D39): a linked recipe's own category when it is not blank,
``(recipe category, "pair")``; otherwise the positive's value unchanged,
blank or not, ``(positive category, "positive")``.
"""
if has_recipe and not _blank(recipe_category):
return recipe_category, "pair"
return positive_category, "positive"
class _OwnCategoryUnread:
"""The own-category read failed: hide the category (#332 DC5)."""
def __repr__(self) -> str:
return "OWN_CATEGORY_UNREAD"
OWN_CATEGORY_UNREAD = _OwnCategoryUnread()
def apply_own_category(library: dict, own) -> dict:
"""*library* with an image's own category applied (#332 §5).
``own`` is the own category (a string, shown with scope ``"image"``
whatever the link status), None (the inherited values stay) or
OWN_CATEGORY_UNREAD (category and scope hidden rather than possibly
wrong). Returns a new dict.
"""
result = dict(library)
if own is OWN_CATEGORY_UNREAD:
result.update(category=None, category_scope=None)
elif isinstance(own, str):
result.update(category=own, category_scope="image")
return result
def library_metadata(
positive_category: str | bytes | None,
positive_tags: list,
@@ -312,13 +355,14 @@ def library_metadata(
the linked positive prompt's, the values Set category writes and the Images
filters match. Tag sets are never merged. ``recipe`` is (category, tags).
"""
category, category_scope = positive_category, "positive"
category, category_scope = inherited_category(
positive_category,
recipe[0] if recipe is not None else None,
recipe is not None,
)
tags, tags_scope = positive_tags, "positive"
if recipe is not None:
if not _blank(recipe[0]):
category, category_scope = recipe[0], "pair"
if recipe[1]:
tags, tags_scope = recipe[1], "pair"
if recipe is not None and recipe[1]:
tags, tags_scope = recipe[1], "pair"
return {
"category": category,
"category_scope": category_scope,
@@ -345,9 +389,13 @@ def _library_projection(library: dict) -> dict:
value = library.get(key)
if value is None or isinstance(value, str):
result[key] = value
for key in ("category_scope", "tags_scope"):
# An image's own category (#332) has the scope "image"; tags never do.
for key, scopes in (
("category_scope", ("pair", "positive", "image", None)),
("tags_scope", ("pair", "positive", None)),
):
scope = library.get(key)
if scope in ("pair", "positive", None):
if scope in scopes:
result[key] = scope
tags = library.get("tags")
result["tags"] = (
+6 -1
View File
@@ -21,6 +21,8 @@ import logging
import threading
from concurrent.futures import ThreadPoolExecutor
from .exit_signal import checked
SHUTDOWN_REFUSAL = "interpreter shutdown"
_logger = logging.getLogger("prompt_manager.parallel")
_noted = set()
@@ -49,8 +51,11 @@ def map_parallel(func, items, workers):
``workers`` of 1 (or fewer) runs on the calling thread with no pool. An
exception in ``func`` propagates to the caller, so callers that must
survive bad input wrap ``func`` themselves. A pool refused at interpreter
exit runs the items it did not start on the calling thread (#375).
exit runs the items it did not start on the calling thread (#375). An item
that has not started when the interpreter starts exiting raises
ExitRequested instead (#376).
"""
func = checked(func)
items = list(items)
workers = max(1, int(workers or 1))
if workers == 1 or len(items) <= 1:
+76 -3
View File
@@ -10,6 +10,7 @@ import threading
import zlib
from dataclasses import dataclass, field as dataclass_field
from pathlib import Path
from collections.abc import Mapping
from typing import Any
from PIL import Image
@@ -461,6 +462,14 @@ def read_inspection_metadata(path: Path) -> dict[str, str]:
return read_inspection_file(path).metadata
# A side whose branch recorded run history named (#386): the save node that
# wrote the file picked the sampler; the text is still that branch's graph walk.
RECORDED_RUN_SOURCE = "recorded-run"
RECORDED_TEXT_REASON = "PromptManager recorded this text when the run executed"
# Outcomes the file alone leaves open; only these ask recorded history.
_UNDECIDED = frozenset({"ambiguous", "unresolved"})
def _side(status: str, text: str | None = None, source: str | None = None) -> dict:
return {"status": status, "text": text, "source": source}
@@ -529,15 +538,60 @@ def _parsed_workflow(value: Any) -> dict | None:
return value if isinstance(value, dict) else None
def _recorded_execution(resolved: dict, texts: Mapping) -> dict | None:
"""Evidence from the texts recorded at one sampler of the file's run (#386).
Only for an outcome that names one sampler (a single stage stopped at a
run-time value, such as a language model's output feeding PromptManager)
and only when that sampler's positive text was recorded and is not empty.
"""
sampler = resolved.get("sampler_id")
roles = texts.get(sampler) if isinstance(sampler, str) else None
if not isinstance(roles, Mapping):
return None
positive = roles.get("positive")
if not isinstance(positive, str) or not positive.strip():
return None
negative = roles.get("negative")
prompts = {
"positive": _side("resolved", positive, "execution"),
"negative": (
_side("resolved", negative, "execution")
if isinstance(negative, str)
else _side("unresolved")
),
}
resolution = {
"state": "resolved",
"sampler_id": sampler,
"reason": RECORDED_TEXT_REASON,
"reason_code": None,
"node": None,
}
return {"prompts": prompts, "resolution": resolution}
def inspect_prompt_evidence(
metadata: dict,
saved: dict | None = None,
*,
profile: ResolverProfile = INSPECTION_PROFILE,
file_name: str | None = None,
recorded_saver_id: str | None = None,
recorded_texts: Mapping[str, Mapping[str, str | None]] | None = None,
) -> dict:
"""Keep uncertainty visible; a saved association never replaces file evidence.
``recorded_saver_id`` is the save node recorded run history says wrote the
file (#386). It is asked only when the file alone leaves the prompt open
(ambiguous or unresolved), and the anchored result is kept only when it
names the positive (resolved, or a guider with no negative); its
graph-inferred sides are then labelled ``recorded-run``. ``recorded_texts``
(sampler id to the effective texts PromptManager recorded at execution in
that run) is used last: when the prompt is still open at one named sampler
whose positive text was recorded, the recorded texts are shown as
``execution``.
Inspection routes resolve with ``INSPECTION_PROFILE``: bounded work, not
graph size. Ingestion (scan, live monitor, admin import) calls the resolver
directly with its default profile. ``resolution`` carries the closed
@@ -599,6 +653,24 @@ def inspect_prompt_evidence(
# Only to recognize overwritten PromptManager text (R1); never a role.
canonical["workflow"] = workflow
resolved = resolve_prompt_pair(canonical, profile=profile, file_name=file_name)
recorded = False
if resolved["state"] in _UNDECIDED and recorded_saver_id is not None:
anchored = resolve_prompt_pair(
canonical,
profile=profile,
file_name=file_name,
recorded_saver_id=recorded_saver_id,
)
# Only a result whose samplers history's saver chose is about this file.
if anchored.get("recorded_anchor") is True:
if anchored["state"] not in _UNDECIDED and isinstance(
anchored["positive"], dict
):
resolved, recorded = anchored, True
elif recorded_texts:
executed = _recorded_execution(anchored, recorded_texts)
if executed is not None:
return executed
state = resolved["state"]
result = {}
for role in ("positive", "negative"):
@@ -608,9 +680,10 @@ def inspect_prompt_evidence(
and isinstance(side, dict)
and (role == "negative" or side["text"].strip())
):
result[role] = _side(
"resolved", side["text"], side["provenance"]["text_source"]
)
source = side["provenance"]["text_source"]
if recorded and source == "graph-inferred":
source = RECORDED_RUN_SOURCE
result[role] = _side("resolved", side["text"], source)
elif state == "resolved" and isinstance(side, dict):
# A resolved pair withholds only an empty positive.
result[role] = {**_side("unresolved"), "message": _EMPTY_POSITIVE}
+67 -10
View File
@@ -1601,6 +1601,11 @@ class _GraphResolver:
several, an unreadable walk or no candidate sampler is ``None``: the
file is then judged unanchored, never guessed (#381).
"""
match = self.saver_match(file_name)
return self.saver_walk(match) if match is not None else None
def saver_match(self, file_name: Any) -> Optional[str]:
"""The one node whose ``filename_prefix`` names ``file_name``, or None."""
stem = _file_stem(file_name)
if stem is None:
return None
@@ -1611,15 +1616,31 @@ class _GraphResolver:
for node_id, node in self.graph.items()
if self.writes(node, stem)
]
if len(matches) != 1:
return None
_, inputs = self.node(matches[0])
except (_Unresolved, _Diverging):
return None
return matches[0] if len(matches) == 1 else None
def recorded_saver_stages(self, node_id: str) -> Optional[list]:
"""The candidate samplers of the save node recorded history names (#386).
The node is walked exactly like a file-name match. A node the graph
lacks, a failing walk or no candidate sampler is ``None``: unanchored.
"""
if not isinstance(self.graph.get(node_id), dict):
return None
self.budget.reset(self.per_candidate_visits)
return self.saver_walk(node_id)
def saver_walk(self, saver_id: str) -> Optional[list]:
"""The listed samplers in one saver's linked inputs' ancestry, or None."""
try:
_, inputs = self.node(saver_id)
values = [
value
for name, value in inputs.items()
if name != "filename_prefix" and _link(value)
]
found = self.stages_from(matches[0], values)
found = self.stages_from(saver_id, values)
except (_Unresolved, _Diverging):
return None
candidates = [
@@ -2256,6 +2277,7 @@ def resolve_prompt_pair(
max_depth: Optional[int] = None,
max_text_length: Optional[int] = None,
file_name: Optional[str] = None,
recorded_saver_id: Optional[str] = None,
) -> Dict[str, Any]:
"""Resolve parsed metadata; a present API graph is authoritative over text.
@@ -2268,9 +2290,14 @@ def resolve_prompt_pair(
``profile`` selects the bounds; explicit ``max_*`` keywords override it.
``file_name`` (inspection only, unanchored) anchors on the one saver that
wrote that file (``_GraphResolver.saver_stages``) when exactly one matches.
``recorded_saver_id`` (inspection only, unanchored) is the save node that
recorded history says wrote the file (#386); it is tried before the file
name, and a node that yields no candidate sampler anchors nothing.
"""
profile = _profile_with(profile, max_visits, max_depth, max_text_length)
result = _resolve(metadata, output_node_id, adapters, profile, file_name)
result = _resolve(
metadata, output_node_id, adapters, profile, file_name, recorded_saver_id
)
return _legacy_result(result) if profile.legacy else result
@@ -2324,6 +2351,7 @@ def _resolve(
adapters: Optional[Mapping[Tuple[str, int], OutputAdapter]],
profile: ResolverProfile,
file_name: Optional[str] = None,
recorded_saver_id: Optional[str] = None,
) -> Dict[str, Any]:
if not isinstance(metadata, dict) or not metadata:
return _result("absent", reason="No parsed prompt metadata")
@@ -2379,11 +2407,12 @@ def _resolve(
"unresolved", reason=str(error), reason_code=error.code, node=error.node
)
else:
anchored = (
resolver.saver_stages(file_name)
if semantics is not None and file_name is not None
else None
)
anchored, by_record = None, False
if semantics is not None and recorded_saver_id is not None:
anchored = _recorded_anchor(resolver, str(recorded_saver_id), file_name)
by_record = anchored is not None
elif semantics is not None and file_name is not None:
anchored = resolver.saver_stages(file_name)
sampler_ids = anchored or [
node_id
for node_id, node in normalized.items()
@@ -2393,6 +2422,34 @@ def _resolve(
]
if semantics is not None and anchored is None:
sampler_ids = resolver.final_stages(sampler_ids)
if semantics is not None and recorded_saver_id is not None:
# The caller can tell whether history's saver chose the samplers.
return {
**_judge(resolver, sampler_ids, normalized, profile),
"recorded_anchor": by_record,
}
return _judge(resolver, sampler_ids, normalized, profile)
def _recorded_anchor(
resolver: "_GraphResolver", saver_id: str, file_name: Optional[str]
) -> Optional[list]:
"""The samplers of the saver history recorded, or None (#386).
A file name that names a different saver contradicts history: no anchor.
"""
named = resolver.saver_match(file_name) if file_name is not None else None
if named is not None and named != saver_id:
return None
return resolver.recorded_saver_stages(saver_id)
def _judge(
resolver: "_GraphResolver",
sampler_ids: list,
normalized: dict,
profile: ResolverProfile,
) -> Dict[str, Any]:
if not sampler_ids:
return _no_sampler(normalized, profile)
if len(sampler_ids) > 1:
+50 -14
View File
@@ -51,6 +51,53 @@ def _links(inputs: Any) -> List[str]:
return [value[0] for value in inputs.values() if is_link(value)]
def upstream_closure(
graph: Dict[str, Any], starts: List[str], limit: Optional[int] = None
) -> Optional[Dict[str, Any]]:
"""``{node_id: [class_type, inputs]}`` of ``starts`` and everything upstream.
None when a node or an upstream link's node is missing, or the closure
exceeds ``limit`` nodes.
"""
closure: Dict[str, Any] = {}
stack = list(starts)
while stack:
current = stack.pop()
if current in closure:
continue
node = graph.get(current)
if not isinstance(node, dict):
return None
if limit is not None and len(closure) >= limit:
return None
inputs = node.get("inputs")
closure[current] = [node.get("class_type"), inputs]
stack.extend(_links(inputs))
return closure
def upstream_nodes(graph: Any, starts: List[str]) -> List[str]:
"""Ids of ``starts`` and every node upstream of them that exists in ``graph``.
Used to tell which nodes a job's outputs to execute need ("run selected
output" jobs leave the rest of the graph out). Missing nodes are skipped.
"""
if not isinstance(graph, dict):
return []
found: Dict[str, None] = {}
stack = [str(node) for node in starts]
while stack:
current = stack.pop()
if current in found:
continue
node = graph.get(current)
if not isinstance(node, dict):
continue
found[current] = None
stack.extend(_links(node.get("inputs")))
return list(found)
def prompt_node_key(graph: Any, node_id: Any) -> Optional[str]:
"""Digest of ``node_id`` and its whole upstream closure, or None.
@@ -62,20 +109,9 @@ def prompt_node_key(graph: Any, node_id: Any) -> Optional[str]:
if not isinstance(graph, dict) or node_id is None:
return None
start = str(node_id)
closure: Dict[str, Any] = {}
stack = [start]
while stack:
current = stack.pop()
if current in closure:
continue
node = graph.get(current)
if not isinstance(node, dict):
return None
if len(closure) >= MAX_VISITS:
return None
inputs = node.get("inputs")
closure[current] = [node.get("class_type"), inputs]
stack.extend(_links(inputs))
closure = upstream_closure(graph, [start], MAX_VISITS)
if closure is None:
return None
try:
encoded = json.dumps(
[RUN_KEY_VERSION, start, closure],
+18 -7
View File
@@ -16,7 +16,9 @@ it never counts. Servers without a ``prompt_queue`` fall back to counting at
submit time.
Text from known pure string nodes is resolved at queue time too. Anything else
(PromptSearchList batches, unknown nodes) is counted by the node as it executes.
(PromptSearchList batches, unknown nodes) is counted by the node as it executes,
and when ComfyUI serves such a node from its cache, the job replays what that
execution counted (``cached_runs``, #347).
"""
import threading
@@ -167,13 +169,22 @@ def queued_pair_runs(graph: Any) -> List[Tuple[str, str, str]]:
def find_pair_id(db: Any, positive: str, negative: str) -> Optional[int]:
"""Find a saved pair without creating or overwriting its metadata."""
positive_row = db.get_prompt_by_hash(generate_prompt_hash(positive))
negative_row = (
db.get_prompt_by_hash(generate_prompt_hash(negative))
if negative.strip()
else None
return find_pair_id_by_hashes(
db,
generate_prompt_hash(positive),
generate_prompt_hash(negative) if negative.strip() else None,
)
if not positive_row or (negative.strip() and not negative_row):
def find_pair_id_by_hashes(
db: Any, positive_hash: str, negative_hash: Optional[str]
) -> Optional[int]:
"""The saved pair of two prompt hashes; ``None`` negative is the empty one."""
positive_row = db.get_prompt_by_hash(positive_hash)
negative_row = (
db.get_prompt_by_hash(negative_hash) if negative_hash is not None else None
)
if not positive_row or (negative_hash is not None and not negative_row):
return None
row = (
db.model.get_connection()
+21 -320
View File
@@ -164,27 +164,31 @@
class="px-2.5 py-1 bg-pm-surface border border-pm hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
Export
</button>
<button id="settingsBtn"
class="px-2.5 py-1 bg-pm-surface border border-pm hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
<a id="settingsBtn" href="workspace#/system/overview" target="_blank" rel="noopener"
title="Opens System › Overview in a new tab"
class="px-2.5 py-1 bg-pm-surface border border-pm hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
Settings
</button>
</a>
<a id="diagnosticsBtn" href="workspace#/system/diagnostics" target="_blank" rel="noopener"
title="Opens System › Diagnostics &amp; logs in a new tab"
class="px-2.5 py-1 bg-pm-surface border border-pm hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
Diagnostics
</a>
<button id="maintenanceBtn"
class="px-2.5 py-1 bg-pm-surface border border-pm hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
<a id="maintenanceBtn" href="workspace#/system/maintenance" target="_blank" rel="noopener"
title="Opens System › Maintenance in a new tab"
class="px-2.5 py-1 bg-pm-surface border border-pm hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
Maintenance
</button>
<button id="backupBtn"
class="px-2.5 py-1 bg-pm-surface border border-pm hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
</a>
<a id="backupBtn" href="workspace#/system/backups" target="_blank" rel="noopener"
title="Opens System › Backup &amp; export in a new tab"
class="px-2.5 py-1 bg-pm-surface border border-pm hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
Backup DB
</button>
<button id="restoreBtn"
class="px-2.5 py-1 bg-pm-surface border border-pm hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
</a>
<a id="restoreBtn" href="workspace#/system/backups" target="_blank" rel="noopener"
title="Opens System › Backup &amp; export in a new tab"
class="px-2.5 py-1 bg-pm-surface border border-pm hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
Restore DB
</button>
</a>
<button id="scanBtn"
class="px-2.5 py-1 bg-pm-surface border border-pm hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
Scan Images
@@ -231,11 +235,12 @@
</select>
<span class="text-xs text-pm-secondary">per page</span>
</div>
<!-- Session-only toggle; the saved default lives in Settings -->
<!-- Session-only toggle; the saved default lives in Page options -->
<label class="flex items-center gap-1.5 text-xs text-pm-secondary cursor-pointer" title="Load the next page as you scroll (this session)">
<input type="checkbox" id="infiniteScrollSession" class="w-3.5 h-3.5 rounded border-pm bg-pm-input text-pm-accent focus:ring-pm-accent">
Infinite scroll
</label>
<button id="pageOptionsBtn" type="button" class="text-xs text-pm-secondary hover:text-pm-accent underline">Page options</button>
</div>
<span class="text-xs font-medium text-pm-secondary bg-pm-input px-2 py-0.5 rounded-full" id="resultsCount">Loading...</span>
</div>
@@ -369,61 +374,13 @@
<!-- Settings Modal -->
<div id="settingsModal" class="fixed inset-0 bg-black/60 backdrop-blur-[2px] hidden items-center justify-center z-50">
<div class="bg-pm-surface rounded-pm-lg p-4 max-w-lg w-full mx-4 border border-pm max-h-[90vh] overflow-y-auto shadow-pm">
<h3 class="text-sm font-semibold text-pm mb-4">Settings</h3>
<h3 class="text-sm font-semibold text-pm mb-4">Page options</h3>
<div class="space-y-4">
<!-- Gallery Settings Section -->
<div class="border-b border-pm pb-3">
<h4 class="text-xs font-semibold text-pm-secondary mb-3 uppercase tracking-wide">Gallery Settings</h4>
<h4 class="text-xs font-semibold text-pm-secondary mb-3 uppercase tracking-wide">Prompt list</h4>
<div class="mb-3">
<label class="block text-xs font-medium text-pm-secondary mb-1">Image Scan Directories</label>
<div id="scanPathsList" class="space-y-2">
<!-- Populated dynamically by JS -->
</div>
<button id="addScanPathBtn" type="button"
class="mt-2 flex items-center gap-1 text-xs text-pm-accent hover:text-pm-accent-hover">
<svg class="w-3.5 h-3.5" fill="none" stroke="currentColor" viewBox="0 0 24 24">
<path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M12 4v16m8-8H4"/>
</svg>
Add Directory
</button>
<p class="text-xs text-pm-muted mt-1">Directories to scan for images. Leave empty to use ComfyUI default output directory.</p>
</div>
<div class="mb-3">
<div class="flex items-center justify-between">
<label class="block text-xs font-medium text-pm-secondary">Currently Monitoring</label>
<button id="refreshMonitoringStatus" class="text-xs text-pm-accent hover:text-pm-accent-hover">Refresh</button>
</div>
<div id="monitoringStatus" class="mt-1.5 px-2.5 py-1.5 bg-pm-input rounded-pm-sm text-xs text-pm-secondary font-mono">
Loading...
</div>
</div>
</div>
<!-- UI Settings Section -->
<div class="border-b border-pm pb-3">
<h4 class="text-xs font-semibold text-pm-secondary mb-3 uppercase tracking-wide">UI Settings</h4>
<div class="mb-3">
<label class="block text-xs font-medium text-pm-secondary mb-1">Search Results Auto-Hide (seconds)</label>
<input type="number" id="resultTimeout" min="0" max="300" value="30"
class="w-full px-2.5 py-1.5 bg-pm-input border border-pm rounded-pm-sm text-pm text-[13px] focus:outline-none">
<p class="text-xs text-pm-muted mt-1">Set to 0 to disable auto-hide.</p>
</div>
<div>
<label class="block text-xs font-medium text-pm-secondary mb-1">Web UI Display Mode</label>
<select id="webuiDisplayMode"
class="w-full px-2.5 py-1.5 bg-pm-input border border-pm rounded-pm-sm text-pm text-[13px] focus:outline-none">
<option value="popup">Popup Window (Current)</option>
<option value="newtab">New Tab</option>
</select>
<p class="text-xs text-pm-muted mt-1">Choose how the Web UI opens from ComfyUI nodes.</p>
</div>
<div class="mt-3 flex items-center justify-between">
<div class="flex items-center justify-between">
<div>
<label for="infiniteScrollDefault" class="block text-xs font-medium text-pm-secondary">Infinite Scrolling</label>
<p class="text-xs text-pm-muted">Load the next page of prompts as you scroll. The checkbox in the list header overrides this for the current session.</p>
@@ -431,76 +388,6 @@
<input type="checkbox" id="infiniteScrollDefault" class="w-4 h-4 rounded border-pm bg-pm-input text-pm-accent focus:ring-pm-accent ml-3">
</div>
</div>
<!-- Performance Section -->
<div class="border-b border-pm pb-3">
<h4 class="text-xs font-semibold text-pm-secondary mb-3 uppercase tracking-wide">Performance</h4>
<div>
<div class="flex items-center justify-between mb-1.5">
<label for="workerThreads" class="text-xs font-medium text-pm-secondary">Worker threads</label>
<span id="workerThreadsValue" class="text-pm-accent text-xs font-mono">1 / 1 cores</span>
</div>
<input type="range" id="workerThreads" min="1" max="1" step="1" value="1"
class="w-full h-1.5 bg-pm-primary rounded-lg appearance-none cursor-pointer accent-[var(--pm-accent)]">
<p class="text-xs text-pm-muted mt-1">Threads used by the image rescan and thumbnail generation. More threads finish sooner; fewer leave CPU for ComfyUI while it generates.</p>
</div>
</div>
<!-- Integrations Section -->
<div class="border-b border-pm pb-3">
<h4 class="text-xs font-semibold text-pm-secondary mb-3 uppercase tracking-wide">Integrations</h4>
<div class="mb-3">
<div class="flex items-center justify-between mb-2">
<div class="flex items-center gap-2">
<label class="block text-xs font-medium text-pm-secondary">LoraManager</label>
<span id="loraDetectionBadge" class="text-[10px] px-1.5 py-0.5 rounded-full bg-pm-input text-pm-muted">checking...</span>
</div>
<label class="relative inline-flex items-center cursor-pointer">
<input type="checkbox" id="loraEnabled" class="sr-only peer" disabled>
<div class="w-8 h-4 bg-pm-input rounded-full peer peer-checked:bg-pm-accent transition-colors pointer-events-none after:content-[''] after:absolute after:top-0.5 after:left-0.5 after:bg-white after:rounded-full after:h-3 after:w-3 after:transition-all peer-checked:after:translate-x-4"></div>
</label>
</div>
<div id="loraSettings" class="hidden space-y-2 mt-2 pl-2 border-l-2 border-pm">
<div>
<label class="block text-xs font-medium text-pm-secondary mb-1">Path</label>
<input type="text" id="loraManagerPath" readonly
class="w-full px-2.5 py-1.5 bg-pm-input border border-pm rounded-pm-sm text-pm text-[13px] text-pm-muted">
</div>
<div class="flex items-center justify-between">
<label class="block text-xs font-medium text-pm-secondary">Auto-inject trigger words</label>
<label class="relative inline-flex items-center cursor-pointer">
<input type="checkbox" id="loraTriggerWords" class="sr-only peer">
<div class="w-8 h-4 bg-pm-input rounded-full peer peer-checked:bg-pm-accent transition-colors pointer-events-none after:content-[''] after:absolute after:top-0.5 after:left-0.5 after:bg-white after:rounded-full after:h-3 after:w-3 after:transition-all peer-checked:after:translate-x-4"></div>
</label>
</div>
<p class="text-xs text-pm-muted">Automatically append trigger words when <code class="text-pm-accent">&lt;lora:name:weight&gt;</code> is detected in prompts.</p>
<div class="mt-2">
<label class="block text-xs font-medium text-pm-secondary mb-1">CivitAI API Key</label>
<input type="password" id="civitaiApiKey" placeholder="Optional — required for NSFW example images"
class="w-full px-2.5 py-1.5 bg-pm-input border border-pm rounded-pm-sm text-pm text-[13px] placeholder:text-pm-muted focus:outline-none">
<p class="text-xs text-pm-muted mt-1">Get your key from <a href="https://civitai.com/user/account" target="_blank" class="text-pm-accent hover:underline">civitai.com/user/account</a></p>
</div>
<label class="mt-1 flex items-center gap-1.5 text-xs text-pm-secondary cursor-pointer">
<input type="checkbox" id="clearCivitaiKey" class="w-3.5 h-3.5 text-pm-accent bg-pm-input border-pm rounded focus:ring-pm-accent">
<span>Clear the stored key on save</span>
</label>
<button id="loraImportBtn" type="button"
class="mt-1 flex items-center gap-1 text-xs text-pm-accent hover:text-pm-accent-hover">
<svg class="w-3.5 h-3.5" fill="none" stroke="currentColor" viewBox="0 0 24 24">
<path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M4 16v1a3 3 0 003 3h10a3 3 0 003-3v-1m-4-8l-4-4m0 0L8 8m4-4v12"/>
</svg>
Import LoRA Data
</button>
</div>
</div>
</div>
</div>
<div class="flex justify-end gap-2 mt-4">
@@ -516,32 +403,6 @@
</div>
</div>
<!-- LoRA Import Progress Modal -->
<div id="loraImportModal" class="fixed inset-0 bg-black/60 backdrop-blur-[2px] hidden items-center justify-center z-50">
<div class="bg-pm-surface rounded-pm-lg p-4 max-w-md w-full mx-4 border border-pm shadow-pm">
<h3 class="text-sm font-semibold text-pm mb-4">Importing LoRA Data</h3>
<div class="mb-3">
<p class="text-pm-secondary text-xs">Status:</p>
<p id="loraImportStatus" class="text-pm text-xs font-medium truncate">Initializing...</p>
</div>
<div class="mb-3">
<div class="flex justify-between text-xs text-pm-secondary mb-1">
<span>Progress</span>
<span id="loraImportPercent">0%</span>
</div>
<div class="w-full bg-pm-input rounded-full h-1.5">
<div id="loraImportBar" class="bg-pm-accent h-1.5 rounded-full transition-all duration-300" style="width: 0%"></div>
</div>
<div class="flex justify-between text-[11px] text-pm-muted mt-1">
<span id="loraImportProcessed">0 processed</span>
<span id="loraImportImported">0 imported</span>
</div>
</div>
</div>
</div>
<!-- Bulk Tag Modal -->
<div id="bulkTagModal" class="fixed inset-0 bg-black/60 backdrop-blur-[2px] hidden items-center justify-center z-50">
<div class="bg-pm-surface rounded-pm-lg p-4 max-w-md w-full mx-4 border border-pm shadow-pm">
@@ -614,136 +475,6 @@
</div>
</div>
<!-- Database Restore Modal -->
<div id="restoreModal" class="fixed inset-0 bg-black/60 backdrop-blur-[2px] hidden items-center justify-center z-50">
<div class="bg-pm-surface rounded-pm-lg p-4 max-w-md w-full mx-4 border border-pm shadow-pm">
<h3 class="text-sm font-semibold text-pm mb-4">Restore Database</h3>
<div class="mb-3">
<div class="bg-pm-warning-tint border border-pm-warning/40 rounded-pm-md p-3 mb-3">
<h4 class="text-pm-warning font-medium text-xs">Warning</h4>
<p class="text-pm text-xs mt-1">This will replace your current database. A backup will be created automatically.</p>
</div>
<label class="block text-xs font-medium text-pm-secondary mb-1">Select Database File (.db)</label>
<input type="file" id="restoreFileInput" accept=".db"
class="w-full px-2.5 py-1.5 bg-pm-input border border-pm rounded-pm-sm text-pm text-[13px] file:mr-3 file:py-1 file:px-3 file:rounded-pm-sm file:border-0 file:text-xs file:font-medium file:bg-pm-accent file:text-pm-accent-fg">
</div>
<div class="flex justify-end gap-2 mt-4">
<button id="cancelRestore"
class="px-3 py-1.5 bg-pm-surface border border-pm hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
Cancel
</button>
<button id="confirmRestore"
class="px-3 py-1.5 bg-pm-warning hover:opacity-80 text-black text-xs font-medium rounded-pm-sm transition-colors">
Restore Database
</button>
</div>
</div>
</div>
<!-- Maintenance Modal -->
<div id="maintenanceModal" class="fixed inset-0 bg-black/60 backdrop-blur-[2px] hidden items-center justify-center z-50">
<div class="bg-pm-surface rounded-pm-lg max-w-4xl w-full max-h-[90vh] overflow-hidden mx-4 border border-pm shadow-pm">
<div class="p-4 border-b border-pm">
<div class="flex items-center justify-between">
<h3 class="text-sm font-semibold text-pm">Database Maintenance</h3>
<button onclick="window.admin.closeMaintenance()" class="p-1 rounded-pm-sm hover:bg-pm-hover transition-colors">
<svg class="w-4 h-4 text-pm-secondary" fill="none" stroke="currentColor" viewBox="0 0 24 24">
<path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M6 18L18 6M6 6l12 12"></path>
</svg>
</button>
</div>
</div>
<div class="p-4 overflow-y-auto max-h-[calc(90vh-120px)]">
<!-- Maintenance Options -->
<div class="space-y-3 mb-4">
<h4 class="text-xs font-semibold text-pm-secondary uppercase tracking-wide mb-3">Select Maintenance Operations:</h4>
<div class="grid grid-cols-1 md:grid-cols-2 gap-3">
<label class="flex items-center gap-2.5 p-3 bg-pm-input rounded-pm-md hover:bg-pm-hover cursor-pointer">
<input type="checkbox" id="cleanup_duplicates" class="maintenance-option w-3.5 h-3.5 accent-[var(--pm-accent)]" checked>
<div class="flex-1">
<div class="text-pm text-xs font-medium">Remove Duplicates</div>
<div class="text-[11px] text-pm-secondary">Find and remove duplicate prompts based on text content</div>
</div>
</label>
<label class="flex items-center gap-2.5 p-3 bg-pm-input rounded-pm-md hover:bg-pm-hover cursor-pointer">
<input type="checkbox" id="vacuum" class="maintenance-option w-3.5 h-3.5 accent-[var(--pm-accent)]" checked>
<div class="flex-1">
<div class="text-pm text-xs font-medium">Vacuum Database</div>
<div class="text-[11px] text-pm-secondary">Optimize database file size and performance</div>
</div>
</label>
<label class="flex items-center gap-2.5 p-3 bg-pm-input rounded-pm-md hover:bg-pm-hover cursor-pointer">
<input type="checkbox" id="cleanup_orphaned_images" class="maintenance-option w-3.5 h-3.5 accent-[var(--pm-accent)]" checked>
<div class="flex-1">
<div class="text-pm text-xs font-medium">Clean Orphaned Images</div>
<div class="text-[11px] text-pm-secondary">Remove image records for files that no longer exist</div>
</div>
</label>
<label class="flex items-center gap-2.5 p-3 bg-pm-input rounded-pm-md hover:bg-pm-hover cursor-pointer">
<input type="checkbox" id="check_hash_duplicates" class="maintenance-option w-3.5 h-3.5 accent-[var(--pm-accent)]">
<div class="flex-1">
<div class="text-pm text-xs font-medium">Check Hash Duplicates</div>
<div class="text-[11px] text-pm-secondary">Identify prompts with duplicate hash values</div>
</div>
</label>
<label class="flex items-center gap-2.5 p-3 bg-pm-input rounded-pm-md hover:bg-pm-hover cursor-pointer">
<input type="checkbox" id="check_consistency" class="maintenance-option w-3.5 h-3.5 accent-[var(--pm-accent)]">
<div class="flex-1">
<div class="text-pm text-xs font-medium">Consistency Check</div>
<div class="text-[11px] text-pm-secondary">Verify database integrity and data consistency</div>
</div>
</label>
<label class="flex items-center gap-2.5 p-3 bg-pm-input rounded-pm-md hover:bg-pm-hover cursor-pointer">
<input type="checkbox" id="statistics" class="maintenance-option w-3.5 h-3.5 accent-[var(--pm-accent)]">
<div class="flex-1">
<div class="text-pm text-xs font-medium">Database Statistics</div>
<div class="text-[11px] text-pm-secondary">Generate detailed database statistics report</div>
</div>
</label>
<label class="flex items-center gap-2.5 p-3 bg-pm-input rounded-pm-md hover:bg-pm-hover cursor-pointer">
<input type="checkbox" id="prune_orphaned_prompts" class="maintenance-option w-3.5 h-3.5 accent-[var(--pm-accent)]">
<div class="flex-1">
<div class="text-pm text-xs font-medium">Prune Orphaned Prompts</div>
<div class="text-[11px] text-pm-secondary">Remove prompts that have no linked images</div>
</div>
</label>
</div>
</div>
<!-- Action Buttons -->
<div class="flex items-center justify-between mb-4">
<div class="flex gap-2">
<button id="runMaintenanceBtn" class="px-3 py-1.5 bg-pm-accent hover:bg-pm-accent-hover text-pm-accent-fg text-xs font-medium rounded-pm-sm transition-colors">
Run Maintenance
</button>
<button id="selectAllMaintenanceBtn" class="px-2.5 py-1.5 bg-pm-surface border border-pm hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
Select All
</button>
<button id="clearAllMaintenanceBtn" class="px-2.5 py-1.5 bg-pm-surface border border-pm hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
Clear All
</button>
</div>
</div>
<!-- Results -->
<div id="maintenanceResults" class="space-y-3">
<div class="text-center py-6 text-pm-secondary text-xs">
Click "Run Maintenance" to perform selected operations
</div>
</div>
</div>
</div>
</div>
<!-- Gallery Modal -->
<div id="galleryModal" class="fixed inset-0 bg-black/75 backdrop-blur-[2px] hidden items-center justify-center z-50">
<div class="bg-pm-surface rounded-pm-lg max-w-6xl w-full max-h-[90vh] overflow-hidden mx-4 border border-pm shadow-pm">
@@ -1257,36 +988,6 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr
</div>
</div>
<!-- Auto Tag Download Modal -->
<div id="autoTagDownloadModal" class="fixed inset-0 bg-black/60 backdrop-blur-[2px] hidden items-center justify-center z-50">
<div class="bg-pm-surface rounded-pm-lg p-4 max-w-md w-full mx-4 border border-pm shadow-pm">
<h3 class="text-sm font-semibold text-pm mb-4">Downloading Model</h3>
<div class="mb-3">
<p class="text-pm-secondary text-xs">Downloading:</p>
<p id="downloadModelName" class="text-pm text-xs font-medium">GGUF Model</p>
</div>
<div class="mb-4">
<div class="flex justify-between text-xs text-pm-secondary mb-1">
<span id="downloadStatus">Preparing...</span>
<span id="downloadProgressPercent">0%</span>
</div>
<div class="w-full bg-pm-input rounded-full h-1.5">
<div id="downloadProgressBar" class="bg-pm-accent h-1.5 rounded-full transition-all duration-300" style="width: 0%"></div>
</div>
<p class="text-pm-secondary text-xs mt-2">The download continues in Activity.</p>
</div>
<div class="flex justify-center">
<button id="cancelDownloadBtn"
class="px-3 py-1.5 text-pm-error bg-transparent hover:bg-pm-error-tint border border-pm text-xs font-medium rounded-pm-sm transition-colors">
Close
</button>
</div>
</div>
</div>
<script src="js/comfy-metadata.js"></script>
<script src="js/prompt-list-sort.js"></script>
<script src="js/view-router.js"></script>
+27 -249
View File
@@ -8,8 +8,6 @@
<script src="js/html-escape.js"></script>
<script src="js/image-helpers.js"></script>
<script src="js/gallery-settings.js"></script>
<script src="js/cache-ownership.js"></script>
<script src="js/scan-progress.js"></script>
<script src="js/api-client.js"></script>
<script src="js/png-metadata.js"></script>
<script src="js/list-state.js"></script>
@@ -63,7 +61,7 @@
<button id="thumbMdBtn" class="px-2 py-1 bg-pm-accent text-pm-accent-fg text-xs rounded-pm-sm transition-colors" title="Medium thumbnails">M</button>
<button id="thumbLgBtn" class="px-2 py-1 bg-pm-input hover:bg-pm-hover text-pm text-xs rounded-pm-sm transition-colors" title="Large thumbnails">L</button>
</div>
<div class="flex items-center gap-2">
<div class="flex flex-wrap items-center gap-2">
<button id="refreshBtn"
class="px-3 py-1 bg-pm-accent hover:bg-pm-accent-hover text-pm-accent-fg text-xs font-medium rounded-pm-sm transition-colors">
Refresh
@@ -76,9 +74,14 @@
class="px-3 py-1 bg-pm-input hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
List View
</button>
<button id="settingsBtn"
class="px-3 py-1 bg-pm-input hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
<a id="settingsBtn" href="workspace#/system/appearance" target="_blank" rel="noopener"
title="Opens System › Appearance in a new tab"
class="px-3 py-1 bg-pm-input hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
Settings
</a>
<button id="pageOptionsBtn" type="button"
class="px-3 py-1 bg-pm-input hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
Page options
</button>
<button id="autoTagBtn"
class="px-3 py-1 bg-pm-accent hover:bg-pm-accent-hover text-pm-accent-fg text-xs font-medium rounded-pm-sm transition-colors">
@@ -176,7 +179,7 @@
<div class="bg-pm-surface rounded-pm-md max-w-2xl w-full max-h-[90vh] overflow-hidden mx-4 border border-pm">
<div class="p-3 border-b border-pm">
<div class="flex items-center justify-between">
<h3 class="text-sm font-semibold text-pm">Gallery Settings</h3>
<h3 class="text-sm font-semibold text-pm">Page options</h3>
<button id="closeSettingsBtn" class="p-2 rounded-pm-sm hover:bg-pm-hover transition-colors">
<svg class="w-4 h-4 text-pm-secondary" fill="none" stroke="currentColor" viewBox="0 0 24 24">
<path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M6 18L18 6M6 6l12 12"></path>
@@ -190,25 +193,9 @@
<div class="space-y-3">
<h4 class="text-sm font-semibold text-pm">Performance Settings</h4>
<!-- Thumbnail Generation -->
<div class="bg-pm-input rounded-pm-sm p-3">
<div class="flex items-center justify-between mb-2">
<div>
<h5 class="text-pm text-xs font-medium">Generate Thumbnails</h5>
<p class="text-xs text-pm-secondary">Create optimized thumbnails for faster loading</p>
</div>
<button id="generateThumbnailsBtn" class="px-3 py-1 bg-pm-accent hover:bg-pm-accent-hover text-pm-accent-fg text-xs rounded-pm-sm transition-colors">
Generate
</button>
</div>
<div class="text-xs text-pm-muted mb-2">
Status: <span id="thumbnailStatus">Not generated</span>
<div class="text-xs text-pm-muted mt-1">
Thumbnails speed up gallery loading. Original images always used for viewing and metadata.
</div>
</div>
<div class="bg-pm-input rounded-pm-sm p-3 space-y-2">
<!-- Check thumbnails at startup -->
<div class="flex items-center justify-between py-2 border-t border-pm">
<div class="flex items-center justify-between">
<div>
<label class="text-xs text-pm-secondary">Check thumbnails at startup</label>
<p class="text-xs text-pm-muted">Prompt to generate thumbnails if missing when opening gallery</p>
@@ -218,43 +205,9 @@
<div class="w-11 h-6 bg-pm-input peer-focus:outline-none peer-focus:ring-4 peer-focus:ring-pm-accent rounded-full peer peer-checked:after:translate-x-full peer-checked:after:border-white after:content-[''] after:absolute after:top-[2px] after:left-[2px] after:bg-white after:border-pm after:border after:rounded-full after:h-5 after:w-5 after:transition-all peer-checked:bg-pm-accent"></div>
</label>
</div>
<!-- Progress bar (hidden by default) -->
<div id="thumbnailProgress" class="hidden">
<div class="flex items-center justify-between text-xs text-pm-secondary mb-1">
<span id="progressText">Processing...</span>
<div class="flex items-center gap-2">
<span id="progressPercent">0%</span>
<button id="cancelThumbnailsBtn" class="text-pm-error hover:text-pm-error text-xs px-2 py-1 bg-pm-error-tint rounded-pm-sm hidden">
Cancel
</button>
</div>
</div>
<div class="w-full bg-pm-input rounded-full h-2">
<div id="progressBar" class="bg-pm-accent h-2 rounded-full transition-all duration-300" style="width: 0%"></div>
</div>
<div class="flex items-center justify-between text-xs text-pm-muted mt-1">
<span id="progressDetails">Scanning images...</span>
<span id="progressETA" class="text-pm-muted"></span>
</div>
</div>
</div>
<!-- Image Loading -->
<div class="bg-pm-input rounded-pm-sm p-3">
<h5 class="text-pm text-xs font-medium mb-2">Image Loading</h5>
<div class="space-y-2">
<div class="flex items-center justify-between">
<label class="text-xs text-pm-secondary">Lazy Loading</label>
<input type="checkbox" id="lazyLoadingToggle" checked class="w-4 h-4 text-pm-accent bg-pm-input border-pm rounded focus:ring-pm-accent">
</div>
<div class="flex items-center justify-between">
<label class="text-xs text-pm-secondary">Image Quality</label>
<select id="imageQualitySelect" class="px-2 py-1 bg-pm-input border border-pm rounded-pm-sm text-pm text-xs">
<option value="high">High (Original)</option>
<option value="medium" selected>Medium (Optimized)</option>
<option value="low">Low (Fast)</option>
</select>
</div>
<div class="flex items-center justify-between">
<label class="text-xs text-pm-secondary">Lazy Loading</label>
<input type="checkbox" id="lazyLoadingToggle" checked class="w-4 h-4 text-pm-accent bg-pm-input border-pm rounded focus:ring-pm-accent">
</div>
</div>
</div>
@@ -264,13 +217,6 @@
<h4 class="text-sm font-semibold text-pm">Display Settings</h4>
<div class="bg-pm-input rounded-pm-sm p-3 space-y-2">
<div class="flex items-center justify-between">
<label class="text-xs text-pm-secondary">Default View Mode</label>
<select id="defaultViewModeSelect" class="px-2 py-1 bg-pm-input border border-pm rounded-pm-sm text-pm text-xs">
<option value="grid" selected>Grid View</option>
<option value="list">List View</option>
</select>
</div>
<div class="flex items-center justify-between">
<label class="text-xs text-pm-secondary">Items per Page</label>
<select id="defaultLimitSelect" class="px-2 py-1 bg-pm-input border border-pm rounded-pm-sm text-pm text-xs">
@@ -289,14 +235,6 @@
<option value="10">10 Columns</option>
</select>
</div>
<div class="flex items-center justify-between">
<label class="text-xs text-pm-secondary">Show Media Info</label>
<input type="checkbox" id="showImageInfoToggle" checked class="w-4 h-4 text-pm-accent bg-pm-input border-pm rounded focus:ring-pm-accent">
</div>
<div class="flex items-center justify-between">
<label class="text-xs text-pm-secondary">Collapse Sidebar by Default</label>
<input type="checkbox" id="sidebarCollapsedByDefaultToggle" class="w-4 h-4 text-pm-accent bg-pm-input border-pm rounded focus:ring-pm-accent">
</div>
<div class="flex items-center justify-between">
<div>
<label class="text-xs text-pm-secondary">Infinite Scrolling</label>
@@ -305,78 +243,6 @@
<input type="checkbox" id="infiniteScrollToggle" class="w-4 h-4 text-pm-accent bg-pm-input border-pm rounded focus:ring-pm-accent">
</div>
</div>
<!-- Video Settings -->
<div class="bg-pm-input rounded-pm-sm p-3 space-y-2">
<h5 class="text-pm text-xs font-medium mb-2">Video Settings</h5>
<div class="flex items-center justify-between">
<label class="text-xs text-pm-secondary">Autoplay Videos</label>
<input type="checkbox" id="videoAutoplayToggle" class="w-4 h-4 text-pm-accent bg-pm-input border-pm rounded focus:ring-pm-accent">
</div>
<div class="flex items-center justify-between">
<label class="text-xs text-pm-secondary">Mute Videos by Default</label>
<input type="checkbox" id="videoMuteToggle" checked class="w-4 h-4 text-pm-accent bg-pm-input border-pm rounded focus:ring-pm-accent">
</div>
<div class="flex items-center justify-between">
<label class="text-xs text-pm-secondary">Loop Videos</label>
<input type="checkbox" id="videoLoopToggle" checked class="w-4 h-4 text-pm-accent bg-pm-input border-pm rounded focus:ring-pm-accent">
</div>
</div>
</div>
<!-- Metadata Settings -->
<div class="space-y-3">
<h4 class="text-sm font-semibold text-pm">Metadata Settings</h4>
<div class="bg-pm-input rounded-pm-sm p-3 space-y-2">
<div class="flex items-center justify-between">
<label class="text-xs text-pm-secondary">Auto-load Metadata</label>
<input type="checkbox" id="autoLoadMetadataToggle" checked class="w-4 h-4 text-pm-accent bg-pm-input border-pm rounded focus:ring-pm-accent">
</div>
<div class="flex items-center justify-between">
<label class="text-xs text-pm-secondary">Cache Metadata</label>
<input type="checkbox" id="cacheMetadataToggle" checked class="w-4 h-4 text-pm-accent bg-pm-input border-pm rounded focus:ring-pm-accent">
</div>
<div class="flex items-center justify-between">
<label class="text-xs text-pm-secondary">Show File Paths</label>
<input type="checkbox" id="showFilePathsToggle" checked class="w-4 h-4 text-pm-accent bg-pm-input border-pm rounded focus:ring-pm-accent">
</div>
</div>
</div>
<!-- Maintenance -->
<div class="space-y-3">
<h4 class="text-sm font-semibold text-pm">Maintenance</h4>
<div class="bg-pm-input rounded-pm-sm p-3 space-y-2">
<div class="flex items-center justify-between">
<div>
<h5 class="text-pm text-xs font-medium">Clear Cache</h5>
<p class="text-xs text-pm-secondary">Clear cached metadata and thumbnails</p>
</div>
<button id="clearCacheBtn" class="px-3 py-1 bg-pm-error hover:bg-pm-error text-pm text-xs rounded-pm-sm transition-colors">
Clear
</button>
</div>
<div class="flex items-center justify-between">
<div>
<h5 class="text-pm text-xs font-medium">Rescan Output Folder</h5>
<p class="text-xs text-pm-secondary">Refresh the image list from disk</p>
</div>
<button id="rescanFolderBtn" class="px-3 py-1 bg-pm-accent hover:bg-pm-accent-hover text-pm-accent-fg text-xs rounded-pm-sm transition-colors">
Rescan
</button>
</div>
<div class="flex items-center justify-between">
<div>
<h5 class="text-pm text-xs font-medium">Find Duplicate Images</h5>
<p class="text-xs text-pm-secondary">Scan for duplicate images based on file content</p>
</div>
<button id="scanDuplicatesBtn" class="px-3 py-1 bg-pm-warning hover:bg-pm-warning text-pm text-xs rounded-pm-sm transition-colors">
Scan
</button>
</div>
</div>
</div>
<!-- Advanced Settings -->
@@ -395,6 +261,14 @@
</div>
</div>
<p class="text-xs text-pm-secondary">Thumbnails, cache, file listing and duplicates are in System.
<a id="pageOptionsMaintenanceLink" href="workspace#/system/maintenance" target="_blank" rel="noopener"
title="Opens System › Maintenance in a new tab" class="text-pm-accent hover:underline">Maintenance</a>
·
<a id="pageOptionsSourcesLink" href="workspace#/system/sources" target="_blank" rel="noopener"
title="Opens System › Sources &amp; scans in a new tab" class="text-pm-accent hover:underline">Sources &amp; scans</a>
</p>
</div>
<div class="p-3 border-t border-pm flex justify-end gap-2">
<button id="resetSettingsBtn" class="px-3 py-1 bg-pm-input hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
@@ -407,71 +281,6 @@
</div>
</div>
<!-- Minimized duplicate scan: shown while it runs with its modal hidden, or when results are waiting -->
<button id="dupScanPill" type="button" title="Show duplicate scan"
class="hidden fixed bottom-4 right-4 z-40 px-3 py-1.5 bg-pm-surface border border-pm shadow-pm rounded-pm-md text-xs text-pm hover:bg-pm-hover transition-colors">
<span class="inline-block w-1.5 h-1.5 rounded-full bg-pm-accent mr-1.5 align-middle animate-pulse"></span><span id="dupScanPillText" class="align-middle">Scanning...</span>
</button>
<!-- Duplicates Modal -->
<div id="duplicatesModal" class="fixed inset-0 bg-black/50 hidden items-center justify-center z-50">
<div class="bg-pm-surface rounded-pm-md max-w-4xl w-full max-h-[90vh] overflow-hidden mx-4 border border-pm">
<div class="p-3 border-b border-pm">
<div class="flex items-center justify-between">
<h3 class="text-sm font-semibold text-pm">Duplicate Images Found</h3>
<button id="closeDuplicatesBtn" class="p-2 rounded-pm-sm hover:bg-pm-hover transition-colors">
<svg class="w-4 h-4 text-pm-secondary" fill="none" stroke="currentColor" viewBox="0 0 24 24">
<path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M6 18L18 6M6 6l12 12"></path>
</svg>
</button>
</div>
</div>
<div class="p-3 overflow-y-auto max-h-[calc(90vh-200px)]">
<div id="duplicatesScanStatus" class="py-6">
<div class="max-w-md mx-auto bg-pm-input rounded-pm-md p-3">
<div class="flex items-center justify-between mb-1.5">
<span class="text-pm text-xs font-medium">Scanning for duplicate images</span>
<span id="dupProgressText" class="text-pm-secondary text-xs">Starting...</span>
</div>
<div class="w-full bg-pm-primary rounded-full h-1.5">
<div id="dupProgressBar" class="bg-pm-accent h-1.5 rounded-full transition-all duration-300" style="width: 0%"></div>
</div>
<p class="text-[11px] text-pm-muted mt-2">Every file is hashed, so this takes a while on a large library. You can close this window; the scan keeps running and the results will be waiting.</p>
</div>
</div>
<div id="duplicatesContent" class="hidden">
<div class="mb-3 p-3 bg-pm-error-tint border border-pm-accent rounded-pm-sm">
<div class="flex items-start">
<svg class="w-4 h-4 text-pm-error mt-0.5 mr-2 flex-shrink-0" fill="currentColor" viewBox="0 0 20 20">
<path fill-rule="evenodd" d="M8.257 3.099c.765-1.36 2.722-1.36 3.486 0l5.58 9.92c.75 1.334-.213 2.98-1.742 2.98H4.42c-1.53 0-2.493-1.646-1.743-2.98l5.58-9.92zM11 13a1 1 0 11-2 0 1 1 0 012 0zm-1-8a1 1 0 00-1 1v3a1 1 0 002 0V6a1 1 0 00-1-1z" clip-rule="evenodd"></path>
</svg>
<div>
<h4 class="text-pm-error text-xs font-medium">Warning: Removing duplicate images cannot be undone!</h4>
<p class="text-pm-secondary text-xs mt-1">When you remove a duplicate image, the selected file will be permanently deleted from disk. Choose carefully which images to keep.</p>
</div>
</div>
</div>
<div id="duplicatesList" class="space-y-3">
<!-- Duplicates will be populated here -->
</div>
</div>
</div>
<div id="duplicatesFooter" class="hidden p-3 border-t border-pm flex justify-between items-center">
<div class="text-xs text-pm-secondary">
<span id="duplicatesCount">0</span> duplicate image groups found
</div>
<div class="flex gap-2">
<button id="selectAllBtn" class="px-3 py-1 bg-pm-input hover:bg-pm-hover text-pm text-xs font-medium rounded-pm-sm transition-colors">
Select All to Delete
</button>
<button id="removeDuplicatesBtn" class="px-3 py-1 bg-pm-error hover:bg-pm-error text-pm text-xs font-medium rounded-pm-sm transition-colors">
Delete Selected
</button>
</div>
</div>
</div>
</div>
<!-- Auto Tag Modal -->
<div id="autoTagModal" class="fixed inset-0 bg-black/50 hidden items-center justify-center z-50">
<div class="bg-pm-surface rounded-pm-md p-3 max-w-lg w-full mx-4 border border-pm">
@@ -549,9 +358,11 @@
<span class="text-pm-secondary text-xs ml-2">(Faster processing, uses ~4-16GB VRAM)</span>
</div>
</div>
<button id="unloadModelBtn" class="px-3 py-1 bg-pm-error hover:bg-pm-error text-pm text-xs rounded-pm-sm transition-colors hidden">
Unload Model
</button>
<a id="unloadModelBtn" href="workspace#/system/tagging" target="_blank" rel="noopener"
title="Unload the model in System › Auto-tagging (opens a new tab)"
class="px-3 py-1 bg-pm-error hover:bg-pm-error text-pm text-xs rounded-pm-sm transition-colors hidden">
Manage in System
</a>
</label>
<p id="modelLoadedStatus" class="text-pm-muted text-xs mt-1 hidden">
<span class="text-pm-success">●</span> Model currently loaded: <span id="loadedModelType">none</span>
@@ -711,39 +522,6 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr
</div>
</div>
<!-- Auto Tag Download Modal -->
<div id="autoTagDownloadModal" class="fixed inset-0 bg-black/50 hidden items-center justify-center z-50">
<div class="bg-pm-surface rounded-pm-md p-3 max-w-md w-full mx-4 border border-pm">
<h3 class="text-sm font-semibold text-pm mb-3">Downloading Model</h3>
<!-- Model Name -->
<div class="mb-3">
<p class="text-pm-secondary text-xs">Downloading:</p>
<p id="downloadModelName" class="text-pm text-xs font-medium">GGUF Model</p>
</div>
<!-- Progress Bar -->
<div class="mb-3">
<div class="flex justify-between text-xs text-pm-secondary mb-1">
<span id="downloadStatus">Preparing...</span>
<span id="downloadProgressPercent">0%</span>
</div>
<div class="w-full bg-pm-input rounded-full h-2">
<div id="downloadProgressBar" class="bg-pm-accent h-2 rounded-full transition-all duration-300" style="width: 0%"></div>
</div>
<p class="text-pm-secondary text-xs mt-2">The download continues in Activity.</p>
</div>
<!-- Cancel Button -->
<div class="flex justify-center">
<button id="cancelDownloadBtn"
class="px-3 py-1 bg-pm-error hover:bg-pm-error text-pm text-xs font-medium rounded-pm-sm transition-colors">
Close
</button>
</div>
</div>
</div>
<!-- Add Prompt Modal -->
<div id="addPromptModal" class="fixed inset-0 bg-black/50 hidden items-center justify-center z-50">
<div class="bg-pm-surface rounded-pm-md p-3 max-w-2xl w-full mx-4 border border-pm max-h-[90vh] overflow-y-auto">
+16 -671
View File
@@ -3,11 +3,8 @@
this.prompts = [];
this.listMode = "recent"; // "recent" or "search"; a sort change reloads it
this.selectedPrompts = new Set();
// The page's own options (Page options); every other setting lives in System.
this.settings = {
resultTimeout: 5,
webuiDisplayMode: 'popup',
workerThreads: 1,
cpuCount: 1,
infiniteScroll: false,
};
this.categories = [];
@@ -125,11 +122,8 @@
document.getElementById("bulkTagBtn").addEventListener("click", () => this.showBulkTagModal());
document.getElementById("bulkCategoryBtn").addEventListener("click", () => this.showBulkCategoryModal());
document.getElementById("exportBtn").addEventListener("click", () => this.exportPrompts());
document.getElementById("settingsBtn").addEventListener("click", () => this.showSettingsModal());
document.getElementById("pageOptionsBtn").addEventListener("click", () => this.showSettingsModal());
document.getElementById("statsBtn").addEventListener("click", () => this.showStats());
document.getElementById("maintenanceBtn").addEventListener("click", () => this.showMaintenanceModal());
document.getElementById("backupBtn").addEventListener("click", () => this.backupDatabase());
document.getElementById("restoreBtn").addEventListener("click", () => this.showRestoreModal());
document.getElementById("scanBtn").addEventListener("click", () => this.showScanModal());
document.getElementById("autoTagBtn").addEventListener("click", () => this.showAutoTagModal());
document.getElementById("addPromptBtn").addEventListener("click", () => this.showAddPromptModal());
@@ -193,8 +187,6 @@
// Settings modal
document.getElementById("saveSettings").addEventListener("click", () => this.saveSettings());
document.getElementById("cancelSettings").addEventListener("click", () => this.hideModal("settingsModal"));
document.getElementById("refreshMonitoringStatus").addEventListener("click", () => this.updateMonitoringStatus());
document.getElementById("addScanPathBtn").addEventListener("click", () => this.addScanPath());
// Bulk tag modal
document.getElementById("confirmBulkTag").addEventListener("click", () => this.confirmBulkTag());
@@ -208,14 +200,6 @@
document.getElementById("confirmBulkCategory").addEventListener("click", () => this.confirmBulkCategory());
document.getElementById("cancelBulkCategory").addEventListener("click", () => this.hideModal("bulkCategoryModal"));
// Restore modal
document.getElementById("confirmRestore").addEventListener("click", () => this.confirmRestore());
document.getElementById("cancelRestore").addEventListener("click", () => this.hideModal("restoreModal"));
// Maintenance modal
document.getElementById("runMaintenanceBtn").addEventListener("click", () => this.runMaintenance());
document.getElementById("selectAllMaintenanceBtn").addEventListener("click", () => this.selectAllMaintenance());
document.getElementById("clearAllMaintenanceBtn").addEventListener("click", () => this.clearAllMaintenance());
// Scan modal
document.getElementById("startScan").addEventListener("click", () => this.startScan());
@@ -229,7 +213,6 @@
document.getElementById("skipReviewBtn").addEventListener("click", () => this.skipReviewImage());
document.getElementById("applyReviewBtn").addEventListener("click", () => this.applyReviewTags());
document.getElementById("cancelAutoTagBtn").addEventListener("click", () => this.cancelAutoTag());
document.getElementById("cancelDownloadBtn").addEventListener("click", () => this.cancelDownload());
// Re-tag confirmation modal buttons
document.getElementById("retagSkipBtn").addEventListener("click", () => this.handleRetagChoice('skip'));
@@ -249,7 +232,6 @@
document.getElementById('wd14GeneralThreshold').addEventListener('input', (e) => {
document.getElementById('wd14GeneralThresholdValue').textContent = parseFloat(e.target.value).toFixed(2);
});
document.getElementById("workerThreads").addEventListener("input", () => this.renderWorkerThreads());
document.getElementById('wd14CharacterThreshold').addEventListener('input', (e) => {
document.getElementById('wd14CharacterThresholdValue').textContent = parseFloat(e.target.value).toFixed(2);
});
@@ -298,21 +280,9 @@
if (response.ok) {
const data = await response.json();
if (data.success && data.settings) {
this.settings.resultTimeout = data.settings.result_timeout ?? 5;
this.settings.webuiDisplayMode = data.settings.webui_display_mode || 'popup';
this.settings.infiniteScroll = data.settings.infinite_scroll === true;
// The saved default seeds the session until the header checkbox is used
if (!this.infiniteScrollChosen) this.setInfiniteScroll(this.settings.infiniteScroll, { session: false });
this.settings.cpuCount = Math.max(1, parseInt(data.settings.cpu_count, 10) || 1);
this.settings.workerThreads = Math.min(
this.settings.cpuCount, Math.max(1, parseInt(data.settings.worker_threads, 10) || 1)
);
this.settings.monitoredDirectories = data.settings.monitored_directories || [];
this.settings.galleryRootPaths = data.settings.gallery_root_paths || [];
// Backward compat: if server only returned old field
if (!this.settings.galleryRootPaths.length && data.settings.gallery_root_path) {
this.settings.galleryRootPaths = [data.settings.gallery_root_path];
}
}
}
} catch (error) {
@@ -996,255 +966,27 @@
});
}
renderWorkerThreads() {
const threads = document.getElementById("workerThreads");
document.getElementById("workerThreadsValue").textContent =
`${threads.value} / ${this.settings.cpuCount} cores`;
}
/** Page options: the prompt list's saved Infinite Scrolling default (scan folders and the rest are in System). */
showSettingsModal() {
document.getElementById("resultTimeout").value = this.settings.resultTimeout;
document.getElementById("webuiDisplayMode").value = this.settings.webuiDisplayMode;
document.getElementById("infiniteScrollDefault").checked = this.settings.infiniteScroll;
const threads = document.getElementById("workerThreads");
threads.max = this.settings.cpuCount;
threads.value = this.settings.workerThreads;
this.renderWorkerThreads();
this.renderScanPaths();
this.updateMonitoringStatus();
this.detectLoraManager();
this.showModal("settingsModal");
}
async updateMonitoringStatus() {
const statusEl = document.getElementById("monitoringStatus");
try {
const response = await ApiClient.fetchUrl("/prompt_manager/settings");
if (response.ok) {
const data = await response.json();
const dirs = data.settings?.monitored_directories || [];
if (dirs.length > 0) {
statusEl.innerHTML = dirs.map(d => `<div class="truncate" title="${escapeHtml(d)}">✓ ${escapeHtml(d)}</div>`).join('');
} else {
statusEl.textContent = 'No directories being monitored (auto-detect on restart)';
}
}
} catch (error) {
statusEl.textContent = 'Unable to fetch status';
}
}
// ── LoraManager integration ──────────────────────────────
async detectLoraManager() {
const badge = document.getElementById("loraDetectionBadge");
const toggle = document.getElementById("loraEnabled");
const settings = document.getElementById("loraSettings");
try {
// First check current status
const statusRes = await ApiClient.fetchUrl("/prompt_manager/lora/status");
if (statusRes.ok) {
const status = await statusRes.json();
if (status.success && status.enabled) {
badge.textContent = "enabled";
badge.className = "text-[10px] px-1.5 py-0.5 rounded-full bg-green-500/20 text-green-400";
toggle.checked = true;
toggle.disabled = false;
document.getElementById("loraManagerPath").value = status.path;
document.getElementById("loraTriggerWords").checked = status.trigger_words_enabled;
const keyInput = document.getElementById("civitaiApiKey");
keyInput.value = "";
keyInput.placeholder = status.has_civitai_api_key
? "A key is stored — leave empty to keep it"
: "Optional — required for NSFW example images";
document.getElementById("clearCivitaiKey").checked = false;
settings.classList.remove("hidden");
this._loraPath = status.path;
this._bindLoraEvents();
return;
}
}
// Try auto-detection
const detectRes = await ApiClient.fetchUrl("/prompt_manager/lora/detect");
if (!detectRes.ok) return;
const data = await detectRes.json();
if (data.detected) {
badge.textContent = "detected";
badge.className = "text-[10px] px-1.5 py-0.5 rounded-full bg-blue-500/20 text-blue-400";
toggle.disabled = false;
document.getElementById("loraManagerPath").value = data.path;
this._loraPath = data.path;
} else {
badge.textContent = "not installed";
badge.className = "text-[10px] px-1.5 py-0.5 rounded-full bg-pm-input text-pm-muted";
toggle.disabled = true;
}
this._bindLoraEvents();
} catch (e) {
badge.textContent = "error";
badge.className = "text-[10px] px-1.5 py-0.5 rounded-full bg-red-500/20 text-red-400";
}
}
_bindLoraEvents() {
if (this._loraBound) return;
this._loraBound = true;
const toggle = document.getElementById("loraEnabled");
const settings = document.getElementById("loraSettings");
toggle.addEventListener("change", () => {
if (toggle.checked) {
settings.classList.remove("hidden");
} else {
settings.classList.add("hidden");
}
});
document.getElementById("loraImportBtn").addEventListener("click", () => {
this.runLoraImport();
});
}
async saveLoraSettings() {
const enabled = document.getElementById("loraEnabled").checked;
const triggerWords = document.getElementById("loraTriggerWords").checked;
const civitaiKey = document.getElementById("civitaiApiKey").value.trim();
const clearKey = document.getElementById("clearCivitaiKey").checked;
const path = this._loraPath || "";
// The server never returns the stored key: an empty field keeps it, the
// checkbox removes it, and a non-empty field replaces it
const payload = { enabled, path, trigger_words_enabled: triggerWords };
if (clearKey) payload.clear_civitai_api_key = true;
else if (civitaiKey) payload.civitai_api_key = civitaiKey;
try {
const res = await ApiClient.fetchUrl("/prompt_manager/lora/enable", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(payload),
});
const data = await res.json();
if (!data.success) {
this.showNotification(data.error || "Failed to save LoRA settings", "error");
return false;
}
return true;
} catch (e) {
this.showNotification("Failed to save LoRA settings", "error");
return false;
}
}
async runLoraImport() {
const saved = await this.saveLoraSettings();
if (!saved) return;
// Close settings and show progress modal
this.hideModal("settingsModal");
document.getElementById("loraImportStatus").textContent = "Initializing...";
document.getElementById("loraImportBar").style.width = "0%";
document.getElementById("loraImportPercent").textContent = "0%";
document.getElementById("loraImportProcessed").textContent = "0 processed";
document.getElementById("loraImportImported").textContent = "0 imported";
this.showModal("loraImportModal");
try {
const res = await ApiClient.fetchUrl("/prompt_manager/lora/scan", { method: "POST" });
const reader = res.body.getReader();
const decoder = new TextDecoder();
let buffer = "";
while (true) {
const { done, value } = await reader.read();
if (done) break;
buffer += decoder.decode(value, { stream: true });
const lines = buffer.split("\n\n");
buffer = lines.pop();
for (const line of lines) {
if (!line.startsWith("data: ")) continue;
try {
const data = JSON.parse(line.slice(6));
const pct = data.progress || 0;
document.getElementById("loraImportBar").style.width = `${pct}%`;
document.getElementById("loraImportPercent").textContent = `${pct}%`;
if (data.status) {
document.getElementById("loraImportStatus").textContent = data.status;
}
if (data.processed !== undefined) {
document.getElementById("loraImportProcessed").textContent = `${data.processed} processed`;
}
if (data.imported !== undefined) {
document.getElementById("loraImportImported").textContent = `${data.imported} imported`;
}
if (data.type === "complete") {
document.getElementById("loraImportStatus").textContent =
`Done — ${data.imported} imported, ${data.skipped} skipped`;
this.loadStatistics();
setTimeout(() => this.hideModal("loraImportModal"), 2000);
} else if (data.type === "error") {
// A refusal (another job runs) or a failed import: name it and close the modal.
document.getElementById("loraImportStatus").textContent = `Error: ${data.message}`;
this.showNotification(`LoRA import failed: ${data.message}`, "error");
setTimeout(() => this.hideModal("loraImportModal"), 3000);
}
} catch (e) { /* skip malformed SSE */ }
}
}
} catch (e) {
document.getElementById("loraImportStatus").textContent = `Error: ${e.message}`;
setTimeout(() => this.hideModal("loraImportModal"), 3000);
}
}
/** Saves Page options: only `infinite_scroll`, so a stale page never overwrites what System saved (#330 S13-9). */
async saveSettings() {
const timeout = parseInt(document.getElementById("resultTimeout").value);
const displayMode = document.getElementById("webuiDisplayMode").value;
const workerThreads = parseInt(document.getElementById("workerThreads").value, 10) || 1;
const infiniteScroll = document.getElementById("infiniteScrollDefault").checked;
const galleryPaths = this._collectScanPaths().filter(p => p !== '');
this.settings.resultTimeout = timeout;
this.settings.webuiDisplayMode = displayMode;
this.settings.workerThreads = workerThreads;
this.settings.infiniteScroll = infiniteScroll;
this.settings.galleryRootPaths = galleryPaths;
try {
const response = await ApiClient.fetchUrl("/prompt_manager/settings", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
result_timeout: timeout,
webui_display_mode: displayMode,
worker_threads: workerThreads,
infinite_scroll: infiniteScroll,
gallery_root_paths: galleryPaths
}),
body: JSON.stringify({ infinite_scroll: infiniteScroll }),
});
if (response.ok) {
const data = await response.json();
// Save LoRA integration settings (fire-and-forget)
await this.saveLoraSettings();
this.settings.infiniteScroll = infiniteScroll;
this.setInfiniteScroll(infiniteScroll, { session: false });
if (data.restart_required) {
this.showNotification("Settings saved. Restart ComfyUI for gallery path changes to take effect.", "warning");
} else {
this.showNotification("Settings saved successfully", "success");
}
this.showNotification("Page options saved", "success");
this.hideModal("settingsModal");
} else {
// A refusal names its reason (a corrupt config.json, a changed file): show it.
@@ -1257,57 +999,6 @@
}
}
renderScanPaths() {
const container = document.getElementById("scanPathsList");
const paths = this.settings.galleryRootPaths || [];
container.textContent = ''; // Clear safely
const toRender = paths.length > 0 ? paths : [''];
toRender.forEach((p, i) => {
const row = document.createElement('div');
row.className = 'flex gap-2 items-center';
row.dataset.pathIndex = i;
const input = document.createElement('input');
input.type = 'text';
input.value = p;
input.placeholder = 'Leave empty for auto-detect';
input.className = 'flex-1 px-2.5 py-1.5 bg-pm-input border border-pm rounded-pm-sm text-pm text-[13px] focus:outline-none scan-path-input';
const removeBtn = document.createElement('button');
removeBtn.type = 'button';
removeBtn.className = 'text-pm-muted hover:text-red-400 p-1 text-lg leading-none';
removeBtn.title = 'Remove';
removeBtn.textContent = '\u00d7'; // multiplication sign as X
removeBtn.addEventListener('click', () => this.removeScanPath(i));
row.appendChild(input);
row.appendChild(removeBtn);
container.appendChild(row);
});
}
addScanPath() {
const paths = this._collectScanPaths();
paths.push('');
this.settings.galleryRootPaths = paths;
this.renderScanPaths();
const inputs = document.querySelectorAll('.scan-path-input');
if (inputs.length) inputs[inputs.length - 1].focus();
}
removeScanPath(index) {
const paths = this._collectScanPaths();
paths.splice(index, 1);
this.settings.galleryRootPaths = paths;
this.renderScanPaths();
}
_collectScanPaths() {
return Array.from(document.querySelectorAll('.scan-path-input'))
.map(input => input.value.trim());
}
toggleSelectAll(checked) {
document.querySelectorAll(".prompt-checkbox").forEach((checkbox) => {
checkbox.checked = checked;
@@ -1788,287 +1479,6 @@
return `${size.toFixed(1)} ${units[unitIndex]}`;
}
// Maintenance functionality
showMaintenanceModal() {
this.showModal("maintenanceModal");
}
closeMaintenance() {
this.hideModal("maintenanceModal");
}
selectAllMaintenance() {
document.querySelectorAll('.maintenance-option').forEach(checkbox => {
checkbox.checked = true;
});
}
clearAllMaintenance() {
document.querySelectorAll('.maintenance-option').forEach(checkbox => {
checkbox.checked = false;
});
}
async runMaintenance() {
const resultsContainer = document.getElementById("maintenanceResults");
const runButton = document.getElementById("runMaintenanceBtn");
// Get selected operations
const selectedOps = Array.from(document.querySelectorAll('.maintenance-option:checked'))
.map(checkbox => checkbox.id);
if (selectedOps.length === 0) {
this.showNotification('Please select at least one maintenance operation', 'warning');
return;
}
// Show loading state
runButton.disabled = true;
runButton.innerHTML = '🔄 Running...';
resultsContainer.innerHTML = `
<div class="text-center py-4">
<div class="w-12 h-12 border-4 border-teal-500/30 border-t-teal-500 rounded-full animate-spin mx-auto mb-4"></div>
<p class="text-pm-secondary">Running maintenance operations...</p>
<p class="text-sm text-pm-muted mt-2">Operations: ${selectedOps.map(op => op.replace(/_/g, ' ')).join(', ')}</p>
</div>
`;
try {
const response = await ApiClient.fetchUrl('/prompt_manager/maintenance', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ operations: selectedOps })
});
if (response.ok) {
const data = await response.json();
if (data.success) {
this.renderMaintenanceResults(data);
// Show overall success notification
if (data.all_successful) {
this.showNotification('✅ All maintenance operations completed successfully!', 'success');
} else {
this.showNotification('⚠️ Some maintenance operations had issues', 'warning');
}
// Refresh stats after maintenance
setTimeout(() => this.loadStatistics(), 1000);
} else {
throw new Error(data.error || 'Maintenance failed');
}
} else {
// A refusal names the job that runs ("Scan for metadata already running").
const data = await response.json().catch(() => ({}));
throw new Error(data.error || `HTTP ${response.status}: ${response.statusText}`);
}
} catch (error) {
console.error('Maintenance error:', error);
resultsContainer.innerHTML = `
<div class="bg-pm-error-tint border border-pm-error rounded-pm-sm p-4">
<h4 class="text-pm-error font-medium mb-2">❌ Maintenance Failed</h4>
<p class="text-pm-secondary">${escapeHtml(error.message)}</p>
</div>
`;
this.showNotification('❌ Maintenance failed', 'error');
} finally {
runButton.disabled = false;
runButton.innerHTML = '🔧 Run Maintenance';
}
}
renderMaintenanceResults(data) {
const resultsContainer = document.getElementById("maintenanceResults");
let html = '<div class="space-y-4">';
// Summary
html += '<div class="bg-pm-surface rounded-pm-sm p-4">';
html += '<h4 class="text-pm font-medium mb-3">📋 Maintenance Summary</h4>';
html += '<div class="grid grid-cols-1 md:grid-cols-2 gap-3">';
html += `
<div class="flex justify-between items-center p-2 bg-pm-input rounded">
<span class="text-pm-secondary">Operations Completed</span>
<span class="text-pm-success font-mono text-sm">${escapeHtml(data.operations_completed)}</span>
</div>
`;
html += `
<div class="flex justify-between items-center p-2 bg-pm-input rounded">
<span class="text-pm-secondary">Overall Success</span>
<span class="${data.all_successful ? 'text-pm-success' : 'text-pm-warning'} font-mono text-sm">
${data.all_successful ? '✅ ALL PASSED' : '⚠️ SOME ISSUES'}
</span>
</div>
`;
html += '</div></div>';
// Detailed results
for (const [operation, result] of Object.entries(data.results)) {
const bgColor = result.success ? 'bg-pm-success-tint border-pm-success' : 'bg-pm-error-tint border-pm-error';
const statusIcon = result.success ? '✅' : '❌';
const statusText = result.success ? 'SUCCESS' : 'FAILED';
html += `<div class="${bgColor} border rounded-pm-sm p-4">`;
html += `<h4 class="text-pm font-medium mb-2 flex items-center space-x-2">`;
html += `<span>${statusIcon}</span>`;
html += `<span class="capitalize">${escapeHtml(operation.replace(/_/g, ' '))}</span>`;
html += `<span class="text-sm font-mono ${result.success ? 'text-pm-success' : 'text-pm-error'}">${statusText}</span>`;
html += `</h4>`;
if (result.message) {
html += `<p class="text-pm-secondary mb-2">${escapeHtml(result.message)}</p>`;
}
// Show specific details
if (result.removed_count !== undefined) {
html += `<div class="text-sm text-pm-secondary">Items removed: ${escapeHtml(result.removed_count)}</div>`;
}
if (result.duplicate_hashes !== undefined) {
html += `<div class="text-sm text-pm-secondary">Duplicate hash groups found: ${escapeHtml(result.duplicate_hashes)}</div>`;
}
if (result.issues_found !== undefined) {
html += `<div class="text-sm text-pm-secondary">Issues found: ${escapeHtml(result.issues_found)}</div>`;
if (result.issues && result.issues.length > 0) {
html += `<ul class="ml-4 list-disc text-xs text-pm-muted mt-1">`;
result.issues.forEach(issue => {
html += `<li>${escapeHtml(issue)}</li>`;
});
html += `</ul>`;
}
}
if (result.info) {
html += `<div class="text-sm text-pm-secondary mt-2">`;
html += `<p>Total prompts: ${escapeHtml(result.info.total_prompts || 'N/A')}</p>`;
html += `<p>Database size: ${escapeHtml(result.info.database_size_bytes ? this.formatFileSize(result.info.database_size_bytes) : 'N/A')}</p>`;
html += `</div>`;
}
if (result.error) {
html += `<div class="text-sm text-pm-error mt-2 font-mono bg-pm-error-tint p-2 rounded">`;
html += `Error: ${escapeHtml(result.error)}`;
html += `</div>`;
}
html += '</div>';
}
html += '</div>';
resultsContainer.innerHTML = html;
}
// Database backup and restore functionality
async backupDatabase() {
try {
const response = await ApiClient.fetchUrl('/prompt_manager/backup');
if (response.ok) {
// Get the filename from the Content-Disposition header
const contentDisposition = response.headers.get('Content-Disposition');
let filename = 'prompts_backup.db';
if (contentDisposition) {
const filenameMatch = contentDisposition.match(/filename="(.+)"/);
if (filenameMatch) {
filename = filenameMatch[1];
}
}
// Create blob and download
const blob = await response.blob();
const url = window.URL.createObjectURL(blob);
const a = document.createElement('a');
a.href = url;
a.download = filename;
document.body.appendChild(a);
a.click();
window.URL.revokeObjectURL(url);
document.body.removeChild(a);
this.showNotification('💾 Database backup downloaded successfully!', 'success');
} else {
const errorData = await response.json();
throw new Error(errorData.error || 'Backup failed');
}
} catch (error) {
console.error('Backup error:', error);
this.showNotification(`❌ Backup failed: ${error.message}`, 'error');
}
}
showRestoreModal() {
// Clear previous file selection
document.getElementById('restoreFileInput').value = '';
this.showModal('restoreModal');
}
async confirmRestore() {
const fileInput = document.getElementById('restoreFileInput');
const file = fileInput.files[0];
if (!file) {
this.showNotification('Please select a database file to restore', 'warning');
return;
}
if (!file.name.endsWith('.db')) {
this.showNotification('Please select a valid .db file', 'warning');
return;
}
// Show confirmation
const confirmed = confirm(
`Are you sure you want to restore from "${file.name}"?\n\n` +
'This will replace your current database. A backup will be created automatically.\n\n' +
'This action cannot be undone.'
);
if (!confirmed) {
return;
}
try {
// Disable the restore button and show loading
const confirmBtn = document.getElementById('confirmRestore');
confirmBtn.disabled = true;
confirmBtn.innerHTML = '🔄 Restoring...';
// Create FormData for file upload
const formData = new FormData();
formData.append('database_file', file);
const response = await ApiClient.fetchUrl('/prompt_manager/restore', {
method: 'POST',
body: formData
});
const data = await response.json();
if (response.ok && data.success) {
this.showNotification(
`✅ Database restored successfully! Found ${data.prompt_count} prompts.`,
'success'
);
this.hideModal('restoreModal');
// Refresh the interface to show the restored data
await this.loadInitialData();
} else {
throw new Error(data.error || 'Restore failed');
}
} catch (error) {
console.error('Restore error:', error);
this.showNotification(`❌ Restore failed: ${error.message}`, 'error');
} finally {
// Re-enable the restore button
const confirmBtn = document.getElementById('confirmRestore');
confirmBtn.disabled = false;
confirmBtn.innerHTML = 'Restore Database';
}
}
// Copy prompt to clipboard functionality
async copyPromptToClipboard(promptId) {
// Find the prompt text from the prompts array
@@ -2773,7 +2183,6 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
// State for AutoTag feature
autoTagState = {
eventSource: null,
downloadEventSource: null,
reviewImages: [],
reviewIndex: 0,
currentTags: [],
@@ -3187,11 +2596,15 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
span.textContent = '\u2713 Downloaded';
el.appendChild(span);
} else {
const btn = document.createElement('button');
btn.className = 'px-3 py-1 bg-pm-accent hover:bg-pm-accent-hover text-pm text-xs rounded transition-colors';
btn.textContent = 'Download';
btn.addEventListener('click', () => window.admin.downloadModel(modelKey));
el.appendChild(btn);
// Download lives in System › Auto-tagging (#330); the batch refuses a model that is not downloaded.
const link = document.createElement('a');
link.href = 'workspace#/system/tagging';
link.target = '_blank';
link.rel = 'noopener';
link.title = 'Opens System › Auto-tagging in a new tab';
link.className = 'px-3 py-1 bg-pm-accent hover:bg-pm-accent-hover text-pm text-xs rounded transition-colors';
link.textContent = 'Manage in System';
el.appendChild(link);
}
}
@@ -3234,74 +2647,6 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
}
}
async downloadModel(modelType) {
const modelNames = {
'gguf': 'GGUF Model',
'hf': 'HuggingFace Model',
'wd14-swinv2': 'WD14 SwinV2',
'wd14-vit': 'WD14 ViT'
};
const modelName = modelNames[modelType] || modelType;
document.getElementById('downloadModelName').textContent = modelName;
document.getElementById('downloadStatus').textContent = 'Preparing...';
document.getElementById('downloadProgressPercent').textContent = '0%';
document.getElementById('downloadProgressBar').style.width = '0%';
document.getElementById('downloadProgressBar').classList.remove('animate-pulse');
this.hideModal("autoTagModal");
this.showModal("autoTagDownloadModal");
try {
this.autoTagState.downloadEventSource = SseStream.connect(`/prompt_manager/autotag/download/${modelType}`, { method: 'POST' });
this.autoTagState.downloadEventSource.onmessage = (event) => {
const data = JSON.parse(event.data);
// The server sends `type` (progress / complete / error); the download is an
// Activity job with no byte progress, so the bar is indeterminate (#328).
if (data.type === 'progress') {
document.getElementById('downloadStatus').textContent = 'Downloading\u2026';
document.getElementById('downloadProgressPercent').textContent = '';
const bar = document.getElementById('downloadProgressBar');
bar.style.width = '100%';
bar.classList.add('animate-pulse');
} else if (data.type === 'complete') {
this.autoTagState.downloadEventSource.close();
this.hideModal("autoTagDownloadModal");
this.showModal("autoTagModal");
this.checkAutoTagModels();
this.showNotification('Model downloaded successfully!', 'success');
} else if (data.type === 'error') {
this.autoTagState.downloadEventSource.close();
this.hideModal("autoTagDownloadModal");
this.showModal("autoTagModal");
this.showNotification(`Download failed: ${data.message}`, 'error');
}
};
this.autoTagState.downloadEventSource.onerror = () => {
this.autoTagState.downloadEventSource.close();
this.hideModal("autoTagDownloadModal");
this.showModal("autoTagModal");
this.showNotification('Download connection lost', 'error');
};
} catch (error) {
console.error('Download error:', error);
this.hideModal("autoTagDownloadModal");
this.showModal("autoTagModal");
this.showNotification('Failed to start download', 'error');
}
}
cancelDownload() {
if (this.autoTagState.downloadEventSource) {
this.autoTagState.downloadEventSource.close();
this.autoTagState.downloadEventSource = null;
}
this.hideModal("autoTagDownloadModal");
this.showModal("autoTagModal");
}
async startAutoTag() {
const modelType = document.querySelector('input[name="autoTagModel"]:checked').value;
const prompt = document.getElementById('autoTagPrompt').value;
+54 -803
View File
File diff suppressed because it is too large Load Diff
+10 -3
View File
@@ -20,6 +20,8 @@ const REPRESENTATIONS = {
limited: 'Record exceeds inspection limits; text is unavailable',
};
const TEXTLESS = new Set(['invalid', 'limited']);
// #332: the hidden prefix that names where a category chip's value came from.
const CATEGORY_PREFIX = {image: 'Image category: ', positive: 'Positive prompt category: ', pair: 'Recipe category: '};
const REVEAL_URL = '/prompt_manager/images/reveal';
// D38 Show in folder: the route's answers, in the maintainer's words.
const REVEAL_OPENED = 'Opened the folder on this computer.';
@@ -158,9 +160,11 @@ function generationFacts(record, listClass, noteClass) {
}
// D21: the panel badge names the strongest resolved prompt source; any other source shows no badge.
// A prompt that left out a run-time PromptManager part (D125) or a later stage's added conditioning (D126) is still
// inferred from the workflow.
// inferred from the workflow. "Recorded run" (#386): recorded run history named the save node that wrote the file; it
// ranks below any workflow-inferred side, so a branch with an omitted or added part still reads "Workflow inferred".
const EVIDENCE_BADGES = [['execution', 'Execution recorded'], ['graph-inferred', 'Workflow inferred'], ['runtime-parts-omitted', 'Workflow inferred'],
['added-conditioning', 'Workflow inferred'], ['runtime-parts-omitted-and-added-conditioning', 'Workflow inferred'], ['embedded-parameters', 'Embedded parameters']];
['added-conditioning', 'Workflow inferred'], ['runtime-parts-omitted-and-added-conditioning', 'Workflow inferred'], ['recorded-run', 'Recorded run'],
['embedded-parameters', 'Embedded parameters']];
function evidenceBadge(prompts) {
const sources = new Set(Object.values(prompts).filter(role => role?.status === 'resolved').map(role => role.source));
return EVIDENCE_BADGES.find(([source]) => sources.has(source))?.[1] ?? null;
@@ -263,7 +267,10 @@ export function mountFileInspector(container, {api, sender, parentScope, variant
inner.append(heading, el('p', filename, 'small'));
const badges = el('div', undefined, 'tags'), badge = evidenceBadge(record.prompts);
if (badge) badges.append(el('span', badge, 'chip blue'));
if (linked && library.category) badges.append(prefixedChip(library.category_scope === 'pair' ? 'Recipe category: ' : 'Positive prompt category: ', library.category));
// #332: an image's own category shows for any link status; an inherited one only when linked.
if (library?.category && (library.category_scope === 'image' || linked)) {
badges.append(prefixedChip(CATEGORY_PREFIX[library.category_scope] ?? CATEGORY_PREFIX.positive, library.category));
}
if (badges.childNodes.length) inner.append(badges);
const prompts = el('div'); inner.append(prompts);
const blocks = mountPromptBlocks(prompts, {sender, parentScope:contentScope, includePairAction:false, contextLabel:promptContext, title});
+3 -1
View File
@@ -201,7 +201,9 @@ function resolutionFor(value) {
function libraryFor(value) {
if (!object(value) || !['linked', 'none', 'ambiguous', 'unavailable', 'limited'].includes(value.status) || value.scope !== 'saved-library' ||
!Array.isArray(value.tags) || value.tags.some(tag => typeof tag !== 'string') ||
![value.category_scope, value.tags_scope].every(scope => [null, 'pair', 'positive'].includes(scope))) throw new TypeError('Invalid inspection library');
![null, 'pair', 'positive', 'image'].includes(value.category_scope) ||
![null, 'pair', 'positive'].includes(value.tags_scope)) throw new TypeError('Invalid inspection library');
// #332: an image's own category (scope "image") may sit on any link status; tags never have that scope.
// D39: each field names its own source, so a recipe's category can sit beside its positive's tags.
const library = { status: value.status, scope: value.scope, tags: [...value.tags], category_scope: value.category_scope, tags_scope: value.tags_scope };
for (const key of ['positive_id', 'pair_id', 'title', 'category']) {