Compare commits
30
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
112d66f4e1 | ||
|
|
cdb771d722 | ||
|
|
55fa2c4870 | ||
|
|
6325d33831 | ||
|
|
11422d4584 | ||
|
|
8aebfd4aaa | ||
|
|
b5a56bee0b | ||
|
|
95586e6395 | ||
|
|
5583159497 | ||
|
|
3f085d0724 | ||
|
|
1701e927f7 | ||
|
|
894b48964c | ||
|
|
e769845417 | ||
|
|
463bda0e38 | ||
|
|
74e7cdaa62 | ||
|
|
12c27ae2b0 | ||
|
|
9a90c7677f | ||
|
|
94089d452f | ||
|
|
718d4a5cf1 | ||
|
|
1110ccbd33 | ||
|
|
20b6514882 | ||
|
|
e1ae60d5eb | ||
|
|
0de5c0a4e7 | ||
|
|
df6f538283 | ||
|
|
0237fba77a | ||
|
|
cd573ccf1a | ||
|
|
1db6088b2f | ||
|
|
cb874d7826 | ||
|
|
b543b7db89 | ||
|
|
2455ab21dd |
@@ -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
|
||||
|
||||

|
||||
|
||||
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
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
@@ -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
@@ -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]:
|
||||
|
||||
@@ -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
@@ -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
@@ -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,
|
||||
|
||||
@@ -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
@@ -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
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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')} "
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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"},
|
||||
)
|
||||
@@ -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
@@ -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 = {
|
||||
|
||||
@@ -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
@@ -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"
|
||||
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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');
|
||||
}
|
||||
});
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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, []);
|
||||
});
|
||||
|
||||
@@ -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, []);
|
||||
|
||||
@@ -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']]);
|
||||
});
|
||||
|
||||
|
||||
@@ -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');
|
||||
}
|
||||
});
|
||||
|
||||
@@ -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>' },
|
||||
});
|
||||
|
||||
|
||||
@@ -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 = []
|
||||
|
||||
@@ -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))
|
||||
@@ -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):
|
||||
|
||||
@@ -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"], "")
|
||||
|
||||
@@ -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(), [])
|
||||
@@ -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))
|
||||
@@ -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"])
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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"))
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
@@ -49,6 +49,7 @@ SUMMARY_KEYS = {
|
||||
"image_links",
|
||||
"runs",
|
||||
"saved_views",
|
||||
"collections",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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__":
|
||||
|
||||
Vendored
+6
-4
@@ -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/);
|
||||
});
|
||||
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
@@ -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
@@ -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:
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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 & 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 & 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 & 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"><lora:name:weight></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
@@ -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 & scans in a new tab" class="text-pm-accent hover:underline">Sources & 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
@@ -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
File diff suppressed because it is too large
Load Diff
@@ -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});
|
||||
|
||||
@@ -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']) {
|
||||
|
||||
Reference in New Issue
Block a user