Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Vito Sansevero
|
f8bec865cc | ||
|
Vito Sansevero
|
399c0cb63f | ||
|
Vito Sansevero
|
1e2f448646 | ||
|
Vito Sansevero
|
4a615a0c8d | ||
|
Vito Sansevero
|
c9f78fac58 | ||
|
Vito Sansevero
|
2b5a945c4d | ||
|
Vito Sansevero
|
50c6f261eb | ||
|
Vito Sansevero
|
58ededde6d | ||
|
Vito Sansevero
|
237a97d7cb | ||
|
Vito Sansevero
|
7df509e6fb | ||
|
Vito Sansevero
|
f09d2250df |
@@ -2,15 +2,18 @@
|
||||
"required_status_checks": {
|
||||
"strict": true,
|
||||
"contexts": [
|
||||
"py3.10 on ubuntu-latest",
|
||||
"py3.11 on ubuntu-latest",
|
||||
"py3.12 on ubuntu-latest",
|
||||
"py3.13 on ubuntu-latest",
|
||||
"py3.14 on ubuntu-latest",
|
||||
"py3.10 on windows-latest",
|
||||
"py3.11 on windows-latest",
|
||||
"py3.12 on windows-latest",
|
||||
"py3.13 on windows-latest",
|
||||
"py3.14 on windows-latest",
|
||||
"py3.10 on macos-latest",
|
||||
"py3.11 on macos-latest",
|
||||
"py3.12 on macos-latest",
|
||||
"py3.13 on macos-latest",
|
||||
"py3.14 on macos-latest",
|
||||
"coverage gate (85%)",
|
||||
"black and flake8",
|
||||
"frontend modules",
|
||||
|
||||
@@ -19,17 +19,14 @@ jobs:
|
||||
security:
|
||||
name: bandit (high severity)
|
||||
runs-on: ubuntu-latest
|
||||
# ceil(1.5 x slowest of three green main runs) + 3, never under 5 (test.yml): runs 40257217665,
|
||||
# 40254137380, 40240706583 took 23s / 16s / 23s -> 5 min.
|
||||
timeout-minutes: 5
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
|
||||
- name: Set up Python 3.12
|
||||
- name: Set up Python 3.10
|
||||
uses: actions/setup-python@v7
|
||||
with:
|
||||
python-version: "3.12"
|
||||
python-version: "3.10"
|
||||
|
||||
- name: Install bandit
|
||||
run: pip install bandit[toml]
|
||||
|
||||
@@ -19,28 +19,15 @@ concurrency:
|
||||
env:
|
||||
PYTHONFAULTHANDLER: "1"
|
||||
|
||||
# Every job has a time limit: ceil(1.5 x the slowest of three recent green runs on main) + 3
|
||||
# minutes, never under 5. The 3 minutes leave room for the Python runner's watchdog (120 s, see
|
||||
# tests/run_suite.py) to print its stack dump before GitHub cancels the job. A limit changes only
|
||||
# with new measurements written here (TESTING.md "Time limits and hangs").
|
||||
# Measured 2026-10-26 on runs 40257217681 (3be0ada), 40240706568 (fb9134f), 40228167328 (f9835aa):
|
||||
# py matrix, slowest py3.10 on windows-latest: 12m43s / 15m54s / 14m53s -> 27 min
|
||||
# (ubuntu 5m27s-6m27s, macos 6m14s-8m30s, windows 11m42s-15m54s)
|
||||
# 3.10 and 3.11 left the matrix and 3.14 joined it in #356; re-measure after three green main runs.
|
||||
# real ComfyUI compatibility: 13m32s / 8m55s / 11m46s -> 24 min
|
||||
# coverage gate (85%): 7m26s / 7m22s / 7m54s -> 15 min
|
||||
# perf budget 30s / 29s / 27s, black and flake8 39s / 37s / 39s,
|
||||
# frontend modules 23s / 19s / 29s -> 5 min (the floor)
|
||||
jobs:
|
||||
test:
|
||||
name: py${{ matrix.python-version }} on ${{ matrix.os }}
|
||||
runs-on: ${{ matrix.os }}
|
||||
timeout-minutes: 27
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os: [ubuntu-latest, windows-latest, macos-latest]
|
||||
python-version: ["3.12", "3.13", "3.14"]
|
||||
python-version: ["3.10", "3.11", "3.12", "3.13"]
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
@@ -58,12 +45,12 @@ jobs:
|
||||
pip install -r requirements-dev.txt
|
||||
|
||||
- name: Run tests
|
||||
run: python -m tests.run_suite -v
|
||||
run: python -m unittest discover -s tests -v
|
||||
|
||||
comfyui:
|
||||
name: real ComfyUI compatibility
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 24
|
||||
timeout-minutes: 25
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
with:
|
||||
@@ -140,7 +127,6 @@ jobs:
|
||||
coverage:
|
||||
name: coverage gate (85%)
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
|
||||
@@ -157,7 +143,7 @@ jobs:
|
||||
|
||||
- name: Measure coverage
|
||||
run: |
|
||||
python -m coverage run -m tests.run_suite
|
||||
python -m coverage run -m unittest discover -s tests
|
||||
python -m coverage combine
|
||||
python -m coverage report --fail-under=85
|
||||
|
||||
@@ -168,7 +154,6 @@ jobs:
|
||||
name: perf budget
|
||||
if: github.event_name == 'push' || github.event_name == 'workflow_dispatch'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
|
||||
@@ -191,7 +176,6 @@ jobs:
|
||||
lint:
|
||||
name: black and flake8
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
- uses: actions/setup-python@v7
|
||||
@@ -204,7 +188,6 @@ jobs:
|
||||
js-test:
|
||||
name: frontend modules
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
|
||||
|
||||
@@ -211,6 +211,3 @@ coverage/
|
||||
|
||||
# Disposable upstream checkout for real-host CI
|
||||
.ci-comfyui/
|
||||
|
||||
# Old library files a Move library renamed (never deleted; #249 S6b)
|
||||
*.moved-to-*
|
||||
|
||||
@@ -22,7 +22,7 @@ repos:
|
||||
|
||||
- id: tests
|
||||
name: tests
|
||||
entry: python -m tests.run_suite -q
|
||||
entry: python -m unittest discover tests/ -v
|
||||
language: system
|
||||
pass_filenames: false
|
||||
stages: [pre-push]
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# ComfyUI Prompt Manager
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -483,15 +483,11 @@ If you use [ComfyUI-Lora-Manager](https://github.com/willmiao/ComfyUI-Lora-Manag
|
||||
|
||||
Click **Import LoRA Data** in the settings panel to start the import. A progress popup shows real-time status as each LoRA is processed:
|
||||
|
||||
- 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)
|
||||
- Links LoRA Manager's preview (including `.preview.webp` and video previews) and every downloaded example image and video where they are, without copying them
|
||||
- Downloads example images from CivitAI as 512px thumbnails only for LoRAs that have no local image
|
||||
- Scans all LoRA directories (including `extra_model_paths.yaml` paths) for metadata
|
||||
- Downloads example images from CivitAI as 512px thumbnails
|
||||
- Creates a prompt entry for each LoRA with trigger words, example prompts, and preview images
|
||||
- All imported LoRAs are tagged with `lora-manager` category for easy filtering
|
||||
|
||||
The System › Integrations review names the settings it read (source, storage mode and LoRA Manager library) and counts metadata files, LoRAs and LoRAs with local images. If LoRA Manager's settings cannot be used, it says why and nothing is changed.
|
||||
|
||||
|
||||
|
||||
After import, filter by the `lora-manager` category to browse your LoRA collection.
|
||||
@@ -505,8 +501,6 @@ To add your key:
|
||||
2. Paste it into the **CivitAI API Key** field in the Integrations settings
|
||||
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).
|
||||
|
||||
> **Note:** Re-importing is safe — PromptManager skips LoRAs that were already imported. To do a fresh import, the previous `lora-manager` data is cleared automatically before re-scanning.
|
||||
|
||||
### 🌐 Web Interface Features
|
||||
@@ -757,51 +751,9 @@ ComfyUI_PromptManager/
|
||||
|
||||
## Configuration
|
||||
|
||||
### Where your library lives
|
||||
|
||||
New installs keep the library (`prompts.db`, `config.json` with the CivitAI key, the
|
||||
database backups and the logs) in ComfyUI's system user folder
|
||||
`ComfyUI/user/__promptmanager/`. An update, a reinstall or ComfyUI-Manager replacing
|
||||
the node folder no longer touches it, and ComfyUI's `/userdata` web API does not
|
||||
serve `__` system folders. Thumbnails stay in `output/thumbnails`.
|
||||
|
||||
- **Upgrading from an earlier version:** your library stays where it is, in the node
|
||||
folder, until you move it; nothing is moved on its own. The workspace shows a notice
|
||||
above every page and System › Overview a "Library location" row. **Move library**
|
||||
copies the database with SQLite's backup API, verifies the copy (integrity check, every
|
||||
workflow document, the same rows in every table), copies `config.json` and the backups,
|
||||
switches to the new folder without a restart and only then renames the old files
|
||||
`prompts.db.moved-to-__promptmanager` and `config.json.moved-to-__promptmanager` (never
|
||||
deleted; the old backups stay too). If anything fails, the library stays in the node
|
||||
folder, unchanged. **Keep it here** stops the notice; the Overview row stays. Take a
|
||||
backup first (System › Backup & export or the dashboard's Backup) as with any change.
|
||||
- **Your own folder:** set `PROMPT_MANAGER_DATA_DIR` to an absolute folder outside
|
||||
ComfyUI's output, input and temp folders and your image sources. If it is empty and a
|
||||
node-folder library exists, that library stays in use until you move it there.
|
||||
- **An explicit `database.default_path`** in `config.json` and `PROMPT_MANAGER_CONFIG_PATH`
|
||||
keep working as before (described below); such a library is not offered a move.
|
||||
|
||||
### Thumbnails
|
||||
|
||||
Previews in `output/thumbnails` are written as **WebP, quality 80, 512 px** by default
|
||||
(System › Appearance › Thumbnails sets the format, quality and size limit for every
|
||||
preview). Previews from older versions keep their old format (usually PNG) until you
|
||||
run **Convert outdated** (System › Overview or Appearance): it rewrites each one from
|
||||
its original and then removes the old file. Nothing converts on its own.
|
||||
|
||||
**JPEG XL** (smaller still) can be chosen once the server has an encoder. Install one
|
||||
of these, then restart ComfyUI:
|
||||
|
||||
- `pip install pillow-jxl-plugin` in ComfyUI's Python environment (simplest), or
|
||||
- an `ffmpeg` on PATH built with `libjxl`, or
|
||||
- `cjxl` from [libjxl](https://github.com/libjxl/libjxl) on PATH.
|
||||
|
||||
Browsers that cannot show JPEG XL see the original images in the grid instead.
|
||||
|
||||
### Database Settings
|
||||
|
||||
You can customize the database path by creating a `config.json` file in the library's
|
||||
folder (the node folder for a library that was not moved):
|
||||
You can customize the database path by creating a `config.json` file in the extension root:
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -928,12 +880,12 @@ db.model.backup_database("backup_prompts.db")
|
||||
```bash
|
||||
cd ComfyUI_PromptManager
|
||||
pip install -r requirements-dev.txt
|
||||
python -m tests.run_suite # Python (what CI runs: unittest discovery plus a hang watchdog)
|
||||
python -m unittest discover -s tests # Python (what CI runs)
|
||||
node --test tests/js/*.test.js # frontend pure modules
|
||||
python -m coverage run -m tests.run_suite && python -m coverage combine && python -m coverage report
|
||||
python -m coverage run -m unittest discover -s tests && python -m coverage combine && python -m coverage report
|
||||
```
|
||||
|
||||
CI runs the suite on Linux, Windows and macOS for Python 3.12 to 3.14 with an 85% coverage gate. `python -m unittest discover -s tests` still works; see `TESTING.md` ("Time limits and hangs") for the watchdog. Tests that delete a database file first check that nothing in the process still holds it open (`tests/open_handles.py`), so a connection leak fails on Linux the way it would on Windows.
|
||||
CI runs the suite on Linux, Windows and macOS for Python 3.10 to 3.12 with an 85% coverage gate. Tests that delete a database file first check that nothing in the process still holds it open (`tests/open_handles.py`), so a connection leak fails on Linux the way it would on Windows.
|
||||
|
||||
### Code Style
|
||||
|
||||
@@ -1018,15 +970,6 @@ MIT License - see LICENSE file for details.
|
||||
|
||||
## Changelog
|
||||
|
||||
### Unreleased (Library location)
|
||||
|
||||
- **📁 Library outside the node folder**: new installs keep `prompts.db`, `config.json`, the database backups and the logs in `ComfyUI/user/__promptmanager/`, where an update or reinstall of the node cannot delete them and ComfyUI's `/userdata` API does not reach. An existing library stays in the node folder until you press **Move library** (a verified copy, then a live switch; the old files are renamed `*.moved-to-__promptmanager`, never deleted) or **Keep it here**. `PROMPT_MANAGER_DATA_DIR` names your own folder. After moving, restart ComfyUI once so the logs follow. Going back to an older version after a move: an older version does not look in `user/__promptmanager`, so stop ComfyUI and copy `user/__promptmanager/prompts.db` (with its `-wal` file, if one is left) and `config.json` into the node folder; the `*.moved-to-__promptmanager` files are the library as it was at the move, without later changes.
|
||||
|
||||
### Unreleased (Workflows stored once)
|
||||
|
||||
- **🗜️ Smaller database**: each distinct workflow and prompt graph is stored once, compressed, instead of again for every image. Images recorded by this version use it from the first start; the maintenance job "Store each workflow once" rewrites the existing records and then reclaims the space. A 20,000-image test library went from 1.4 GB to 60 MB. The job always takes a verified backup first: the first time, `prompts.db.pre-workflow-storage_<date time>` beside the database, kept until you delete it; on later runs, a normal pooled safety backup. No image or video file is opened or changed.
|
||||
- **⚠️ Downgrading**: images recorded by this version (and, after the job, all images) are stored in the new form. An older Prompt Manager version still opens and even restores such a database, but shows those images without workflows or prompt graphs until you update again. To go back to an older version, stop ComfyUI, delete `prompts.db-wal` and `prompts.db-shm`, and copy the **earliest** `prompts.db.pre-workflow-storage_<date time>` over `prompts.db`, or, if the job never ran, a backup taken before this update (records added since that backup are lost). Otherwise stay on this version.
|
||||
|
||||
### v3.3.1 (Proxy Routing Fix)
|
||||
|
||||
- **🌐 Proxy support**: the dashboard, gallery, tag page and metadata viewer now work when ComfyUI is reached through a proxy mount such as SwarmUI's `/ComfyBackendDirect/`. Scripts, images, API calls, progress streams and gallery folder filters stay under the mount, and the canvas node's buttons use ComfyUI's `fileURL()`. Direct ComfyUI URLs are unchanged. Thanks to @ghosttigerllc-bit ([#164](https://github.com/ComfyAssets/ComfyUI_PromptManager/pull/164))
|
||||
|
||||
+51
-1331
File diff suppressed because it is too large
Load Diff
+1
-26
@@ -87,24 +87,10 @@ try:
|
||||
routes = config.routes
|
||||
api = PromptManagerAPI()
|
||||
api.add_routes(routes)
|
||||
# Temp files of preview writes a crash interrupted (#317), off the
|
||||
# import path: the walk may cover a large thumbnails folder.
|
||||
import threading
|
||||
|
||||
threading.Thread(
|
||||
target=api.sweep_thumbnail_temps,
|
||||
name="PromptManagerPreviewTempSweep",
|
||||
daemon=True,
|
||||
).start()
|
||||
host_app = getattr(config.server_instance, "app", None)
|
||||
if host_app is not None and hasattr(host_app, "on_cleanup"):
|
||||
host_app.on_cleanup.append(api.cleanup_history_reads)
|
||||
host_app.on_cleanup.append(api.cleanup_annotation_captures)
|
||||
host_app.on_cleanup.append(api.cleanup_gallery_listings)
|
||||
if host_app is not None and hasattr(host_app, "on_startup"):
|
||||
# Build each output root's Images listing once the server's loop runs,
|
||||
# so the first open does not wait for the walk (#348).
|
||||
host_app.on_startup.append(api.warm_gallery_listings)
|
||||
if host_app is not None and hasattr(host_app, "on_shutdown"):
|
||||
# Open Needs review change streams end before the server waits on them.
|
||||
host_app.on_shutdown.append(api.close_review_events)
|
||||
@@ -146,22 +132,11 @@ try:
|
||||
_review_api = globals().get("api")
|
||||
if _review_api is not None:
|
||||
_global_image_monitor.set_review_recorder(_review_api.record_review_file)
|
||||
# Every settled output's preview is queued for the API's background
|
||||
# worker (#249): the workspace never asks to generate thumbnails.
|
||||
_global_image_monitor.set_thumbnail_queue(_review_api.queue_capture_thumbnail)
|
||||
# Every settled output walks its Images listing again (#348): the
|
||||
# API hops from the monitor's thread onto the server's loop.
|
||||
_global_image_monitor.set_output_listener(_review_api.notify_output_settled)
|
||||
_review_host = getattr(config.server_instance, "app", None)
|
||||
if _review_host is not None and hasattr(_review_host, "on_cleanup"):
|
||||
|
||||
async def _stop_review_recording(
|
||||
app, monitor=_global_image_monitor, api=_review_api
|
||||
):
|
||||
async def _stop_review_recording(app, monitor=_global_image_monitor):
|
||||
monitor.set_review_recorder(None)
|
||||
monitor.set_thumbnail_queue(None)
|
||||
monitor.set_output_listener(None)
|
||||
api.close_auto_thumbnails(final=True)
|
||||
|
||||
_review_host.on_cleanup.insert(0, _stop_review_recording)
|
||||
|
||||
|
||||
@@ -1,201 +0,0 @@
|
||||
"""Atomic replacement of one imported category (#221, #249 System S10).
|
||||
|
||||
The LoRA reference import used to delete its whole category and then save
|
||||
each record in its own transaction, so a failure part way left the previous
|
||||
collection deleted and the new one partial. ``replace_category_prompts``
|
||||
does both in one ``BEGIN IMMEDIATE`` transaction: the category's current
|
||||
prompts (with their tags and image links) are replaced by the staged
|
||||
records, or, on any error, nothing changes.
|
||||
|
||||
A staged record is ``{text, hash, tags, images}``. A record whose hash a
|
||||
prompt outside the category already has links its images to that prompt
|
||||
(the legacy import did the same); a second record with the hash of an
|
||||
earlier one links its images to that one. Neither is a new prompt.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Sequence
|
||||
|
||||
try:
|
||||
from .maintenance import PlanChanged
|
||||
from .pairs import PairDependencyError
|
||||
from .workflow_documents import documents_of_rows, release_documents
|
||||
except ImportError: # pragma: no cover - flat test imports
|
||||
from database.maintenance import PlanChanged
|
||||
from database.pairs import PairDependencyError
|
||||
from database.workflow_documents import documents_of_rows, release_documents
|
||||
|
||||
|
||||
class EmptyReplacement(ValueError):
|
||||
"""The replacement has no record and would remove every current import."""
|
||||
|
||||
|
||||
# The category's prompt ids, as a subquery (one bound parameter).
|
||||
_IN_CATEGORY = "SELECT id FROM prompts WHERE category = ?"
|
||||
|
||||
# Each record's outcome and the result counter it adds to.
|
||||
_COUNTERS = {"imported": "imported", "linked": "linked", "duplicate": "duplicates"}
|
||||
|
||||
|
||||
def _placeholders(values: Sequence[Any]) -> str:
|
||||
return ",".join("?" * len(values))
|
||||
|
||||
|
||||
class CategoryImportOperations:
|
||||
"""Read the state of an imported category and replace it atomically."""
|
||||
|
||||
def category_import_state(self, category: str) -> Dict[str, Any]:
|
||||
"""``{ids, paired}``: the category's prompt ids (ascending) and how
|
||||
many of them a saved pair uses (those block a replacement)."""
|
||||
with self.model.get_connection() as conn:
|
||||
ids = [
|
||||
row[0]
|
||||
for row in conn.execute(
|
||||
"SELECT id FROM prompts WHERE category = ? ORDER BY id",
|
||||
(category,),
|
||||
).fetchall()
|
||||
]
|
||||
paired = conn.execute(
|
||||
"SELECT COUNT(*) FROM prompts WHERE category = ? AND id IN "
|
||||
"(SELECT positive_prompt_id FROM prompt_pairs UNION "
|
||||
"SELECT negative_prompt_id FROM prompt_pairs)",
|
||||
(category,),
|
||||
).fetchone()[0]
|
||||
return {"ids": ids, "paired": paired}
|
||||
|
||||
def existing_prompt_hashes(
|
||||
self, hashes: Sequence[str], exclude_category: str
|
||||
) -> set:
|
||||
"""The given hashes a prompt outside ``exclude_category`` already has."""
|
||||
found = set()
|
||||
wanted = list(dict.fromkeys(hashes))
|
||||
with self.model.get_connection() as conn:
|
||||
for start in range(0, len(wanted), 500):
|
||||
chunk = wanted[start : start + 500]
|
||||
rows = conn.execute(
|
||||
f"SELECT hash FROM prompts WHERE hash IN ({_placeholders(chunk)}) "
|
||||
"AND (category IS NULL OR category != ?)",
|
||||
[*chunk, exclude_category],
|
||||
).fetchall()
|
||||
found.update(row[0] for row in rows)
|
||||
return found
|
||||
|
||||
def replace_category_prompts(
|
||||
self,
|
||||
category: str,
|
||||
records: List[Dict[str, Any]],
|
||||
*,
|
||||
expected_ids: Optional[List[int]] = None,
|
||||
allow_empty: bool = False,
|
||||
) -> Dict[str, Any]:
|
||||
"""Replace every prompt of ``category`` with ``records`` in one transaction.
|
||||
|
||||
Raises (and changes nothing): ``PlanChanged`` when ``expected_ids`` is
|
||||
given and the category's prompt ids differ from it;
|
||||
``PairDependencyError`` when a saved pair uses one of its prompts;
|
||||
``EmptyReplacement`` when ``records`` is empty while the category has
|
||||
prompts, unless ``allow_empty`` (an empty replacement that was
|
||||
reviewed). Any other error rolls everything back too.
|
||||
|
||||
Returns ``{removed, imported, linked, duplicates, images, rows}``:
|
||||
prompts removed and created, records linked to an existing prompt or
|
||||
to an earlier record with the same text, image links added, and per
|
||||
record ``{name, outcome}`` (``imported``/``linked``/``duplicate``).
|
||||
"""
|
||||
with self._pair_transaction() as conn:
|
||||
ids = [
|
||||
row[0]
|
||||
for row in conn.execute(
|
||||
"SELECT id FROM prompts WHERE category = ? ORDER BY id",
|
||||
(category,),
|
||||
).fetchall()
|
||||
]
|
||||
if expected_ids is not None and ids != list(expected_ids):
|
||||
raise PlanChanged("The imported prompts changed since the review")
|
||||
if not records and ids and not allow_empty:
|
||||
raise EmptyReplacement(
|
||||
"Nothing to import; the previous imports were kept"
|
||||
)
|
||||
self._ensure_category_unpaired(conn, category)
|
||||
removed = self._delete_category_rows(conn, category, ids)
|
||||
result = {
|
||||
"removed": removed,
|
||||
"imported": 0,
|
||||
"linked": 0,
|
||||
"duplicates": 0,
|
||||
"images": 0,
|
||||
"rows": [],
|
||||
}
|
||||
created = {}
|
||||
for record in records:
|
||||
self._import_record(conn, category, record, created, result)
|
||||
return result
|
||||
|
||||
def _delete_category_rows(self, conn, category: str, ids: List[int]) -> int:
|
||||
"""Delete the category's prompts with their image links and tags.
|
||||
|
||||
No commit. The rows are selected by category (just verified against
|
||||
``ids`` under the write lock), so no statement binds one parameter per
|
||||
prompt and SQLite's variable limit never caps a collection.
|
||||
"""
|
||||
if not ids:
|
||||
return 0
|
||||
documents = documents_of_rows(conn, "prompt_id", ids) # chunked
|
||||
conn.execute(
|
||||
f"DELETE FROM generated_images WHERE prompt_id IN ({_IN_CATEGORY})",
|
||||
(category,),
|
||||
)
|
||||
release_documents(conn, documents)
|
||||
conn.execute(
|
||||
f"DELETE FROM prompt_tags WHERE prompt_id IN ({_IN_CATEGORY})",
|
||||
(category,),
|
||||
)
|
||||
return conn.execute(
|
||||
"DELETE FROM prompts WHERE category = ?", (category,)
|
||||
).rowcount
|
||||
|
||||
@staticmethod
|
||||
def _ensure_category_unpaired(conn, category: str) -> None:
|
||||
"""Raise PairDependencyError when a saved pair uses a category prompt."""
|
||||
row = conn.execute(
|
||||
f"SELECT id FROM prompt_pairs WHERE positive_prompt_id IN ({_IN_CATEGORY}) "
|
||||
f"OR negative_prompt_id IN ({_IN_CATEGORY}) LIMIT 1",
|
||||
(category, category),
|
||||
).fetchone()
|
||||
if row:
|
||||
raise PairDependencyError(
|
||||
f"Prompt belongs to saved pair {row[0]}; delete the pair first "
|
||||
"or save changed text as a new prompt"
|
||||
)
|
||||
|
||||
def _import_record(self, conn, category, record, created, result) -> None:
|
||||
"""Save one staged record on ``conn``, or link it to the prompt it repeats."""
|
||||
prompt_hash = record["hash"]
|
||||
name = record.get("name") or record["text"]
|
||||
if prompt_hash in created:
|
||||
prompt_id, outcome = created[prompt_hash], "duplicate"
|
||||
else:
|
||||
row = conn.execute(
|
||||
"SELECT id FROM prompts WHERE hash = ?", (prompt_hash,)
|
||||
).fetchone()
|
||||
if row is not None:
|
||||
prompt_id, outcome = row[0], "linked"
|
||||
else:
|
||||
prompt_id = self._insert_prompt(
|
||||
conn,
|
||||
record["text"],
|
||||
category,
|
||||
list(record.get("tags") or []),
|
||||
None,
|
||||
None,
|
||||
prompt_hash,
|
||||
)
|
||||
outcome = "imported"
|
||||
created[prompt_hash] = prompt_id
|
||||
result[_COUNTERS[outcome]] += 1
|
||||
# Linked last to first: every workspace list shows a prompt's media
|
||||
# newest link first, so the record's first image (the LoRA cover,
|
||||
# #355) is the one a Library card shows.
|
||||
for image in reversed(record.get("images") or []):
|
||||
if self._insert_image_link(conn, prompt_id, image):
|
||||
result["images"] += 1
|
||||
result["rows"].append({"name": name, "outcome": outcome})
|
||||
@@ -25,7 +25,6 @@ from .annotations import _BUMP_INDEX
|
||||
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
|
||||
from .workflow_documents import documents_of_rows, release_documents
|
||||
|
||||
LINKS_BY_FILE_SQL = "SELECT id, prompt_id FROM generated_images WHERE file_path = ?"
|
||||
LEGACY_LINKS_SQL = (
|
||||
@@ -156,9 +155,7 @@ class DuplicateCopyOperations:
|
||||
break
|
||||
if existing is not None:
|
||||
conn.execute(MOVE_PAIR_LINKS_SQL, (existing[0], image_id))
|
||||
documents = documents_of_rows(conn, "id", [image_id])
|
||||
conn.execute(DELETE_LINK_SQL, (image_id,))
|
||||
release_documents(conn, documents)
|
||||
counts["merged"] += 1
|
||||
continue
|
||||
conn.execute(
|
||||
|
||||
@@ -196,14 +196,3 @@ class ExecutionHistoryOperations:
|
||||
.fetchone()
|
||||
)
|
||||
return dict(row) if row is not None else None
|
||||
|
||||
|
||||
def rekey_history_restore(old_path: str, new_path: str) -> None:
|
||||
"""Move the observers of ``old_path`` to ``new_path`` (Move library)."""
|
||||
old, new = _database_identity(old_path), _database_identity(new_path)
|
||||
if old == new:
|
||||
return
|
||||
with _restore_lock:
|
||||
observers = _restore_observers.pop(old, None)
|
||||
if observers:
|
||||
_restore_observers.setdefault(new, {}).update(observers)
|
||||
|
||||
@@ -4,8 +4,6 @@ import sqlite3
|
||||
from contextlib import ExitStack
|
||||
from typing import Callable, ContextManager
|
||||
|
||||
from .run_attribution import notify_outputs_sealed
|
||||
|
||||
try:
|
||||
from ..utils.execution_history_events import RunFinish, encode_command
|
||||
from ..utils.execution_outputs import OutputBundle, validate_output_bundle
|
||||
@@ -115,10 +113,6 @@ class ExecutionOutputOperations:
|
||||
"VALUES(?,?,?,?,?,?,?,?)",
|
||||
[(run_id, *reference) for reference in prepared],
|
||||
)
|
||||
# Committed: images waiting for this run may now be attributed (#342)
|
||||
notify_outputs_sealed(
|
||||
run_id, [dict(zip(_REFERENCE_FIELDS, row)) for row in prepared]
|
||||
)
|
||||
return True
|
||||
|
||||
def get_execution_outputs(self, run_id: str) -> dict:
|
||||
|
||||
@@ -1,805 +0,0 @@
|
||||
"""Move library (#249 S6b): copy the node-folder library, verify it, switch.
|
||||
|
||||
``plan_move`` checks a ``utils.library_location`` location can move and
|
||||
sizes the copy; ``move_library`` does it, in this order:
|
||||
|
||||
1. **Disk check** on the target's volume (the copy, the backups and
|
||||
config.json, plus a margin). The source is never renamed across folders,
|
||||
so another drive is simply a copy.
|
||||
2. **Copy** with SQLite's backup API from a private connection holding one
|
||||
read transaction: the per-table counts, the workflow documents' digest
|
||||
and ``PRAGMA data_version`` describe exactly the copied snapshot. The copy
|
||||
is a temp file in the target folder, made self-contained
|
||||
(``journal_mode=DELETE``), fsynced and owner-only.
|
||||
3. **Verify**: ``verify_database_file`` (integrity check, schema, every
|
||||
workflow document decompressed against its hash, #315), the same row
|
||||
count in every table and the same documents.
|
||||
4. **Copy files**: every pooled safety backup and pre-migration backup,
|
||||
checked by size and SHA-256 after the copy (their times are kept, so the
|
||||
pool prunes in the same order).
|
||||
5. **Publish** under the old database's write lock (``BEGIN IMMEDIATE``) and
|
||||
the config.json lock: if any other connection wrote since the snapshot
|
||||
(``data_version``), steps 2–4 run again (``attempts``, then
|
||||
``library-busy``). Else config.json is copied (owner-only, checked), a
|
||||
``pending`` marker is written, the temp files take their final names
|
||||
(``os.replace`` inside the target folder), the marker turns ``moved``
|
||||
and ``switch`` points the process at the new library while the old file
|
||||
is still locked, so no write can land in it.
|
||||
6. **Retire**: the old file's WAL is checkpointed and the old ``prompts.db``
|
||||
and ``config.json`` are renamed ``*.moved-to-<folder>`` (never deleted;
|
||||
the old backups stay where they are). A file still open elsewhere is left
|
||||
in place and reported (``old-in-use``): the marker already makes a restart
|
||||
use the new library.
|
||||
|
||||
Any failure before the marker says ``moved`` removes everything the move
|
||||
wrote in the target folder and leaves the old library in use, untouched.
|
||||
A crash at any point leaves either the old library in use (the marker is
|
||||
absent or ``pending``; the next move replaces its own leftovers) or the new
|
||||
one (``moved``). ``MoveFailed`` carries a code and a message without paths.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sqlite3
|
||||
import time
|
||||
import types
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Callable, Dict, List, Optional, Tuple
|
||||
|
||||
from .models import BUSY_TIMEOUT_SECONDS, PromptModel
|
||||
|
||||
try:
|
||||
from ..utils import library_location as ll
|
||||
from ..utils.logging_config import get_logger
|
||||
except ImportError:
|
||||
from utils import library_location as ll
|
||||
from utils.logging_config import get_logger
|
||||
|
||||
logger = get_logger("prompt_manager.database.library_move")
|
||||
|
||||
TEMP_PREFIX = ".pm-move-"
|
||||
MOVED_INFIX = ".moved-to-"
|
||||
DISK_MARGIN_BYTES = 16 * 2**20
|
||||
DEFAULT_ATTEMPTS = 3
|
||||
_COPY_CHUNK = 4 * 2**20
|
||||
# Backups that live beside the database and move with it.
|
||||
BACKUP_INFIXES = (".backup_", ".pre-workflow-storage_")
|
||||
|
||||
MESSAGES = {
|
||||
"not-movable": (
|
||||
"This library is not in the node folder, so there is nothing to move."
|
||||
),
|
||||
"same-folder": "The new folder is the node folder itself.",
|
||||
"target-invalid": "The new folder cannot hold the library.",
|
||||
"target-has-library": (
|
||||
"The new folder already holds a Prompt Manager library. Nothing was moved."
|
||||
),
|
||||
"target-conflict": (
|
||||
"The new folder already holds files with the same names. Nothing was moved."
|
||||
),
|
||||
"insufficient-space": (
|
||||
"There is not enough free space for the copy. Nothing was moved."
|
||||
),
|
||||
"copy-failed": "The copy could not be written or did not match. Nothing was moved.",
|
||||
"verify-failed": "The copy did not pass verification. Nothing was moved.",
|
||||
"library-busy": (
|
||||
"Prompt Manager kept writing to the library while it was copied. "
|
||||
"Nothing was moved; try again when no workflow is running."
|
||||
),
|
||||
"switch-failed": (
|
||||
"Prompt Manager could not switch to the verified copy, so the copy was "
|
||||
"removed. Nothing was moved."
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
class MoveFailed(Exception):
|
||||
"""A move that did not happen (or did not finish), with a coded reason."""
|
||||
|
||||
def __init__(self, code: str, message: Optional[str] = None):
|
||||
self.code = code
|
||||
self.message = message or MESSAGES.get(code, "The library was not moved.")
|
||||
super().__init__(self.message)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MovePlan:
|
||||
location: "ll.LibraryLocation"
|
||||
db_path: str
|
||||
source_dir: str
|
||||
target_dir: str
|
||||
config_path: Optional[str]
|
||||
backups: Tuple[str, ...]
|
||||
db_bytes: int
|
||||
files_bytes: int
|
||||
|
||||
@property
|
||||
def required_bytes(self) -> int:
|
||||
return self.db_bytes + self.files_bytes + DISK_MARGIN_BYTES
|
||||
|
||||
@property
|
||||
def target_location(self) -> "ll.LibraryLocation":
|
||||
"""Where the library is once moved (an explicit config.json stays)."""
|
||||
source = self.location.target_source or ll.SOURCE_DEFAULT
|
||||
config = (
|
||||
self.location.config_path
|
||||
if self.location.config_explicit
|
||||
else os.path.join(self.target_dir, ll.CONFIG_NAME)
|
||||
)
|
||||
return ll.LibraryLocation(
|
||||
source,
|
||||
os.path.join(self.target_dir, ll.DATABASE_NAME),
|
||||
config,
|
||||
self.target_dir,
|
||||
self.target_dir,
|
||||
None,
|
||||
source,
|
||||
self.location.config_explicit,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MoveResult:
|
||||
location: "ll.LibraryLocation"
|
||||
counts: Dict[str, int]
|
||||
files: List[str]
|
||||
renamed: List[str] = field(default_factory=list)
|
||||
retire_problem: Optional[str] = None
|
||||
|
||||
|
||||
def _size(path: str) -> int:
|
||||
try:
|
||||
return os.path.getsize(path)
|
||||
except OSError:
|
||||
return 0
|
||||
|
||||
|
||||
def _free_bytes(path: str) -> int:
|
||||
return shutil.disk_usage(path).free
|
||||
|
||||
|
||||
def _existing_ancestor(path: str) -> str:
|
||||
current = os.path.abspath(path)
|
||||
while not os.path.isdir(current):
|
||||
parent = os.path.dirname(current)
|
||||
if parent == current:
|
||||
break
|
||||
current = parent
|
||||
return current
|
||||
|
||||
|
||||
def backup_names(db_path: str) -> Tuple[str, ...]:
|
||||
"""The pooled and pre-migration backups beside ``db_path`` (names)."""
|
||||
folder = os.path.dirname(os.path.abspath(db_path))
|
||||
base = os.path.basename(db_path)
|
||||
prefixes = tuple(base + infix for infix in BACKUP_INFIXES)
|
||||
try:
|
||||
names = os.listdir(folder)
|
||||
except OSError:
|
||||
return ()
|
||||
return tuple(
|
||||
sorted(
|
||||
name
|
||||
for name in names
|
||||
if name.startswith(prefixes) and os.path.isfile(os.path.join(folder, name))
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def plan_move(location) -> MovePlan:
|
||||
"""Check ``location`` can move to its target; raises MoveFailed."""
|
||||
if not location.movable:
|
||||
raise MoveFailed("not-movable")
|
||||
target = os.path.abspath(location.target_dir)
|
||||
source_dir = os.path.dirname(os.path.abspath(location.db_path))
|
||||
if ll.same_folder(target, source_dir) or os.path.normcase(
|
||||
os.path.realpath(target)
|
||||
) == os.path.normcase(os.path.realpath(source_dir)):
|
||||
raise MoveFailed("same-folder")
|
||||
if os.path.exists(target) and not os.path.isdir(target):
|
||||
raise MoveFailed("target-invalid")
|
||||
backups = backup_names(location.db_path)
|
||||
# config.json moves with a node-folder library; one named by
|
||||
# PROMPT_MANAGER_CONFIG_PATH is the user's and stays where it is.
|
||||
config = None if location.config_explicit else location.config_path
|
||||
config = config if config and os.path.isfile(config) else None
|
||||
if ll.read_marker(target).get("state") != ll.MARKER_PENDING:
|
||||
if os.path.exists(os.path.join(target, ll.DATABASE_NAME)):
|
||||
raise MoveFailed("target-has-library")
|
||||
names = [ll.CONFIG_NAME] if config else []
|
||||
if any(os.path.exists(os.path.join(target, n)) for n in names + list(backups)):
|
||||
raise MoveFailed("target-conflict")
|
||||
db_bytes = _size(location.db_path) + _size(location.db_path + "-wal")
|
||||
files_bytes = sum(_size(os.path.join(source_dir, n)) for n in backups)
|
||||
files_bytes += _size(config) if config else 0
|
||||
return MovePlan(
|
||||
location,
|
||||
os.path.abspath(location.db_path),
|
||||
source_dir,
|
||||
target,
|
||||
config,
|
||||
backups,
|
||||
db_bytes,
|
||||
files_bytes,
|
||||
)
|
||||
|
||||
|
||||
def space_problem(plan: MovePlan) -> Optional[Tuple[int, int]]:
|
||||
"""``(required, free)`` when the target's volume is short, else None."""
|
||||
free = _free_bytes(_existing_ancestor(plan.target_dir))
|
||||
if free < plan.required_bytes:
|
||||
return plan.required_bytes, free
|
||||
return None
|
||||
|
||||
|
||||
# ── copy and verify ───────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _tables(conn) -> List[str]:
|
||||
return [
|
||||
row[0]
|
||||
for row in conn.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' "
|
||||
"AND name NOT LIKE 'sqlite_%' ORDER BY name"
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def _describe(conn) -> Tuple[Dict[str, int], str]:
|
||||
"""Row counts per table and a digest of the workflow documents."""
|
||||
counts = {}
|
||||
for table in _tables(conn):
|
||||
# A name read from sqlite_master, quoted as an identifier.
|
||||
quoted = '"' + table.replace('"', '""') + '"'
|
||||
sql = f"SELECT COUNT(*) FROM {quoted}" # nosec B608
|
||||
counts[table] = conn.execute(sql).fetchone()[0]
|
||||
digest = hashlib.sha256()
|
||||
if "workflow_documents" in counts:
|
||||
for row in conn.execute(
|
||||
"SELECT id, sha256, byte_length FROM workflow_documents ORDER BY id"
|
||||
):
|
||||
digest.update(repr((row[0], bytes(row[1]), row[2])).encode())
|
||||
return counts, digest.hexdigest()
|
||||
|
||||
|
||||
def _fsync_file(path: str) -> None:
|
||||
with open(path, "rb+") as handle:
|
||||
os.fsync(handle.fileno())
|
||||
|
||||
|
||||
def _fsync_dir(path: str) -> None:
|
||||
if os.name != "posix":
|
||||
return
|
||||
fd = os.open(path, os.O_RDONLY)
|
||||
try:
|
||||
os.fsync(fd)
|
||||
finally:
|
||||
os.close(fd)
|
||||
|
||||
|
||||
def _private(path: str) -> None:
|
||||
if os.name == "posix":
|
||||
os.chmod(path, 0o600)
|
||||
|
||||
|
||||
def _copy_snapshot(source: sqlite3.Connection, temp: str):
|
||||
"""Back up ``source``'s current snapshot to ``temp``; return what it holds."""
|
||||
source.execute("BEGIN")
|
||||
try:
|
||||
version = source.execute("PRAGMA data_version").fetchone()[0]
|
||||
counts, documents = _describe(source)
|
||||
destination = sqlite3.connect(temp)
|
||||
try:
|
||||
source.backup(destination)
|
||||
finally:
|
||||
destination.close()
|
||||
finally:
|
||||
source.execute("COMMIT")
|
||||
with contextlib.closing(sqlite3.connect(temp)) as conn:
|
||||
conn.execute("PRAGMA journal_mode=DELETE").fetchone()
|
||||
_fsync_file(temp)
|
||||
_private(temp)
|
||||
return version, counts, documents
|
||||
|
||||
|
||||
def _verify_copy(temp: str, counts, documents) -> None:
|
||||
ok, reason = PromptModel.verify_database_file(temp)
|
||||
if not ok:
|
||||
logger.error(f"Library move: the copy did not verify: {reason}")
|
||||
detail = reason if reason.startswith("Workflow document") else None
|
||||
raise MoveFailed(
|
||||
"verify-failed",
|
||||
f"{MESSAGES['verify-failed']} {detail}" if detail else None,
|
||||
)
|
||||
uri = Path(os.path.abspath(temp)).as_uri() + "?mode=ro"
|
||||
try:
|
||||
with contextlib.closing(sqlite3.connect(uri, uri=True)) as conn:
|
||||
copied = _describe(conn)
|
||||
except sqlite3.Error as error:
|
||||
logger.error(f"Library move: the copy could not be read: {error}")
|
||||
raise MoveFailed("verify-failed") from error
|
||||
if copied != (counts, documents):
|
||||
logger.error("Library move: the copy's rows or documents differ")
|
||||
raise MoveFailed("verify-failed")
|
||||
|
||||
|
||||
def _hash_copy(source: str, destination: str) -> Tuple[int, str]:
|
||||
"""Copy ``source`` to ``destination`` (fsynced, owner-only, times kept)."""
|
||||
digest = hashlib.sha256()
|
||||
size = 0
|
||||
with open(source, "rb") as reader, open(destination, "wb") as writer:
|
||||
while True:
|
||||
chunk = reader.read(_COPY_CHUNK)
|
||||
if not chunk:
|
||||
break
|
||||
digest.update(chunk)
|
||||
size += len(chunk)
|
||||
writer.write(chunk)
|
||||
writer.flush()
|
||||
os.fsync(writer.fileno())
|
||||
stat = os.stat(source)
|
||||
os.utime(destination, ns=(stat.st_atime_ns, stat.st_mtime_ns))
|
||||
_private(destination)
|
||||
return size, digest.hexdigest()
|
||||
|
||||
|
||||
def _file_digest(path: str) -> Tuple[int, str]:
|
||||
digest = hashlib.sha256()
|
||||
size = 0
|
||||
with open(path, "rb") as reader:
|
||||
while True:
|
||||
chunk = reader.read(_COPY_CHUNK)
|
||||
if not chunk:
|
||||
return size, digest.hexdigest()
|
||||
digest.update(chunk)
|
||||
size += len(chunk)
|
||||
|
||||
|
||||
def _check_copies(copies) -> None:
|
||||
for source, temp, expected in copies:
|
||||
if _file_digest(temp) != expected or _file_digest(source) != expected:
|
||||
logger.error(
|
||||
f"Library move: {os.path.basename(source)} did not copy unchanged"
|
||||
)
|
||||
raise MoveFailed("copy-failed")
|
||||
|
||||
|
||||
def write_marker(
|
||||
directory: str, state: str, source_db: str, source_config: Optional[str] = None
|
||||
) -> str:
|
||||
"""Atomically write the move marker; return its path.
|
||||
|
||||
``from`` and ``config_from`` name the old files, so a restart after a
|
||||
crash between the switch and the rename can finish the rename.
|
||||
"""
|
||||
path = os.path.join(directory, ll.MARKER_NAME)
|
||||
temp = os.path.join(directory, f"{TEMP_PREFIX}{uuid.uuid4().hex}.marker")
|
||||
payload = json.dumps(
|
||||
{
|
||||
"version": 1,
|
||||
"state": state,
|
||||
"from": source_db,
|
||||
"config_from": source_config,
|
||||
"updated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
||||
},
|
||||
indent=2,
|
||||
).encode("utf-8")
|
||||
fd = os.open(temp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
|
||||
try:
|
||||
with os.fdopen(fd, "wb") as handle:
|
||||
handle.write(payload)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
_private(temp)
|
||||
os.replace(temp, path)
|
||||
except BaseException:
|
||||
_remove(temp)
|
||||
raise
|
||||
_fsync_dir(directory)
|
||||
return path
|
||||
|
||||
|
||||
def _remove(*paths) -> None:
|
||||
for path in paths:
|
||||
try:
|
||||
if path and os.path.lexists(path):
|
||||
os.remove(path)
|
||||
except OSError:
|
||||
logger.warning(f"Library move: could not remove {os.path.basename(path)}")
|
||||
|
||||
|
||||
def sweep_move_temps(directory: str) -> List[str]:
|
||||
"""Remove the temp files an interrupted move left in ``directory``."""
|
||||
removed = []
|
||||
try:
|
||||
names = os.listdir(directory)
|
||||
except OSError:
|
||||
return removed
|
||||
for name in names:
|
||||
if name.startswith(TEMP_PREFIX):
|
||||
_remove(os.path.join(directory, name))
|
||||
removed.append(name)
|
||||
return removed
|
||||
|
||||
|
||||
# ── the move ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _noop(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
|
||||
def switch_process(old_db: str, location, detached: list) -> None:
|
||||
"""Point this process at ``location``: resolver, every model, observers.
|
||||
|
||||
Adds the old file's connections, detached and still open, to
|
||||
``detached`` as soon as they are detached, so the mover retires them
|
||||
(``PromptModel.retire_detached``) whether or not the rest succeeds.
|
||||
"""
|
||||
from .execution_history import rekey_history_restore
|
||||
from .review_changes import rekey_review_changes
|
||||
|
||||
ll.switch_to(location)
|
||||
detached.extend(PromptModel.repoint_instances(old_db, location.db_path))
|
||||
rekey_review_changes(old_db, location.db_path)
|
||||
rekey_history_restore(old_db, location.db_path)
|
||||
|
||||
|
||||
def unswitch_process(old_db: str, new_db: str) -> list:
|
||||
"""Undo ``switch_process`` (any part of it); returns detached connections."""
|
||||
from .execution_history import rekey_history_restore
|
||||
from .review_changes import rekey_review_changes
|
||||
|
||||
ll.clear_override()
|
||||
detached = PromptModel.repoint_instances(new_db, old_db)
|
||||
rekey_review_changes(new_db, old_db)
|
||||
rekey_history_restore(new_db, old_db)
|
||||
return detached
|
||||
|
||||
|
||||
class _Move:
|
||||
def __init__(self, plan, *, switch, config_lock, progress, checkpoint):
|
||||
self.plan = plan
|
||||
self.switch = switch
|
||||
self.config_lock = config_lock
|
||||
self.progress = progress
|
||||
self.checkpoint = checkpoint
|
||||
self.created_dir = False
|
||||
self.temps: List[str] = []
|
||||
self.published: List[str] = []
|
||||
self.marker: Optional[str] = None
|
||||
self.moved = False
|
||||
# Retired connections still in use; awaited once the locks are free.
|
||||
self.retiring: list = []
|
||||
|
||||
def temp(self, suffix: str) -> str:
|
||||
path = os.path.join(
|
||||
self.plan.target_dir, f"{TEMP_PREFIX}{uuid.uuid4().hex}{suffix}"
|
||||
)
|
||||
self.temps.append(path)
|
||||
return path
|
||||
|
||||
def prepare_target(self) -> None:
|
||||
target = self.plan.target_dir
|
||||
self.created_dir = not os.path.isdir(target)
|
||||
ll.ensure_private_dir(target)
|
||||
sweep_move_temps(target)
|
||||
if ll.read_marker(target).get("state") == ll.MARKER_PENDING:
|
||||
# An earlier move stopped before it switched: its files are ours.
|
||||
for name in (ll.DATABASE_NAME, ll.CONFIG_NAME) + self.plan.backups:
|
||||
_remove(os.path.join(target, name))
|
||||
for suffix in ("-wal", "-shm", "-journal"):
|
||||
_remove(os.path.join(target, ll.DATABASE_NAME + suffix))
|
||||
|
||||
def copy_round(self, source):
|
||||
"""Steps 2–4 once: the verified temp copies and the snapshot version."""
|
||||
self.progress("Copying the database")
|
||||
temp_db = self.temp(".db")
|
||||
version, counts, documents = _copy_snapshot(source, temp_db)
|
||||
self.checkpoint("copied")
|
||||
self.progress("Verifying the copy")
|
||||
_verify_copy(temp_db, counts, documents)
|
||||
self.checkpoint("verified")
|
||||
copies = []
|
||||
for name in self.plan.backups:
|
||||
src = os.path.join(self.plan.source_dir, name)
|
||||
temp = self.temp(".bak")
|
||||
try:
|
||||
expected = _hash_copy(src, temp)
|
||||
except OSError as error:
|
||||
logger.error(f"Library move: copying {name} failed: {error}")
|
||||
raise MoveFailed("copy-failed") from error
|
||||
copies.append((src, temp, expected, name))
|
||||
if copies:
|
||||
self.progress(f"Copied {len(copies)} backup file(s)")
|
||||
self.checkpoint("files")
|
||||
_check_copies([(s, t, e) for s, t, e, _ in copies])
|
||||
return version, counts, temp_db, copies
|
||||
|
||||
def discard_round(self, temp_db, copies) -> None:
|
||||
_remove(temp_db, *[temp for _s, temp, _e, _n in copies])
|
||||
|
||||
def publish(self, temp_db, copies) -> List[str]:
|
||||
"""Under the locks: config.json, marker, final names, marker moved."""
|
||||
target = self.plan.target_dir
|
||||
files = []
|
||||
if self.plan.config_path:
|
||||
temp = self.temp(".config")
|
||||
try:
|
||||
expected = _hash_copy(self.plan.config_path, temp)
|
||||
except OSError as error:
|
||||
raise MoveFailed("copy-failed") from error
|
||||
_check_copies([(self.plan.config_path, temp, expected)])
|
||||
copies = copies + [(self.plan.config_path, temp, expected, ll.CONFIG_NAME)]
|
||||
self.marker = write_marker(
|
||||
target, ll.MARKER_PENDING, self.plan.db_path, self.plan.config_path
|
||||
)
|
||||
self.checkpoint("pending")
|
||||
for _src, temp, _expected, name in copies:
|
||||
final = os.path.join(target, name)
|
||||
os.replace(temp, final)
|
||||
self.published.append(final)
|
||||
files.append(name)
|
||||
final_db = os.path.join(target, ll.DATABASE_NAME)
|
||||
os.replace(temp_db, final_db)
|
||||
self.published.append(final_db)
|
||||
files.append(ll.DATABASE_NAME)
|
||||
_fsync_dir(target)
|
||||
self.checkpoint("published")
|
||||
write_marker(target, ll.MARKER_MOVED, self.plan.db_path, self.plan.config_path)
|
||||
self.moved = True
|
||||
self.checkpoint("marked")
|
||||
return files
|
||||
|
||||
def cleanup(self) -> None:
|
||||
if self.moved:
|
||||
return
|
||||
_remove(*self.temps, *self.published, self.marker)
|
||||
if self.created_dir:
|
||||
try:
|
||||
os.rmdir(self.plan.target_dir)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _unique(path: str) -> str:
|
||||
candidate, number = path, 1
|
||||
while os.path.lexists(candidate):
|
||||
candidate = f"{path}_{number}"
|
||||
number += 1
|
||||
return candidate
|
||||
|
||||
|
||||
def retire_old_files(plan: MovePlan) -> Tuple[List[str], Optional[str]]:
|
||||
"""Rename the old database and config.json; ``(renamed, problem code)``.
|
||||
|
||||
Every connection to the old file must be closed. The WAL is checkpointed
|
||||
first; when it cannot be emptied (another process still has the file
|
||||
open) nothing is renamed (``old-in-use``).
|
||||
"""
|
||||
old = plan.db_path
|
||||
folder = os.path.basename(os.path.normpath(plan.target_dir))
|
||||
try:
|
||||
with contextlib.closing(sqlite3.connect(old, timeout=1)) as conn:
|
||||
busy = conn.execute("PRAGMA wal_checkpoint(TRUNCATE)").fetchone()[0]
|
||||
except sqlite3.Error as error:
|
||||
logger.warning(
|
||||
f"Library move: the old database could not be checkpointed: {error}"
|
||||
)
|
||||
return [], "old-in-use"
|
||||
if busy or _size(old + "-wal") > 0:
|
||||
return [], "old-in-use"
|
||||
renamed = []
|
||||
try:
|
||||
for path in (old, plan.config_path):
|
||||
if not path or not os.path.isfile(path):
|
||||
continue
|
||||
new_name = _unique(path + MOVED_INFIX + folder)
|
||||
os.rename(path, new_name)
|
||||
renamed.append(os.path.basename(new_name))
|
||||
for suffix in ("-wal", "-shm"):
|
||||
if os.path.isfile(old + suffix) and _size(old + suffix) == 0:
|
||||
os.remove(old + suffix)
|
||||
except OSError as error:
|
||||
logger.warning(f"Library move: the old files could not be renamed: {error}")
|
||||
return renamed, "old-locked"
|
||||
return renamed, None
|
||||
|
||||
|
||||
def finish_interrupted_move(location) -> List[str]:
|
||||
"""Rename old files a crash left between the switch and the rename.
|
||||
|
||||
The marker of the folder in use says ``moved`` and names the old files;
|
||||
an old ``prompts.db`` (and config.json) still there under its own name is
|
||||
renamed ``*.moved-to-<folder>`` now, never deleted. Returns the new
|
||||
names (nothing when there is nothing to finish or the file is in use).
|
||||
"""
|
||||
if location.source not in (ll.SOURCE_DEFAULT, ll.SOURCE_ENVIRONMENT):
|
||||
return []
|
||||
marker = ll.read_marker(location.data_dir)
|
||||
old = marker.get("from")
|
||||
if marker.get("state") != ll.MARKER_MOVED or not isinstance(old, str):
|
||||
return []
|
||||
if not os.path.isfile(old) or os.path.basename(old) != ll.DATABASE_NAME:
|
||||
return []
|
||||
if ll.same_folder(os.path.dirname(old), location.data_dir):
|
||||
return []
|
||||
config = marker.get("config_from")
|
||||
if not (
|
||||
isinstance(config, str)
|
||||
and os.path.basename(config) == ll.CONFIG_NAME
|
||||
and os.path.isfile(config)
|
||||
and os.path.normcase(os.path.abspath(config))
|
||||
!= os.path.normcase(os.path.abspath(location.config_path))
|
||||
):
|
||||
config = None
|
||||
pending = types.SimpleNamespace(
|
||||
db_path=old, config_path=config, target_dir=location.data_dir
|
||||
)
|
||||
renamed, problem = retire_old_files(pending)
|
||||
if problem:
|
||||
logger.warning(f"Library move: the old files are still in use ({problem})")
|
||||
elif renamed:
|
||||
logger.info(f"Library move finished: renamed {', '.join(renamed)}")
|
||||
return renamed
|
||||
|
||||
|
||||
def move_library(
|
||||
plan: MovePlan,
|
||||
*,
|
||||
switch: Callable = switch_process,
|
||||
unswitch: Callable = unswitch_process,
|
||||
after_switch: Callable = _noop,
|
||||
config_lock=None,
|
||||
progress: Callable[[str], None] = _noop,
|
||||
checkpoint: Callable[[str], None] = _noop,
|
||||
attempts: int = DEFAULT_ATTEMPTS,
|
||||
) -> MoveResult:
|
||||
"""Move the library as planned (blocking; see the module docstring).
|
||||
|
||||
``switch(old_db, location, detached)`` points the process at the copy
|
||||
and adds the old connections it detached to ``detached``; if it raises,
|
||||
``unswitch(old_db, new_db)`` undoes it while both files are still
|
||||
locked, every detached connection is closed, the copy is removed and
|
||||
the move fails (``switch-failed``). ``after_switch()`` runs
|
||||
once the switch is final; it is best effort (logged, never fails the move).
|
||||
"""
|
||||
short = space_problem(plan)
|
||||
if short is not None:
|
||||
raise MoveFailed("insufficient-space")
|
||||
lock = config_lock if config_lock is not None else contextlib.nullcontext()
|
||||
move = _Move(
|
||||
plan,
|
||||
switch=switch,
|
||||
config_lock=lock,
|
||||
progress=progress,
|
||||
checkpoint=checkpoint,
|
||||
)
|
||||
move.unswitch = unswitch
|
||||
location = plan.target_location
|
||||
try:
|
||||
move.prepare_target()
|
||||
files, counts = _copy_and_switch(move, location, attempts)
|
||||
except Exception:
|
||||
move.cleanup()
|
||||
raise
|
||||
try:
|
||||
after_switch()
|
||||
except Exception:
|
||||
logger.exception("Library move: a follow-up after the switch failed")
|
||||
from .execution_history import notify_history_restore
|
||||
|
||||
notify_history_restore(location.db_path)
|
||||
progress("Renaming the old files")
|
||||
renamed, problem = retire_old_files(plan)
|
||||
checkpoint("retired")
|
||||
return MoveResult(location, counts, files, renamed, problem)
|
||||
|
||||
|
||||
def _switch_guarded(move: _Move, location) -> list:
|
||||
"""Switch while the copy is locked too; undo it all if any part fails.
|
||||
|
||||
The copy's write lock keeps every repointed thread from writing to it
|
||||
until the switch is whole, so undoing a partial switch loses nothing.
|
||||
Returns the old file's detached connections (still open). On failure
|
||||
every detached connection is retired; those still in use are added to
|
||||
``move.retiring`` for the caller to await once the locks are released.
|
||||
"""
|
||||
plan = move.plan
|
||||
guard = sqlite3.connect(
|
||||
location.db_path, timeout=BUSY_TIMEOUT_SECONDS, isolation_level=None
|
||||
)
|
||||
try:
|
||||
guard.execute("BEGIN IMMEDIATE")
|
||||
detached: list = []
|
||||
try:
|
||||
move.switch(plan.db_path, location, detached)
|
||||
except Exception as error:
|
||||
logger.exception("Library move: switching failed; undoing it")
|
||||
move.moved = False # the copy and its marker are removed
|
||||
# Retire, never wait, here: a connection still in use may be
|
||||
# waiting on the very locks this thread holds.
|
||||
try:
|
||||
undone = move.unswitch(plan.db_path, location.db_path)
|
||||
except Exception:
|
||||
logger.exception("Library move: undoing the switch failed")
|
||||
else:
|
||||
_retire(move, undone)
|
||||
# The old file's connections a partial switch detached: the
|
||||
# models reopen it on next use, so close them now, not at GC.
|
||||
_retire(move, detached)
|
||||
raise MoveFailed("switch-failed") from error
|
||||
move.checkpoint("switched")
|
||||
return detached
|
||||
finally:
|
||||
if guard.in_transaction:
|
||||
guard.execute("ROLLBACK")
|
||||
guard.close()
|
||||
|
||||
|
||||
def _copy_and_switch(move: _Move, location, attempts: int):
|
||||
"""Copy, then switch under both locks; retired connections are awaited.
|
||||
|
||||
A connection that was inside a statement when it was retired is closed
|
||||
by its own thread once the statement ends, which may need the locks
|
||||
released first; the old files are renamed (and a failed copy removed)
|
||||
only after that, since Windows refuses to rename or delete an open file.
|
||||
"""
|
||||
try:
|
||||
return _copy_and_switch_locked(move, location, attempts)
|
||||
finally:
|
||||
try:
|
||||
PromptModel.await_detached(move.retiring, logger)
|
||||
except Exception:
|
||||
logger.exception("Library move: waiting for old connections failed")
|
||||
|
||||
|
||||
def _retire(move: _Move, connections: list) -> None:
|
||||
"""Retire ``connections``, adding those still in use to ``move.retiring``.
|
||||
|
||||
Never raises: a switch that completed must not be reported as failed
|
||||
because one of the old connections could not be retired (logged; the
|
||||
rename then reports the old files as still in use).
|
||||
"""
|
||||
try:
|
||||
move.retiring.extend(PromptModel.retire_detached(connections, logger))
|
||||
except Exception:
|
||||
logger.exception("Library move: retiring the old connections failed")
|
||||
|
||||
|
||||
def _copy_and_switch_locked(move: _Move, location, attempts: int):
|
||||
plan = move.plan
|
||||
for attempt in range(1, attempts + 1):
|
||||
source = sqlite3.connect(
|
||||
plan.db_path, timeout=BUSY_TIMEOUT_SECONDS, isolation_level=None
|
||||
)
|
||||
try:
|
||||
source.execute(f"PRAGMA busy_timeout = {BUSY_TIMEOUT_SECONDS * 1000}")
|
||||
version, counts, temp_db, copies = move.copy_round(source)
|
||||
source.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
if source.execute("PRAGMA data_version").fetchone()[0] != version:
|
||||
move.progress("The library changed during the copy; copying again")
|
||||
move.discard_round(temp_db, copies)
|
||||
continue
|
||||
move.checkpoint("locked")
|
||||
with move.config_lock:
|
||||
files = move.publish(temp_db, copies)
|
||||
detached = _switch_guarded(move, location)
|
||||
# Still holding the old file's lock: a write that was waiting
|
||||
# on it is interrupted and cannot commit, so it fails (and
|
||||
# runs again on the new library) instead of landing in the
|
||||
# old file. Its own thread closes the connection.
|
||||
_retire(move, detached)
|
||||
finally:
|
||||
if source.in_transaction:
|
||||
source.execute("ROLLBACK")
|
||||
return files, counts
|
||||
finally:
|
||||
source.close()
|
||||
raise MoveFailed("library-busy")
|
||||
+2
-18
@@ -42,15 +42,12 @@ from pathlib import Path
|
||||
from .execution_contexts import LIVE_HINT_CTES
|
||||
from .health import BATCH_ROWS, link_states, readable_roots, readonly_uri
|
||||
from .pairs import PairDependencyError
|
||||
from .workflow_documents import documents_of_rows, release_documents
|
||||
|
||||
MAINTENANCE_OPERATIONS = (
|
||||
"duplicate_prompts",
|
||||
"optimize",
|
||||
"unused_components",
|
||||
"missing_links",
|
||||
# Applied as its own batched job (database/workflow_storage.py).
|
||||
"workflow_storage",
|
||||
)
|
||||
MAX_DETAILS = 50 # detail rows a plan lists
|
||||
MAX_RESULT_ROWS = 200 # changed records an apply describes (the job lists 200)
|
||||
@@ -263,24 +260,16 @@ def _orphaned_ids(conn):
|
||||
|
||||
|
||||
def _delete_ids(conn, table, ids):
|
||||
"""Delete ``table`` rows by id in batches; returns the rows removed.
|
||||
|
||||
The workflow documents only those rows referenced go with them (a
|
||||
prompt's images go by cascade), in the same transaction.
|
||||
"""
|
||||
"""Delete ``table`` rows by id in batches; returns the rows removed."""
|
||||
if table not in ("prompts", "generated_images"):
|
||||
raise ValueError(f"Unsupported table {table!r}")
|
||||
removed = 0
|
||||
for batch in _batches(ids):
|
||||
documents = documents_of_rows(
|
||||
conn, "id" if table == "generated_images" else "prompt_id", batch
|
||||
)
|
||||
cursor = conn.execute(
|
||||
f"DELETE FROM {table} WHERE id IN ({_marks(batch)})", # nosec B608
|
||||
batch,
|
||||
)
|
||||
removed += cursor.rowcount
|
||||
release_documents(conn, documents)
|
||||
return removed
|
||||
|
||||
|
||||
@@ -390,13 +379,10 @@ class MaintenanceOperations:
|
||||
conn, primary_id, merged_metadata, strict=strict
|
||||
)
|
||||
removed = 0
|
||||
# Images the lenient merge could not move go with their prompt.
|
||||
documents = documents_of_rows(conn, "prompt_id", duplicate_ids)
|
||||
for duplicate_id in duplicate_ids:
|
||||
conn.execute("DELETE FROM prompts WHERE id = ?", (duplicate_id,))
|
||||
removed += 1
|
||||
self.logger.debug(f"Removed duplicate prompt {duplicate_id}")
|
||||
release_documents(conn, documents)
|
||||
return removed, moved
|
||||
|
||||
@staticmethod
|
||||
@@ -422,8 +408,6 @@ class MaintenanceOperations:
|
||||
roots = [str(root) for root in roots]
|
||||
if operation == "optimize":
|
||||
return self._plan_optimize()
|
||||
if operation == "workflow_storage":
|
||||
return self._plan_workflow_storage()
|
||||
if operation == "missing_links":
|
||||
rows = _read_snapshot(self.model.db_path, self._link_rows)
|
||||
return self._plan_missing_links(rows, roots)
|
||||
@@ -672,7 +656,7 @@ class MaintenanceOperations:
|
||||
records changed and up to ``MAX_RESULT_ROWS`` descriptions of them
|
||||
(``{prompt_id[, kept_id]}`` or ``{image_path}``).
|
||||
"""
|
||||
if operation not in MAINTENANCE_OPERATIONS or operation == "workflow_storage":
|
||||
if operation not in MAINTENANCE_OPERATIONS:
|
||||
raise ValueError(f"Unknown maintenance operation {operation!r}")
|
||||
if operation == "optimize":
|
||||
# VACUUM cannot run in a transaction; the plan binds the database.
|
||||
|
||||
@@ -1,353 +0,0 @@
|
||||
"""A SQLite connection that is never closed under the thread using it.
|
||||
|
||||
``PromptModel`` gives each thread its own connection, but Move library and
|
||||
``close_all`` must retire every thread's connection from one thread. Closing
|
||||
a ``sqlite3.Connection`` while its owner is inside a statement (Python runs
|
||||
``sqlite3_step`` without the GIL) frees memory the owner then reads: on
|
||||
Windows the process dies with "access violation", elsewhere the owner sees
|
||||
garbage such as a ``MemoryError``.
|
||||
|
||||
``ManagedConnection`` counts the calls its owner is inside. ``retire()`` is
|
||||
safe from any thread and never raises: an idle connection is closed at once;
|
||||
one in use is closed by its owner as the call ends. A retired connection
|
||||
refuses new work (cursors, statements, commits) with the error a closed
|
||||
connection raises.
|
||||
|
||||
What retiring a connection in use does and does not do:
|
||||
|
||||
- It calls ``sqlite3_interrupt`` (thread-safe). That ends a running
|
||||
statement at its next step, but it does not wake a statement that is
|
||||
waiting for another connection's lock: that one waits until the lock is
|
||||
released (or the busy timeout ends it), and only then is it interrupted.
|
||||
An interrupt that arrives between two statements, for example between the
|
||||
implicit ``BEGIN`` and the DML or between ``executemany`` rows, is lost.
|
||||
- Writes are kept out of a file the process has switched away from because
|
||||
a retired connection never commits: legacy-mode DML runs inside a
|
||||
transaction that only ``commit()`` (refused) or ``close()`` (rolls back)
|
||||
ends. A statement that commits by itself is not covered: DDL or a PRAGMA
|
||||
write outside a transaction, or an explicit ``COMMIT`` passed to
|
||||
``execute``/``executescript`` (``executescript`` also commits first, in C,
|
||||
on Python before 3.12). Managed connections must never be used that way
|
||||
for library writes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
from typing import Iterable, Iterator, List, Optional
|
||||
|
||||
import sqlite3
|
||||
|
||||
CLOSED_MESSAGE = "Cannot operate on a closed database."
|
||||
# Rows fetched per step when a cursor is iterated: the connection is marked
|
||||
# in use once per chunk instead of once per row (iteration costs about what
|
||||
# a plain cursor costs). A loop that stops early leaves up to this many rows
|
||||
# fetched; ``fetchone``/``fetchmany`` after a partial loop skip them.
|
||||
ITER_CHUNK = 128
|
||||
|
||||
|
||||
class ManagedCursor(sqlite3.Cursor):
|
||||
"""A cursor whose calls into SQLite count as its connection being in use."""
|
||||
|
||||
def execute(self, sql, parameters=(), /):
|
||||
conn = self.connection
|
||||
conn._enter()
|
||||
try:
|
||||
return super().execute(sql, parameters)
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
def executemany(self, sql, seq_of_parameters, /):
|
||||
conn = self.connection
|
||||
conn._enter()
|
||||
try:
|
||||
return super().executemany(sql, seq_of_parameters)
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
def executescript(self, sql_script, /):
|
||||
conn = self.connection
|
||||
conn._enter()
|
||||
try:
|
||||
return super().executescript(sql_script)
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
def fetchone(self):
|
||||
conn = self.connection
|
||||
conn._enter()
|
||||
try:
|
||||
return super().fetchone()
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
def fetchmany(self, *args, **kwargs):
|
||||
conn = self.connection
|
||||
conn._enter()
|
||||
try:
|
||||
return super().fetchmany(*args, **kwargs)
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
def fetchall(self):
|
||||
conn = self.connection
|
||||
conn._enter()
|
||||
try:
|
||||
return super().fetchall()
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
def __iter__(self) -> Iterator:
|
||||
return self._chunks()
|
||||
|
||||
def _chunks(self) -> Iterator:
|
||||
while True:
|
||||
rows = self.fetchmany(ITER_CHUNK)
|
||||
if not rows:
|
||||
return
|
||||
yield from rows
|
||||
|
||||
def __next__(self):
|
||||
conn = self.connection
|
||||
conn._enter()
|
||||
try:
|
||||
return super().__next__()
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
def close(self):
|
||||
conn = self.connection
|
||||
try:
|
||||
conn._enter()
|
||||
except sqlite3.ProgrammingError:
|
||||
return super().close() # the connection is closed or closing
|
||||
try:
|
||||
return super().close()
|
||||
finally:
|
||||
conn._leave()
|
||||
|
||||
|
||||
class ManagedConnection(sqlite3.Connection):
|
||||
"""A connection only its owner closes while a call into it is running."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
# Guards _depth, _retired, _closing and _cursors. Re-entrant so that
|
||||
# code running in the closing thread while it holds the lock (a
|
||||
# cursor() during the cursor-set copy) is refused instead of hanging.
|
||||
self._use_lock = threading.RLock()
|
||||
self._depth = 0
|
||||
self._retired = False
|
||||
self._closing = False
|
||||
self._released = threading.Event()
|
||||
# Python before 3.12 keeps the file open after close() while a cursor
|
||||
# still holds a statement; these are closed first (see _close_now).
|
||||
self._cursors: "weakref.WeakSet[sqlite3.Cursor]" = weakref.WeakSet()
|
||||
self.close_logger: Optional[logging.Logger] = None
|
||||
|
||||
# ── in use ───────────────────────────────────────────────────────
|
||||
|
||||
def _enter(self) -> None:
|
||||
with self._use_lock:
|
||||
if self._retired:
|
||||
raise sqlite3.ProgrammingError(CLOSED_MESSAGE)
|
||||
self._depth += 1
|
||||
|
||||
def _leave(self) -> None:
|
||||
with self._use_lock:
|
||||
self._depth -= 1
|
||||
close_now = self._retired and self._depth == 0
|
||||
if close_now:
|
||||
self._close_now()
|
||||
|
||||
@contextlib.contextmanager
|
||||
def using(self):
|
||||
"""Mark the connection in use for a call this class cannot wrap.
|
||||
|
||||
For example as the *target* of another connection's ``backup()``.
|
||||
"""
|
||||
self._enter()
|
||||
try:
|
||||
yield self
|
||||
finally:
|
||||
self._leave()
|
||||
|
||||
def cursor(self, factory=None):
|
||||
with self._use_lock:
|
||||
if self._retired:
|
||||
raise sqlite3.ProgrammingError(CLOSED_MESSAGE)
|
||||
cursor = super().cursor(factory or ManagedCursor)
|
||||
self._cursors.add(cursor)
|
||||
return cursor
|
||||
|
||||
def execute(self, sql, parameters=(), /):
|
||||
return self.cursor().execute(sql, parameters)
|
||||
|
||||
def executemany(self, sql, seq_of_parameters, /):
|
||||
return self.cursor().executemany(sql, seq_of_parameters)
|
||||
|
||||
def executescript(self, sql_script, /):
|
||||
return self.cursor().executescript(sql_script)
|
||||
|
||||
def commit(self):
|
||||
self._enter()
|
||||
try:
|
||||
return super().commit()
|
||||
finally:
|
||||
self._leave()
|
||||
|
||||
def rollback(self):
|
||||
self._enter() # a closed connection raises, as sqlite3's does
|
||||
try:
|
||||
return super().rollback()
|
||||
finally:
|
||||
self._leave()
|
||||
|
||||
def backup(self, target, **kwargs):
|
||||
self._enter()
|
||||
try:
|
||||
return super().backup(target, **kwargs)
|
||||
finally:
|
||||
self._leave()
|
||||
|
||||
def create_function(self, *args, **kwargs):
|
||||
self._enter()
|
||||
try:
|
||||
return super().create_function(*args, **kwargs)
|
||||
finally:
|
||||
self._leave()
|
||||
|
||||
@property
|
||||
def in_transaction(self) -> bool:
|
||||
self._enter()
|
||||
try:
|
||||
return super().in_transaction
|
||||
finally:
|
||||
self._leave()
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
# The C implementation calls its own commit/rollback, not these.
|
||||
if exc_type is None:
|
||||
try:
|
||||
self.commit()
|
||||
except BaseException:
|
||||
with contextlib.suppress(sqlite3.Error):
|
||||
self.rollback()
|
||||
raise
|
||||
else:
|
||||
self.rollback()
|
||||
return False
|
||||
|
||||
# ── closing ──────────────────────────────────────────────────────
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close now if idle, else as soon as the running call ends.
|
||||
|
||||
Unlike ``sqlite3.Connection.close``, a call from another thread while
|
||||
the owner is inside a call returns before the file is closed. A caller
|
||||
that must know the file is closed (before renaming or deleting it)
|
||||
waits for ``wait_closed`` (``PromptModel.close_detached`` does).
|
||||
"""
|
||||
self.retire()
|
||||
|
||||
def retire(self) -> bool:
|
||||
"""Refuse new work and close; True when the connection is closed already.
|
||||
|
||||
Safe from any thread, and never raises: failures are logged.
|
||||
"""
|
||||
try:
|
||||
with self._use_lock:
|
||||
if not self._retired:
|
||||
self._retired = True
|
||||
if self._depth:
|
||||
# Ends the running statement (see the module docstring
|
||||
# for what it does not end). The owner cannot start
|
||||
# closing the connection while this lock is held.
|
||||
self.interrupt()
|
||||
return False
|
||||
elif self._depth:
|
||||
return False
|
||||
else:
|
||||
return self._released.wait(0)
|
||||
except Exception:
|
||||
self._log_exception("Error retiring database connection")
|
||||
return self._released.wait(0)
|
||||
self._close_now()
|
||||
return True
|
||||
|
||||
def wait_closed(self, timeout: float) -> bool:
|
||||
return self._released.wait(timeout)
|
||||
|
||||
def _close_now(self) -> None:
|
||||
# Only ever called with no call into the connection running, and a
|
||||
# retired connection refuses new calls and cursors: its cursors are
|
||||
# idle, and the set is copied under the lock that guards it.
|
||||
with self._use_lock:
|
||||
if self._closing:
|
||||
return # another thread closes it and sets _released
|
||||
self._closing = True
|
||||
cursors = list(self._cursors)
|
||||
try:
|
||||
for cursor in cursors:
|
||||
with contextlib.suppress(sqlite3.Error):
|
||||
sqlite3.Cursor.close(cursor)
|
||||
super().close()
|
||||
except Exception:
|
||||
self._log_exception("Error closing database connection")
|
||||
finally:
|
||||
self._released.set()
|
||||
|
||||
def _log_exception(self, message: str) -> None:
|
||||
if self.close_logger is not None:
|
||||
self.close_logger.exception(message)
|
||||
|
||||
|
||||
def retire_connections(
|
||||
connections: Iterable[sqlite3.Connection], logger=None
|
||||
) -> List[ManagedConnection]:
|
||||
"""Retire every connection; return those whose owner has yet to close them.
|
||||
|
||||
Safe while holding a lock the owners wait on: nothing here waits for them.
|
||||
One connection failing to retire is logged and never skips the rest.
|
||||
"""
|
||||
pending = []
|
||||
for conn in connections:
|
||||
try:
|
||||
if isinstance(conn, ManagedConnection):
|
||||
if logger is not None:
|
||||
conn.close_logger = logger
|
||||
if not conn.retire():
|
||||
pending.append(conn)
|
||||
else:
|
||||
conn.close()
|
||||
except Exception:
|
||||
if logger is not None:
|
||||
logger.exception("Error closing database connection")
|
||||
return pending
|
||||
|
||||
|
||||
def await_closed(
|
||||
connections: Iterable[ManagedConnection], timeout: float, logger=None
|
||||
) -> bool:
|
||||
"""Wait until the owners closed ``connections`` (retired); False on timeout.
|
||||
|
||||
A connection still open after ``timeout`` (its owner is blocked, for
|
||||
example on another process's lock after a lost interrupt) is logged and
|
||||
left to close when its call ends; a Move library then reports the old
|
||||
files as still in use instead of renaming them.
|
||||
"""
|
||||
deadline = time.monotonic() + timeout
|
||||
still_open = 0
|
||||
for conn in connections:
|
||||
if not conn.wait_closed(max(0.0, deadline - time.monotonic())):
|
||||
still_open += 1
|
||||
if still_open and logger is not None:
|
||||
logger.warning(
|
||||
f"{still_open} database connection(s) still in use after "
|
||||
f"{timeout:g}s; they close when their statement ends"
|
||||
)
|
||||
return not still_open
|
||||
+81
-183
@@ -9,17 +9,12 @@ import math
|
||||
import sqlite3
|
||||
import os
|
||||
import threading
|
||||
import tempfile
|
||||
import weakref
|
||||
from contextlib import closing
|
||||
from pathlib import Path
|
||||
from typing import Optional, Set, Tuple
|
||||
|
||||
from .managed_connection import (
|
||||
ManagedConnection,
|
||||
await_closed,
|
||||
retire_connections,
|
||||
)
|
||||
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 .annotation_schema import annotation_schema_error, migrate_annotation_state
|
||||
@@ -29,16 +24,10 @@ from .review_schema import (
|
||||
review_without_annotations_error,
|
||||
)
|
||||
from .saved_views import migrate_saved_views, saved_views_schema_error
|
||||
from .workflow_documents import (
|
||||
DOCUMENT_ID_COLUMNS,
|
||||
migrate_workflow_documents,
|
||||
workflow_documents_schema_error,
|
||||
)
|
||||
|
||||
|
||||
from .execution_history_schema import (
|
||||
history_schema_error,
|
||||
migrate_execution_history,
|
||||
rotate_history_epoch,
|
||||
)
|
||||
|
||||
# Import logging system
|
||||
@@ -61,19 +50,6 @@ MAX_PAIR_PROVENANCE_LENGTH = 4_000_000
|
||||
MAX_PAIR_PROVENANCE_DEPTH = 64
|
||||
|
||||
|
||||
def _refuse_rebuild_with_documents(columns) -> None:
|
||||
"""The legacy generated_images rebuilds copy a fixed column list.
|
||||
|
||||
A file that needs one (``prompt_id`` TEXT, no ``file_path``) but already
|
||||
has the workflow document ids would lose them, so it is refused.
|
||||
"""
|
||||
if any(column in columns for column in DOCUMENT_ID_COLUMNS):
|
||||
raise sqlite3.DatabaseError(
|
||||
"This database needs an old image table rebuild but already has "
|
||||
"workflow document references; the rebuild would drop them"
|
||||
)
|
||||
|
||||
|
||||
def pair_provenance_error(encoded: object) -> Optional[str]:
|
||||
"""Validate the same bounded JSON object for live writes and restores."""
|
||||
if not isinstance(encoded, str):
|
||||
@@ -340,8 +316,6 @@ class PromptModel:
|
||||
raise sqlite3.DatabaseError(review_error)
|
||||
migrate_annotation_state(conn)
|
||||
migrate_review_state(conn)
|
||||
# Last: the older generated_images rebuilds run only before file_path.
|
||||
migrate_workflow_documents(conn)
|
||||
|
||||
def _migrate_prompt_pairs(self, conn: sqlite3.Connection) -> None:
|
||||
conn.execute("SAVEPOINT prompt_pair_schema")
|
||||
@@ -554,10 +528,7 @@ class PromptModel:
|
||||
"WHERE pp.positive_prompt_id IS NOT gi.prompt_id LIMIT 1"
|
||||
).fetchone():
|
||||
return "Image pair association has a different positive prompt owner"
|
||||
# fetchone, not iteration: a managed connection's cursor prefetches
|
||||
# ITER_CHUNK rows when iterated, and each provenance may be ~4 MB.
|
||||
cursor = conn.execute("SELECT provenance FROM image_prompt_pairs")
|
||||
while (row := cursor.fetchone()) is not None:
|
||||
for row in conn.execute("SELECT provenance FROM image_prompt_pairs"):
|
||||
error = pair_provenance_error(row[0])
|
||||
if error:
|
||||
return error
|
||||
@@ -616,42 +587,19 @@ class PromptModel:
|
||||
if conn is not None and self._local.generation == self._generation:
|
||||
return conn
|
||||
|
||||
while True:
|
||||
# The path and generation are read together, and the new
|
||||
# connection is kept only if no Move library repointed the model
|
||||
# meanwhile: otherwise it would write to the old file.
|
||||
with self._conn_lock:
|
||||
path, generation = self.db_path, self._generation
|
||||
# check_same_thread=False only so close_all and Move library can
|
||||
# retire other threads' connections; each connection is still
|
||||
# used by one thread, and ManagedConnection never closes it under
|
||||
# a running statement (see managed_connection).
|
||||
conn = sqlite3.connect(
|
||||
path,
|
||||
check_same_thread=False,
|
||||
timeout=BUSY_TIMEOUT_SECONDS,
|
||||
factory=ManagedConnection,
|
||||
)
|
||||
try:
|
||||
conn.row_factory = sqlite3.Row
|
||||
configure_connection(conn)
|
||||
except BaseException:
|
||||
conn.close()
|
||||
raise
|
||||
with self._conn_lock:
|
||||
if generation == self._generation:
|
||||
self._connections.add(conn)
|
||||
break
|
||||
conn.close()
|
||||
# check_same_thread=False only so close_all can clean up other
|
||||
# threads' connections; each connection is still used by one thread.
|
||||
conn = sqlite3.connect(
|
||||
self.db_path, check_same_thread=False, timeout=BUSY_TIMEOUT_SECONDS
|
||||
)
|
||||
conn.row_factory = sqlite3.Row
|
||||
configure_connection(conn)
|
||||
self._local.conn = conn
|
||||
self._local.generation = generation
|
||||
self._local.generation = self._generation
|
||||
with self._conn_lock:
|
||||
self._connections.add(conn)
|
||||
return conn
|
||||
|
||||
@property
|
||||
def generation(self) -> int:
|
||||
"""Bumped whenever every connection is closed or the model repointed."""
|
||||
return self._generation
|
||||
|
||||
@classmethod
|
||||
def close_all_instances(cls) -> None:
|
||||
"""Close every connection of every live model in this process."""
|
||||
@@ -660,76 +608,22 @@ class PromptModel:
|
||||
for model in models:
|
||||
model.close_all()
|
||||
|
||||
@classmethod
|
||||
def repoint_instances(cls, old_path: str, new_path: str) -> list:
|
||||
"""Point every model of ``old_path`` at ``new_path`` (Move library).
|
||||
|
||||
Every thread opens the new file on its next ``get_connection``. The
|
||||
old connections are detached, not closed, and returned: the caller
|
||||
retires them (``retire_detached``) while it still holds the old
|
||||
file's write lock. A write waiting on that lock is then interrupted
|
||||
and can never commit, so it fails (and, through
|
||||
``PromptDatabase.survive_library_switch``, runs again on the new file).
|
||||
"""
|
||||
old_key = os.path.normcase(os.path.realpath(old_path))
|
||||
with cls._instances_lock:
|
||||
models = list(cls._instances)
|
||||
with cls._init_lock:
|
||||
cls._initialized_paths.add(os.path.normcase(os.path.realpath(new_path)))
|
||||
detached = []
|
||||
for model in models:
|
||||
if model._schema_key() == old_key:
|
||||
detached.extend(model._detach_connections(new_path))
|
||||
return detached
|
||||
|
||||
@staticmethod
|
||||
def retire_detached(connections, logger=None) -> list:
|
||||
"""Retire detached connections without waiting for their owners.
|
||||
|
||||
Idle ones close now. One whose owner thread is inside a statement is
|
||||
never closed from here (that frees memory the owner is using: an
|
||||
access violation on Windows); it is interrupted, refuses to commit and
|
||||
is closed by its owner as the statement ends. Returns those still
|
||||
open, for ``await_detached`` once no lock the owners wait on is held.
|
||||
"""
|
||||
return retire_connections(connections, logger)
|
||||
|
||||
@staticmethod
|
||||
def await_detached(connections, logger=None) -> bool:
|
||||
"""Wait until retired connections are closed; False if some are not."""
|
||||
return await_closed(connections, BUSY_TIMEOUT_SECONDS, logger)
|
||||
|
||||
@staticmethod
|
||||
def close_detached(connections, logger=None) -> bool:
|
||||
"""Retire connections ``repoint_instances`` detached and await them.
|
||||
|
||||
Do not call it while holding a lock their owners may wait on; retire
|
||||
them there (``retire_detached``) and await them after releasing it.
|
||||
"""
|
||||
pending = retire_connections(connections, logger)
|
||||
return await_closed(pending, BUSY_TIMEOUT_SECONDS, logger)
|
||||
|
||||
def close_all(self) -> None:
|
||||
"""Close every thread's connection; they reopen lazily on next use.
|
||||
|
||||
Needed before deleting the database file (Windows keeps open files
|
||||
locked) and by tests that ran queries through executor threads.
|
||||
"""
|
||||
self.close_detached(self._detach_connections(), self.logger)
|
||||
|
||||
def _detach_connections(self, new_path: Optional[str] = None) -> list:
|
||||
"""Forget every connection (optionally pointing the model elsewhere).
|
||||
|
||||
Returns them for the caller to close; threads reopen on next use.
|
||||
"""
|
||||
with self._conn_lock:
|
||||
if new_path is not None:
|
||||
self.db_path = new_path
|
||||
connections = list(self._connections)
|
||||
self._connections.clear()
|
||||
self._generation += 1
|
||||
self._local.conn = None
|
||||
return connections
|
||||
for conn in connections:
|
||||
try:
|
||||
conn.close()
|
||||
except sqlite3.Error as e:
|
||||
self.logger.warning(f"Error closing database connection: {e}")
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the calling thread's connection, if it has one.
|
||||
@@ -743,7 +637,10 @@ class PromptModel:
|
||||
self._local.conn = None
|
||||
with self._conn_lock:
|
||||
self._connections.discard(conn)
|
||||
self.close_detached([conn], self.logger)
|
||||
try:
|
||||
conn.close()
|
||||
except sqlite3.Error as e:
|
||||
self.logger.warning(f"Error closing database connection: {e}")
|
||||
|
||||
def _migrate_workflow_name_removal(self, conn: sqlite3.Connection) -> None:
|
||||
"""
|
||||
@@ -814,7 +711,6 @@ class PromptModel:
|
||||
columns = {column[1]: column[2] for column in cursor.fetchall()}
|
||||
if columns.get("prompt_id") != "TEXT":
|
||||
return
|
||||
_refuse_rebuild_with_documents(columns)
|
||||
|
||||
self.logger.info("Migrating foreign key types: prompt_id TEXT -> INTEGER")
|
||||
with self._table_rebuild(conn):
|
||||
@@ -980,7 +876,6 @@ class PromptModel:
|
||||
}
|
||||
if "file_path" in columns:
|
||||
return
|
||||
_refuse_rebuild_with_documents(columns)
|
||||
|
||||
self.logger.info("Migrating database: image uniqueness by full path")
|
||||
with self._table_rebuild(conn):
|
||||
@@ -1134,7 +1029,7 @@ class PromptModel:
|
||||
|
||||
@staticmethod
|
||||
def verify_database_file(
|
||||
path: str, *, require_library_state: bool = False, documents: bool = True
|
||||
path: str, *, require_library_state: bool = False
|
||||
) -> Tuple[bool, str]:
|
||||
"""
|
||||
Check that a file is a healthy prompts database before restoring it.
|
||||
@@ -1148,9 +1043,6 @@ class PromptModel:
|
||||
database has (Library state, image favorites, saved views,
|
||||
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
|
||||
damaged documents on the stage (``database/restore_stage.py``).
|
||||
|
||||
Returns:
|
||||
(ok, reason): reason is "ok" when the file is acceptable, otherwise
|
||||
@@ -1225,11 +1117,6 @@ class PromptModel:
|
||||
)
|
||||
if review_error:
|
||||
return False, review_error
|
||||
documents_error = workflow_documents_schema_error(
|
||||
conn, allow_missing=not require_library_state, documents=documents
|
||||
)
|
||||
if documents_error:
|
||||
return False, documents_error
|
||||
return True, "ok"
|
||||
except sqlite3.Error as e:
|
||||
return False, f"Invalid SQLite database: {e}"
|
||||
@@ -1252,51 +1139,21 @@ class PromptModel:
|
||||
Path(src_path).resolve().as_uri() + "?mode=ro", uri=True
|
||||
)
|
||||
try:
|
||||
target = self.get_connection()
|
||||
# The target is written under backup's step too: mark it in use
|
||||
# so a concurrent close_all retires it instead of closing it.
|
||||
with target.using():
|
||||
source.backup(target)
|
||||
source.backup(self.get_connection())
|
||||
finally:
|
||||
source.close()
|
||||
|
||||
def stage_restore(self, src_path: str):
|
||||
"""Stage ``src_path`` beside this database for a restore (blocking).
|
||||
|
||||
Returns a ``StagedRestore`` (``database/restore_stage.py``): a
|
||||
migrated, repaired and verified copy in the library's folder, its
|
||||
digest and what it holds. Neither the live database nor the source
|
||||
is touched. Raises ValueError when either does not verify.
|
||||
"""
|
||||
return restore_stage.stage_restore(PromptModel, self.db_path, src_path)
|
||||
|
||||
def publish_staged(self, staged) -> str:
|
||||
"""Replace the live database with ``staged``; return the safety backup.
|
||||
|
||||
See ``restore_stage.publish_staged``: the staged bytes are rechecked,
|
||||
the current database is kept as ``<db_path>.backup_<stamp>`` and the
|
||||
WAL is truncated afterwards.
|
||||
"""
|
||||
backup_path = restore_stage.publish_staged(self, staged)
|
||||
self.logger.info("Database restored from a staged copy")
|
||||
from .execution_history import notify_history_restore
|
||||
|
||||
notify_history_restore(self.db_path)
|
||||
return backup_path
|
||||
|
||||
def restore_from_file(self, src_path: str) -> str:
|
||||
"""
|
||||
Replace the live database with the file at src_path.
|
||||
|
||||
``stage_restore`` then ``publish_staged``: a disposable copy of the
|
||||
source is migrated, repaired and verified in the library's folder
|
||||
first. Only then is the current database backed up next to itself as
|
||||
``<db_path>.backup_<YYYYmmdd_HHMMSS_ffffff>`` (WAL content included),
|
||||
and the staged pages replace the live pages through SQLite's backup
|
||||
API. Staging failures never replace live content or modify the
|
||||
source. An online replacement failure leaves the live database intact
|
||||
and propagates; the completed safety backup remains available. The
|
||||
stage is removed either way.
|
||||
A disposable SQLite backup of the source is migrated and verified first.
|
||||
Only then is the current database backed up next to itself as
|
||||
``<db_path>.backup_<YYYYmmdd_HHMMSS>`` (WAL content included), and the
|
||||
staged pages replace the live pages through SQLite's backup API.
|
||||
Migration failures never replace live content or modify the source.
|
||||
An online replacement failure leaves the live database intact and
|
||||
propagates; the completed safety backup remains available.
|
||||
|
||||
Args:
|
||||
src_path: SQLite file to restore
|
||||
@@ -1306,14 +1163,55 @@ class PromptModel:
|
||||
there was no database to back up.
|
||||
|
||||
Raises:
|
||||
ValueError: If the source or its migrated copy does not verify
|
||||
SafetyBackupFailed: If the previous database could not be backed up
|
||||
RuntimeError: If the previous database could not be backed up
|
||||
sqlite3.OperationalError: If SQLite refuses online replacement
|
||||
"""
|
||||
staged = self.stage_restore(src_path)
|
||||
try:
|
||||
backup_path = self.publish_staged(staged)
|
||||
finally:
|
||||
staged.discard()
|
||||
self.logger.info(f"Database restored from {src_path}")
|
||||
return backup_path
|
||||
ok, reason = self.verify_database_file(src_path)
|
||||
if not ok:
|
||||
raise ValueError(reason)
|
||||
with tempfile.TemporaryDirectory(prefix="prompt-manager-restore-") as directory:
|
||||
staged_path = str(Path(directory) / "staged.db")
|
||||
source = sqlite3.connect(
|
||||
Path(src_path).resolve().as_uri() + "?mode=ro", uri=True
|
||||
)
|
||||
try:
|
||||
with closing(sqlite3.connect(staged_path)) as destination:
|
||||
source.backup(destination)
|
||||
finally:
|
||||
source.close()
|
||||
staged = None
|
||||
try:
|
||||
staged = PromptModel(staged_path)
|
||||
rotate_history_epoch(staged.get_connection())
|
||||
ok, reason = self.verify_database_file(
|
||||
staged_path, require_library_state=True
|
||||
)
|
||||
if not ok:
|
||||
raise ValueError(reason)
|
||||
staged.close_all()
|
||||
backup_path = ""
|
||||
if os.path.exists(self.db_path):
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
backup_path = f"{self.db_path}.backup_{timestamp}"
|
||||
if not self.backup_database(backup_path):
|
||||
raise RuntimeError("Could not back up the current database")
|
||||
self.logger.info(f"Current database backed up to: {backup_path}")
|
||||
try:
|
||||
self._replace_contents_from(staged_path)
|
||||
except sqlite3.OperationalError as exc:
|
||||
exc._prompt_manager_restore_backup_created = bool(backup_path)
|
||||
raise
|
||||
with PromptModel._init_lock:
|
||||
PromptModel._initialized_paths.add(self._schema_key())
|
||||
self.logger.info(f"Database restored from {src_path}")
|
||||
from .execution_history import notify_history_restore
|
||||
|
||||
notify_history_restore(self.db_path)
|
||||
return backup_path
|
||||
finally:
|
||||
if staged is not None:
|
||||
staged.close_all()
|
||||
with PromptModel._init_lock:
|
||||
PromptModel._initialized_paths.discard(
|
||||
os.path.normcase(os.path.realpath(staged_path))
|
||||
)
|
||||
|
||||
+131
-300
@@ -25,15 +25,6 @@ from .execution_outputs import ExecutionOutputOperations
|
||||
from .history_reads import HistoryReadOperations
|
||||
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 .workflow_documents import (
|
||||
DocumentReader,
|
||||
documents_of_rows,
|
||||
release_documents,
|
||||
resolve_image_row,
|
||||
stored_values,
|
||||
)
|
||||
from .maintenance import ( # noqa: F401 (ORPHANED_PROMPTS_SQL is re-exported)
|
||||
ORPHANED_PROMPTS_SQL,
|
||||
_LINK_ROWS_SQL,
|
||||
@@ -47,7 +38,6 @@ try:
|
||||
from ..utils.logging_config import get_logger
|
||||
from ..utils.hashing import generate_prompt_hash
|
||||
from ..utils.json_strict import dumps_strict
|
||||
from ..utils import library_location
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
@@ -56,7 +46,6 @@ except ImportError:
|
||||
from utils.logging_config import get_logger
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.json_strict import dumps_strict
|
||||
from utils import library_location
|
||||
|
||||
|
||||
# Subquery to fetch tags from junction table, embedded in SELECT statements
|
||||
@@ -102,56 +91,40 @@ def escape_like(term: str) -> str:
|
||||
)
|
||||
|
||||
|
||||
def _configured_default_path():
|
||||
"""config.json's ``database.default_path``, or None.
|
||||
|
||||
The package-qualified config wins over a standalone ``py.config`` (a
|
||||
custom-node package uses its own configuration).
|
||||
"""
|
||||
try:
|
||||
try:
|
||||
from ..py.config import PromptManagerConfig
|
||||
except ImportError:
|
||||
from py.config import PromptManagerConfig
|
||||
except Exception:
|
||||
return None
|
||||
return getattr(PromptManagerConfig, "DEFAULT_DB_PATH", None)
|
||||
|
||||
|
||||
def resolve_library_location():
|
||||
"""Where the library lives now (``utils.library_location``, #249 S6b).
|
||||
|
||||
The image sources that rule out a ``PROMPT_MANAGER_DATA_DIR`` are read by
|
||||
the resolver itself, exactly as ``config_location`` reads them.
|
||||
"""
|
||||
return library_location.resolve_library(_configured_default_path())
|
||||
|
||||
|
||||
def _resolve_db_path(db_path: Optional[str] = None) -> str:
|
||||
"""Resolve the database path through the one library resolver.
|
||||
"""Resolve the database path from config, falling back to defaults.
|
||||
|
||||
With no explicit path, ``utils.library_location.resolve_library`` decides
|
||||
(``PROMPT_MANAGER_DATA_DIR``, an explicit ``database.default_path``, a
|
||||
node-folder library not moved yet, else ComfyUI's ``user/__promptmanager``).
|
||||
An explicit relative path is resolved against the extension root, as
|
||||
before, so the location does not depend on the working directory.
|
||||
When no explicit path is given, reads PromptManagerConfig.DEFAULT_DB_PATH
|
||||
(set via config.json). Relative paths are resolved against the extension
|
||||
root directory so the database location is predictable regardless of the
|
||||
process working directory (important for Docker deployments).
|
||||
|
||||
Args:
|
||||
db_path: Explicit database path, or None to use the resolver.
|
||||
db_path: Explicit database path, or None to use config/default.
|
||||
|
||||
Returns:
|
||||
Absolute path to the SQLite database file; its folder exists (a new
|
||||
data folder is created owner-only).
|
||||
Absolute path to the SQLite database file.
|
||||
"""
|
||||
extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
if db_path is None:
|
||||
db_path = resolve_library_location().db_path
|
||||
elif not os.path.isabs(db_path):
|
||||
extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
try:
|
||||
try:
|
||||
from ..py.config import PromptManagerConfig
|
||||
except ImportError:
|
||||
from py.config import PromptManagerConfig
|
||||
|
||||
db_path = PromptManagerConfig.DEFAULT_DB_PATH
|
||||
except Exception:
|
||||
db_path = "prompts.db"
|
||||
|
||||
if not os.path.isabs(db_path):
|
||||
db_path = os.path.normpath(os.path.join(extension_root, db_path))
|
||||
|
||||
# Ensure parent directory exists for custom paths
|
||||
parent = os.path.dirname(db_path)
|
||||
if parent:
|
||||
library_location.ensure_private_dir(parent)
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
|
||||
return db_path
|
||||
|
||||
@@ -366,8 +339,6 @@ class PromptDatabase(
|
||||
HealthOperations,
|
||||
MaintenanceOperations,
|
||||
DuplicateCopyOperations,
|
||||
WorkflowStorageOperations,
|
||||
CategoryImportOperations,
|
||||
):
|
||||
"""Database operations class for managing prompts."""
|
||||
|
||||
@@ -385,25 +356,6 @@ class PromptDatabase(
|
||||
self.model = PromptModel(db_path)
|
||||
self.logger.debug("Database operations initialized successfully")
|
||||
|
||||
def survive_library_switch(self, write):
|
||||
"""Run ``write()``; once more if it failed while the library moved.
|
||||
|
||||
A background write (run counting, the image monitor) that was waiting
|
||||
on the old file when Move library switched fails when its connection
|
||||
is closed. When the model's generation changed during the attempt,
|
||||
the write runs again on a fresh connection, which is the new library;
|
||||
any other failure is raised as it was. ``write`` must open its own
|
||||
transaction (every PromptDatabase method does).
|
||||
"""
|
||||
generation = self.model.generation
|
||||
try:
|
||||
return write()
|
||||
except Exception:
|
||||
if self.model.generation == generation:
|
||||
raise
|
||||
self.logger.info("Running a write again on the moved library")
|
||||
return write()
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the calling thread's database connection (see PromptModel.close)."""
|
||||
self.model.close()
|
||||
@@ -457,9 +409,26 @@ class PromptDatabase(
|
||||
|
||||
try:
|
||||
with self.model.get_connection() as conn:
|
||||
prompt_id = self._insert_prompt(
|
||||
conn, text, category, tags, rating, notes, prompt_hash
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
INSERT INTO prompts (
|
||||
text, category, tags, rating, notes, hash, created_at,
|
||||
updated_at
|
||||
) VALUES (?, ?, NULL, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
text.strip(),
|
||||
category,
|
||||
rating,
|
||||
notes,
|
||||
prompt_hash,
|
||||
datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
||||
datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
||||
),
|
||||
)
|
||||
prompt_id = cursor.lastrowid
|
||||
if tags:
|
||||
self._sync_prompt_tags(conn, prompt_id, tags)
|
||||
conn.commit()
|
||||
self.logger.debug(f"Successfully saved prompt with ID: {prompt_id}")
|
||||
return prompt_id
|
||||
@@ -475,32 +444,6 @@ class PromptDatabase(
|
||||
)
|
||||
return existing["id"]
|
||||
|
||||
def _insert_prompt(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
text: str,
|
||||
category: Optional[str],
|
||||
tags: Optional[List[str]],
|
||||
rating: Optional[int],
|
||||
notes: Optional[str],
|
||||
prompt_hash: Optional[str],
|
||||
) -> int:
|
||||
"""Insert one prompt row and its tags on ``conn`` (no commit)."""
|
||||
now = datetime.datetime.now(datetime.timezone.utc).isoformat()
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
INSERT INTO prompts (
|
||||
text, category, tags, rating, notes, hash, created_at,
|
||||
updated_at
|
||||
) VALUES (?, ?, NULL, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(text.strip(), category, rating, notes, prompt_hash, now, now),
|
||||
)
|
||||
prompt_id = cursor.lastrowid
|
||||
if tags:
|
||||
self._sync_prompt_tags(conn, prompt_id, tags)
|
||||
return prompt_id
|
||||
|
||||
def record_prompt_use(self, prompt_id: int, times: int = 1) -> bool:
|
||||
"""
|
||||
Count runs of a prompt and mark it as most recently used.
|
||||
@@ -833,12 +776,10 @@ class PromptDatabase(
|
||||
"""
|
||||
with self._pair_transaction() as conn:
|
||||
self._ensure_prompts_unpaired(conn, [prompt_id])
|
||||
documents = documents_of_rows(conn, "prompt_id", [prompt_id])
|
||||
# First delete related images to avoid foreign key constraint
|
||||
conn.execute(
|
||||
"DELETE FROM generated_images WHERE prompt_id = ?", (prompt_id,)
|
||||
)
|
||||
release_documents(conn, documents)
|
||||
# Then delete the prompt
|
||||
cursor = conn.execute("DELETE FROM prompts WHERE id = ?", (prompt_id,))
|
||||
conn.commit()
|
||||
@@ -862,12 +803,10 @@ class PromptDatabase(
|
||||
return 0
|
||||
self._ensure_prompts_unpaired(conn, ids)
|
||||
placeholders = ",".join("?" * len(ids))
|
||||
documents = documents_of_rows(conn, "prompt_id", ids)
|
||||
conn.execute(
|
||||
f"DELETE FROM generated_images WHERE prompt_id IN ({placeholders})",
|
||||
ids,
|
||||
)
|
||||
release_documents(conn, documents)
|
||||
conn.execute(
|
||||
f"DELETE FROM prompt_tags WHERE prompt_id IN ({placeholders})",
|
||||
ids,
|
||||
@@ -1628,137 +1567,84 @@ class PromptDatabase(
|
||||
metadata: Optional metadata about the image
|
||||
|
||||
Returns:
|
||||
int: The ID of the created image record, or 0 when nothing was
|
||||
linked (invalid or unknown prompt, already linked, or an error,
|
||||
which is logged)
|
||||
int: The ID of the created image record
|
||||
"""
|
||||
# Validate that prompt_id is a valid integer and exists in prompts table
|
||||
try:
|
||||
return self.link_image_to_prompt_or_raise(prompt_id, image_path, metadata)
|
||||
# Convert prompt_id to integer if it's a string number
|
||||
if isinstance(prompt_id, str) and prompt_id.isdigit():
|
||||
prompt_id_int = int(prompt_id)
|
||||
elif isinstance(prompt_id, int):
|
||||
prompt_id_int = prompt_id
|
||||
else:
|
||||
# Handle temporary IDs like "temp_123456" - skip linking
|
||||
if isinstance(prompt_id, str) and prompt_id.startswith("temp_"):
|
||||
self.logger.debug(
|
||||
f"Skipping image linking for temporary prompt ID: {prompt_id}"
|
||||
)
|
||||
return 0
|
||||
else:
|
||||
self.logger.warning(
|
||||
f"Invalid prompt_id format: {prompt_id}, skipping image linking"
|
||||
)
|
||||
return 0
|
||||
|
||||
# Verify the prompt exists in the database
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT id FROM prompts WHERE id = ?", (prompt_id_int,)
|
||||
)
|
||||
if not cursor.fetchone():
|
||||
self.logger.warning(
|
||||
f"Prompt ID {prompt_id_int} not found in database, skipping "
|
||||
"image linking"
|
||||
)
|
||||
return 0
|
||||
|
||||
# Proceed with linking
|
||||
if not isinstance(metadata, dict):
|
||||
metadata = {}
|
||||
filename = os.path.basename(image_path)
|
||||
file_path = normalize_image_path(image_path)
|
||||
file_size, width, height, image_format = _coerce_file_info(metadata)
|
||||
|
||||
# INSERT OR IGNORE skips a file already linked to this prompt
|
||||
# (same prompt_id + normalised full path)
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
INSERT OR IGNORE INTO generated_images
|
||||
(prompt_id, image_path, filename, file_path, file_size, width,
|
||||
height, format, workflow_data, prompt_metadata, parameters)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
prompt_id_int,
|
||||
image_path,
|
||||
filename,
|
||||
file_path,
|
||||
file_size,
|
||||
width,
|
||||
height,
|
||||
image_format,
|
||||
dumps_strict(metadata.get("workflow", {})),
|
||||
dumps_strict(metadata.get("prompt", {})),
|
||||
dumps_strict(metadata.get("parameters", {})),
|
||||
),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
if cursor.rowcount == 0:
|
||||
self.logger.debug(
|
||||
f"Image {file_path} already linked to prompt {prompt_id_int}"
|
||||
)
|
||||
return 0
|
||||
|
||||
return cursor.lastrowid
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error linking image to prompt {prompt_id}: {e}")
|
||||
return 0
|
||||
|
||||
def link_image_to_prompt_or_raise(
|
||||
self, prompt_id: str, image_path: str, metadata: Optional[Dict[str, Any]] = None
|
||||
) -> int:
|
||||
"""``link_image_to_prompt`` that raises database errors.
|
||||
|
||||
The image monitor links through this so a write that failed because
|
||||
Move library switched reaches ``survive_library_switch`` and runs
|
||||
again on the new library (#249 S6b). Each call is one transaction:
|
||||
a failed attempt rolls back, and a committed one does not raise.
|
||||
Returns 0 for an invalid or unknown prompt or an image already linked.
|
||||
"""
|
||||
# Validate that prompt_id is a valid integer and exists in prompts table
|
||||
# Convert prompt_id to integer if it's a string number
|
||||
if isinstance(prompt_id, str) and prompt_id.isdigit():
|
||||
prompt_id_int = int(prompt_id)
|
||||
elif isinstance(prompt_id, int):
|
||||
prompt_id_int = prompt_id
|
||||
else:
|
||||
# Handle temporary IDs like "temp_123456" - skip linking
|
||||
if isinstance(prompt_id, str) and prompt_id.startswith("temp_"):
|
||||
self.logger.debug(
|
||||
f"Skipping image linking for temporary prompt ID: {prompt_id}"
|
||||
)
|
||||
return 0
|
||||
else:
|
||||
self.logger.warning(
|
||||
f"Invalid prompt_id format: {prompt_id}, skipping image linking"
|
||||
)
|
||||
return 0
|
||||
|
||||
# Verify the prompt exists in the database
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT id FROM prompts WHERE id = ?", (prompt_id_int,)
|
||||
)
|
||||
if not cursor.fetchone():
|
||||
self.logger.warning(
|
||||
f"Prompt ID {prompt_id_int} not found in database, skipping "
|
||||
"image linking"
|
||||
)
|
||||
return 0
|
||||
|
||||
# Proceed with linking
|
||||
image_id = self._insert_image_link(
|
||||
conn, prompt_id_int, image_path, metadata
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
if not image_id:
|
||||
self.logger.debug(
|
||||
f"Image {normalize_image_path(image_path)} already linked to "
|
||||
f"prompt {prompt_id_int}"
|
||||
)
|
||||
return 0
|
||||
|
||||
return image_id
|
||||
|
||||
def _insert_image_link(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
prompt_id_int: int,
|
||||
image_path: str,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> int:
|
||||
"""Insert one image link on ``conn`` (no commit).
|
||||
|
||||
Returns the new image row's id, or 0 when the file is already linked
|
||||
to this prompt (``INSERT OR IGNORE``).
|
||||
"""
|
||||
if not isinstance(metadata, dict):
|
||||
metadata = {}
|
||||
filename = os.path.basename(image_path)
|
||||
file_path = normalize_image_path(image_path)
|
||||
file_size, width, height, image_format = _coerce_file_info(metadata)
|
||||
# Each workflow and prompt graph is stored once (new form).
|
||||
documents = stored_values(
|
||||
conn,
|
||||
{
|
||||
"workflow_data": dumps_strict(metadata.get("workflow", {})),
|
||||
"prompt_metadata": dumps_strict(metadata.get("prompt", {})),
|
||||
},
|
||||
)
|
||||
|
||||
# INSERT OR IGNORE skips a file already linked to this prompt
|
||||
# (same prompt_id + normalised full path)
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
INSERT OR IGNORE INTO generated_images
|
||||
(prompt_id, image_path, filename, file_path, file_size, width,
|
||||
height, format, workflow_data, prompt_metadata, parameters,
|
||||
workflow_document_id, prompt_document_id)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
prompt_id_int,
|
||||
image_path,
|
||||
filename,
|
||||
file_path,
|
||||
file_size,
|
||||
width,
|
||||
height,
|
||||
image_format,
|
||||
documents["workflow_data"],
|
||||
documents["prompt_metadata"],
|
||||
dumps_strict(metadata.get("parameters", {})),
|
||||
documents["workflow_document_id"],
|
||||
documents["prompt_document_id"],
|
||||
),
|
||||
)
|
||||
if cursor.rowcount == 0:
|
||||
# Already linked: a document made only for this row goes.
|
||||
release_documents(
|
||||
conn,
|
||||
(
|
||||
documents["workflow_document_id"],
|
||||
documents["prompt_document_id"],
|
||||
),
|
||||
)
|
||||
return 0
|
||||
return cursor.lastrowid
|
||||
|
||||
def get_prompt_images(self, prompt_id: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get all images associated with a prompt.
|
||||
@@ -1778,8 +1664,7 @@ class PromptDatabase(
|
||||
""",
|
||||
(prompt_id,),
|
||||
)
|
||||
reader = DocumentReader(conn)
|
||||
return [self._image_row_to_dict(row, reader) for row in cursor]
|
||||
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
|
||||
|
||||
def get_recent_images(
|
||||
self, limit: int = 50, offset: int = 0
|
||||
@@ -1805,39 +1690,7 @@ class PromptDatabase(
|
||||
""",
|
||||
(int(limit), max(0, int(offset))),
|
||||
)
|
||||
reader = DocumentReader(conn)
|
||||
return [self._image_row_to_dict(row, reader) for row in cursor]
|
||||
|
||||
def autotag_batch_candidates(self) -> List[Dict[str, Any]]:
|
||||
"""The images a Tag output images batch looks at (#328).
|
||||
|
||||
``{id, image_path, prompt_id, prompt_tags}`` for every image with a
|
||||
stored path that is linked to a prompt, newest first (the
|
||||
``get_all_images`` order, ties newest row first). Never reads
|
||||
``workflow_data`` or the other documents: a batch needs none of them.
|
||||
"""
|
||||
sql = (
|
||||
"SELECT gi.id, gi.image_path, gi.prompt_id, "
|
||||
"(SELECT GROUP_CONCAT(t.name, '|||') FROM prompt_tags pt "
|
||||
"JOIN tags t ON pt.tag_id = t.id WHERE pt.prompt_id = gi.prompt_id) "
|
||||
"AS prompt_tags "
|
||||
"FROM generated_images gi "
|
||||
"INNER JOIN prompts p ON gi.prompt_id = p.id "
|
||||
"WHERE gi.image_path IS NOT NULL AND gi.image_path != '' "
|
||||
"ORDER BY gi.generation_time DESC, gi.id DESC"
|
||||
)
|
||||
with self.model.get_connection() as conn:
|
||||
return [
|
||||
{
|
||||
"id": row["id"],
|
||||
"image_path": row["image_path"],
|
||||
"prompt_id": row["prompt_id"],
|
||||
"prompt_tags": (
|
||||
row["prompt_tags"].split("|||") if row["prompt_tags"] else []
|
||||
),
|
||||
}
|
||||
for row in conn.execute(sql)
|
||||
]
|
||||
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
|
||||
|
||||
def get_all_images(
|
||||
self,
|
||||
@@ -1871,10 +1724,10 @@ class PromptDatabase(
|
||||
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(sql, params)
|
||||
reader = DocumentReader(conn)
|
||||
rows = cursor.fetchall()
|
||||
result = []
|
||||
for row in cursor:
|
||||
data = self._image_row_to_dict(row, reader)
|
||||
for row in rows:
|
||||
data = self._image_row_to_dict(row)
|
||||
if row["_prompt_tags_list"]:
|
||||
data["prompt_tags"] = row["_prompt_tags_list"].split("|||")
|
||||
else:
|
||||
@@ -1912,8 +1765,7 @@ class PromptDatabase(
|
||||
max(0, int(offset)),
|
||||
),
|
||||
)
|
||||
reader = DocumentReader(conn)
|
||||
return [self._image_row_to_dict(row, reader) for row in cursor]
|
||||
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
|
||||
|
||||
def get_image_by_id(self, image_id: int) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
@@ -1930,7 +1782,7 @@ class PromptDatabase(
|
||||
"SELECT * FROM generated_images WHERE id = ?", (image_id,)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
return self._image_row_to_dict(row, DocumentReader(conn)) if row else None
|
||||
return self._image_row_to_dict(row) if row else None
|
||||
|
||||
def delete_image(self, image_id: int) -> bool:
|
||||
"""
|
||||
@@ -1943,11 +1795,9 @@ class PromptDatabase(
|
||||
bool: True if deletion was successful
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
documents = documents_of_rows(conn, "id", [image_id])
|
||||
cursor = conn.execute(
|
||||
"DELETE FROM generated_images WHERE id = ?", (image_id,)
|
||||
)
|
||||
release_documents(conn, documents)
|
||||
conn.commit()
|
||||
return cursor.rowcount > 0
|
||||
|
||||
@@ -1981,16 +1831,12 @@ class PromptDatabase(
|
||||
cursor = conn.execute("SELECT id, image_path FROM generated_images")
|
||||
images = cursor.fetchall()
|
||||
|
||||
removed = [
|
||||
image["id"]
|
||||
for image in images
|
||||
if not os.path.exists(image["image_path"])
|
||||
]
|
||||
documents = documents_of_rows(conn, "id", removed)
|
||||
for image_id in removed:
|
||||
conn.execute("DELETE FROM generated_images WHERE id = ?", (image_id,))
|
||||
removed_count += 1
|
||||
release_documents(conn, documents)
|
||||
for image in images:
|
||||
if not os.path.exists(image["image_path"]):
|
||||
conn.execute(
|
||||
"DELETE FROM generated_images WHERE id = ?", (image["id"],)
|
||||
)
|
||||
removed_count += 1
|
||||
|
||||
conn.commit()
|
||||
|
||||
@@ -2281,13 +2127,11 @@ class PromptDatabase(
|
||||
count = 0
|
||||
with self._pair_transaction() as conn:
|
||||
self._ensure_prompts_unpaired(conn, prompt_ids)
|
||||
documents = documents_of_rows(conn, "prompt_id", prompt_ids)
|
||||
for pid in prompt_ids:
|
||||
conn.execute("DELETE FROM generated_images WHERE prompt_id = ?", (pid,))
|
||||
cursor = conn.execute("DELETE FROM prompts WHERE id = ?", (pid,))
|
||||
if cursor.rowcount > 0:
|
||||
count += 1
|
||||
release_documents(conn, documents)
|
||||
conn.commit()
|
||||
return count
|
||||
|
||||
@@ -2361,8 +2205,7 @@ class PromptDatabase(
|
||||
cursor = conn.execute(
|
||||
"""SELECT gi.prompt_id, p.text, p.category, p.rating, p.notes,
|
||||
gi.workflow_data, gi.prompt_metadata, gi.generation_time,
|
||||
gi.parameters, gi.id, gi.image_path, gi.file_path,
|
||||
gi.workflow_document_id, gi.prompt_document_id
|
||||
gi.parameters, gi.id, gi.image_path, gi.file_path
|
||||
FROM generated_images gi
|
||||
JOIN prompts p ON gi.prompt_id = p.id
|
||||
WHERE gi.image_path = ? OR gi.file_path = ?
|
||||
@@ -2442,21 +2285,12 @@ class PromptDatabase(
|
||||
"text_source": source,
|
||||
}
|
||||
tags = self._get_prompt_tags(conn, row[0])
|
||||
texts = resolve_image_row(DocumentReader(conn), row)
|
||||
try:
|
||||
workflow = (
|
||||
json.loads(texts["workflow_data"])
|
||||
if texts["workflow_data"]
|
||||
else None
|
||||
)
|
||||
workflow = json.loads(row[5]) if row[5] else None
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
workflow = None
|
||||
try:
|
||||
prompt_meta = (
|
||||
json.loads(texts["prompt_metadata"])
|
||||
if texts["prompt_metadata"]
|
||||
else None
|
||||
)
|
||||
prompt_meta = json.loads(row[6]) if row[6] else None
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
prompt_meta = None
|
||||
return {
|
||||
@@ -2706,20 +2540,17 @@ class PromptDatabase(
|
||||
with self.model.get_connection() as conn:
|
||||
return list(iter_consistency_issues(conn))
|
||||
|
||||
def _image_row_to_dict(
|
||||
self, row: sqlite3.Row, reader: DocumentReader
|
||||
) -> Dict[str, Any]:
|
||||
def _image_row_to_dict(self, row: sqlite3.Row) -> Dict[str, Any]:
|
||||
"""
|
||||
Convert an image database row to a dictionary with parsed JSON fields.
|
||||
|
||||
Args:
|
||||
row: SQLite row object
|
||||
reader: The workflow document reader of the row's connection
|
||||
|
||||
Returns:
|
||||
Dictionary representation of the row
|
||||
"""
|
||||
data = resolve_image_row(reader, row)
|
||||
data = dict(row)
|
||||
|
||||
# Parse JSON fields
|
||||
for field in ["workflow_data", "prompt_metadata", "parameters"]:
|
||||
|
||||
@@ -5,13 +5,6 @@ from typing import Any, Dict
|
||||
|
||||
from .models import normalize_image_path
|
||||
from .operations import _coerce_file_info
|
||||
from .workflow_documents import (
|
||||
TEXT_COLUMNS,
|
||||
DocumentReader,
|
||||
release_documents,
|
||||
resolve_image_row,
|
||||
stored_values,
|
||||
)
|
||||
|
||||
try:
|
||||
from ..utils.hashing import generate_prompt_hash
|
||||
@@ -23,26 +16,10 @@ except ImportError:
|
||||
from utils.prompt_pair_resolver import RESOLVER_VERSION
|
||||
|
||||
|
||||
def blank_positive(resolution: Dict[str, Any]) -> bool:
|
||||
"""A resolved pair whose stored positive text is empty (#342).
|
||||
|
||||
A PromptManager node whose editor is empty encodes only its prepend and
|
||||
append text, which is never stored as a prompt: such a result names no
|
||||
prompt to link.
|
||||
"""
|
||||
if resolution.get("state") != "resolved":
|
||||
return False
|
||||
positive = resolution.get("positive") or {}
|
||||
text = positive.get("storage_text", positive.get("text"))
|
||||
return not (isinstance(text, str) and text.strip())
|
||||
|
||||
|
||||
def ingest_resolved_pair(
|
||||
db: Any, image_path: str, metadata: Dict[str, Any], resolution: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
state = resolution.get("state", "unresolved")
|
||||
if state == "resolved" and blank_positive(resolution):
|
||||
state = "unresolved"
|
||||
if state != "resolved":
|
||||
return {"status": state, "pairs_added": 0, "links_added": 0}
|
||||
with db._pair_transaction() as conn:
|
||||
@@ -139,35 +116,25 @@ def _ingest_resolved_pair_in(
|
||||
changed = False
|
||||
if existing:
|
||||
image_id = existing["id"]
|
||||
# Compared as stored text, whichever form the row is in.
|
||||
current = resolve_image_row(DocumentReader(conn), existing)
|
||||
updates = {
|
||||
key: value for key, value in snapshots.items() if current[key] != value
|
||||
key: value for key, value in snapshots.items() if existing[key] != value
|
||||
}
|
||||
if updates:
|
||||
replaced = [
|
||||
existing[reference]
|
||||
for column, reference in TEXT_COLUMNS
|
||||
if column in updates
|
||||
]
|
||||
updates = _document_form(conn, updates)
|
||||
conn.execute(
|
||||
"UPDATE generated_images SET "
|
||||
+ ",".join(f"{key}=?" for key in updates)
|
||||
+ " WHERE id=?",
|
||||
[*updates.values(), image_id],
|
||||
)
|
||||
release_documents(conn, replaced)
|
||||
changed = True
|
||||
else:
|
||||
stored = _document_form(conn, snapshots)
|
||||
columns = ["prompt_id", "image_path", "filename", "file_path", *stored]
|
||||
columns = ["prompt_id", "image_path", "filename", "file_path", *snapshots]
|
||||
values = [
|
||||
pair["positive_prompt_id"],
|
||||
image_path,
|
||||
os.path.basename(image_path),
|
||||
normalized,
|
||||
*stored.values(),
|
||||
*snapshots.values(),
|
||||
]
|
||||
image_id = conn.execute(
|
||||
"INSERT INTO generated_images("
|
||||
@@ -195,18 +162,3 @@ def _ingest_resolved_pair_in(
|
||||
"pair_id": pair_id,
|
||||
"image_id": image_id,
|
||||
}
|
||||
|
||||
|
||||
def _document_form(conn: Any, values: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""``values`` with the workflow and prompt graph texts stored once.
|
||||
|
||||
Each of those columns becomes NULL beside its document id (new form);
|
||||
every other column is unchanged.
|
||||
"""
|
||||
texts = {
|
||||
column: values[column]
|
||||
for column, _reference in TEXT_COLUMNS
|
||||
if column in values
|
||||
}
|
||||
rest = {key: value for key, value in values.items() if key not in texts}
|
||||
return {**rest, **stored_values(conn, texts)}
|
||||
|
||||
@@ -1,365 +0,0 @@
|
||||
"""A restore in two steps: stage the upload, then publish it (#249 System S8).
|
||||
|
||||
**Stage** (``stage_restore``) never touches the live database or the source.
|
||||
In a new folder ``.pm-restore-*`` in the library's own folder (#318: never
|
||||
the system temp folder, which may be another disk), it:
|
||||
|
||||
1. checks the source (``verify_database_file`` without reading every
|
||||
workflow document: a damaged one is repaired below, not refused);
|
||||
2. copies it with SQLite's backup API (a rollback-journal file and a WAL file
|
||||
alike) and opens the copy as a ``PromptModel``, which migrates an older
|
||||
schema;
|
||||
3. repairs damaged workflow documents (``repair_damaged_documents``),
|
||||
rotates the history epoch and counts what the copy holds;
|
||||
4. makes the copy one self-contained file and verifies it fully (every
|
||||
current table, every document);
|
||||
5. records the copy's SHA-256, so publishing binds exactly these bytes.
|
||||
|
||||
**Publish** (``publish_staged``) re-hashes the stage (``StageChanged`` if it
|
||||
differs), keeps the current database as a uniquely named, standalone safety
|
||||
backup beside it (``database/safety_backup.py``; not verified, because the
|
||||
database being replaced may be the damaged one), copies the stage over the
|
||||
live pages through the backup API (no file is swapped under open handles; an
|
||||
``OperationalError`` there leaves the live data intact, #226), then
|
||||
checkpoints the WAL with TRUNCATE so the copied pages do not leave a
|
||||
database-sized ``-wal`` behind. A reader that keeps a snapshot past the busy
|
||||
wait is logged, not fatal: the restore itself is done.
|
||||
|
||||
Every handle on the stage is closed before its folder is removed (Windows).
|
||||
``sweep_restore_temps`` removes what an interrupted restore left (startup).
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
import sqlite3
|
||||
import tempfile
|
||||
from contextlib import closing
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional
|
||||
|
||||
from . import workflow_documents as wd
|
||||
from .execution_history_schema import rotate_history_epoch
|
||||
from .safety_backup import take_safety_backup
|
||||
|
||||
logger = logging.getLogger("prompt_manager.database.restore_stage")
|
||||
|
||||
STAGE_PREFIX = ".pm-restore-"
|
||||
UPLOAD_PREFIX = ".pm-restore-upload-"
|
||||
STAGED_NAME = "staged.db"
|
||||
# How long the post-restore TRUNCATE checkpoint waits for readers to leave.
|
||||
CHECKPOINT_BUSY_MS = 30_000
|
||||
_HASH_CHUNK = 1024 * 1024
|
||||
SUMMARY_KEYS = (
|
||||
"prompts",
|
||||
"positives",
|
||||
"pairs",
|
||||
"archived_pairs",
|
||||
"image_links",
|
||||
"runs",
|
||||
"saved_views",
|
||||
)
|
||||
|
||||
|
||||
class StageChanged(ValueError):
|
||||
"""The staged file is no longer the bytes that were reviewed."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class StagedRestore:
|
||||
"""A verified, migrated copy of an upload, ready to publish."""
|
||||
|
||||
path: str
|
||||
digest: str
|
||||
size: int
|
||||
summary: Dict[str, int]
|
||||
schema: Dict[str, bool]
|
||||
documents: Dict[str, int] = field(
|
||||
default_factory=lambda: {"repaired": 0, "cleared": 0, "images": 0}
|
||||
)
|
||||
|
||||
@property
|
||||
def directory(self) -> str:
|
||||
return os.path.dirname(self.path)
|
||||
|
||||
def discard(self) -> None:
|
||||
"""Remove the stage folder (idempotent; its handles are all closed)."""
|
||||
remove_stage_folder(self.directory)
|
||||
|
||||
|
||||
def remove_stage_folder(directory: str) -> None:
|
||||
"""Remove a stage folder and its files, logging what cannot be removed."""
|
||||
if not directory or not os.path.isdir(directory):
|
||||
return
|
||||
for name in os.listdir(directory):
|
||||
try:
|
||||
os.remove(os.path.join(directory, name))
|
||||
except OSError as error:
|
||||
logger.warning(f"Could not remove restore stage file {name}: {error}")
|
||||
try:
|
||||
os.rmdir(directory)
|
||||
except OSError as error:
|
||||
logger.warning(f"Could not remove the restore stage folder: {error}")
|
||||
|
||||
|
||||
def file_digest(path: str) -> str:
|
||||
"""``sha256:<hex>`` of the file's bytes (blocking)."""
|
||||
digest = hashlib.sha256()
|
||||
with open(path, "rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(_HASH_CHUNK), b""):
|
||||
digest.update(chunk)
|
||||
return "sha256:" + digest.hexdigest()
|
||||
|
||||
|
||||
def _readonly(path: str) -> sqlite3.Connection:
|
||||
return sqlite3.connect(Path(path).resolve().as_uri() + "?mode=ro", uri=True)
|
||||
|
||||
|
||||
def _schema(path: str) -> list:
|
||||
"""The file's schema objects, to tell whether staging migrated it."""
|
||||
with closing(_readonly(path)) as conn:
|
||||
return sorted(
|
||||
conn.execute(
|
||||
"SELECT type, name, tbl_name, sql FROM sqlite_master "
|
||||
"WHERE name NOT LIKE 'sqlite_%'"
|
||||
).fetchall(),
|
||||
key=repr,
|
||||
)
|
||||
|
||||
|
||||
def _count(conn: sqlite3.Connection, table: str, where: str = "") -> int:
|
||||
exists = conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (table,)
|
||||
).fetchone()
|
||||
if not exists:
|
||||
return 0
|
||||
# ``table`` and ``where`` are constants of library_summary, never input.
|
||||
return conn.execute(f"SELECT COUNT(*) FROM {table} {where}").fetchone()[0] # nosec
|
||||
|
||||
|
||||
def library_summary(conn: sqlite3.Connection) -> Dict[str, int]:
|
||||
"""What a library holds, for a restore's review (reads only).
|
||||
|
||||
``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.
|
||||
"""
|
||||
from .library import _ELIGIBLE
|
||||
|
||||
has_state = conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND "
|
||||
"name='prompt_library_state'"
|
||||
).fetchone()
|
||||
has_pairs = conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND name='prompt_pairs'"
|
||||
).fetchone()
|
||||
if has_state and has_pairs:
|
||||
positives = conn.execute(
|
||||
"SELECT COUNT(*) FROM prompts p WHERE " + _ELIGIBLE
|
||||
).fetchone()[0]
|
||||
else:
|
||||
positives = _count(conn, "prompts")
|
||||
return {
|
||||
"prompts": _count(conn, "prompts"),
|
||||
"positives": positives,
|
||||
"pairs": _count(conn, "prompt_pairs", "WHERE archived=0"),
|
||||
"archived_pairs": _count(conn, "prompt_pairs", "WHERE archived<>0"),
|
||||
"image_links": _count(conn, "generated_images"),
|
||||
"runs": _count(conn, "execution_runs"),
|
||||
"saved_views": _count(conn, "saved_views"),
|
||||
}
|
||||
|
||||
|
||||
# The newest row of each table a restore would overwrite: with the counts,
|
||||
# a write after the review (a new run, image, prompt or edit) changes it.
|
||||
_FINGERPRINT_SQL = (
|
||||
("prompts", "SELECT MAX(id), MAX(updated_at), MAX(last_used_at) FROM prompts"),
|
||||
("generated_images", "SELECT MAX(id) FROM generated_images"),
|
||||
("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"),
|
||||
)
|
||||
|
||||
|
||||
def live_fingerprint(conn: sqlite3.Connection):
|
||||
"""``(summary, fingerprint)`` of a live library, read in one snapshot.
|
||||
|
||||
The fingerprint changes when rows are added to or removed from the
|
||||
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.
|
||||
"""
|
||||
import json
|
||||
|
||||
opened = not conn.in_transaction
|
||||
if opened:
|
||||
conn.execute("BEGIN") # one read snapshot for every count
|
||||
try:
|
||||
summary = library_summary(conn)
|
||||
newest = []
|
||||
for table, sql in _FINGERPRINT_SQL:
|
||||
try:
|
||||
newest.append([table, list(conn.execute(sql).fetchone())])
|
||||
except sqlite3.OperationalError:
|
||||
newest.append([table, None]) # no such table or column
|
||||
finally:
|
||||
if opened:
|
||||
conn.rollback() # it only read
|
||||
text = json.dumps([summary, newest], sort_keys=True, default=str)
|
||||
return summary, "sha256:" + hashlib.sha256(text.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _copy(src_path: str, staged_path: str) -> None:
|
||||
source = _readonly(src_path)
|
||||
try:
|
||||
with closing(sqlite3.connect(staged_path)) as destination:
|
||||
source.backup(destination)
|
||||
finally:
|
||||
source.close()
|
||||
|
||||
|
||||
def _standalone(path: str) -> None:
|
||||
with closing(sqlite3.connect(path)) as conn:
|
||||
conn.execute("PRAGMA journal_mode=DELETE").fetchone()
|
||||
|
||||
|
||||
def stage_restore(model_cls, db_path: str, src_path: str) -> StagedRestore:
|
||||
"""Stage ``src_path`` beside ``db_path`` (blocking); see the module docstring.
|
||||
|
||||
Raises ValueError with a readable reason when the source or the migrated
|
||||
stage does not verify; the stage folder is then already removed.
|
||||
"""
|
||||
ok, reason = model_cls.verify_database_file(src_path, documents=False)
|
||||
if not ok:
|
||||
raise ValueError(reason)
|
||||
before = _schema(src_path)
|
||||
with closing(_readonly(src_path)) as conn:
|
||||
library_state = bool(
|
||||
conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND "
|
||||
"name='prompt_library_state'"
|
||||
).fetchone()
|
||||
)
|
||||
folder = os.path.dirname(os.path.abspath(db_path))
|
||||
directory = tempfile.mkdtemp(prefix=STAGE_PREFIX, dir=folder)
|
||||
staged_path = os.path.join(directory, STAGED_NAME)
|
||||
staged_key = os.path.normcase(os.path.realpath(staged_path))
|
||||
try:
|
||||
_copy(src_path, staged_path)
|
||||
staged = None
|
||||
try:
|
||||
staged = model_cls(staged_path)
|
||||
conn = staged.get_connection()
|
||||
documents = wd.repair_damaged_documents(conn)
|
||||
rotate_history_epoch(conn)
|
||||
summary = library_summary(conn)
|
||||
finally:
|
||||
if staged is not None:
|
||||
staged.close_all()
|
||||
_standalone(staged_path)
|
||||
ok, reason = model_cls.verify_database_file(
|
||||
staged_path, require_library_state=True
|
||||
)
|
||||
if not ok:
|
||||
raise ValueError(reason)
|
||||
result = StagedRestore(
|
||||
path=staged_path,
|
||||
digest=file_digest(staged_path),
|
||||
size=os.path.getsize(staged_path),
|
||||
summary=summary,
|
||||
schema={
|
||||
"migrated": _schema(staged_path) != before,
|
||||
"library_state": library_state,
|
||||
},
|
||||
documents=documents,
|
||||
)
|
||||
except BaseException:
|
||||
remove_stage_folder(directory)
|
||||
raise
|
||||
finally:
|
||||
with model_cls._init_lock:
|
||||
model_cls._initialized_paths.discard(staged_key)
|
||||
return result
|
||||
|
||||
|
||||
def checkpoint_truncate(conn: sqlite3.Connection) -> bool:
|
||||
"""Checkpoint the WAL and truncate it to 0 bytes; False when a reader held it."""
|
||||
previous = conn.execute("PRAGMA busy_timeout").fetchone()[0]
|
||||
try:
|
||||
conn.execute(f"PRAGMA busy_timeout = {int(CHECKPOINT_BUSY_MS)}")
|
||||
busy = conn.execute("PRAGMA wal_checkpoint(TRUNCATE)").fetchone()[0]
|
||||
finally:
|
||||
conn.execute(f"PRAGMA busy_timeout = {int(previous)}")
|
||||
return not busy
|
||||
|
||||
|
||||
def publish_staged(model, staged: StagedRestore) -> str:
|
||||
"""Publish ``staged`` over ``model``'s database; return the safety backup path.
|
||||
|
||||
"" when there was no database to back up. Raises StageChanged before
|
||||
anything is written when the stage's bytes changed since staging, and
|
||||
SafetyBackupFailed when the current database could not be kept. An
|
||||
``OperationalError`` from the page copy propagates with
|
||||
``_prompt_manager_restore_backup_created`` set; the live data is intact.
|
||||
"""
|
||||
if not os.path.isfile(staged.path) or file_digest(staged.path) != staged.digest:
|
||||
raise StageChanged("The staged restore changed since it was reviewed")
|
||||
backup_path = ""
|
||||
if os.path.exists(model.db_path):
|
||||
backup_path = take_safety_backup(model, verify=False)
|
||||
logger.info(f"Current database backed up to: {os.path.basename(backup_path)}")
|
||||
try:
|
||||
model._replace_contents_from(staged.path)
|
||||
except sqlite3.OperationalError as exc:
|
||||
exc._prompt_manager_restore_backup_created = bool(backup_path)
|
||||
raise
|
||||
with model._init_lock:
|
||||
model._initialized_paths.add(model._schema_key())
|
||||
_truncate_wal(model)
|
||||
return backup_path
|
||||
|
||||
|
||||
def _truncate_wal(model) -> Optional[bool]:
|
||||
"""TRUNCATE-checkpoint the live WAL after a publish (never raises)."""
|
||||
try:
|
||||
conn = model.get_connection()
|
||||
with conn.using():
|
||||
truncated = checkpoint_truncate(conn)
|
||||
except sqlite3.Error as error:
|
||||
logger.warning(f"The database log could not be truncated: {error}")
|
||||
return None
|
||||
if not truncated:
|
||||
logger.warning(
|
||||
"The database log could not be truncated after the restore: a "
|
||||
"reader kept it open. It shrinks at the next checkpoint with no "
|
||||
"reader."
|
||||
)
|
||||
return truncated
|
||||
|
||||
|
||||
def sweep_restore_temps(db_path: str) -> list:
|
||||
"""Remove the stages and uploads an interrupted restore left (startup)."""
|
||||
folder = os.path.dirname(os.path.abspath(db_path))
|
||||
removed = []
|
||||
try:
|
||||
names = os.listdir(folder)
|
||||
except OSError:
|
||||
return removed
|
||||
for name in names:
|
||||
path = os.path.join(folder, name)
|
||||
if name.startswith(UPLOAD_PREFIX) and os.path.isfile(path):
|
||||
try:
|
||||
os.remove(path)
|
||||
removed.append(name)
|
||||
except OSError:
|
||||
continue
|
||||
elif name.startswith(STAGE_PREFIX) and os.path.isdir(path):
|
||||
remove_stage_folder(path)
|
||||
if not os.path.exists(path):
|
||||
removed.append(name)
|
||||
return removed
|
||||
@@ -62,14 +62,3 @@ def notify_review_change(db_path: str) -> None:
|
||||
callback()
|
||||
except Exception:
|
||||
_logger.warning("Needs review change observer failed", exc_info=True)
|
||||
|
||||
|
||||
def rekey_review_changes(old_path: str, new_path: str) -> None:
|
||||
"""Move the observers of ``old_path`` to ``new_path`` (Move library)."""
|
||||
old, new = _identity(old_path), _identity(new_path)
|
||||
if old == new:
|
||||
return
|
||||
with _lock:
|
||||
observers = _observers.pop(old, None)
|
||||
if observers:
|
||||
_observers.setdefault(new, {}).update(observers)
|
||||
|
||||
@@ -380,21 +380,15 @@ def _structure_error(conn: sqlite3.Connection) -> Optional[str]:
|
||||
|
||||
|
||||
def _evidence_error(conn: sqlite3.Connection) -> Optional[str]:
|
||||
"""Stream every evidence document once; one row is held at a time.
|
||||
|
||||
Rows are read with ``fetchone``: iterating a managed connection's cursor
|
||||
would prefetch ``ITER_CHUNK`` (128) rows (``database/managed_connection.py``).
|
||||
"""
|
||||
items = conn.execute(
|
||||
"""Stream every evidence document once; one row is held at a time."""
|
||||
for evidence_hash, text in conn.execute(
|
||||
f"SELECT evidence_hash, evidence FROM {ITEMS_TABLE}" # nosec B608
|
||||
)
|
||||
while (row := items.fetchone()) is not None:
|
||||
evidence_hash, text = row[0], row[1]
|
||||
):
|
||||
if not _is_canonical(text) or evidence_digest(text) != evidence_hash:
|
||||
return f"Review queue table {ITEMS_TABLE} contains invalid evidence"
|
||||
decisions = conn.execute(f"SELECT evidence FROM {DECISIONS_TABLE}") # nosec B608
|
||||
while (row := decisions.fetchone()) is not None:
|
||||
text = row[0]
|
||||
for (text,) in conn.execute(
|
||||
f"SELECT evidence FROM {DECISIONS_TABLE}" # nosec B608
|
||||
):
|
||||
if not _is_canonical(text):
|
||||
return f"Review queue table {DECISIONS_TABLE} contains invalid evidence"
|
||||
return None
|
||||
|
||||
@@ -1,244 +0,0 @@
|
||||
"""Attribute an output file to a stored prompt through recorded history (#342).
|
||||
|
||||
Execution history records, per run, the files ComfyUI reported
|
||||
(``execution_run_outputs``) and the PromptManager texts the run used
|
||||
(``execution_run_contexts``, whose ``positive_prompt_id`` hint is set only
|
||||
while the raw text matches exactly one stored prompt). A file is attributed
|
||||
only when all of this holds:
|
||||
|
||||
- exactly one recorded run reported the file's output location (same output
|
||||
root, subfolder and name);
|
||||
- the file was last written while that run was executing, from its dequeue
|
||||
to its terminal event (else the finish callback): a later run that history
|
||||
did not record may have reused the name;
|
||||
- every positive context of that run names the same stored prompt.
|
||||
|
||||
Anything else returns None: the caller leaves the file unlinked. Read-only.
|
||||
|
||||
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
|
||||
still open the library.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sqlite3
|
||||
import threading
|
||||
import weakref
|
||||
from typing import Any, Callable, Dict, Iterable, List, Optional
|
||||
|
||||
try:
|
||||
from ..utils.execution_outputs import snapshot_output_roots
|
||||
from ..utils.library_location import import_folder_paths
|
||||
except ImportError:
|
||||
from utils.execution_outputs import snapshot_output_roots
|
||||
from utils.library_location import import_folder_paths
|
||||
|
||||
# File times may trail or lead the recorded run times a little (coarse file
|
||||
# system clocks, the write finishing just after the host reported it).
|
||||
MTIME_SLACK_MS = 2000
|
||||
|
||||
# Distinct file names per lookup query: well under SQLite's bound-parameter
|
||||
# limit (999 on old builds), so a scan batch costs one query per chunk.
|
||||
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) "
|
||||
"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})"
|
||||
)
|
||||
|
||||
|
||||
def host_output_directory() -> Optional[str]:
|
||||
"""ComfyUI's output directory, or None outside ComfyUI."""
|
||||
folder_paths = import_folder_paths()
|
||||
if folder_paths is None:
|
||||
return None
|
||||
try:
|
||||
value = folder_paths.get_output_directory()
|
||||
except Exception:
|
||||
return None
|
||||
return value if isinstance(value, str) and value else None
|
||||
|
||||
|
||||
def _inside(image_path: str, output_root: str) -> Optional[str]:
|
||||
"""``image_path`` relative to ``output_root``, or None when outside it."""
|
||||
try:
|
||||
relative = os.path.relpath(image_path, output_root)
|
||||
except ValueError: # another drive on Windows
|
||||
return None
|
||||
if os.path.isabs(relative) or relative == os.pardir:
|
||||
return None
|
||||
if relative.startswith(os.pardir + os.sep):
|
||||
return None
|
||||
return relative
|
||||
|
||||
|
||||
def _resolved(path: str) -> str:
|
||||
try:
|
||||
return os.path.realpath(path)
|
||||
except (OSError, ValueError):
|
||||
return os.path.abspath(path)
|
||||
|
||||
|
||||
def _location(image_path: str, output_root: str, path_style: str):
|
||||
"""(root key, subfolder, filename) as history records them, or None.
|
||||
|
||||
The root key is always the host's own spelling of its output folder,
|
||||
because history recorded that one. The file may be spelled through the
|
||||
same folder or through its resolved path (a symlinked or junction output
|
||||
folder, as scan roots are): the lexical spelling is tried first, then
|
||||
both sides resolved.
|
||||
"""
|
||||
root_key = snapshot_output_roots(path_style, output_root, output_root).output_key
|
||||
if root_key is None:
|
||||
return None
|
||||
relative = _inside(os.path.abspath(image_path), os.path.abspath(output_root))
|
||||
if relative is None:
|
||||
relative = _inside(_resolved(image_path), _resolved(output_root))
|
||||
if relative is None:
|
||||
return None
|
||||
subfolder, filename = os.path.split(relative)
|
||||
if not filename:
|
||||
return None
|
||||
return root_key, subfolder, filename
|
||||
|
||||
|
||||
def in_output_folder(image_path: str, output_root: Optional[str]) -> bool:
|
||||
"""Whether history could ever attribute ``image_path`` (inside the host's
|
||||
output folder, in either spelling)."""
|
||||
if not image_path or not output_root:
|
||||
return False
|
||||
style = "windows" if os.name == "nt" else "posix"
|
||||
return _location(image_path, output_root, style) is not None
|
||||
|
||||
|
||||
def _written_during(image_path: str, started: Any, finished: Any) -> bool:
|
||||
if not isinstance(started, int) or not isinstance(finished, int):
|
||||
return False
|
||||
try:
|
||||
written = os.stat(image_path).st_mtime_ns // 1_000_000
|
||||
except OSError:
|
||||
return False
|
||||
return started - MTIME_SLACK_MS <= written <= finished + MTIME_SLACK_MS
|
||||
|
||||
|
||||
def stored_run_prompt(
|
||||
db: Any,
|
||||
image_path: str,
|
||||
output_root: Optional[str],
|
||||
path_style: Optional[str] = None,
|
||||
) -> Optional[int]:
|
||||
"""The stored prompt id history attributes to ``image_path``, else None."""
|
||||
if not image_path:
|
||||
return None
|
||||
return stored_run_prompts(db, [image_path], output_root, path_style).get(image_path)
|
||||
|
||||
|
||||
def _runs_by_location(db: Any, names: List[str]) -> Optional[Dict[tuple, dict]]:
|
||||
"""``{(filename, subfolder, root_key, style): {run_id: (start, end)}}``.
|
||||
|
||||
One query per ``LOOKUP_CHUNK`` names; None without history tables.
|
||||
"""
|
||||
found: Dict[tuple, dict] = {}
|
||||
conn = db.model.get_connection()
|
||||
for start in range(0, len(names), LOOKUP_CHUNK):
|
||||
chunk = names[start : start + LOOKUP_CHUNK]
|
||||
query = _RUNS_FOR_NAMES.format(marks=",".join("?" for _ in chunk))
|
||||
try:
|
||||
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:
|
||||
runs = found.setdefault((name, subfolder, root_key, style), {})
|
||||
runs[run_id] = (began, ended)
|
||||
return found
|
||||
|
||||
|
||||
def _named_prompt(db: Any, run_id: str) -> Optional[int]:
|
||||
contexts = db.get_execution_contexts(run_id).get("contexts") or []
|
||||
named = {
|
||||
context.get("positive_prompt_id")
|
||||
for context in contexts
|
||||
if context.get("positive") is not None
|
||||
}
|
||||
if len(named) != 1:
|
||||
return None
|
||||
(prompt_id,) = named
|
||||
return prompt_id if isinstance(prompt_id, int) else None
|
||||
|
||||
|
||||
def stored_run_prompts(
|
||||
db: Any,
|
||||
image_paths: Iterable[str],
|
||||
output_root: Optional[str],
|
||||
path_style: Optional[str] = None,
|
||||
) -> Dict[str, Optional[int]]:
|
||||
"""``{path: stored prompt id or None}`` for many files in few queries.
|
||||
|
||||
The output lookup is one query per ``LOOKUP_CHUNK`` distinct names; each
|
||||
run that alone reported a file is then read once for its contexts.
|
||||
"""
|
||||
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
|
||||
if run_id not in prompts:
|
||||
prompts[run_id] = _named_prompt(db, run_id)
|
||||
result[path] = prompts[run_id]
|
||||
return result
|
||||
|
||||
|
||||
_LOGGER = logging.getLogger("prompt_manager.run_attribution")
|
||||
_SEALED_LISTENERS: List[weakref.WeakMethod] = []
|
||||
_SEALED_LOCK = threading.Lock()
|
||||
|
||||
|
||||
def subscribe_outputs_sealed(method: Callable[[str, List[dict]], None]) -> None:
|
||||
"""Call bound ``method(run_id, references)`` after a run's outputs are recorded.
|
||||
|
||||
Held weakly: a listener that is garbage collected is dropped.
|
||||
"""
|
||||
with _SEALED_LOCK:
|
||||
_SEALED_LISTENERS.append(weakref.WeakMethod(method))
|
||||
|
||||
|
||||
def notify_outputs_sealed(run_id: str, references: Iterable[dict]) -> None:
|
||||
"""Tell listeners a run's output references were just committed.
|
||||
|
||||
A listener failure is logged and never reaches the history writer.
|
||||
"""
|
||||
references = list(references)
|
||||
with _SEALED_LOCK:
|
||||
alive = [ref for ref in _SEALED_LISTENERS if ref() is not None]
|
||||
_SEALED_LISTENERS[:] = alive
|
||||
listeners = [ref() for ref in alive]
|
||||
for listener in listeners:
|
||||
if listener is None:
|
||||
continue
|
||||
try:
|
||||
listener(run_id, references)
|
||||
except Exception:
|
||||
_LOGGER.warning("Output seal listener failed", exc_info=True)
|
||||
@@ -1,185 +0,0 @@
|
||||
"""Safety backups beside the live database: one writer for every caller.
|
||||
|
||||
A safety backup is ``<db>.<label>_<YYYYmmdd_HHMMSS_ffffff>[_n]`` in the
|
||||
library's own folder. It is written by SQLite's backup API into a temp file
|
||||
beside the database (``SAFETY_TEMP_PREFIX``, swept at startup after a crash),
|
||||
made one self-contained file (rollback journal), optionally verified, fsynced,
|
||||
made private (``0o600`` on POSIX) and only then given its unique name, so a
|
||||
backup name never carries a half-written file. The names are unique within
|
||||
one second (microseconds, then a counter): on Windows two backups in the same
|
||||
second must not overwrite each other.
|
||||
|
||||
Maintenance applies take a verified backup (``py/api/maintenance.py``); a
|
||||
restore takes an unverified one, because the database it replaces may be the
|
||||
damaged one the user is recovering from (``database/restore_stage.py``).
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import os
|
||||
import sqlite3
|
||||
import tempfile
|
||||
from contextlib import closing
|
||||
|
||||
SAFETY_TEMP_PREFIX = ".pm-safety-"
|
||||
_NAME_ATTEMPTS = 1000
|
||||
|
||||
|
||||
class SafetyBackupFailed(Exception):
|
||||
"""The safety backup could not be written, verified or put in place.
|
||||
|
||||
``reason`` is the verify reason when the copy did not verify.
|
||||
"""
|
||||
|
||||
def __init__(self, message, reason=None):
|
||||
super().__init__(message)
|
||||
self.reason = reason
|
||||
|
||||
|
||||
def _now():
|
||||
return datetime.datetime.now()
|
||||
|
||||
|
||||
def remove_quietly(*paths):
|
||||
"""Remove each existing path, ignoring errors (cleanup after a failure)."""
|
||||
for path in paths:
|
||||
try:
|
||||
if path and os.path.exists(path):
|
||||
os.remove(path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def backup_names(db_path, label="backup"):
|
||||
"""Candidate names, unique within one second (microseconds, then ``_n``)."""
|
||||
base = f"{db_path}.{label}_{_now().strftime('%Y%m%d_%H%M%S_%f')}"
|
||||
for attempt in range(_NAME_ATTEMPTS):
|
||||
yield base if attempt == 0 else f"{base}_{attempt}"
|
||||
|
||||
|
||||
def reserve_backup_name(db_path, label="backup"):
|
||||
"""Create and return an unused ``<db>.<label>_<stamp>[_n]`` (O_EXCL)."""
|
||||
for candidate in backup_names(db_path, label):
|
||||
try:
|
||||
fd = os.open(candidate, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
|
||||
except FileExistsError:
|
||||
continue
|
||||
os.close(fd)
|
||||
return candidate
|
||||
raise SafetyBackupFailed("No free backup name")
|
||||
|
||||
|
||||
def place_backup(temp, db_path, label="backup"):
|
||||
"""Give the finished ``temp`` its unique backup name; return that path.
|
||||
|
||||
``os.link`` is atomic and refuses an existing name, so no empty or half
|
||||
file ever carries a backup name, even across a crash. A filesystem
|
||||
without hard links falls back to reserving the name (``O_EXCL``) and
|
||||
``os.replace``. The temp file is removed by the caller.
|
||||
"""
|
||||
for candidate in backup_names(db_path, label):
|
||||
try:
|
||||
os.link(temp, candidate)
|
||||
except FileExistsError:
|
||||
continue
|
||||
except (OSError, NotImplementedError):
|
||||
break # no hard links here
|
||||
else:
|
||||
return candidate
|
||||
else:
|
||||
raise SafetyBackupFailed("No free backup name")
|
||||
final = reserve_backup_name(db_path, label)
|
||||
try:
|
||||
os.replace(temp, final)
|
||||
except OSError:
|
||||
remove_quietly(final)
|
||||
raise
|
||||
return final
|
||||
|
||||
|
||||
def sweep_safety_temps(db_path):
|
||||
"""Remove the temp copies an interrupted safety backup left (blocking).
|
||||
|
||||
Called once at startup, before any job can take a backup. Returns the
|
||||
names removed.
|
||||
"""
|
||||
folder = os.path.dirname(os.path.abspath(db_path))
|
||||
removed = []
|
||||
try:
|
||||
names = os.listdir(folder)
|
||||
except OSError:
|
||||
return removed
|
||||
for name in names:
|
||||
if not (name.startswith(SAFETY_TEMP_PREFIX) and ".tmp" in name):
|
||||
continue
|
||||
try:
|
||||
os.remove(os.path.join(folder, name))
|
||||
removed.append(name)
|
||||
except OSError:
|
||||
continue
|
||||
return removed
|
||||
|
||||
|
||||
def make_standalone(path):
|
||||
"""Make the copy one self-contained file (rollback journal, not WAL).
|
||||
|
||||
The backup API copies the live database's WAL flag; a WAL file opened
|
||||
read-only later (a verify, a restore preview) would leave ``-wal`` and
|
||||
``-shm`` sidecars beside it.
|
||||
"""
|
||||
try:
|
||||
with closing(sqlite3.connect(path)) as conn:
|
||||
conn.execute("PRAGMA journal_mode=DELETE").fetchone()
|
||||
except sqlite3.Error as exc:
|
||||
raise SafetyBackupFailed("The backup could not be finished") from exc
|
||||
|
||||
|
||||
def fsync_file(path):
|
||||
with open(path, "rb+") as handle:
|
||||
os.fsync(handle.fileno())
|
||||
|
||||
|
||||
def safety_temp(db_path):
|
||||
"""A new, closed temp file beside the database (``SAFETY_TEMP_PREFIX``)."""
|
||||
fd, temp = tempfile.mkstemp(
|
||||
dir=os.path.dirname(os.path.abspath(db_path)),
|
||||
prefix=SAFETY_TEMP_PREFIX,
|
||||
suffix=".tmp",
|
||||
)
|
||||
os.close(fd)
|
||||
return temp
|
||||
|
||||
|
||||
def take_safety_backup(model, label="backup", *, verify=True):
|
||||
"""Write (and by default verify) a safety backup; return its path (blocking).
|
||||
|
||||
The copy is made in a temp file beside the database and only renamed
|
||||
onto its final, unique name ``<db>.<label>_<stamp>`` once it is complete
|
||||
(and verified), so a backup file on disk is always a whole one. ``label``
|
||||
"backup" joins the pruned pool; another label (the workflow storage
|
||||
migration's ``pre-workflow-storage``) is kept until the user removes it.
|
||||
Raises SafetyBackupFailed, leaving no file.
|
||||
"""
|
||||
db_path = os.path.abspath(model.db_path)
|
||||
temp = safety_temp(db_path)
|
||||
final = None
|
||||
try:
|
||||
if not model.backup_database(temp):
|
||||
raise SafetyBackupFailed("The backup could not be written")
|
||||
make_standalone(temp)
|
||||
if verify:
|
||||
ok, reason = model.verify_database_file(temp)
|
||||
if not ok:
|
||||
raise SafetyBackupFailed("The backup did not verify", reason)
|
||||
fsync_file(temp)
|
||||
if os.name != "nt":
|
||||
os.chmod(temp, 0o600)
|
||||
final = place_backup(temp, db_path, label)
|
||||
return final
|
||||
except OSError as exc:
|
||||
remove_quietly(final)
|
||||
raise SafetyBackupFailed("The backup could not be put in place") from exc
|
||||
except SafetyBackupFailed:
|
||||
remove_quietly(final)
|
||||
raise
|
||||
finally:
|
||||
remove_quietly(temp, temp + "-journal", temp + "-wal", temp + "-shm")
|
||||
@@ -1,652 +0,0 @@
|
||||
"""Each workflow and prompt graph stored once, compressed (#249 S6 carry-over).
|
||||
|
||||
``generated_images.workflow_data`` (the UI workflow) and ``prompt_metadata``
|
||||
(the API prompt graph) used to hold the same large JSON text again for every
|
||||
image of a run. Now each distinct text is one ``workflow_documents`` row:
|
||||
|
||||
- ``sha256``: the SHA-256 of the text's exact UTF-8 bytes (unique key);
|
||||
- ``byte_length``: the length of those bytes;
|
||||
- ``data``: ``zlib.compress`` of those bytes.
|
||||
|
||||
An image row references its texts through ``workflow_document_id`` and
|
||||
``prompt_document_id``. Two forms coexist until the migration has rewritten
|
||||
every row, and both stay readable forever:
|
||||
|
||||
- old form: the text inline in ``workflow_data`` / ``prompt_metadata``, no id;
|
||||
- new form: the inline column NULL and the document id set.
|
||||
|
||||
The text is stored and hashed exactly as written, never parsed and
|
||||
re-serialized, because a JSON round trip can change bytes (key order,
|
||||
spacing, number spelling, duplicate keys, ComfyUI's ``NaN`` tokens).
|
||||
|
||||
``read_stored_text`` is the one decoder. Inline text, when present, wins (a
|
||||
row an older app version rewrote inline stays right); otherwise the document
|
||||
is read by its rowid, decompressed within its recorded length and checked
|
||||
against its length and hash. A document that fails a check reads as
|
||||
unreadable (None, logged), never as other text. ``resolve_image_row`` puts
|
||||
the texts back into a ``generated_images`` row and drops the id keys, so
|
||||
every payload keeps its shape.
|
||||
|
||||
No trigger: every restore validator refuses triggers in a restored file. A
|
||||
document is removed in code: ``release_documents`` in the same transaction
|
||||
that drops a row's last reference, and ``sweep_documents`` (at startup, after
|
||||
the migration and before Reclaim space) for anything a path missed.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import re
|
||||
import sqlite3
|
||||
import zlib
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Dict, Iterable, Optional, Set
|
||||
|
||||
logger = logging.getLogger("prompt_manager.database.workflow_documents")
|
||||
|
||||
TABLE = "workflow_documents"
|
||||
# (inline text column, document id column) on generated_images.
|
||||
TEXT_COLUMNS = (
|
||||
("workflow_data", "workflow_document_id"),
|
||||
("prompt_metadata", "prompt_document_id"),
|
||||
)
|
||||
DOCUMENT_ID_COLUMNS = tuple(reference for _column, reference in TEXT_COLUMNS)
|
||||
INDEXES = {
|
||||
"workflow_document_id": "idx_image_workflow_document",
|
||||
"prompt_document_id": "idx_image_prompt_document",
|
||||
}
|
||||
COMPRESSION_LEVEL = 6
|
||||
# A text larger than this stays inline: a document is decompressed in one
|
||||
# piece, and its recorded length bounds that (a restored file cannot make a
|
||||
# small row expand without limit).
|
||||
MAX_DOCUMENT_BYTES = 256 * 2**20
|
||||
ID_BATCH = 400 # ids bound per statement (under SQLite's old 999 limit)
|
||||
|
||||
_DDL = (
|
||||
f"CREATE TABLE {TABLE} (id INTEGER PRIMARY KEY, sha256 BLOB NOT NULL UNIQUE, "
|
||||
"byte_length INTEGER NOT NULL, data BLOB NOT NULL)"
|
||||
)
|
||||
_AUTOINDEX = f"sqlite_autoindex_{TABLE}_1"
|
||||
_REFERENCE_TYPE = f"INTEGER REFERENCES {TABLE}(id)"
|
||||
|
||||
_UNREFERENCED = (
|
||||
"NOT EXISTS (SELECT 1 FROM generated_images g "
|
||||
"WHERE g.workflow_document_id=workflow_documents.id) "
|
||||
"AND NOT EXISTS (SELECT 1 FROM generated_images g "
|
||||
"WHERE g.prompt_document_id=workflow_documents.id)"
|
||||
)
|
||||
LOOKUP_SQL = (
|
||||
"SELECT id, sha256, byte_length, data FROM workflow_documents WHERE sha256=?"
|
||||
)
|
||||
REPAIR_SQL = "UPDATE workflow_documents SET byte_length=?, data=? WHERE id=?"
|
||||
ALL_DOCUMENTS_SQL = "SELECT sha256, byte_length, data, id FROM workflow_documents"
|
||||
READ_SQL = "SELECT sha256, byte_length, data FROM workflow_documents WHERE id=?"
|
||||
INSERT_SQL = (
|
||||
"INSERT INTO workflow_documents (sha256, byte_length, data) VALUES (?, ?, ?)"
|
||||
)
|
||||
# Built from constants only; the ids are bound as "?" marks.
|
||||
RELEASE_SQL = "DELETE FROM workflow_documents WHERE id IN ({marks}) AND "
|
||||
RELEASE_SQL += _UNREFERENCED
|
||||
# The sweep walks the ids through the small unique index, not the blobs.
|
||||
SWEEP_SQL = (
|
||||
"DELETE FROM workflow_documents WHERE id IN (SELECT id FROM workflow_documents "
|
||||
"INDEXED BY sqlite_autoindex_workflow_documents_1 WHERE "
|
||||
)
|
||||
SWEEP_SQL += _UNREFERENCED + ")"
|
||||
IMAGE_DOCUMENTS_SQL = (
|
||||
"SELECT workflow_document_id, prompt_document_id FROM generated_images WHERE "
|
||||
)
|
||||
|
||||
# Statements whose query plans tests/test_workflow_documents.py records.
|
||||
QUERY_PLAN_STATEMENTS = {
|
||||
"lookup by hash": (LOOKUP_SQL, (b"\0" * 32,)),
|
||||
"read by id": (READ_SQL, (1,)),
|
||||
"release": (RELEASE_SQL.format(marks="?,?"), (1, 2)),
|
||||
"sweep": (SWEEP_SQL, ()),
|
||||
"documents of images": (IMAGE_DOCUMENTS_SQL + "id IN (?,?)", (1, 2)),
|
||||
"documents of prompts": (IMAGE_DOCUMENTS_SQL + "prompt_id IN (?,?)", (1, 2)),
|
||||
}
|
||||
|
||||
|
||||
class DocumentMismatch(Exception):
|
||||
"""A document would not give back exactly the text it was made from."""
|
||||
|
||||
|
||||
def _compact(sql: str) -> str:
|
||||
return re.sub(r"\s+", "", sql or "").lower()
|
||||
|
||||
|
||||
# ── schema ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _reference_columns(conn: sqlite3.Connection) -> Dict[str, tuple]:
|
||||
return {
|
||||
row[1]: tuple(row)
|
||||
for row in conn.execute("PRAGMA table_info(generated_images)")
|
||||
if row[1] in DOCUMENT_ID_COLUMNS
|
||||
}
|
||||
|
||||
|
||||
def _reference_error(conn: sqlite3.Connection, columns: Dict[str, tuple]):
|
||||
for name, info in columns.items():
|
||||
# (cid, name, type, notnull, default, pk)
|
||||
if (str(info[2]).upper(), info[3], info[4], info[5]) != ("INTEGER", 0, None, 0):
|
||||
return "Workflow document references have an invalid column"
|
||||
keys = {
|
||||
(row[3], row[2], row[4])
|
||||
for row in conn.execute("PRAGMA foreign_key_list(generated_images)")
|
||||
if row[3] in DOCUMENT_ID_COLUMNS
|
||||
}
|
||||
if keys != {(name, TABLE, "id") for name in DOCUMENT_ID_COLUMNS}:
|
||||
return "Workflow document references have invalid foreign keys"
|
||||
for column, index in INDEXES.items():
|
||||
info = conn.execute(f'PRAGMA index_info("{index}")').fetchall()
|
||||
owner = conn.execute(
|
||||
"SELECT tbl_name FROM sqlite_master WHERE type='index' AND name=?",
|
||||
(index,),
|
||||
).fetchone()
|
||||
if [row[2] for row in info] != [column] or not owner:
|
||||
return "Workflow document references are missing their index"
|
||||
if owner[0] != "generated_images":
|
||||
return "Workflow document references are missing their index"
|
||||
return None
|
||||
|
||||
|
||||
def _value_error(conn: sqlite3.Connection):
|
||||
if conn.execute(
|
||||
"SELECT 1 FROM workflow_documents WHERE typeof(sha256)!='blob' "
|
||||
"OR length(sha256)!=32 OR typeof(byte_length)!='integer' OR byte_length<0 "
|
||||
"OR byte_length>? OR typeof(data)!='blob' LIMIT 1",
|
||||
(MAX_DOCUMENT_BYTES,),
|
||||
).fetchone():
|
||||
return "Workflow documents contain invalid values"
|
||||
for column in DOCUMENT_ID_COLUMNS:
|
||||
# typeof() reads no overflow page; a dangling id is a rowid probe.
|
||||
# ``column`` is one of DOCUMENT_ID_COLUMNS, never caller data.
|
||||
sql = (
|
||||
f"SELECT 1 FROM generated_images g WHERE typeof(g.{column}) " # nosec B608
|
||||
f"NOT IN ('null','integer') OR (g.{column} IS NOT NULL AND NOT EXISTS "
|
||||
f"(SELECT 1 FROM workflow_documents d WHERE d.id=g.{column})) LIMIT 1"
|
||||
)
|
||||
if conn.execute(sql).fetchone():
|
||||
return "Workflow document references point at missing documents"
|
||||
return None
|
||||
|
||||
|
||||
def damaged_message(document_id: int) -> str:
|
||||
"""The verify reason for a damaged document: which one, and what to do."""
|
||||
return (
|
||||
f"Workflow document {document_id} is damaged: its text no longer matches "
|
||||
"its checksum, and the images that use it show no workflow. To fix it, "
|
||||
"restore a backup from before the damage, or store a new image with the "
|
||||
"same workflow (that repairs it)."
|
||||
)
|
||||
|
||||
|
||||
def _damaged_document(conn: sqlite3.Connection) -> Optional[int]:
|
||||
"""The first document failing its length or hash check, or None."""
|
||||
for row in conn.execute(ALL_DOCUMENTS_SQL):
|
||||
if _decode(row[:3], row[3], quiet=True) is None:
|
||||
return row[3]
|
||||
return None
|
||||
|
||||
|
||||
def workflow_documents_schema_error(
|
||||
conn: sqlite3.Connection, *, allow_missing: bool = True, documents: bool = False
|
||||
) -> Optional[str]:
|
||||
"""Reject an unsupported document schema or invalid stored values (reads only).
|
||||
|
||||
A database from before the dedupe has neither the table nor the id
|
||||
columns; that is acceptable unless ``allow_missing`` is false. Half of
|
||||
the schema never is.
|
||||
"""
|
||||
objects = conn.execute(
|
||||
"SELECT type,name,sql FROM sqlite_master WHERE name=? "
|
||||
"OR (type IN ('index','trigger') AND tbl_name=?)",
|
||||
(TABLE, TABLE),
|
||||
).fetchall()
|
||||
has_images = conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND name='generated_images'"
|
||||
).fetchone()
|
||||
columns = _reference_columns(conn) if has_images else {}
|
||||
if not objects and not columns:
|
||||
if allow_missing or not has_images:
|
||||
return None
|
||||
return "Missing workflow document schema"
|
||||
if not objects or set(columns) != set(DOCUMENT_ID_COLUMNS):
|
||||
return "Incomplete workflow document schema"
|
||||
tables = [sql for kind, _name, sql in objects if kind == "table"]
|
||||
others = {(kind, name, sql) for kind, name, sql in objects if kind != "table"}
|
||||
if len(tables) != 1 or others != {("index", _AUTOINDEX, None)}:
|
||||
return "Workflow documents name is reserved by an unsupported schema object"
|
||||
# The exact definition fixes the columns, their types and constraints.
|
||||
if _compact(tables[0]) != _compact(_DDL):
|
||||
return "Workflow documents have an invalid table definition"
|
||||
error = _reference_error(conn, columns) or _value_error(conn)
|
||||
if error is None and documents:
|
||||
# Restore and backup verification only: every document is read.
|
||||
damaged = _damaged_document(conn)
|
||||
if damaged is not None:
|
||||
return damaged_message(damaged)
|
||||
return error
|
||||
|
||||
|
||||
def migrate_workflow_documents(conn: sqlite3.Connection) -> None:
|
||||
"""Install the table, the id columns and their indexes atomically.
|
||||
|
||||
Runs after every older migration (the table rebuilds only run on
|
||||
databases without ``file_path``, before this). Validates before and
|
||||
after, then sweeps documents no row references.
|
||||
"""
|
||||
conn.execute("SAVEPOINT workflow_documents_schema")
|
||||
try:
|
||||
error = workflow_documents_schema_error(conn)
|
||||
if error:
|
||||
raise sqlite3.DatabaseError(error)
|
||||
conn.execute(_DDL.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS", 1))
|
||||
present = _reference_columns(conn)
|
||||
for column in DOCUMENT_ID_COLUMNS:
|
||||
if column not in present:
|
||||
conn.execute(
|
||||
f"ALTER TABLE generated_images ADD COLUMN {column} "
|
||||
f"{_REFERENCE_TYPE}"
|
||||
)
|
||||
for column, index in INDEXES.items():
|
||||
conn.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS {index} ON generated_images({column})"
|
||||
)
|
||||
error = workflow_documents_schema_error(conn, allow_missing=False)
|
||||
if error:
|
||||
raise sqlite3.DatabaseError(error)
|
||||
sweep_documents(conn)
|
||||
conn.execute("RELEASE workflow_documents_schema")
|
||||
except BaseException:
|
||||
conn.execute("ROLLBACK TO workflow_documents_schema")
|
||||
conn.execute("RELEASE workflow_documents_schema")
|
||||
raise
|
||||
|
||||
|
||||
# ── write ─────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def storable(text: Any) -> bool:
|
||||
"""Whether ``text`` can be a document (text within MAX_DOCUMENT_BYTES)."""
|
||||
if not isinstance(text, str):
|
||||
return False
|
||||
# UTF-8 takes at most 4 bytes a character: most texts need no encoding.
|
||||
if len(text) * 4 <= MAX_DOCUMENT_BYTES:
|
||||
return True
|
||||
return len(text.encode("utf-8")) <= MAX_DOCUMENT_BYTES
|
||||
|
||||
|
||||
def store_text(
|
||||
conn: sqlite3.Connection, text: Optional[str], verified: Optional[set] = None
|
||||
) -> Optional[int]:
|
||||
"""The id of the document holding exactly ``text``; None for None.
|
||||
|
||||
Reuses the document with the same hash, or inserts a new one whose
|
||||
compressed bytes were checked to decompress to ``text``'s bytes. A reused
|
||||
document is decompressed and checked first; a damaged one (a restored
|
||||
file's corruption under a valid hash) is repaired from ``text``, whose
|
||||
hash is its key, so every row that references it reads right again.
|
||||
``verified``, when given, holds the ids already checked (a migration
|
||||
batch), which are then not read again. Opens a write transaction first
|
||||
when none is open, so a concurrent ``release_documents`` cannot remove a
|
||||
reused document before the caller references it. Raises TypeError for a
|
||||
non-text value and ValueError for a text over MAX_DOCUMENT_BYTES.
|
||||
"""
|
||||
if text is None:
|
||||
return None
|
||||
if not isinstance(text, str):
|
||||
raise TypeError("Only text is stored as a workflow document")
|
||||
raw = text.encode("utf-8")
|
||||
if len(raw) > MAX_DOCUMENT_BYTES:
|
||||
raise ValueError("Text is too large for a workflow document")
|
||||
digest = hashlib.sha256(raw).digest()
|
||||
if not conn.in_transaction:
|
||||
conn.execute("BEGIN IMMEDIATE")
|
||||
row = conn.execute(LOOKUP_SQL, (digest,)).fetchone()
|
||||
if row is not None and verified is not None and row[0] in verified:
|
||||
return row[0]
|
||||
data = zlib.compress(raw, COMPRESSION_LEVEL)
|
||||
if zlib.decompress(data) != raw: # pragma: no cover - zlib is exact
|
||||
raise DocumentMismatch("Compression did not round-trip")
|
||||
if row is None:
|
||||
document_id = conn.execute(INSERT_SQL, (digest, len(raw), data)).lastrowid
|
||||
else:
|
||||
document_id = row[0]
|
||||
if _decode(row[1:], document_id, quiet=True) != text:
|
||||
logger.warning(f"Workflow document {document_id} was damaged; repaired")
|
||||
conn.execute(REPAIR_SQL, (len(raw), data, document_id))
|
||||
if verified is not None:
|
||||
verified.add(document_id)
|
||||
return document_id
|
||||
|
||||
|
||||
def stored_values(
|
||||
conn: sqlite3.Connection, texts: Dict[str, Optional[str]]
|
||||
) -> Dict[str, Any]:
|
||||
"""New-form column values for ``{text column: text}`` (writers).
|
||||
|
||||
Each text becomes ``{column: None, id column: document id}``; None
|
||||
stays ``{column: None, id column: None}``; a value that cannot be a
|
||||
document (not text, or too large) stays inline with no id.
|
||||
"""
|
||||
references = dict(TEXT_COLUMNS)
|
||||
values: Dict[str, Any] = {}
|
||||
for column, text in texts.items():
|
||||
reference = references[column]
|
||||
if text is None or not storable(text):
|
||||
values[column], values[reference] = text, None
|
||||
else:
|
||||
values[column], values[reference] = None, store_text(conn, text)
|
||||
return values
|
||||
|
||||
|
||||
# ── read ──────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _decode(row, document_id, quiet=False) -> Optional[str]:
|
||||
"""``(sha256, byte_length, data)`` as checked text, or None (logged)."""
|
||||
if row is None:
|
||||
logger.warning(f"Workflow document {document_id} is missing")
|
||||
return None
|
||||
digest, length, data = bytes(row[0]), row[1], row[2]
|
||||
try:
|
||||
if not isinstance(length, int) or not 0 <= length <= MAX_DOCUMENT_BYTES:
|
||||
raise ValueError("invalid length")
|
||||
inflater = zlib.decompressobj()
|
||||
raw = inflater.decompress(bytes(data), length + 1)
|
||||
if len(raw) != length or inflater.unconsumed_tail or not inflater.eof:
|
||||
raise ValueError("length differs")
|
||||
if hashlib.sha256(raw).digest() != digest:
|
||||
raise ValueError("hash differs")
|
||||
return raw.decode("utf-8")
|
||||
except (zlib.error, ValueError, TypeError) as error:
|
||||
if not quiet:
|
||||
logger.warning(f"Workflow document {document_id} is unreadable: {error}")
|
||||
return None
|
||||
|
||||
|
||||
def read_stored_text(
|
||||
conn: sqlite3.Connection, inline: Any, document_id: Optional[int]
|
||||
) -> Any:
|
||||
"""The stored text of one column: inline when present, else the document.
|
||||
|
||||
Returns None when there is neither, or when the document is missing or
|
||||
fails its checks (logged). An inline value is returned unchanged.
|
||||
"""
|
||||
if inline is not None:
|
||||
return inline
|
||||
if document_id is None:
|
||||
return None
|
||||
return _decode(conn.execute(READ_SQL, (document_id,)).fetchone(), document_id)
|
||||
|
||||
|
||||
class DocumentReader:
|
||||
"""``read_stored_text`` with a small memo for list reads on one connection."""
|
||||
|
||||
MEMO = 32
|
||||
|
||||
def __init__(self, conn: sqlite3.Connection):
|
||||
self.conn = conn
|
||||
self._memo: "OrderedDict[int, Optional[str]]" = OrderedDict()
|
||||
|
||||
def text(self, inline: Any, document_id: Optional[int]) -> Any:
|
||||
if inline is not None or document_id is None:
|
||||
return read_stored_text(self.conn, inline, document_id)
|
||||
if document_id in self._memo:
|
||||
self._memo.move_to_end(document_id)
|
||||
return self._memo[document_id]
|
||||
text = read_stored_text(self.conn, None, document_id)
|
||||
self._memo[document_id] = text
|
||||
if len(self._memo) > self.MEMO:
|
||||
self._memo.popitem(last=False)
|
||||
return text
|
||||
|
||||
|
||||
def resolve_image_row(reader: DocumentReader, row) -> Dict[str, Any]:
|
||||
"""A ``generated_images`` row as a dict with its texts read back.
|
||||
|
||||
The id columns are dropped, so a payload built from it has the shape it
|
||||
had before the dedupe. Columns the row did not select stay absent.
|
||||
"""
|
||||
data = dict(row)
|
||||
for column, reference in TEXT_COLUMNS:
|
||||
has_reference = reference in data
|
||||
document_id = data.pop(reference, None)
|
||||
if column in data or has_reference:
|
||||
data[column] = reader.text(data.get(column), document_id)
|
||||
return data
|
||||
|
||||
|
||||
# ── garbage ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _chunks(values):
|
||||
values = list(values)
|
||||
for start in range(0, len(values), ID_BATCH):
|
||||
yield values[start : start + ID_BATCH]
|
||||
|
||||
|
||||
def image_documents(conn: sqlite3.Connection, where: str, params) -> Set[int]:
|
||||
"""The document ids the image rows matching ``where`` reference.
|
||||
|
||||
``where`` is a fixed condition written in this package (``id IN (...)``,
|
||||
``prompt_id IN (...)``), never caller data; values go in ``params``.
|
||||
"""
|
||||
ids: Set[int] = set()
|
||||
sql = IMAGE_DOCUMENTS_SQL + where
|
||||
for row in conn.execute(sql, list(params)): # nosec B608
|
||||
ids.update(value for value in row if value is not None)
|
||||
return ids
|
||||
|
||||
|
||||
def documents_of_rows(
|
||||
conn: sqlite3.Connection, column: str, values: Iterable[Any]
|
||||
) -> Set[int]:
|
||||
"""Document ids of the image rows whose ``column`` is in ``values``.
|
||||
|
||||
``column`` is ``id`` or ``prompt_id`` (the rows a prompt delete cascades to).
|
||||
"""
|
||||
if column not in ("id", "prompt_id"):
|
||||
raise ValueError(f"Unsupported column {column!r}")
|
||||
ids: Set[int] = set()
|
||||
for chunk in _chunks(values):
|
||||
marks = ",".join("?" for _ in chunk)
|
||||
ids |= image_documents(conn, f"{column} IN ({marks})", chunk)
|
||||
return ids
|
||||
|
||||
|
||||
def release_documents(conn: sqlite3.Connection, ids: Iterable[Optional[int]]) -> int:
|
||||
"""Delete the documents in ``ids`` that no image row references now.
|
||||
|
||||
Runs in the caller's transaction, right after the delete or update that
|
||||
dropped the references. Returns how many documents were removed.
|
||||
"""
|
||||
removed = 0
|
||||
for chunk in _chunks(sorted({i for i in ids if i is not None})):
|
||||
marks = ",".join("?" for _ in chunk)
|
||||
# Only "?" marks are interpolated.
|
||||
sql = RELEASE_SQL.format(marks=marks)
|
||||
removed += conn.execute(sql, chunk).rowcount # nosec B608
|
||||
return removed
|
||||
|
||||
|
||||
def sweep_documents(conn: sqlite3.Connection) -> int:
|
||||
"""Delete every document no image row references; returns the count."""
|
||||
return conn.execute(SWEEP_SQL).rowcount
|
||||
|
||||
|
||||
# ── repair (a restore's stage) ────────────────────────────────────────
|
||||
|
||||
INLINE_TEXTS_SQL = (
|
||||
"SELECT workflow_data, prompt_metadata FROM generated_images "
|
||||
"WHERE workflow_data IS NOT NULL OR prompt_metadata IS NOT NULL"
|
||||
)
|
||||
REKEY_SQL = "UPDATE workflow_documents SET sha256=?, byte_length=? WHERE id=?"
|
||||
|
||||
|
||||
def _rows(conn: sqlite3.Connection, sql: str, params=()):
|
||||
"""Each row of ``sql``, read one at a time.
|
||||
|
||||
Iterating a managed connection's cursor prefetches ``ITER_CHUNK`` (128)
|
||||
rows, and a document or inline text can be megabytes.
|
||||
"""
|
||||
cursor = conn.execute(sql, params)
|
||||
while (row := cursor.fetchone()) is not None:
|
||||
yield row
|
||||
|
||||
|
||||
def _inflated(data) -> Optional[bytes]:
|
||||
"""The complete text bytes ``data`` inflates to, or None.
|
||||
|
||||
Complete means the zlib stream ends (its Adler-32 checked) within
|
||||
MAX_DOCUMENT_BYTES with nothing after it: the body is intact even when
|
||||
the row's recorded hash or length is not.
|
||||
"""
|
||||
try:
|
||||
inflater = zlib.decompressobj()
|
||||
raw = inflater.decompress(bytes(data), MAX_DOCUMENT_BYTES + 1)
|
||||
if not inflater.eof or inflater.unconsumed_tail or inflater.unused_data:
|
||||
return None
|
||||
if len(raw) > MAX_DOCUMENT_BYTES:
|
||||
return None
|
||||
raw.decode("utf-8")
|
||||
return raw
|
||||
except (zlib.error, TypeError, UnicodeDecodeError):
|
||||
return None
|
||||
|
||||
|
||||
def _one_column_off(row, raw: bytes) -> bool:
|
||||
"""Whether only the recorded hash or only the recorded length is wrong.
|
||||
|
||||
The two columns check the body independently. When the body inflates
|
||||
completely and one of them still agrees with it, that column vouches for
|
||||
the body and the other one is the damaged value. When both disagree, the
|
||||
body itself is the suspect (another valid stream), so it is not trusted.
|
||||
"""
|
||||
hash_ok = hashlib.sha256(raw).digest() == bytes(row[0])
|
||||
length_ok = row[1] == len(raw)
|
||||
return hash_ok != length_ok
|
||||
|
||||
|
||||
def _inline_copies(conn: sqlite3.Connection, wanted: Dict[bytes, int]):
|
||||
"""``{document id: text}`` for each wanted hash an inline row still holds."""
|
||||
found: Dict[int, str] = {}
|
||||
for row in _rows(conn, INLINE_TEXTS_SQL):
|
||||
for text in row:
|
||||
if not isinstance(text, str) or not storable(text):
|
||||
continue
|
||||
document_id = wanted.get(hashlib.sha256(text.encode("utf-8")).digest())
|
||||
if document_id is not None and document_id not in found:
|
||||
found[document_id] = text
|
||||
if len(found) == len(wanted):
|
||||
break
|
||||
return found
|
||||
|
||||
|
||||
def _rekey(conn: sqlite3.Connection, document_id: int, raw: bytes) -> Optional[int]:
|
||||
"""Give an intact body its true key, or merge it into the document that has it.
|
||||
|
||||
Returns the holder's id when its own body was damaged and has just been
|
||||
rewritten from ``raw`` (it is repaired, never to be cleared), else None.
|
||||
"""
|
||||
digest = hashlib.sha256(raw).digest()
|
||||
holder = conn.execute(LOOKUP_SQL, (digest,)).fetchone()
|
||||
if holder is None or holder[0] == document_id:
|
||||
conn.execute(REKEY_SQL, (digest, len(raw), document_id))
|
||||
return None
|
||||
for column in DOCUMENT_ID_COLUMNS:
|
||||
# ``column`` is one of DOCUMENT_ID_COLUMNS, never caller data.
|
||||
conn.execute(
|
||||
f"UPDATE generated_images SET {column}=? WHERE {column}=?", # nosec B608
|
||||
(holder[0], document_id),
|
||||
)
|
||||
conn.execute("DELETE FROM workflow_documents WHERE id=?", (document_id,))
|
||||
if _decode(holder[1:], holder[0], quiet=True) is None:
|
||||
conn.execute(
|
||||
REPAIR_SQL, (len(raw), zlib.compress(raw, COMPRESSION_LEVEL), holder[0])
|
||||
)
|
||||
return holder[0]
|
||||
return None
|
||||
|
||||
|
||||
def repair_damaged_documents(conn: sqlite3.Connection) -> Dict[str, int]:
|
||||
"""Repair or clear every document failing its checks; commit (writes).
|
||||
|
||||
For a restore's stage only, never the live library. In order:
|
||||
|
||||
1. A document whose body inflates completely and whose recorded hash or
|
||||
length (not both) disagrees is re-keyed from its text, or merged
|
||||
into the document that already holds that text (``_one_column_off``).
|
||||
2. A damaged body whose exact text an inline row still holds (the same
|
||||
SHA-256, its key) is stored again from that text.
|
||||
3. Any other damaged document is unreadable for good: its images'
|
||||
references are cleared (they already showed no workflow) and the
|
||||
document is removed, so the stage verifies.
|
||||
|
||||
Returns ``{repaired, cleared, images}``: documents repaired (1 and 2),
|
||||
documents cleared and the distinct image rows whose reference was
|
||||
cleared. Rows are read one at a time.
|
||||
"""
|
||||
damaged: Dict[bytes, int] = {}
|
||||
intact: Dict[int, bytes] = {}
|
||||
for row in _rows(conn, ALL_DOCUMENTS_SQL):
|
||||
if _decode(row[:3], row[3], quiet=True) is None:
|
||||
raw = _inflated(row[2])
|
||||
if raw is not None and _one_column_off(row, raw):
|
||||
intact[row[3]] = raw
|
||||
else:
|
||||
damaged[bytes(row[0])] = row[3]
|
||||
result = {"repaired": 0, "cleared": 0, "images": 0}
|
||||
if not damaged and not intact:
|
||||
return result
|
||||
with conn:
|
||||
healed = []
|
||||
for document_id, raw in intact.items():
|
||||
holder = _rekey(conn, document_id, raw)
|
||||
if holder is not None and holder in damaged.values():
|
||||
# Rewritten from the merged body: repaired, not to be cleared
|
||||
# (it now also holds the merged images' references).
|
||||
damaged = {k: v for k, v in damaged.items() if v != holder}
|
||||
healed.append(holder)
|
||||
copies = _inline_copies(conn, damaged) if damaged else {}
|
||||
for document_id, text in copies.items():
|
||||
raw = text.encode("utf-8")
|
||||
data = zlib.compress(raw, COMPRESSION_LEVEL)
|
||||
conn.execute(REPAIR_SQL, (len(raw), data, document_id))
|
||||
lost = sorted(set(damaged.values()) - set(copies))
|
||||
images: Set[int] = set()
|
||||
for chunk in _chunks(lost):
|
||||
marks = ",".join("?" * len(chunk))
|
||||
# ``marks`` is only "?" placeholders; the ids are bound.
|
||||
for (image_id,) in _rows(
|
||||
conn,
|
||||
"SELECT id FROM generated_images WHERE " # nosec B608
|
||||
f"workflow_document_id IN ({marks}) OR prompt_document_id IN ({marks})",
|
||||
chunk + chunk,
|
||||
):
|
||||
images.add(image_id)
|
||||
for column in DOCUMENT_ID_COLUMNS:
|
||||
conn.execute(
|
||||
f"UPDATE generated_images SET {column}=NULL " # nosec B608
|
||||
f"WHERE {column} IN ({marks})",
|
||||
chunk,
|
||||
)
|
||||
conn.execute(
|
||||
f"DELETE FROM workflow_documents WHERE id IN ({marks})", # nosec B608
|
||||
chunk,
|
||||
)
|
||||
for document_id in list(intact) + healed + list(copies):
|
||||
logger.warning(f"Workflow document {document_id} was damaged; repaired")
|
||||
if lost:
|
||||
logger.warning(
|
||||
f"{len(lost)} damaged workflow document(s) could not be repaired; "
|
||||
f"{len(images)} image(s) show no workflow"
|
||||
)
|
||||
result.update(
|
||||
repaired=len(intact) + len(healed) + len(copies),
|
||||
cleared=len(lost),
|
||||
images=len(images),
|
||||
)
|
||||
return result
|
||||
@@ -1,455 +0,0 @@
|
||||
"""The workflow storage migration and Reclaim space (#249 S6 carry-over).
|
||||
|
||||
``migrate_workflow_storage`` rewrites old-form image rows (workflow and
|
||||
prompt graph text inline) to the new form (a ``workflow_documents`` id,
|
||||
``database/workflow_documents.py``) in batches:
|
||||
|
||||
- each batch is one ``BEGIN IMMEDIATE`` transaction over at most
|
||||
``BATCH_ROWS`` rows or about ``BATCH_BYTES`` of text, found by a rowid
|
||||
range search (``NEXT_IDS_SQL``) and read one row at a time, so memory is
|
||||
bounded by one row plus the batch's distinct documents;
|
||||
- each text is stored (or its document reused) and read back through the
|
||||
one accessor before its row changes; a text that would not come back
|
||||
identical fails the batch (``DocumentMismatch``) and rolls it back whole;
|
||||
- a value that is not text (or is over ``MAX_DOCUMENT_BYTES``) stays inline
|
||||
and is counted as protected;
|
||||
- each row runs in its own savepoint: one that cannot be rewritten (text
|
||||
SQLite cannot decode, a text that would not come back) stays inline and
|
||||
is listed, and the rest of the batch goes on;
|
||||
- a batch ends after ``BATCH_SECONDS``, and the write lock is then left
|
||||
free for ``BATCH_PAUSE``, so other writers are never starved.
|
||||
|
||||
A crash or a stop between batches leaves committed batches in the new form
|
||||
and the rest in the old form; both read the same, and a rerun continues
|
||||
from the first old-form row. Only database rows change: no image or video
|
||||
file is opened.
|
||||
|
||||
``reclaim_space`` is the VACUUM afterwards: it sweeps unreferenced
|
||||
documents, rebuilds the file and checkpoints the WAL with TRUNCATE, so the
|
||||
freed pages go back to the disk.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sqlite3
|
||||
import time
|
||||
from contextlib import closing
|
||||
|
||||
from .workflow_documents import (
|
||||
TEXT_COLUMNS,
|
||||
DocumentMismatch,
|
||||
read_stored_text,
|
||||
release_documents,
|
||||
storable,
|
||||
store_text,
|
||||
sweep_documents,
|
||||
)
|
||||
|
||||
OPERATION = "workflow_storage"
|
||||
# The first apply's backup: ``<db>.pre-workflow-storage_<stamp>``, outside the
|
||||
# pruned ``<db>.backup_*`` pool. Only one is ever made; later applies take a
|
||||
# pooled safety backup.
|
||||
PRE_MIGRATION_LABEL = "pre-workflow-storage"
|
||||
BATCH_ROWS = 200
|
||||
BATCH_BYTES = 16 * 2**20 # characters of text rewritten per transaction
|
||||
# A batch holds the write lock: it ends after this long, and the loop then
|
||||
# leaves the lock free for BATCH_PAUSE, longer than SQLite's longest busy
|
||||
# poll (100 ms), so the image monitor and the run counting get in between.
|
||||
BATCH_SECONDS = 0.5
|
||||
BATCH_PAUSE = 0.15
|
||||
|
||||
|
||||
def _row_error(error) -> bool:
|
||||
"""Whether ``error`` is one row's own cause, which keeps only that row inline.
|
||||
|
||||
Text SQLite cannot decode for Python, and a text that would not come
|
||||
back identical. Anything else (a full disk, an I/O error, a malformed
|
||||
database, a bug) fails the whole batch. The decode error is matched by
|
||||
its message: ``sqlite_errorcode`` needs Python 3.11.
|
||||
"""
|
||||
if isinstance(error, (DocumentMismatch, UnicodeError)):
|
||||
return True
|
||||
return isinstance(error, sqlite3.OperationalError) and str(error).startswith(
|
||||
"Could not decode"
|
||||
)
|
||||
|
||||
|
||||
logger = logging.getLogger("prompt_manager.database.workflow_storage")
|
||||
_clock = time.monotonic
|
||||
|
||||
# typeof() reads a column's type from the record header: no overflow page.
|
||||
NEXT_IDS_SQL = (
|
||||
"SELECT id FROM generated_images WHERE id>? AND "
|
||||
"(typeof(workflow_data)!='null' OR typeof(prompt_metadata)!='null') "
|
||||
"ORDER BY id LIMIT ?"
|
||||
)
|
||||
ROW_SQL = (
|
||||
"SELECT workflow_data, prompt_metadata, workflow_document_id, "
|
||||
"prompt_document_id FROM generated_images WHERE id=?"
|
||||
)
|
||||
UPDATE_SQL = (
|
||||
"UPDATE generated_images SET {assignments} WHERE id=?" # columns are constants
|
||||
)
|
||||
# One pass over the record headers: rows with inline text, and inline
|
||||
# values that are not text (they stay inline).
|
||||
INLINE_COUNTS_SQL = (
|
||||
"SELECT "
|
||||
"SUM(typeof(workflow_data)='text' OR typeof(prompt_metadata)='text'), "
|
||||
"SUM(typeof(workflow_data) NOT IN ('null','text')) "
|
||||
"+ SUM(typeof(prompt_metadata) NOT IN ('null','text')) "
|
||||
"FROM generated_images"
|
||||
)
|
||||
# octet_length() (SQLite 3.43+) reads no overflow page either.
|
||||
INLINE_BYTES_SQL = (
|
||||
"SELECT SUM(octet_length(workflow_data)), SUM(octet_length(prompt_metadata)) "
|
||||
"FROM generated_images"
|
||||
)
|
||||
|
||||
QUERY_PLAN_STATEMENTS = {
|
||||
"next old-form rows": (NEXT_IDS_SQL, (0, BATCH_ROWS)),
|
||||
"one row": (ROW_SQL, (1,)),
|
||||
}
|
||||
|
||||
|
||||
def _rewrite_row(conn, image_id, row, verified, checked):
|
||||
"""Store the row's inline texts as documents; returns (changed, protected, size).
|
||||
|
||||
``verified`` maps document ids already read back in this batch to their
|
||||
text, and ``checked`` holds the reused ids ``store_text`` already
|
||||
checked, so a document shared by many rows is decompressed once.
|
||||
"""
|
||||
changes, replaced = {}, []
|
||||
protected = size = 0
|
||||
for index, (column, reference) in enumerate(TEXT_COLUMNS):
|
||||
value, current = row[index], row[index + len(TEXT_COLUMNS)]
|
||||
if value is None:
|
||||
continue
|
||||
if not storable(value):
|
||||
protected += 1
|
||||
continue
|
||||
document_id = store_text(conn, value, checked)
|
||||
if document_id not in verified:
|
||||
verified[document_id] = read_stored_text(conn, None, document_id)
|
||||
if verified[document_id] != value:
|
||||
raise DocumentMismatch(f"Image {image_id}: {column} would change")
|
||||
changes[column], changes[reference] = None, document_id
|
||||
if current is not None and current != document_id:
|
||||
replaced.append(current)
|
||||
size += len(value)
|
||||
if changes:
|
||||
assignments = ",".join(f"{key}=?" for key in changes)
|
||||
conn.execute(
|
||||
UPDATE_SQL.format(assignments=assignments),
|
||||
[*changes.values(), image_id],
|
||||
)
|
||||
release_documents(conn, replaced)
|
||||
return bool(changes), protected, size
|
||||
|
||||
|
||||
def migrate_batch(
|
||||
conn,
|
||||
after_id,
|
||||
max_rows=BATCH_ROWS,
|
||||
max_bytes=BATCH_BYTES,
|
||||
max_seconds=BATCH_SECONDS,
|
||||
):
|
||||
"""Rewrite the next batch after ``after_id`` in one transaction.
|
||||
|
||||
The batch ends after ``max_rows`` rows, about ``max_bytes`` of text or
|
||||
``max_seconds``, whichever comes first. Each row runs in its own
|
||||
savepoint: a row that cannot be rewritten (``_row_error``) is rolled
|
||||
back alone, stays inline and is listed in ``errors``; anything else rolls
|
||||
the whole batch back and propagates. Returns ``{last_id, rows,
|
||||
protected, errors, done}``: the last row looked at, rows rewritten,
|
||||
inline values kept (not text), the ids kept by an error and whether
|
||||
there was nothing left to look at.
|
||||
"""
|
||||
if conn.in_transaction:
|
||||
raise RuntimeError("The migration needs its own transaction")
|
||||
started = _clock()
|
||||
conn.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
ids = [row[0] for row in conn.execute(NEXT_IDS_SQL, (after_id, max_rows))]
|
||||
result = {
|
||||
"last_id": after_id,
|
||||
"rows": 0,
|
||||
"protected": 0,
|
||||
"errors": [],
|
||||
"done": not ids,
|
||||
}
|
||||
verified, checked, size = {}, set(), 0
|
||||
for image_id in ids:
|
||||
conn.execute("SAVEPOINT workflow_row")
|
||||
try:
|
||||
row = conn.execute(ROW_SQL, (image_id,)).fetchone()
|
||||
changed, protected, row_size = _rewrite_row(
|
||||
conn, image_id, row, verified, checked
|
||||
)
|
||||
except Exception as error:
|
||||
if not _row_error(error):
|
||||
raise
|
||||
conn.execute("ROLLBACK TO workflow_row")
|
||||
conn.execute("RELEASE workflow_row")
|
||||
verified.clear() # the row's documents were rolled back too
|
||||
checked.clear()
|
||||
logger.warning(f"Image record {image_id} kept inline: {error}")
|
||||
result["errors"].append(image_id)
|
||||
changed, protected, row_size = False, 0, 0
|
||||
else:
|
||||
conn.execute("RELEASE workflow_row")
|
||||
result["rows"] += int(changed)
|
||||
result["protected"] += protected
|
||||
result["last_id"] = image_id
|
||||
size += row_size
|
||||
if size >= max_bytes or _clock() - started >= max_seconds:
|
||||
break
|
||||
conn.commit()
|
||||
return result
|
||||
except BaseException:
|
||||
conn.rollback()
|
||||
raise
|
||||
|
||||
|
||||
def pre_migration_backups(db_path):
|
||||
"""The kept pre-migration backups beside ``db_path``, earliest first."""
|
||||
db_path = os.path.abspath(db_path)
|
||||
prefix = os.path.basename(db_path) + "." + PRE_MIGRATION_LABEL + "_"
|
||||
folder = os.path.dirname(db_path)
|
||||
try:
|
||||
names = os.listdir(folder)
|
||||
except OSError:
|
||||
return []
|
||||
return sorted(
|
||||
name
|
||||
for name in names
|
||||
if name.startswith(prefix) and os.path.isfile(os.path.join(folder, name))
|
||||
)
|
||||
|
||||
|
||||
def _file_bytes(db_path):
|
||||
return sum(
|
||||
os.path.getsize(path)
|
||||
for path in (db_path, db_path + "-wal")
|
||||
if os.path.exists(path)
|
||||
)
|
||||
|
||||
|
||||
class WorkflowStorageOperations:
|
||||
"""PromptDatabase mixin: the workflow storage migration and Reclaim space."""
|
||||
|
||||
def migrate_workflow_storage(
|
||||
self,
|
||||
*,
|
||||
on_batch=None,
|
||||
should_stop=None,
|
||||
max_rows=BATCH_ROWS,
|
||||
max_bytes=BATCH_BYTES,
|
||||
) -> dict:
|
||||
"""Rewrite every old-form row, batch by batch (blocking).
|
||||
|
||||
``on_batch(progress)`` runs after each committed batch with the
|
||||
running totals ``{rows, protected, batches, last_id}``;
|
||||
``should_stop()`` is asked before each batch. Each committed batch is
|
||||
followed by ``BATCH_PAUSE`` with the write lock free. Returns
|
||||
``{rows, protected, batches, stopped, errors, documents_removed}``
|
||||
(``errors``: the ids of rows kept inline by an error). Unreferenced
|
||||
documents are swept once every row was looked at.
|
||||
"""
|
||||
totals = {"rows": 0, "protected": 0, "batches": 0, "last_id": 0}
|
||||
errors = []
|
||||
stopped = False
|
||||
while True:
|
||||
if should_stop is not None and should_stop():
|
||||
stopped = True
|
||||
break
|
||||
batch = self.migrate_workflow_batch(totals["last_id"], max_rows, max_bytes)
|
||||
if batch["done"]:
|
||||
break
|
||||
time.sleep(BATCH_PAUSE) # leave the write lock to other writers
|
||||
errors += batch["errors"]
|
||||
totals["rows"] += batch["rows"]
|
||||
totals["protected"] += batch["protected"]
|
||||
totals["batches"] += 1
|
||||
totals["last_id"] = batch["last_id"]
|
||||
if on_batch is not None:
|
||||
on_batch(dict(totals))
|
||||
removed = 0 if stopped else self.sweep_workflow_documents()
|
||||
return {
|
||||
"rows": totals["rows"],
|
||||
"protected": totals["protected"],
|
||||
"batches": totals["batches"],
|
||||
"stopped": stopped,
|
||||
"errors": errors,
|
||||
"documents_removed": removed,
|
||||
}
|
||||
|
||||
def migrate_workflow_batch(
|
||||
self, after_id, max_rows=BATCH_ROWS, max_bytes=BATCH_BYTES
|
||||
) -> dict:
|
||||
"""One batch on the calling thread's connection (``migrate_batch``)."""
|
||||
return migrate_batch(self.model.get_connection(), after_id, max_rows, max_bytes)
|
||||
|
||||
def sweep_workflow_documents(self) -> int:
|
||||
"""Remove every document no image row references (one transaction)."""
|
||||
conn = self.model.get_connection()
|
||||
with conn:
|
||||
return sweep_documents(conn)
|
||||
|
||||
def workflow_storage_status(self) -> dict:
|
||||
"""``workflow_storage_counts`` on one read-only snapshot (blocking)."""
|
||||
from .maintenance import _read_snapshot
|
||||
|
||||
return _read_snapshot(self.model.db_path, self.workflow_storage_counts)
|
||||
|
||||
def workflow_storage_counts(self, conn) -> dict:
|
||||
"""``{rows, protected, documents, inline_bytes}`` on ``conn`` (reads only).
|
||||
|
||||
``inline_bytes`` is None where SQLite has no ``octet_length``.
|
||||
"""
|
||||
rows, protected = conn.execute(INLINE_COUNTS_SQL).fetchone()
|
||||
documents = conn.execute("SELECT COUNT(*) FROM workflow_documents").fetchone()
|
||||
inline_bytes = None
|
||||
if sqlite3.sqlite_version_info >= (3, 43, 0):
|
||||
row = conn.execute(INLINE_BYTES_SQL).fetchone()
|
||||
inline_bytes = sum(value or 0 for value in row)
|
||||
return {
|
||||
"rows": rows or 0,
|
||||
"protected": protected or 0,
|
||||
"documents": documents[0],
|
||||
"inline_bytes": inline_bytes,
|
||||
}
|
||||
|
||||
def reclaim_space(self) -> dict:
|
||||
"""Sweep, VACUUM and truncate the WAL (blocking); file bytes before/after.
|
||||
|
||||
The caller holds the registry's ``exclusive`` slot. Raises whatever
|
||||
stopped VACUUM; the database is then unchanged.
|
||||
"""
|
||||
db_path = self.model.db_path
|
||||
removed = self.sweep_workflow_documents()
|
||||
before = _file_bytes(db_path)
|
||||
self.model.vacuum_database()
|
||||
with closing(sqlite3.connect(db_path)) as checkpoint:
|
||||
busy = checkpoint.execute("PRAGMA wal_checkpoint(TRUNCATE)").fetchone()[0]
|
||||
return {
|
||||
"before_bytes": before,
|
||||
"after_bytes": _file_bytes(db_path),
|
||||
"documents_removed": removed,
|
||||
"checkpoint_busy": bool(busy),
|
||||
}
|
||||
|
||||
# ── preview ───────────────────────────────────────────────────────
|
||||
|
||||
def _plan_workflow_storage(self):
|
||||
"""The reviewed plan of the migration (one read-only snapshot)."""
|
||||
from .maintenance import (
|
||||
MaintenanceDraft,
|
||||
_Fingerprint,
|
||||
_megabytes,
|
||||
_plural,
|
||||
_read_snapshot,
|
||||
)
|
||||
|
||||
def read(conn):
|
||||
sizes = [
|
||||
conn.execute(f"PRAGMA {name}").fetchone()[0]
|
||||
for name in ("page_size", "page_count", "freelist_count")
|
||||
]
|
||||
return sizes, self.workflow_storage_counts(conn)
|
||||
|
||||
(page_size, page_count, free_pages), counts = _read_snapshot(
|
||||
self.model.db_path, read
|
||||
)
|
||||
size = page_size * page_count
|
||||
details = [
|
||||
{
|
||||
"label": "Database file",
|
||||
"detail": f"{_megabytes(size)}, of which "
|
||||
f"{_megabytes(free_pages * page_size)} is free pages",
|
||||
},
|
||||
{
|
||||
"label": "Image records",
|
||||
"detail": f"{_plural(counts['rows'], 'record')} keep their workflow "
|
||||
"or prompt graph inline; each distinct text is stored once, "
|
||||
"compressed",
|
||||
},
|
||||
]
|
||||
if not counts["rows"]:
|
||||
details[1]["detail"] = NOTHING_TO_DO
|
||||
if counts["inline_bytes"] is not None:
|
||||
details.append(
|
||||
{
|
||||
"label": "Inline text",
|
||||
"detail": f"{_megabytes(counts['inline_bytes'])} now",
|
||||
}
|
||||
)
|
||||
if counts["documents"]:
|
||||
details.append(
|
||||
{
|
||||
"label": "Already stored once",
|
||||
"detail": _plural(counts["documents"], "workflow document"),
|
||||
}
|
||||
)
|
||||
name = os.path.basename(self.model.db_path)
|
||||
kept = pre_migration_backups(self.model.db_path)
|
||||
backup = (
|
||||
POOLED_BACKUP_NOTE.format(
|
||||
name=f"{name}.backup_<date and time>", kept=kept[0]
|
||||
)
|
||||
if kept
|
||||
else BACKUP_NOTE.format(
|
||||
name=f"{name}.{PRE_MIGRATION_LABEL}_<date and time>"
|
||||
)
|
||||
)
|
||||
details += [
|
||||
{"label": "Backup", "detail": backup},
|
||||
{"label": "Reclaim space", "detail": RECLAIM_NOTE},
|
||||
{"label": "Older versions", "detail": OLDER_VERSIONS_NOTE},
|
||||
]
|
||||
protected = counts["protected"]
|
||||
return MaintenanceDraft(
|
||||
operation=OPERATION,
|
||||
fingerprint=_Fingerprint(OPERATION).hexdigest(),
|
||||
targets=[],
|
||||
scope={
|
||||
"count": counts["rows"],
|
||||
"label": (
|
||||
_plural(counts["rows"], "image record") + " to rewrite"
|
||||
if counts["rows"]
|
||||
else "Nothing to rewrite"
|
||||
),
|
||||
},
|
||||
protected={
|
||||
"count": protected,
|
||||
"reasons": (
|
||||
[f"{_plural(protected, 'stored value')} not text; kept as is"]
|
||||
if protected
|
||||
else []
|
||||
),
|
||||
},
|
||||
media={"deletes_files": False, "label": MEDIA_NOTE},
|
||||
details=details,
|
||||
)
|
||||
|
||||
|
||||
MEDIA_NOTE = "No image or video file is opened or changed; only the database."
|
||||
RECLAIM_NOTE = (
|
||||
"When every record is rewritten, Reclaim space runs alone to give the freed "
|
||||
"space back to the disk (VACUUM). The database is locked while it runs."
|
||||
)
|
||||
BACKUP_NOTE = (
|
||||
"Always taken first, whatever the safety backup setting: {name} beside "
|
||||
"the database, verified, and kept until you remove it (it is never pruned)."
|
||||
)
|
||||
POOLED_BACKUP_NOTE = (
|
||||
"A safety backup is taken first, whatever the safety backup setting: {name} "
|
||||
"among the pooled backups; the pre-migration backup {kept} is already kept."
|
||||
)
|
||||
NOTHING_TO_DO = "Every image record already stores its workflow once."
|
||||
OLDER_VERSIONS_NOTE = (
|
||||
"This version stores the workflows of images it records in the new form, "
|
||||
"whether or not this job runs, and this job rewrites the rest. An older "
|
||||
"Prompt Manager version shows those images without their workflow until you "
|
||||
"update again. The way back to an older version is the earliest "
|
||||
"pre-workflow-storage backup."
|
||||
)
|
||||
@@ -238,9 +238,6 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
|
||||
except Exception as e:
|
||||
# Log error but don't fail the encoding
|
||||
self.logger.warning(f"Failed to save prompt to database: {e}")
|
||||
self._mark_untracked(prompt, unique_id)
|
||||
else:
|
||||
self._skip_empty_prompt(prompt, unique_id)
|
||||
|
||||
# Perform standard CLIP text encoding using the combined text
|
||||
self.logger.debug(
|
||||
|
||||
+3
-60
@@ -24,7 +24,6 @@ try:
|
||||
from .utils.prompt_tracker import get_prompt_tracker
|
||||
from .utils.hashing import generate_prompt_hash
|
||||
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
|
||||
except ImportError:
|
||||
import sys
|
||||
@@ -36,7 +35,6 @@ except ImportError:
|
||||
from utils.prompt_tracker import get_prompt_tracker
|
||||
from utils.hashing import generate_prompt_hash
|
||||
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
|
||||
|
||||
|
||||
@@ -170,67 +168,14 @@ class PromptManagerBase:
|
||||
prompt_id = self._save_prompt_to_database(
|
||||
text=text, category=category, tags=tags, count_run=batch_item
|
||||
)
|
||||
if role["is_run_prompt"] and not prompt_id:
|
||||
# The save failed: an older entry of this key must not stand in (#342)
|
||||
self._mark_untracked(prompt_graph, unique_id)
|
||||
elif role["is_run_prompt"]:
|
||||
if prompt_id and role["is_run_prompt"]:
|
||||
self.prompt_tracker.set_current_prompt(
|
||||
prompt_text=encoding_text,
|
||||
additional_data={**additional_data, "prompt_id": prompt_id},
|
||||
push_to_queue=batch_item or not role["known"],
|
||||
**self._run_identity(prompt_graph, unique_id),
|
||||
)
|
||||
return prompt_id
|
||||
|
||||
@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)."""
|
||||
return {
|
||||
"run_key": prompt_node_key(prompt_graph, unique_id),
|
||||
"run_id": executing_run_id(),
|
||||
}
|
||||
|
||||
def _skip_empty_prompt(self, prompt_graph: Any, unique_id: Any) -> None:
|
||||
"""The ``text`` editor is empty: save and track nothing (#342).
|
||||
|
||||
Prepend and append text are deliberately never stored as the prompt, so
|
||||
an empty editor has no prompt to save, count or link. For the run's
|
||||
positive node this is logged at info and the node's run key is marked
|
||||
empty (``_mark_untracked``).
|
||||
"""
|
||||
self._mark_untracked(
|
||||
prompt_graph,
|
||||
unique_id,
|
||||
"Prompt text is empty (prepend and append text are not stored): "
|
||||
"nothing saved or tracked, and this run's images stay unlinked",
|
||||
)
|
||||
|
||||
def _mark_untracked(
|
||||
self, prompt_graph: Any, unique_id: Any, message: Optional[str] = None
|
||||
) -> None:
|
||||
"""This execution of a positive node has no prompt to link (#342).
|
||||
|
||||
An empty editor, or a save or tracking failure. The node's run key is
|
||||
marked with a placeholder, so no earlier prompt of that key can be
|
||||
linked to this run's images, and a batch item keeps its queue place.
|
||||
``message`` is logged at info for a positive node.
|
||||
"""
|
||||
try:
|
||||
role = self._prompt_role(prompt_graph, unique_id, "")
|
||||
if not role["is_run_prompt"]:
|
||||
if message:
|
||||
self.logger.debug("Empty text on a non-positive node; not saved")
|
||||
return
|
||||
if message:
|
||||
self.logger.info(message)
|
||||
self.prompt_tracker.mark_empty(
|
||||
push_to_queue=role["known"],
|
||||
**self._run_identity(prompt_graph, unique_id),
|
||||
)
|
||||
except Exception as error:
|
||||
# Tracking must never block generation
|
||||
self.logger.warning(f"Could not mark the untracked prompt: {error}")
|
||||
|
||||
def _record_use(self, prompt_id: int, times: int = 1) -> None:
|
||||
"""Count runs; a tracking failure must never block generation."""
|
||||
try:
|
||||
@@ -268,10 +213,8 @@ class PromptManagerBase:
|
||||
|
||||
cache = get_trigger_cache()
|
||||
|
||||
# Lazy-load cache on first use; a failed load waits before the next try
|
||||
# (#355: never a rescan of LoraManager's folders on every run).
|
||||
retry_due = getattr(cache, "retry_due", lambda: True)
|
||||
if not cache.is_loaded and IntegrationConfig.LORA_MANAGER_PATH and retry_due():
|
||||
# Lazy-load cache on first use
|
||||
if not cache.is_loaded and IntegrationConfig.LORA_MANAGER_PATH:
|
||||
cache.load(IntegrationConfig.LORA_MANAGER_PATH)
|
||||
|
||||
modified, injected = inject_trigger_words(text, cache)
|
||||
|
||||
+77
-117
@@ -117,29 +117,85 @@ class PromptManagerPair(PromptManagerFocused):
|
||||
)
|
||||
is_positive = "text" in used_fields
|
||||
pair_used = node_id in intrinsic_pair_nodes(prompt) if known else True
|
||||
if not text.strip():
|
||||
# Nothing to store: prepend/append are never the stored prompt (#342)
|
||||
self._skip_empty_prompt(prompt, unique_id)
|
||||
else:
|
||||
try:
|
||||
pair_id, prompt_id = self._save_and_track_pair(
|
||||
text,
|
||||
try:
|
||||
pair_id = self.db.save_prompt_pair(
|
||||
text,
|
||||
negative_text,
|
||||
category=category.strip() or None,
|
||||
tags=self._parse_tags(tags),
|
||||
)
|
||||
pair = self.db.get_prompt_pair(pair_id)
|
||||
prompt_id = pair["positive_prompt_id"]
|
||||
resolved_positive = (
|
||||
resolve_text_input(prompt, node_id, "text") if known else None
|
||||
)
|
||||
resolved_negative = (
|
||||
resolve_text_input(prompt, node_id, "negative_text") if known else None
|
||||
)
|
||||
positive_seen = resolved_positive is not None and generate_prompt_hash(
|
||||
resolved_positive
|
||||
) == generate_prompt_hash(text)
|
||||
pair_seen = (
|
||||
positive_seen
|
||||
and 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,
|
||||
positive,
|
||||
category=category,
|
||||
tags=tags,
|
||||
prepend_text=prepend_text,
|
||||
append_text=append_text,
|
||||
prompt=prompt,
|
||||
unique_id=unique_id,
|
||||
used_fields=used_fields,
|
||||
is_positive=is_positive,
|
||||
pair_used=pair_used,
|
||||
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
|
||||
):
|
||||
continue
|
||||
counted.add(component_hash)
|
||||
seen_at_queue = (
|
||||
resolved_text is not None
|
||||
and generate_prompt_hash(resolved_text) == component_hash
|
||||
)
|
||||
except Exception as error:
|
||||
self.logger.warning(f"Could not save or track paired prompt: {error}")
|
||||
if is_positive:
|
||||
self._mark_untracked(prompt, unique_id)
|
||||
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
|
||||
)
|
||||
if runs:
|
||||
self.db.record_prompt_pair_use(pair_id, times=runs)
|
||||
if is_positive:
|
||||
self.prompt_tracker.set_current_prompt(
|
||||
prompt_text=positive,
|
||||
additional_data={
|
||||
"prompt_id": prompt_id,
|
||||
"pair_id": pair_id if pair_used else None,
|
||||
"positive_text": positive,
|
||||
"negative_text": negative_text,
|
||||
"positive_storage_text": text.strip(),
|
||||
"negative_storage_text": negative_text.strip(),
|
||||
"prepend_text": prepend_text.strip(),
|
||||
"append_text": append_text.strip(),
|
||||
"final_text": positive,
|
||||
"category": category.strip() or None,
|
||||
"tags": self._parse_tags(tags) or [],
|
||||
},
|
||||
push_to_queue=not positive_seen or (pair_used and not pair_seen),
|
||||
)
|
||||
except Exception as error:
|
||||
self.logger.warning(f"Could not save or track paired prompt: {error}")
|
||||
try:
|
||||
self.comfyui_integration.register_prompt(
|
||||
f"promptmanager_pair_{time.time_ns()}",
|
||||
@@ -167,102 +223,6 @@ class PromptManagerPair(PromptManagerFocused):
|
||||
)
|
||||
return result
|
||||
|
||||
def _save_and_track_pair(
|
||||
self,
|
||||
text: str,
|
||||
negative_text: str,
|
||||
positive: str,
|
||||
*,
|
||||
category: str,
|
||||
tags: str,
|
||||
prepend_text: str,
|
||||
append_text: str,
|
||||
prompt: Optional[Dict[str, Any]],
|
||||
unique_id: Optional[str],
|
||||
used_fields: set,
|
||||
is_positive: bool,
|
||||
pair_used: bool,
|
||||
) -> Tuple[Optional[int], Optional[int]]:
|
||||
"""Save the recipe, count its runs and register the positive for linking."""
|
||||
known = isinstance(prompt, dict) and unique_id is not None
|
||||
node_id = str(unique_id)
|
||||
pair_id = self.db.save_prompt_pair(
|
||||
text,
|
||||
negative_text,
|
||||
category=category.strip() or None,
|
||||
tags=self._parse_tags(tags),
|
||||
)
|
||||
pair = self.db.get_prompt_pair(pair_id)
|
||||
prompt_id = pair["positive_prompt_id"]
|
||||
resolved_positive = (
|
||||
resolve_text_input(prompt, node_id, "text") if known else None
|
||||
)
|
||||
resolved_negative = (
|
||||
resolve_text_input(prompt, node_id, "negative_text") if known else None
|
||||
)
|
||||
positive_seen = resolved_positive is not None and generate_prompt_hash(
|
||||
resolved_positive
|
||||
) == generate_prompt_hash(text)
|
||||
pair_seen = (
|
||||
positive_seen
|
||||
and 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
|
||||
):
|
||||
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(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
|
||||
)
|
||||
if runs:
|
||||
self.db.record_prompt_pair_use(pair_id, times=runs)
|
||||
if is_positive:
|
||||
self.prompt_tracker.set_current_prompt(
|
||||
prompt_text=positive,
|
||||
additional_data={
|
||||
"prompt_id": prompt_id,
|
||||
"pair_id": pair_id if pair_used else None,
|
||||
"positive_text": positive,
|
||||
"negative_text": negative_text,
|
||||
"positive_storage_text": text.strip(),
|
||||
"negative_storage_text": negative_text.strip(),
|
||||
"prepend_text": prepend_text.strip(),
|
||||
"append_text": append_text.strip(),
|
||||
"final_text": positive,
|
||||
"category": category.strip() or None,
|
||||
"tags": self._parse_tags(tags) or [],
|
||||
},
|
||||
push_to_queue=not positive_seen or (pair_used and not pair_seen),
|
||||
**self._run_identity(prompt, unique_id),
|
||||
)
|
||||
return pair_id, prompt_id
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(
|
||||
cls,
|
||||
|
||||
@@ -216,9 +216,6 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC):
|
||||
except Exception as e:
|
||||
# Log error but don't fail the processing
|
||||
self.logger.warning(f"Failed to save prompt to database: {e}")
|
||||
self._mark_untracked(prompt, unique_id)
|
||||
else:
|
||||
self._skip_empty_prompt(prompt, unique_id)
|
||||
|
||||
# Register with ComfyUI integration for standard metadata compatibility
|
||||
node_id = f"promptmanagertext_{int(time.time() * 1000)}"
|
||||
|
||||
+20
-186
@@ -9,17 +9,12 @@ Split into domain-specific mixins for maintainability:
|
||||
- JobRoutesMixin: the session job registry behind Activity (py/api/jobs.py)
|
||||
- ImageTagJobsMixin: bulk Add tags / Auto-tag as image-tags jobs
|
||||
- ThumbnailJobsMixin: thumbnail plan, thumbnail job and the gallery stream on it
|
||||
- ThumbnailConvertMixin: Convert outdated thumbnails (#317), a thumbnail job mode
|
||||
- AutoThumbnailMixin: previews made on first view, on capture and after a scan
|
||||
- ReviewRoutesMixin: the Needs review queue, item detail and decisions
|
||||
- ReviewRecordingMixin: Needs review recording for the scan and the monitor
|
||||
- ReviewEventsMixin: the Needs review change stream an open page re-reads on
|
||||
- SystemRoutesMixin: System health check, Overview and usage statistics
|
||||
- LibraryLocationRoutesMixin: where the library lives, Move library, Keep it here
|
||||
- MaintenanceRoutesMixin: reviewed maintenance plans (preview, apply)
|
||||
- WorkflowStorageRoutesMixin: store each workflow once, then Reclaim space
|
||||
- DuplicateFilesMixin: duplicate image files as find and removal jobs
|
||||
- LoraImportRoutesMixin: the reviewed, staged LoRA reference import job
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -31,7 +26,6 @@ import math
|
||||
import ntpath
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
@@ -48,26 +42,19 @@ from .review_events import ReviewEventsMixin
|
||||
from .saved_views import SavedViewRoutesMixin
|
||||
from .history import HistoryRoutesMixin
|
||||
from .images import ImageRoutesMixin
|
||||
from .gallery_listing import WALK_STOPPED, RootListing, WalkOutcome
|
||||
from .admin import AdminRoutesMixin
|
||||
from .settings_routes import SettingsRoutesMixin
|
||||
from .job_routes import JobRoutesMixin
|
||||
from .image_tag_jobs import ImageTagJobsMixin
|
||||
from .thumbnail_jobs import ThumbnailJobsMixin
|
||||
from .thumbnail_convert import ThumbnailConvertMixin
|
||||
from .thumbnail_auto import AutoThumbnailMixin
|
||||
from .system import SystemRoutesMixin
|
||||
from .library_location import LibraryLocationRoutesMixin
|
||||
from .maintenance import MaintenanceRoutesMixin
|
||||
from .backup_restore import BackupRestoreRoutesMixin
|
||||
from .workflow_storage import WorkflowStorageRoutesMixin
|
||||
from .duplicate_files import DuplicateFilesMixin
|
||||
from .maintenance_plans import PlanStore
|
||||
from .jobs import JobRegistry
|
||||
from .logging_routes import LoggingRoutesMixin
|
||||
from .autotag_routes import AutotagRoutesMixin
|
||||
from .lora_integration import LoraIntegrationMixin
|
||||
from .lora_import_routes import LoraImportRoutesMixin
|
||||
|
||||
try:
|
||||
from ...database.operations import PromptDatabase
|
||||
@@ -89,18 +76,6 @@ def _get_project_root():
|
||||
return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
|
||||
# How long an executor thread waits for the loop to drop the listings.
|
||||
_OFF_LOOP_DROP_WAIT = 30
|
||||
|
||||
|
||||
def _running_loop():
|
||||
"""The event loop running in this thread, or None."""
|
||||
try:
|
||||
return asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return None
|
||||
|
||||
|
||||
def _public_path(path) -> str:
|
||||
"""Render a server path for an API response without revealing the layout.
|
||||
|
||||
@@ -314,18 +289,12 @@ class PromptManagerAPI(
|
||||
JobRoutesMixin,
|
||||
ImageTagJobsMixin,
|
||||
ThumbnailJobsMixin,
|
||||
ThumbnailConvertMixin,
|
||||
AutoThumbnailMixin,
|
||||
SystemRoutesMixin,
|
||||
LibraryLocationRoutesMixin,
|
||||
MaintenanceRoutesMixin,
|
||||
BackupRestoreRoutesMixin,
|
||||
WorkflowStorageRoutesMixin,
|
||||
DuplicateFilesMixin,
|
||||
LoggingRoutesMixin,
|
||||
AutotagRoutesMixin,
|
||||
LoraIntegrationMixin,
|
||||
LoraImportRoutesMixin,
|
||||
):
|
||||
"""REST API handler for PromptManager operations and web interface.
|
||||
|
||||
@@ -358,20 +327,8 @@ class PromptManagerAPI(
|
||||
self._captures_closed = False
|
||||
self._cached_output_dir = None # Lazy-cached by _find_comfyui_output_dir()
|
||||
self._html_cache = {} # Cached HTML file contents keyed by path
|
||||
# One shared listing per output root, keyed by str(root), and the
|
||||
# background walks refreshing them (py/api/gallery_listing.py, #348).
|
||||
self._root_listings: dict[str, RootListing] = {}
|
||||
self._listing_tasks: set = set()
|
||||
# The loop the walks run on, set by the startup warm-up or the first
|
||||
# walk; the monitor's thread hops onto it (notify_output_settled).
|
||||
self._listing_loop = None
|
||||
# Set by cleanup_gallery_listings: no background walk starts after it.
|
||||
self._listings_closed = False
|
||||
# The workspace caches rebuilt behind a stored walk that changed the
|
||||
# listing, once the host's startup warm-up has run (#348).
|
||||
self._workspace_warm_enabled = False
|
||||
self._workspace_warm_task = None
|
||||
self._workspace_warm_again = False
|
||||
self._gallery_cache = {} # dict: path_str -> (files, timestamp)
|
||||
self._gallery_cache_ttl = 30 # Cache TTL in seconds
|
||||
# Workspace listing values and sort orders derived from a cached scan
|
||||
# (py/api/images.py _workspace_listed/_workspace_orders).
|
||||
self._workspace_listed_cache = {}
|
||||
@@ -407,19 +364,12 @@ class PromptManagerAPI(
|
||||
try:
|
||||
from .maintenance import sweep_safety_temps
|
||||
|
||||
from .backup_restore import restore_stage
|
||||
|
||||
swept = sweep_safety_temps(self.db.model.db_path)
|
||||
swept += restore_stage.sweep_restore_temps(self.db.model.db_path)
|
||||
if swept:
|
||||
self.logger.info(f"Removed {len(swept)} interrupted backup copies")
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Could not sweep interrupted backup copies: {e}")
|
||||
|
||||
# A crash between Move library's switch and its rename left the old
|
||||
# node-folder files under their own names (#249 S6b): rename them now.
|
||||
self.finish_interrupted_library_move()
|
||||
|
||||
# Run cleanup on initialization to remove any existing duplicates
|
||||
try:
|
||||
removed = self.db.cleanup_duplicates()
|
||||
@@ -491,122 +441,23 @@ class PromptManagerAPI(
|
||||
return await loop.run_in_executor(None, call)
|
||||
return await loop.run_in_executor(None, func, *args)
|
||||
|
||||
def _root_listing(self, root) -> RootListing:
|
||||
"""The shared listing record of one output root (created empty)."""
|
||||
return self._root_listings.setdefault(str(root), RootListing())
|
||||
def invalidate_gallery_cache(self):
|
||||
"""Invalidate the cached gallery file listing.
|
||||
|
||||
def invalidate_gallery_cache(self, retire=False):
|
||||
"""Drop every root's listing, so the next read waits for a new walk.
|
||||
|
||||
Called after files are deleted through the app (duplicate cleanup and
|
||||
the legacy duplicate delete), never for thumbnail writes. A walk
|
||||
running now is not stored (its generation moved on). The workspace
|
||||
keeps its per-file values and card facts, so the next request
|
||||
rebuilds only new or changed files (see ImageRoutesMixin.
|
||||
_workspace_listed and _card_facts); its sort orders are dropped, as
|
||||
the walk replaces them anyway, and so is the favorites lookup, which
|
||||
holds the old list.
|
||||
|
||||
``retire`` (the roots changed): no background walk of a root starts
|
||||
again until a reader or the warm-up asks for it.
|
||||
Called by the image monitor when new files are detected. The workspace
|
||||
keeps its per-file values and card facts, so the next request rebuilds
|
||||
only new or changed files (see ImageRoutesMixin._workspace_listed and
|
||||
_card_facts); its sort orders
|
||||
are dropped, as the rescan replaces them anyway, and so is the
|
||||
favorites lookup, which holds the old scan list.
|
||||
"""
|
||||
loop = self._listing_loop
|
||||
if loop is not None and loop.is_running() and _running_loop() is not loop:
|
||||
# Called from an executor thread (a settings save, a test seed):
|
||||
# listing records are touched only on the loop, and the caller
|
||||
# expects the drop done when this returns.
|
||||
dropped = threading.Event()
|
||||
|
||||
def drop():
|
||||
try:
|
||||
self._drop_listings(retire)
|
||||
finally:
|
||||
dropped.set()
|
||||
|
||||
loop.call_soon_threadsafe(drop)
|
||||
if not dropped.wait(_OFF_LOOP_DROP_WAIT):
|
||||
self.logger.warning("The Images listing drop is still queued")
|
||||
return
|
||||
self._drop_listings(retire)
|
||||
|
||||
def _drop_listings(self, retire=False):
|
||||
for record in list(self._root_listings.values()):
|
||||
record.drop()
|
||||
if retire:
|
||||
# The roots changed: a removed root must not come back through
|
||||
# a pending timer, a trailing walk or a monitor mark.
|
||||
record.retired = True
|
||||
record.dirty = False
|
||||
handle, record.rewalk_handle = record.rewalk_handle, None
|
||||
if handle is not None:
|
||||
handle.cancel()
|
||||
self._gallery_cache = {}
|
||||
self._workspace_order_cache = {}
|
||||
self._favorite_listing_cache = {}
|
||||
|
||||
async def warm_gallery_listings(self, app=None):
|
||||
"""Start one background walk per output root (aiohttp on_startup).
|
||||
|
||||
Registered by the root ``__init__.py``, so the first Images open after
|
||||
ComfyUI starts finds the listing built. It never fails startup: a
|
||||
root that cannot be walked is retried by the first read.
|
||||
"""
|
||||
# Each stored walk that changes a listing then rebuilds the workspace's
|
||||
# per-file values, orders and favorites behind it (_warm_workspace).
|
||||
self._workspace_warm_enabled = True
|
||||
self._listing_loop = asyncio.get_running_loop()
|
||||
self._listings_closed = False
|
||||
try:
|
||||
for root in self._get_all_output_dirs():
|
||||
self._root_listing(root).retired = False
|
||||
self._schedule_walk(root)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Could not warm the Images listing: {e}")
|
||||
|
||||
def notify_output_settled(self, path):
|
||||
"""The image monitor settled an output at *path* (any thread).
|
||||
|
||||
Called on the monitor's worker thread with the resolved path. It only
|
||||
hands the path to the event loop (``output_settled``) and never
|
||||
touches a listing from this thread; before the loop is known, or once
|
||||
the listings are cleaned up, it does nothing.
|
||||
"""
|
||||
loop = self._listing_loop
|
||||
if loop is None or self._listings_closed or not loop.is_running():
|
||||
return
|
||||
try:
|
||||
loop.call_soon_threadsafe(self.output_settled, path)
|
||||
except RuntimeError:
|
||||
pass # the loop closed meanwhile (shutdown)
|
||||
|
||||
async def cleanup_gallery_listings(self, app=None):
|
||||
"""Stop the background listing walks (aiohttp on_cleanup)."""
|
||||
self._workspace_warm_enabled = False
|
||||
# No mark, deferred walk or trailing walk starts a walk from now on.
|
||||
self._listings_closed = True
|
||||
# Before any await, so no delayed refresh can start a walk meanwhile.
|
||||
for record in self._root_listings.values():
|
||||
handle, record.rewalk_handle = record.rewalk_handle, None
|
||||
if handle is not None:
|
||||
handle.cancel()
|
||||
warm, self._workspace_warm_task = self._workspace_warm_task, None
|
||||
if warm is not None:
|
||||
warm.cancel()
|
||||
await asyncio.gather(warm, return_exceptions=True)
|
||||
tasks = list(self._listing_tasks)
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
self._listing_tasks.clear()
|
||||
# A walk cancelled before its first step never ran its own cleanup.
|
||||
for record in self._root_listings.values():
|
||||
flight, record.flight = record.flight, None
|
||||
if flight is not None and not flight.done():
|
||||
flight.set_result(WalkOutcome.failed(RuntimeError(WALK_STOPPED)))
|
||||
|
||||
def reset_gallery_roots_cache(self):
|
||||
"""Drop every cached listing and derived workspace value (roots changed)."""
|
||||
self.invalidate_gallery_cache(retire=True)
|
||||
self.invalidate_gallery_cache()
|
||||
self._workspace_listed_cache = {}
|
||||
self._workspace_card_cache = {}
|
||||
|
||||
@@ -678,14 +529,11 @@ class PromptManagerAPI(
|
||||
self._register_settings_routes(routes)
|
||||
self._register_job_routes(routes)
|
||||
self._register_system_routes(routes)
|
||||
self._register_library_location_routes(routes)
|
||||
self._register_maintenance_routes(routes)
|
||||
self._register_backup_restore_routes(routes)
|
||||
self._register_duplicate_routes(routes)
|
||||
self._register_logging_routes(routes)
|
||||
self._register_autotag_routes(routes)
|
||||
self._register_lora_routes(routes)
|
||||
self._register_lora_import_routes(routes)
|
||||
|
||||
# Register gzip compression middleware (once)
|
||||
global _gzip_registered
|
||||
@@ -750,12 +598,7 @@ class PromptManagerAPI(
|
||||
# ── Shared utilities used by multiple mixins ──────────────────────
|
||||
|
||||
def _enrich_images(self, images):
|
||||
"""Add url and thumbnail_url to a flat list of image dicts.
|
||||
|
||||
Each dict without ``is_video`` also gets ``is_video`` and
|
||||
``media_type`` from its extension (#355: a linked LoraManager video
|
||||
outside the output folders is shown as a video, not a picture).
|
||||
"""
|
||||
"""Add url and thumbnail_url to a flat list of image dicts."""
|
||||
from urllib.parse import quote as url_quote
|
||||
|
||||
output_dirs = self._get_all_output_dirs()
|
||||
@@ -765,11 +608,6 @@ class PromptManagerAPI(
|
||||
if not image_path_str:
|
||||
continue
|
||||
|
||||
if "is_video" not in image:
|
||||
video = is_video_path(image_path_str)
|
||||
image["is_video"] = video
|
||||
image["media_type"] = "video" if video else "image"
|
||||
|
||||
img_path = Path(image_path_str)
|
||||
|
||||
# Set fallback url via image ID
|
||||
@@ -786,20 +624,16 @@ class PromptManagerAPI(
|
||||
f"{url_quote(rel_path.as_posix(), safe='/')}"
|
||||
)
|
||||
|
||||
# Any preview the file owns, the current format first (#317)
|
||||
from . import thumbnail_settings
|
||||
# Check for thumbnail
|
||||
rel_no_ext = rel_path.with_suffix("")
|
||||
thumb_rel = (
|
||||
f"thumbnails/{rel_no_ext.as_posix()}_thumb{rel_path.suffix}"
|
||||
)
|
||||
from .prompts import thumbnail_url_for
|
||||
|
||||
existing = thumbnail_settings.existing_thumbnail(
|
||||
output_path, rel_path
|
||||
thumbnail_url = thumbnail_url_for(
|
||||
thumb_rel, output_path / thumb_rel
|
||||
)
|
||||
thumbnail_url = None
|
||||
if existing is not None:
|
||||
thumb_rel = (
|
||||
thumbnail_settings.thumbnail_stem_rel(rel_path)
|
||||
+ existing.suffix
|
||||
)
|
||||
thumbnail_url = thumbnail_url_for(thumb_rel, existing)
|
||||
if thumbnail_url:
|
||||
image["thumbnail_url"] = thumbnail_url
|
||||
break # Found matching root, stop searching
|
||||
|
||||
+252
-176
@@ -3,21 +3,21 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
from contextlib import closing
|
||||
import datetime
|
||||
import functools
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import urllib.parse
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from .prompts import thumbnail_url_for
|
||||
from . import thumbnail_settings
|
||||
from .jobs import (
|
||||
DEFAULT_SCAN_SCOPE,
|
||||
SCAN_GROUP_NAMES,
|
||||
@@ -26,6 +26,7 @@ from .jobs import (
|
||||
DUPLICATE_CLEANUP_JOB_KIND,
|
||||
DUPLICATES_JOB_KIND,
|
||||
MAINTENANCE_JOB_KIND,
|
||||
RESTORE_JOB_KIND,
|
||||
JobBusy,
|
||||
)
|
||||
|
||||
@@ -35,12 +36,8 @@ except ImportError:
|
||||
from utils.parallel import map_parallel
|
||||
|
||||
try:
|
||||
from ...database.pair_ingestion import blank_positive, ingest_resolved_pair
|
||||
from ...database.run_attribution import (
|
||||
host_output_directory,
|
||||
stored_run_prompt,
|
||||
stored_run_prompts,
|
||||
)
|
||||
from ...database.operations import PromptDatabase
|
||||
from ...database.pair_ingestion import ingest_resolved_pair
|
||||
from ...utils.generation_params import (
|
||||
PLACEHOLDERS,
|
||||
GenerationParamsError,
|
||||
@@ -50,12 +47,8 @@ try:
|
||||
from ...utils.prompt_pair_resolver import resolve_prompt_pair
|
||||
from ...utils.review_capture import reviewable
|
||||
except ImportError:
|
||||
from database.pair_ingestion import blank_positive, ingest_resolved_pair
|
||||
from database.run_attribution import (
|
||||
host_output_directory,
|
||||
stored_run_prompt,
|
||||
stored_run_prompts,
|
||||
)
|
||||
from database.operations import PromptDatabase
|
||||
from database.pair_ingestion import ingest_resolved_pair
|
||||
from utils.generation_params import (
|
||||
PLACEHOLDERS,
|
||||
GenerationParamsError,
|
||||
@@ -153,48 +146,39 @@ def _physical_file_key(path, stat):
|
||||
return ("path", os.path.normcase(os.path.realpath(str(path))))
|
||||
|
||||
|
||||
def _other_sources(output_path, rel_path):
|
||||
"""Whether another media file shares *rel_path*'s stem in its folder."""
|
||||
folder = Path(output_path) / Path(rel_path).parent
|
||||
stem = Path(rel_path).stem
|
||||
own = Path(rel_path).suffix.lower()
|
||||
for extension in IMAGE_EXTENSIONS + VIDEO_EXTENSIONS:
|
||||
if extension == own:
|
||||
continue
|
||||
for name in (stem + extension, stem + extension.upper()):
|
||||
if (folder / name).is_file():
|
||||
return True
|
||||
return False
|
||||
def _thumbnail_rel_path(rel_path, thumbnail_ext):
|
||||
"""Relative path of the thumbnail the gallery would generate for ``rel_path``."""
|
||||
return f"thumbnails/{rel_path.with_suffix('').as_posix()}_thumb{thumbnail_ext}"
|
||||
|
||||
|
||||
def _media_thumbnail_rel_path(rel_path):
|
||||
"""The thumbnail of a media file: a video's preview is a ``.jpg``."""
|
||||
extension = rel_path.suffix.lower()
|
||||
return _thumbnail_rel_path(
|
||||
rel_path, ".jpg" if extension in VIDEO_EXTENSIONS else rel_path.suffix
|
||||
)
|
||||
|
||||
|
||||
def _remove_thumbnail(output_path, rel_path, logger=None):
|
||||
"""Remove the previews of ``rel_path`` under ``output_path``, if any.
|
||||
"""Remove the thumbnail of ``rel_path`` under ``output_path``, if any.
|
||||
|
||||
Shared by the legacy duplicate delete and the reviewed removal job.
|
||||
Every preview the file owns goes, in any format (#317); a preview name
|
||||
another same-stem file still shares (``x_thumb.webp`` of ``x.png`` and
|
||||
``x.mp4``, ruling R2) stays for it. A preview that is a link, or one that
|
||||
cannot be removed, is left alone (a stale preview never fails a removal).
|
||||
Returns whether one was removed.
|
||||
Shared by the legacy duplicate delete and the reviewed removal job. A
|
||||
thumbnail that is a link, or one that cannot be removed, is left alone
|
||||
(a stale preview never fails a removal). Returns whether one was removed.
|
||||
"""
|
||||
output_path = Path(output_path)
|
||||
rel_path = Path(rel_path)
|
||||
shared = _other_sources(output_path, rel_path)
|
||||
removed = False
|
||||
for thumbnail_path in thumbnail_settings.owned_thumbnails(output_path, rel_path):
|
||||
extension = thumbnail_path.suffix.lower()
|
||||
if shared and extension in thumbnail_settings.WRITTEN_EXTENSIONS:
|
||||
continue
|
||||
try:
|
||||
os.remove(thumbnail_path)
|
||||
except OSError as e:
|
||||
if logger is not None:
|
||||
logger.warning(f"Could not delete thumbnail for {rel_path}: {e}")
|
||||
continue
|
||||
removed = True
|
||||
thumbnail_path = output_path / _media_thumbnail_rel_path(Path(rel_path))
|
||||
try:
|
||||
if thumbnail_path.is_symlink() or not thumbnail_path.is_file():
|
||||
return False
|
||||
os.remove(thumbnail_path)
|
||||
except OSError as e:
|
||||
if logger is not None:
|
||||
logger.debug(f"Deleted associated thumbnail: {thumbnail_path}")
|
||||
return removed
|
||||
logger.warning(f"Could not delete thumbnail for {rel_path}: {e}")
|
||||
return False
|
||||
if logger is not None:
|
||||
logger.debug(f"Deleted associated thumbnail: {thumbnail_path}")
|
||||
return True
|
||||
|
||||
|
||||
async def _read_json(request):
|
||||
@@ -558,24 +542,6 @@ def _current_image_monitor():
|
||||
return getattr(_image_monitor_module(), "_monitor_instance", None)
|
||||
|
||||
|
||||
class _BatchHistory:
|
||||
"""A scan batch's recorded-history lookup, made on the first file needing it.
|
||||
|
||||
Most files link through their metadata and never ask; the first one that
|
||||
does triggers one lookup for the whole batch (``_recorded_runs_for``).
|
||||
"""
|
||||
|
||||
def __init__(self, api, media_files):
|
||||
self._api = api
|
||||
self._files = media_files
|
||||
self._found = None
|
||||
|
||||
def get(self, path):
|
||||
if self._found is None:
|
||||
self._found = self._api._recorded_runs_for(self._files)
|
||||
return self._found.get(path)
|
||||
|
||||
|
||||
class AdminRoutesMixin:
|
||||
"""Mixin providing admin, diagnostics, and maintenance API endpoints."""
|
||||
|
||||
@@ -758,20 +724,9 @@ class AdminRoutesMixin:
|
||||
|
||||
Only the files under ``output_dir`` count (paths are relative to it);
|
||||
a group needs two of them. A file that vanished since it was hashed
|
||||
is left out. Each ``url`` names the source ``output_dir`` is
|
||||
(``?root=``) when it is one, so View Full never opens another
|
||||
source's file at the same path.
|
||||
is left out.
|
||||
"""
|
||||
output_path = Path(os.path.abspath(output_dir))
|
||||
identity = os.path.normcase(os.path.realpath(str(output_path)))
|
||||
root_index = next(
|
||||
(
|
||||
index
|
||||
for index, root in enumerate(self._get_all_output_dirs())
|
||||
if os.path.normcase(os.path.realpath(str(root))) == identity
|
||||
),
|
||||
None,
|
||||
)
|
||||
base = os.path.normcase(str(output_path)).rstrip("\\/") + os.sep
|
||||
duplicates = []
|
||||
for group in groups:
|
||||
@@ -782,10 +737,7 @@ class AdminRoutesMixin:
|
||||
try:
|
||||
images.append(
|
||||
self._describe_media_file(
|
||||
Path(record["abs"]),
|
||||
output_path,
|
||||
file_hash=group["hash"],
|
||||
root_index=root_index,
|
||||
Path(record["abs"]), output_path, file_hash=group["hash"]
|
||||
)
|
||||
)
|
||||
except (OSError, ValueError) as e:
|
||||
@@ -885,36 +837,23 @@ class AdminRoutesMixin:
|
||||
)
|
||||
return groups
|
||||
|
||||
def _describe_media_file(
|
||||
self, media_path, output_path, file_hash=None, root_index=None
|
||||
):
|
||||
def _describe_media_file(self, media_path, output_path, file_hash=None):
|
||||
"""Describe one media file for the duplicates response (hashing it
|
||||
unless ``file_hash`` is given). With ``root_index`` the ``url`` names
|
||||
that source and its path is URL-encoded."""
|
||||
unless ``file_hash`` is given)."""
|
||||
stat = media_path.stat()
|
||||
rel_path = media_path.relative_to(output_path)
|
||||
extension = media_path.suffix.lower()
|
||||
is_video = extension in VIDEO_EXTENSIONS
|
||||
|
||||
existing = thumbnail_settings.existing_thumbnail(output_path, rel_path)
|
||||
thumbnail_url = None
|
||||
if existing is not None:
|
||||
thumb_rel = (
|
||||
thumbnail_settings.thumbnail_stem_rel(rel_path) + existing.suffix
|
||||
)
|
||||
thumbnail_url = thumbnail_url_for(thumb_rel, existing)
|
||||
if root_index is None:
|
||||
url = f"/prompt_manager/images/serve/{rel_path.as_posix()}"
|
||||
else:
|
||||
encoded = urllib.parse.quote(rel_path.as_posix(), safe="/")
|
||||
url = f"/prompt_manager/images/serve/{encoded}?root={root_index}"
|
||||
thumb_rel = _thumbnail_rel_path(rel_path, ".jpg" if is_video else extension)
|
||||
thumbnail_url = thumbnail_url_for(thumb_rel, output_path / thumb_rel)
|
||||
|
||||
return {
|
||||
"id": str(hash(str(media_path))),
|
||||
"filename": media_path.name,
|
||||
"path": str(rel_path),
|
||||
"relative_path": str(rel_path),
|
||||
"url": url,
|
||||
"url": f"/prompt_manager/images/serve/{rel_path.as_posix()}",
|
||||
"thumbnail_url": thumbnail_url,
|
||||
"size": stat.st_size,
|
||||
"modified_time": stat.st_mtime,
|
||||
@@ -1282,6 +1221,216 @@ class AdminRoutesMixin:
|
||||
"message": f"Found {len(issues)} consistency issues",
|
||||
}
|
||||
|
||||
async def backup_database(self, request):
|
||||
"""Download a consistent copy of the prompts database (WAL included)."""
|
||||
try:
|
||||
model = self.db.model
|
||||
db_path = model.db_path
|
||||
|
||||
if not os.path.exists(db_path):
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Database file not found"}, status=404
|
||||
)
|
||||
|
||||
fd, temp_path = tempfile.mkstemp(suffix=".db")
|
||||
os.close(fd)
|
||||
try:
|
||||
ok = await self._run_in_executor(model.backup_database, temp_path)
|
||||
if not ok:
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Failed to create database backup"},
|
||||
status=500,
|
||||
)
|
||||
# Read once (the file is deleted right after) rather than
|
||||
# streaming so the temp file never outlives the request.
|
||||
with open(temp_path, "rb") as fh:
|
||||
file_data = fh.read()
|
||||
finally:
|
||||
if os.path.exists(temp_path):
|
||||
os.unlink(temp_path)
|
||||
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"prompts_backup_{timestamp}.db"
|
||||
|
||||
return web.Response(
|
||||
body=file_data,
|
||||
content_type="application/octet-stream",
|
||||
headers={
|
||||
"Content-Disposition": f'attachment; filename="{filename}"',
|
||||
"Content-Length": str(len(file_data)),
|
||||
},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Backup error: {e}", exc_info=True)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"error": f"Failed to backup database: {self._public_error(e)}",
|
||||
},
|
||||
status=500,
|
||||
)
|
||||
|
||||
async def restore_database(self, request):
|
||||
"""Restore the prompts database from an uploaded SQLite file.
|
||||
|
||||
The upload is streamed to a temp file. Free disk space decides how
|
||||
large it may be (``maintenance.restore_room``), under a 64 GiB ceiling
|
||||
(``restore_max_bytes``): a request whose Content-Length does not fit
|
||||
is refused before anything is written, and one without is refused as
|
||||
soon as its running size does not. The file is
|
||||
verified with PromptModel.verify_database_file and only then copied
|
||||
over the live database, which is backed up first. The whole request
|
||||
holds the registry (records): it is refused while a scan or another
|
||||
record-changing job runs, and those are refused while it runs.
|
||||
"""
|
||||
return await self._run_held(
|
||||
RESTORE_JOB_KIND,
|
||||
"Database restore (admin page)",
|
||||
"records",
|
||||
lambda: self._restore_database(request),
|
||||
)
|
||||
|
||||
def _restore_refusal(self, size, ceiling, room, *, known):
|
||||
"""``400 too-large`` past the ceiling, else ``507 insufficient-space``."""
|
||||
from .maintenance import restore_space_message
|
||||
|
||||
if size > ceiling:
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "too-large",
|
||||
"error": (
|
||||
"File too large. Maximum size is "
|
||||
f"{ceiling // (1024 * 1024)}MB"
|
||||
),
|
||||
},
|
||||
status=400,
|
||||
)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "insufficient-space",
|
||||
"error": restore_space_message(size, room, known=known),
|
||||
},
|
||||
status=507,
|
||||
)
|
||||
|
||||
async def _restore_database(self, request):
|
||||
from .maintenance import RESTORE_MAX_BYTES, restore_room
|
||||
|
||||
ceiling = getattr(self, "restore_max_bytes", RESTORE_MAX_BYTES)
|
||||
temp_path = None
|
||||
try:
|
||||
# Free disk space decides how large a backup may be (#313): the
|
||||
# upload, its staged copy and the swap must all fit.
|
||||
room = await self._run_in_executor(restore_room, self.db.model.db_path)
|
||||
max_bytes = min(ceiling, room)
|
||||
length = request.content_length
|
||||
if length is not None and length > max_bytes:
|
||||
return self._restore_refusal(length, ceiling, room, known=True)
|
||||
reader = await request.multipart()
|
||||
field = await reader.next()
|
||||
|
||||
if not field or field.name != "database_file":
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"error": (
|
||||
"No database file uploaded. "
|
||||
"Expected field name: database_file"
|
||||
),
|
||||
},
|
||||
status=400,
|
||||
)
|
||||
|
||||
fd, temp_path = tempfile.mkstemp(suffix=".db")
|
||||
os.close(fd)
|
||||
size = 0
|
||||
with open(temp_path, "wb") as out:
|
||||
while True:
|
||||
chunk = await field.read_chunk()
|
||||
if not chunk:
|
||||
break
|
||||
size += len(chunk)
|
||||
if size > max_bytes:
|
||||
return self._restore_refusal(size, ceiling, room, known=False)
|
||||
out.write(chunk)
|
||||
|
||||
if size == 0:
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Uploaded file is empty"}, status=400
|
||||
)
|
||||
|
||||
model = self.db.model
|
||||
ok, reason = await self._run_in_executor(
|
||||
model.verify_database_file, temp_path
|
||||
)
|
||||
if ok:
|
||||
reason = await self._run_in_executor(_inspect_restore_upload, temp_path)
|
||||
ok = reason is None
|
||||
if not ok:
|
||||
return web.json_response(
|
||||
{"success": False, "error": reason}, status=400
|
||||
)
|
||||
|
||||
try:
|
||||
backup_path = await self._run_in_executor(
|
||||
model.restore_from_file, temp_path
|
||||
)
|
||||
except sqlite3.OperationalError as exc:
|
||||
if (
|
||||
getattr(exc, "_prompt_manager_restore_backup_created", False)
|
||||
is True
|
||||
):
|
||||
try:
|
||||
await self._run_in_executor(
|
||||
_prune_safety_backups, model.db_path
|
||||
)
|
||||
except Exception:
|
||||
self.logger.warning(
|
||||
"Could not prune safety backups after failed restore"
|
||||
)
|
||||
raise
|
||||
removed = await self._run_in_executor(_prune_safety_backups, model.db_path)
|
||||
if removed:
|
||||
self.logger.info(f"Removed {len(removed)} old safety backups")
|
||||
previous = self.db
|
||||
self.db = PromptDatabase(model.db_path)
|
||||
self._database_replaced() # a health report described the old one
|
||||
# Its connections would otherwise linger (a file lock on Windows)
|
||||
previous.close_all()
|
||||
prompt_count = self.db.model.get_database_info().get("total_prompts", 0)
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"message": (
|
||||
"Database restored successfully. "
|
||||
f"Found {prompt_count} prompts."
|
||||
),
|
||||
"prompt_count": prompt_count,
|
||||
"backup_created": (
|
||||
self._public_path(backup_path) if backup_path else None
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Restore error: {e}", exc_info=True)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"error": f"Failed to restore database: {self._public_error(e)}",
|
||||
},
|
||||
status=500,
|
||||
)
|
||||
finally:
|
||||
if temp_path:
|
||||
for suffix in ("", "-wal", "-shm", "-journal"):
|
||||
if os.path.exists(temp_path + suffix):
|
||||
os.unlink(temp_path + suffix)
|
||||
|
||||
async def scan_images(self, request):
|
||||
"""Scan ComfyUI output images for prompt metadata and add them to the database.
|
||||
|
||||
@@ -1355,9 +1504,6 @@ class AdminRoutesMixin:
|
||||
verb = "cancelled" if state == "cancelled" else "finished"
|
||||
job.log_line(f"Scan {verb} · {_plural(job.processed, 'file')} processed")
|
||||
job.finish(state)
|
||||
if state == "finished":
|
||||
# After the scan, so the guarded thumbnail job is admitted (#249).
|
||||
await self._hand_off_scan_thumbnails(job, output_dirs)
|
||||
|
||||
async def _stream_scan_job(self, request, job):
|
||||
"""Stream ``job`` as SSE; stop quietly if the client goes away."""
|
||||
@@ -1592,8 +1738,6 @@ class AdminRoutesMixin:
|
||||
if outcomes is not None:
|
||||
outcomes.append(_scan_outcome(media_file, roots, status, title))
|
||||
|
||||
batch_results = list(batch_results)
|
||||
history = _BatchHistory(self, [media for media, _meta in batch_results])
|
||||
for media_file, metadata in batch_results:
|
||||
if self._scan_job and self._scan_job.cancelled:
|
||||
break
|
||||
@@ -1607,7 +1751,6 @@ class AdminRoutesMixin:
|
||||
ambiguous=ambiguous,
|
||||
untraced=untraced,
|
||||
records=records,
|
||||
history=history,
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error processing {media_file.name}: {e}")
|
||||
@@ -1666,7 +1809,6 @@ class AdminRoutesMixin:
|
||||
ambiguous=None,
|
||||
untraced=None,
|
||||
records=None,
|
||||
history=None,
|
||||
):
|
||||
"""Enrich exact pairs; retain legacy positive-only import for absent pairs.
|
||||
|
||||
@@ -1684,9 +1826,7 @@ class AdminRoutesMixin:
|
||||
a dict.
|
||||
"""
|
||||
if not metadata:
|
||||
# No prompt metadata (--disable-metadata, bare savers): only the
|
||||
# recorded run can say which prompt made it (#342)
|
||||
return self._link_recorded_run(media_file, {}, detailed, history)
|
||||
return None
|
||||
canonical = dict(metadata)
|
||||
for key in ("prompt", "workflow"):
|
||||
if isinstance(canonical.get(key), str):
|
||||
@@ -1696,12 +1836,9 @@ class AdminRoutesMixin:
|
||||
# Preserve malformed metadata for conservative resolver rejection.
|
||||
pass
|
||||
if not canonical.get("prompt") and not canonical.get("parameters"):
|
||||
return self._link_recorded_run(media_file, canonical, detailed, history)
|
||||
return None
|
||||
_with_file_size(canonical, media_file)
|
||||
resolution = resolve_prompt_pair(canonical)
|
||||
if blank_positive(resolution):
|
||||
# Only prepend/append text: no stored prompt can own it (#342)
|
||||
resolution = {**resolution, "state": "unresolved"}
|
||||
# Ambiguous without a choice: unresolved unless something links it.
|
||||
unchoosable = resolution["state"] == "ambiguous" and not reviewable(resolution)
|
||||
if unchoosable:
|
||||
@@ -1730,11 +1867,6 @@ class AdminRoutesMixin:
|
||||
# Graphs without a recognized sampler keep the legacy extractor below.
|
||||
prompt_text = self._graph_prompt_text(metadata)
|
||||
if not prompt_text:
|
||||
linked = self._link_recorded_run(
|
||||
media_file, canonical, detailed, history
|
||||
)
|
||||
if linked is not None:
|
||||
return linked
|
||||
return "unresolved" if unchoosable else resolution["state"]
|
||||
elif resolution["state"] == "ambiguous":
|
||||
return "unresolved" if unchoosable else "ambiguous"
|
||||
@@ -1758,62 +1890,6 @@ class AdminRoutesMixin:
|
||||
return status
|
||||
return {"status": status, "title": self._scan_title(media_file, prompt_text)}
|
||||
|
||||
def _recorded_runs_for(self, media_files):
|
||||
"""``{path: prompt id or None}`` from recorded history for a scan batch.
|
||||
|
||||
One output lookup per chunk of names (``stored_run_prompts``), so a
|
||||
batch never queries history once per file. An empty dict on failure.
|
||||
"""
|
||||
try:
|
||||
return stored_run_prompts(
|
||||
self.db, [str(path) for path in media_files], host_output_directory()
|
||||
)
|
||||
except Exception:
|
||||
self.logger.warning(
|
||||
"History attribution failed for a scan batch", exc_info=True
|
||||
)
|
||||
return {}
|
||||
|
||||
def _link_recorded_run(self, media_file, canonical, detailed=False, history=None):
|
||||
"""Link a file recorded history attributes to one stored prompt (#342).
|
||||
|
||||
Returns the scan status (with ``detailed``, a stored or linked file's
|
||||
``{"status", "title"}``), or None when history names no single stored
|
||||
prompt for the file (``database.run_attribution``). Never creates a
|
||||
prompt. Works without any prompt metadata in the file. ``history`` is
|
||||
the batch's precomputed lookup; without it the file is looked up alone.
|
||||
"""
|
||||
if history is not None:
|
||||
prompt_id = history.get(str(media_file))
|
||||
else:
|
||||
try:
|
||||
prompt_id = stored_run_prompt(
|
||||
self.db, str(media_file), host_output_directory()
|
||||
)
|
||||
except Exception:
|
||||
self.logger.warning(
|
||||
f"History attribution failed for {media_file.name}",
|
||||
exc_info=True,
|
||||
)
|
||||
return None
|
||||
if prompt_id is None:
|
||||
return None
|
||||
try:
|
||||
image_id = self.db.link_image_to_prompt(
|
||||
prompt_id, str(media_file), metadata=canonical
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to link {media_file.name} from history: {e}")
|
||||
return "found"
|
||||
status = "linked" if image_id else "unchanged"
|
||||
if not detailed or status not in SCAN_TITLED_DETAILS:
|
||||
return status
|
||||
prompt = self.db.get_prompt_by_id(prompt_id) or {}
|
||||
return {
|
||||
"status": status,
|
||||
"title": self._scan_title(media_file, prompt.get("text") or ""),
|
||||
}
|
||||
|
||||
def _scan_title(self, media_file, text):
|
||||
"""The linked prompt's Library title (its stored title, else ``text``'s
|
||||
first line), as the Library and the inspector name it."""
|
||||
|
||||
@@ -10,15 +10,12 @@ filesystem or the database.
|
||||
|
||||
import math
|
||||
|
||||
from ..tagging_defaults import ( # noqa: F401 - re-exported for the route
|
||||
LLM_MODEL_TYPES as LLM_MODELS,
|
||||
MAX_INSTRUCTION_BYTES,
|
||||
MAX_INSTRUCTION_CHARS,
|
||||
MODEL_TYPES,
|
||||
WD14_MODEL_TYPES as WD14_MODELS,
|
||||
validate_template,
|
||||
validate_threshold,
|
||||
)
|
||||
WD14_MODELS = frozenset({"wd14-swinv2", "wd14-vit"})
|
||||
LLM_MODELS = frozenset({"gguf", "hf"})
|
||||
MODEL_TYPES = WD14_MODELS | LLM_MODELS
|
||||
|
||||
MAX_INSTRUCTION_CHARS = 4000
|
||||
MAX_INSTRUCTION_BYTES = 16 * 1024
|
||||
|
||||
SUGGEST_FIELDS = frozenset(
|
||||
{"target", "capture_token", "expected_epoch", "model_type", "options"}
|
||||
@@ -27,7 +24,7 @@ _COMMON = frozenset({"suggestion_limit", "use_gpu"})
|
||||
_WD14_REQUIRED = frozenset({"general_threshold", "character_threshold"})
|
||||
_LLM_REQUIRED = frozenset({"instructions"})
|
||||
|
||||
UNAVAILABLE_HINT = "Download it in System › Auto-tagging."
|
||||
UNAVAILABLE_HINT = "Download it from the legacy Auto-tag page."
|
||||
|
||||
|
||||
class OptionError(ValueError):
|
||||
@@ -35,17 +32,27 @@ class OptionError(ValueError):
|
||||
|
||||
|
||||
def _threshold(name: str, value) -> float:
|
||||
try:
|
||||
return validate_threshold(name, value)
|
||||
except ValueError as error:
|
||||
raise OptionError(str(error)) from None
|
||||
if type(value) not in (int, float) or not math.isfinite(value):
|
||||
raise OptionError(f"{name} must be a finite number")
|
||||
if not 0 <= value <= 1:
|
||||
raise OptionError(f"{name} must be between 0 and 1")
|
||||
return float(value)
|
||||
|
||||
|
||||
def _instructions(value) -> str:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise OptionError("instructions must be non-empty text")
|
||||
if len(value) > MAX_INSTRUCTION_CHARS:
|
||||
raise OptionError(
|
||||
f"instructions must be at most {MAX_INSTRUCTION_CHARS} characters"
|
||||
)
|
||||
try:
|
||||
return validate_template(value)
|
||||
except ValueError as error:
|
||||
raise OptionError(str(error)) from None
|
||||
encoded = value.encode("utf-8")
|
||||
except UnicodeEncodeError:
|
||||
raise OptionError("instructions must be valid Unicode text") from None
|
||||
if len(encoded) > MAX_INSTRUCTION_BYTES:
|
||||
raise OptionError("instructions are too long")
|
||||
return value
|
||||
|
||||
|
||||
def validate_options(model_type, options, *, default_limit: int, max_limit: int):
|
||||
|
||||
+91
-638
@@ -1,131 +1,34 @@
|
||||
"""AutoTag API routes for PromptManager."""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import functools
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from urllib.parse import quote
|
||||
|
||||
from .job_routes import busy_response
|
||||
from .jobs import (
|
||||
AUTOTAG_BATCH_JOB_KIND,
|
||||
AUTOTAG_GROUP_NAMES,
|
||||
MODEL_DOWNLOAD_JOB_KIND,
|
||||
JobBusy,
|
||||
)
|
||||
from .admin import _plural, _scan_location, _scan_roots
|
||||
from .maintenance import disk_plan, short_space_message
|
||||
from .jobs import AUTOTAG_BATCH_JOB_KIND
|
||||
from .prompts import thumbnail_url_for
|
||||
from . import thumbnail_settings
|
||||
from .. import model_gate
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
_INTERNAL_ERROR = "An internal error occurred. Check server logs for details."
|
||||
|
||||
# Tag output images (#328): the native prompt-tag batch.
|
||||
AUTOTAG_JOB_TITLE = "Tag output images"
|
||||
AUTOTAG_PREVIEW_LEASE = "Tag output images review"
|
||||
AUTOTAG_OWNER_KIND = "autotag-job"
|
||||
BATCH_PREVIEW_BODY_MAX = 1024
|
||||
BATCH_JOB_BODY_MAX = 32 * 1024
|
||||
REVIEW_REQUIRED = (
|
||||
"Replacing tags needs a review of each image. Use Preview tagging review, "
|
||||
"or choose Skip or Merge for a batch run."
|
||||
)
|
||||
BATCH_CANCELLED = "Cancelled. Tags already added were kept."
|
||||
_BUSY_ERROR = "Auto-tagging is already running"
|
||||
MODEL_BUSY_ERROR = (
|
||||
"A tagging model is busy with another operation. Try again when it finishes."
|
||||
)
|
||||
|
||||
|
||||
# Who holds the model gate, as a 409 model-busy names it (``operation``). An
|
||||
# owner is ``"<kind>:<uuid>"`` (``model_gate.new_owner``); a download names
|
||||
# its own job, whose title says which model.
|
||||
MODEL_DOWNLOAD_OWNER_KIND = "model-download"
|
||||
OWNER_LABELS = {
|
||||
"autotag-job": "Tag output images",
|
||||
"legacy-batch": "Tag output images (legacy page)",
|
||||
"legacy-single": "Tagging one image (legacy page)",
|
||||
"legacy-unload": "Unloading the tagging model",
|
||||
"image-tag-job": "Auto-tag selected images",
|
||||
"image-annotations": "Generate suggestions",
|
||||
}
|
||||
ANOTHER_OPERATION = "Another tagging operation"
|
||||
_DOWNLOAD_FAILED = "The download failed. Check server logs for details."
|
||||
# POST /jobs/model-download carries only {model_type}.
|
||||
MODEL_DOWNLOAD_BODY_MAX = 1024
|
||||
# The running download's gate owner and job title. Only one download runs at a
|
||||
# time and each needs the gate, so a new start replaces the entry.
|
||||
_download_titles = {}
|
||||
|
||||
|
||||
def _holder_label(holder: Optional[str]) -> str:
|
||||
if not holder:
|
||||
return ANOTHER_OPERATION
|
||||
kind = holder.split(":", 1)[0]
|
||||
if kind == MODEL_DOWNLOAD_OWNER_KIND:
|
||||
return _download_titles.get(holder, ANOTHER_OPERATION)
|
||||
return OWNER_LABELS.get(kind, ANOTHER_OPERATION)
|
||||
|
||||
|
||||
def model_busy_response() -> web.Response:
|
||||
"""``409 model-busy``: another operation holds the shared model gate.
|
||||
|
||||
``operation`` names the holder (``OWNER_LABELS``), or "Another tagging
|
||||
operation" when the gate was freed meanwhile or the kind is unknown.
|
||||
"""
|
||||
"""``409 model-busy``: another operation holds the shared model gate."""
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "model-busy",
|
||||
"error": MODEL_BUSY_ERROR,
|
||||
"operation": _holder_label(model_gate.holder()),
|
||||
},
|
||||
{"success": False, "code": "model-busy", "error": MODEL_BUSY_ERROR},
|
||||
status=409,
|
||||
)
|
||||
|
||||
|
||||
_SIZE = re.compile(r"^(~?[0-9.]+)\s*([KMGT]?B)$")
|
||||
|
||||
|
||||
def size_label(size: str) -> str:
|
||||
"""The engine table's size with a space before the unit (``~400 MB``)."""
|
||||
match = _SIZE.match(str(size or "").strip())
|
||||
return f"{match.group(1)} {match.group(2)}" if match else str(size or "")
|
||||
|
||||
|
||||
def public_models(status: dict, loaded_type: Optional[str]) -> dict:
|
||||
"""Each engine model's public entry: names and sizes from ``MODELS``, the
|
||||
flags from ``status``; never a path (``model_path`` is dropped)."""
|
||||
from ..autotag import MODELS
|
||||
|
||||
models = {}
|
||||
for model_type, config in MODELS.items():
|
||||
flags = status.get(model_type) or {}
|
||||
entry = {
|
||||
"name": config["name"],
|
||||
"label": config["label"],
|
||||
"description": config["description"],
|
||||
"size": config["size"],
|
||||
"size_label": size_label(config["size"]),
|
||||
"downloaded": bool(flags.get("downloaded")),
|
||||
"loaded": loaded_type == model_type,
|
||||
}
|
||||
for key in ("model_exists", "mmproj_exists"):
|
||||
if key in flags:
|
||||
entry[key] = bool(flags[key])
|
||||
models[model_type] = entry
|
||||
return models
|
||||
|
||||
|
||||
# Only these may be handed to the tagging engine or listed by the scanner.
|
||||
_MEDIA_EXTENSIONS = frozenset(
|
||||
{".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".tif"}
|
||||
@@ -161,39 +64,6 @@ def path_is_within(candidate, root) -> bool:
|
||||
return child != parent and child.is_relative_to(parent)
|
||||
|
||||
|
||||
def _thresholds_or_saved(general, character):
|
||||
"""The request's WD14 thresholds as floats; one it omits is the saved
|
||||
default (``autotag.defaults``, OQ4). Raises ValueError/TypeError."""
|
||||
from ..config import AutotagConfig
|
||||
|
||||
return (
|
||||
AutotagConfig.GENERAL_THRESHOLD if general is None else float(general),
|
||||
AutotagConfig.CHARACTER_THRESHOLD if character is None else float(character),
|
||||
)
|
||||
|
||||
|
||||
async def _single(event):
|
||||
"""A stream of one event."""
|
||||
yield event
|
||||
|
||||
|
||||
_monotonic = time.monotonic # the legacy frames' throttle clock (patched in tests)
|
||||
|
||||
# What one image came to in a batch (``on_item``): fixed strings, never an
|
||||
# exception's text.
|
||||
SKIP_DETAILS = {
|
||||
"unlinked": "File not found",
|
||||
"prompt already processed": "Prompt handled by another image",
|
||||
"outside allowed directories": "Outside the output folders",
|
||||
"file missing": "File not found",
|
||||
"already tagged": "Prompt already tagged",
|
||||
}
|
||||
|
||||
|
||||
def _added(count):
|
||||
return f"Added {count} tag" if count == 1 else f"Added {count} tags"
|
||||
|
||||
|
||||
async def _refusal(title):
|
||||
"""The single error event of an auto-tag batch refused by another job."""
|
||||
message = f"{title} already running"
|
||||
@@ -212,10 +82,6 @@ class AutotagRoutesMixin:
|
||||
async def download_autotag_model_route(request):
|
||||
return await self.download_autotag_model(request)
|
||||
|
||||
@routes.post("/prompt_manager/autotag/batch/preview")
|
||||
async def preview_autotag_batch_route(request):
|
||||
return await self.preview_autotag_batch(request)
|
||||
|
||||
@routes.post("/prompt_manager/autotag/start")
|
||||
async def start_autotag_route(request):
|
||||
return await self.start_autotag(request)
|
||||
@@ -243,24 +109,18 @@ class AutotagRoutesMixin:
|
||||
from ..config import AutotagConfig
|
||||
|
||||
service = get_autotag_service()
|
||||
# Off the loop: the HuggingFace check walks the whole HF cache.
|
||||
status = await self._run_in_executor(service.get_models_status)
|
||||
loaded = service.is_model_loaded()
|
||||
loaded_type = service.get_loaded_model_type() if loaded else None
|
||||
models_status = service.get_models_status()
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"models": public_models(status, loaded_type),
|
||||
# The saved tagging defaults (#328, System › Auto-tagging).
|
||||
"default_prompt": AutotagConfig.TEMPLATE,
|
||||
"model_loaded": loaded,
|
||||
"loaded_model_type": loaded_type,
|
||||
"wd14_general_threshold": AutotagConfig.GENERAL_THRESHOLD,
|
||||
"wd14_character_threshold": AutotagConfig.CHARACTER_THRESHOLD,
|
||||
"models": models_status,
|
||||
"default_prompt": service.default_prompt,
|
||||
"model_loaded": service.is_model_loaded(),
|
||||
"loaded_model_type": service.get_loaded_model_type(),
|
||||
"wd14_general_threshold": service.wd14_general_threshold,
|
||||
"wd14_character_threshold": service.wd14_character_threshold,
|
||||
"suggestion_limit": AutotagConfig.suggestion_limit(),
|
||||
"model_type": AutotagConfig.MODEL_TYPE,
|
||||
"existing": AutotagConfig.EXISTING,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -298,156 +158,55 @@ class AutotagRoutesMixin:
|
||||
return response
|
||||
|
||||
async def download_autotag_model(self, request):
|
||||
"""Download a tagging model, streaming the legacy events (POST only).
|
||||
|
||||
The download is the model_download Activity job: this request
|
||||
attaches to a running download of the same model or starts one, and
|
||||
streams the legacy frames until the job ends. A client disconnect
|
||||
ends only the stream; the download goes on. A running download of
|
||||
another model answers 409 model-busy (it holds the gate).
|
||||
"""
|
||||
from ..autotag import MODELS, get_autotag_service
|
||||
|
||||
"""Download an AutoTag model with streaming progress (POST only)."""
|
||||
model_type = request.match_info.get("model_type")
|
||||
if model_type not in MODELS:
|
||||
invalid = {"type": "error", "message": f"Invalid model type: {model_type}"}
|
||||
return await self._stream_sse(request, _single(invalid))
|
||||
try:
|
||||
job = self._jobs.active(MODEL_DOWNLOAD_JOB_KIND)
|
||||
if job is not None and job.scope.get("model_type") != model_type:
|
||||
return model_busy_response()
|
||||
if job is None:
|
||||
job, refusal = self._start_model_download(
|
||||
model_type, get_autotag_service()
|
||||
owner = model_gate.new_owner("legacy-download")
|
||||
|
||||
async def events():
|
||||
try:
|
||||
from ..autotag import get_autotag_service
|
||||
|
||||
service = get_autotag_service()
|
||||
if model_type not in service.models_config:
|
||||
yield {
|
||||
"type": "error",
|
||||
"message": f"Invalid model type: {model_type}",
|
||||
}
|
||||
return
|
||||
|
||||
yield {
|
||||
"type": "progress",
|
||||
"progress": 0,
|
||||
"status": "Starting download...",
|
||||
}
|
||||
|
||||
def progress_callback(status: str, progress: float):
|
||||
pass
|
||||
|
||||
success = await model_gate.run(
|
||||
owner, service.download_model, model_type, progress_callback
|
||||
)
|
||||
if refusal is not None:
|
||||
return refusal
|
||||
except Exception:
|
||||
self.logger.exception("Download model error")
|
||||
failed = {"type": "error", "message": _INTERNAL_ERROR}
|
||||
return await self._stream_sse(request, _single(failed))
|
||||
return await self._stream_sse(request, self._download_frames(job))
|
||||
if success:
|
||||
yield {
|
||||
"type": "complete",
|
||||
"progress": 100,
|
||||
"status": "Download complete",
|
||||
}
|
||||
else:
|
||||
yield {"type": "error", "message": "Download failed"}
|
||||
except Exception:
|
||||
self.logger.exception("Download model error")
|
||||
yield {
|
||||
"type": "error",
|
||||
"message": _INTERNAL_ERROR,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
async def _download_frames(job):
|
||||
"""The legacy download frames for job: start, then its end."""
|
||||
yield {"type": "progress", "progress": 0, "status": "Starting download..."}
|
||||
async with contextlib.aclosing(job.events()) as events:
|
||||
async for _payload in events:
|
||||
pass
|
||||
if job.state == "finished":
|
||||
yield {"type": "complete", "progress": 100, "status": "Download complete"}
|
||||
else:
|
||||
yield {"type": "error", "message": "Download failed"}
|
||||
|
||||
# ── model download job (#328) ───────────────────────────────────────
|
||||
|
||||
async def start_model_download_job(self, request):
|
||||
"""``POST /jobs/model-download {model_type}``: 202 with the job.
|
||||
|
||||
400 for a bad body, 409 ``busy`` while a download (or an exclusive
|
||||
job) runs, 409 ``model-busy`` while another operation holds the gate.
|
||||
"""
|
||||
data, error = await self._read_json_body(
|
||||
request, max_bytes=MODEL_DOWNLOAD_BODY_MAX
|
||||
)
|
||||
if error is not None:
|
||||
return error
|
||||
if set(data) != {"model_type"}:
|
||||
return self._bad_request("The body must be exactly {model_type}")
|
||||
from ..autotag import MODELS, get_autotag_service
|
||||
|
||||
model_type = data["model_type"]
|
||||
if not isinstance(model_type, str) or model_type not in MODELS:
|
||||
return self._bad_request(
|
||||
"model_type must be one of " + ", ".join(sorted(MODELS))
|
||||
)
|
||||
try:
|
||||
job, refusal = self._start_model_download(model_type, get_autotag_service())
|
||||
except Exception as e:
|
||||
self.logger.exception("Model download start error")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e)}, status=500
|
||||
)
|
||||
if refusal is not None:
|
||||
return refusal
|
||||
return web.json_response({"success": True, "job": job.summary()}, status=202)
|
||||
|
||||
def _start_model_download(self, model_type, service):
|
||||
"""Register the download job: ``(job, None)`` or ``(None, refusal)``.
|
||||
|
||||
Synchronous (no await between the checks and the start). The gate is
|
||||
taken before the job registers and released when the runner ends,
|
||||
also when the task is cancelled before the runner ever starts.
|
||||
"""
|
||||
from ..autotag import MODELS
|
||||
|
||||
busy = self._jobs.blocker(MODEL_DOWNLOAD_JOB_KIND, admission="read")
|
||||
if busy is not None:
|
||||
return None, busy_response(JobBusy(busy))
|
||||
owner = model_gate.new_owner(MODEL_DOWNLOAD_OWNER_KIND)
|
||||
if not model_gate.try_acquire(owner):
|
||||
return None, model_busy_response()
|
||||
config = MODELS[model_type]
|
||||
label = config["label"]
|
||||
title = f"Tagging model download · {label}"
|
||||
_download_titles.clear()
|
||||
_download_titles[owner] = title
|
||||
return model_busy_response()
|
||||
try:
|
||||
job = self._jobs.start(
|
||||
MODEL_DOWNLOAD_JOB_KIND,
|
||||
title,
|
||||
admission="read",
|
||||
cancellable=False,
|
||||
scope={
|
||||
"model_type": model_type,
|
||||
"label": label,
|
||||
"size_label": size_label(config["size"]),
|
||||
},
|
||||
runner=functools.partial(
|
||||
self._run_model_download, service, model_type, owner
|
||||
),
|
||||
)
|
||||
except BaseException as refused:
|
||||
# JobBusy, or anything else the registry raised: give the gate back.
|
||||
model_gate.release(owner)
|
||||
_download_titles.pop(owner, None)
|
||||
if isinstance(refused, JobBusy):
|
||||
return None, busy_response(refused)
|
||||
raise
|
||||
|
||||
def ended(_task):
|
||||
model_gate.release(owner)
|
||||
if model_gate.holder() != owner: # not while a thread still runs
|
||||
_download_titles.pop(owner, None)
|
||||
|
||||
job.task.add_done_callback(ended)
|
||||
return job, None
|
||||
|
||||
async def _run_model_download(self, service, model_type, owner, job):
|
||||
label = job.scope["label"]
|
||||
job.update(status_line=f"Downloading {label} ({job.scope['size_label']})")
|
||||
job.log_line(f"Download started: {label}")
|
||||
job.publish({"type": "status", "status": job.status_line})
|
||||
try:
|
||||
ok = await model_gate.run(owner, service.download_model, model_type, None)
|
||||
except Exception:
|
||||
self.logger.exception(f"Tagging model download failed: {model_type}")
|
||||
ok = False
|
||||
return await self._stream_sse(request, events())
|
||||
finally:
|
||||
# Before the job ends, so a reader that sees the end finds the
|
||||
# gate free (a running download thread defers the release).
|
||||
model_gate.release(owner)
|
||||
if ok:
|
||||
job.update(status_line=f"Downloaded {label}")
|
||||
job.log_line(f"Download finished: {label}")
|
||||
job.publish({"type": "status", "status": job.status_line})
|
||||
job.finish("finished")
|
||||
return
|
||||
job.update(status_line=_DOWNLOAD_FAILED)
|
||||
job.log_line(f"Download failed: {label}")
|
||||
job.publish({"type": "status", "status": job.status_line})
|
||||
job.finish("failed")
|
||||
|
||||
async def _read_start_params(self, request) -> Optional[dict]:
|
||||
"""Merge query-string and optional JSON-body parameters for autotag/start.
|
||||
@@ -478,48 +237,35 @@ class AutotagRoutesMixin:
|
||||
and not t.startswith("append:")
|
||||
]
|
||||
|
||||
async def _apply_tags_to_prompt(self, prompt_id, tags) -> int:
|
||||
"""Merge ``tags`` into the prompt; how many new tags were added."""
|
||||
async def _apply_tags_to_prompt(self, prompt_id, tags) -> bool:
|
||||
"""Merge ``tags`` into the prompt; True when something new was added."""
|
||||
existing_prompt = await self._run_in_executor(
|
||||
self.db.get_prompt_by_id, prompt_id
|
||||
)
|
||||
if not existing_prompt:
|
||||
return 0
|
||||
return False
|
||||
existing_tags = existing_prompt.get("tags", [])
|
||||
if isinstance(existing_tags, str):
|
||||
existing_tags = [t.strip() for t in existing_tags.split(",") if t.strip()]
|
||||
new_tags = list(dict.fromkeys(t for t in tags if t not in existing_tags))
|
||||
new_tags = [t for t in tags if t not in existing_tags]
|
||||
if not new_tags:
|
||||
return 0
|
||||
return False
|
||||
await self._run_in_executor(
|
||||
self.db.update_prompt_metadata, prompt_id, tags=existing_tags + new_tags
|
||||
)
|
||||
return len(new_tags)
|
||||
return True
|
||||
|
||||
async def _autotag_batch_events(self, service, opts: dict):
|
||||
"""Yield SSE progress dicts while tagging every image linked to a prompt.
|
||||
"""Yield SSE progress dicts while tagging every image linked to a prompt."""
|
||||
import time as _time
|
||||
|
||||
Both batches run here: the legacy page's (its frames unchanged) and
|
||||
the ``autotag`` job's. ``opts``: ``model_type``, ``keep_in_memory``,
|
||||
``existing`` (``skip`` leaves prompts with real tags alone, ``merge``
|
||||
adds to them), the two thresholds, and optionally ``gate_owner``,
|
||||
``prompt`` (the template, passed to each call; the legacy page sets
|
||||
``custom_prompt`` instead), ``cancelled`` (checked before each image:
|
||||
a cancelled run returns without a ``complete`` frame) and
|
||||
``on_item(index, total, group, image_path, detail)`` (one outcome per
|
||||
image, ``SKIP_DETAILS`` for the skips).
|
||||
"""
|
||||
model_type = opts["model_type"]
|
||||
keep_in_memory = opts["keep_in_memory"]
|
||||
cancelled = opts.get("cancelled") or (lambda: False)
|
||||
on_item = opts.get("on_item") or (lambda *_args: None)
|
||||
status = await self._run_in_executor(service.get_models_status)
|
||||
status = service.get_models_status()
|
||||
if not status.get(model_type, {}).get("downloaded"):
|
||||
yield {"type": "error", "message": f"Model {model_type} not downloaded"}
|
||||
return
|
||||
|
||||
if cancelled():
|
||||
return
|
||||
yield {"type": "progress", "progress": 0, "status": "Loading model..."}
|
||||
run = self._gated_runner(opts.get("gate_owner"))
|
||||
try:
|
||||
@@ -532,21 +278,15 @@ class AutotagRoutesMixin:
|
||||
}
|
||||
return
|
||||
|
||||
if opts.get("custom_prompt"):
|
||||
if opts["custom_prompt"]:
|
||||
service.custom_prompt = opts["custom_prompt"]
|
||||
tag_kwargs = {
|
||||
"general_threshold": opts["general_threshold"],
|
||||
"character_threshold": opts["character_threshold"],
|
||||
}
|
||||
if opts.get("prompt"):
|
||||
tag_kwargs["prompt"] = opts["prompt"]
|
||||
|
||||
yield {
|
||||
"type": "progress",
|
||||
"progress": 5,
|
||||
"status": "Model loaded. Fetching all images from database...",
|
||||
}
|
||||
images = await self._run_in_executor(self.db.autotag_batch_candidates)
|
||||
images = await self._run_in_executor(self.db.get_all_images)
|
||||
total = len(images)
|
||||
if total == 0:
|
||||
yield {
|
||||
@@ -558,8 +298,7 @@ class AutotagRoutesMixin:
|
||||
"skipped_outside_roots": 0,
|
||||
"status": "No images with linked prompts found",
|
||||
}
|
||||
if not keep_in_memory:
|
||||
service.unload_model()
|
||||
service.unload_model()
|
||||
return
|
||||
|
||||
yield {
|
||||
@@ -576,7 +315,7 @@ class AutotagRoutesMixin:
|
||||
"skipped_outside_roots": 0,
|
||||
}
|
||||
tagged_prompt_ids = set()
|
||||
last_update = _monotonic()
|
||||
last_update = _time.monotonic()
|
||||
finished = False
|
||||
|
||||
def progress_event(i, message):
|
||||
@@ -592,8 +331,6 @@ class AutotagRoutesMixin:
|
||||
|
||||
try:
|
||||
for i, image_data in enumerate(images):
|
||||
if cancelled():
|
||||
return
|
||||
image_path = image_data.get("image_path")
|
||||
prompt_id = image_data.get("prompt_id")
|
||||
skip_reason = None
|
||||
@@ -612,7 +349,7 @@ class AutotagRoutesMixin:
|
||||
counts["skipped_outside_roots"] += 1
|
||||
else:
|
||||
skip_reason = "file missing"
|
||||
elif opts["existing"] == "skip" and self._real_tags(
|
||||
elif opts["skip_tagged"] and self._real_tags(
|
||||
image_data.get("prompt_tags", [])
|
||||
):
|
||||
tagged_prompt_ids.add(prompt_id)
|
||||
@@ -620,8 +357,7 @@ class AutotagRoutesMixin:
|
||||
|
||||
if skip_reason is not None:
|
||||
counts["skipped"] += 1
|
||||
on_item(i, total, "skipped", image_path, SKIP_DETAILS[skip_reason])
|
||||
now = _monotonic()
|
||||
now = _time.monotonic()
|
||||
if (now - last_update) >= 0.5 or i == total - 1:
|
||||
yield progress_event(
|
||||
i, f"Skipping {i + 1}/{total} ({skip_reason})..."
|
||||
@@ -631,17 +367,17 @@ class AutotagRoutesMixin:
|
||||
continue
|
||||
|
||||
try:
|
||||
tags = await run(service.generate_tags, allowed_path, **tag_kwargs)
|
||||
tags = await run(
|
||||
service.generate_tags,
|
||||
allowed_path,
|
||||
general_threshold=opts["general_threshold"],
|
||||
character_threshold=opts["character_threshold"],
|
||||
)
|
||||
counts["processed"] += 1
|
||||
added = 0
|
||||
if tags:
|
||||
added = await self._apply_tags_to_prompt(prompt_id, tags)
|
||||
if added:
|
||||
if tags and await self._apply_tags_to_prompt(prompt_id, tags):
|
||||
counts["tagged"] += 1
|
||||
on_item(i, total, "tagged", image_path, _added(added))
|
||||
else:
|
||||
counts["skipped"] += 1
|
||||
on_item(i, total, "unchanged", image_path, "No new tags")
|
||||
tagged_prompt_ids.add(prompt_id)
|
||||
except Exception as img_err:
|
||||
self.logger.error(
|
||||
@@ -650,11 +386,10 @@ class AutotagRoutesMixin:
|
||||
)
|
||||
counts["errors"] += 1
|
||||
counts["processed"] += 1
|
||||
on_item(i, total, "errors", image_path, "Tagging failed")
|
||||
|
||||
yield progress_event(i, f"Processing {i + 1}/{total}...")
|
||||
await asyncio.sleep(0.01)
|
||||
last_update = _monotonic()
|
||||
last_update = _time.monotonic()
|
||||
|
||||
finished = True
|
||||
if not keep_in_memory:
|
||||
@@ -682,22 +417,16 @@ class AutotagRoutesMixin:
|
||||
opts = {
|
||||
"model_type": params.get("model_type", "gguf"),
|
||||
"custom_prompt": params.get("prompt", ""),
|
||||
"existing": (
|
||||
"skip"
|
||||
if str(params.get("skip_tagged", "true")).lower() == "true"
|
||||
else "merge"
|
||||
),
|
||||
"skip_tagged": str(params.get("skip_tagged", "true")).lower() == "true",
|
||||
"keep_in_memory": str(params.get("keep_in_memory", "true")).lower()
|
||||
== "true",
|
||||
"general_threshold": params.get("general_threshold"),
|
||||
"character_threshold": params.get("character_threshold"),
|
||||
}
|
||||
try:
|
||||
opts["general_threshold"], opts["character_threshold"] = (
|
||||
_thresholds_or_saved(
|
||||
opts["general_threshold"], opts["character_threshold"]
|
||||
)
|
||||
)
|
||||
for key in ("general_threshold", "character_threshold"):
|
||||
if opts[key] is not None:
|
||||
opts[key] = float(opts[key])
|
||||
except (ValueError, TypeError):
|
||||
return self._bad_request(
|
||||
"general_threshold and character_threshold must be numeric"
|
||||
@@ -727,9 +456,7 @@ class AutotagRoutesMixin:
|
||||
# never runs beside a scan, restore or the workspace tagging job.
|
||||
busy = self._jobs.blocker(AUTOTAG_BATCH_JOB_KIND, admission="records")
|
||||
if busy is not None:
|
||||
# A second legacy batch keeps its JSON 409; the workspace's Tag
|
||||
# output images job is named in a stream frame.
|
||||
if busy.kind == AUTOTAG_BATCH_JOB_KIND and busy.title != AUTOTAG_JOB_TITLE:
|
||||
if busy.kind == AUTOTAG_BATCH_JOB_KIND: # unchanged payload
|
||||
return web.json_response(
|
||||
{"success": False, "error": _BUSY_ERROR}, status=409
|
||||
)
|
||||
@@ -739,7 +466,7 @@ class AutotagRoutesMixin:
|
||||
return await self._stream_sse(request, _refusal(busy.title))
|
||||
async with self._jobs.hold(
|
||||
AUTOTAG_BATCH_JOB_KIND,
|
||||
OWNER_LABELS["legacy-batch"],
|
||||
"Auto-tag images (legacy page)",
|
||||
admission="records",
|
||||
) as job:
|
||||
response = await self._start_autotag_batch(request, opts, events)
|
||||
@@ -763,278 +490,6 @@ class AutotagRoutesMixin:
|
||||
finally:
|
||||
_autotag_batch_lock.release()
|
||||
|
||||
# ── Tag output images (#328) ────────────────────────────────────────
|
||||
|
||||
async def _batch_body(self, request, keys, max_bytes):
|
||||
"""``(data, None)`` for a JSON object with exactly ``keys``."""
|
||||
data, error = await self._read_json_body(request, max_bytes=max_bytes)
|
||||
if error is not None:
|
||||
return None, error
|
||||
if set(data) != set(keys):
|
||||
return None, self._bad_request(
|
||||
"The body must be exactly {" + ", ".join(keys) + "}"
|
||||
)
|
||||
return data, None
|
||||
|
||||
@staticmethod
|
||||
def _batch_choice(data):
|
||||
"""The model's engine entry and the policy, or a 400 message."""
|
||||
from ..autotag import MODELS
|
||||
from ..config import EXISTING_POLICIES
|
||||
|
||||
model_type, existing = data.get("model_type"), data.get("existing")
|
||||
if not isinstance(model_type, str) or model_type not in MODELS:
|
||||
return None, "model_type must be one of " + ", ".join(sorted(MODELS))
|
||||
if not isinstance(existing, str) or existing not in EXISTING_POLICIES:
|
||||
return None, "existing must be one of " + ", ".join(EXISTING_POLICIES)
|
||||
return MODELS[model_type], None
|
||||
|
||||
async def preview_autotag_batch(self, request):
|
||||
"""Review Tag output images: counts, model, backup and disk; no write.
|
||||
|
||||
Opens no image file. Refused (409 ``busy``) while an exclusive job runs.
|
||||
"""
|
||||
try:
|
||||
async with self._jobs.lease(AUTOTAG_PREVIEW_LEASE):
|
||||
return await self._preview_autotag_batch(request)
|
||||
except JobBusy as busy:
|
||||
return busy_response(busy)
|
||||
|
||||
async def _preview_autotag_batch(self, request):
|
||||
from ..autotag import get_autotag_service
|
||||
|
||||
data, error = await self._batch_body(
|
||||
request, ("model_type", "existing"), BATCH_PREVIEW_BODY_MAX
|
||||
)
|
||||
if error is not None:
|
||||
return error
|
||||
config, message = self._batch_choice(data)
|
||||
if message:
|
||||
return self._bad_request(message)
|
||||
model_type, existing = data["model_type"], data["existing"]
|
||||
policy = self._safety_backup_policy()
|
||||
try:
|
||||
candidates = await self._run_in_executor(self.db.autotag_batch_candidates)
|
||||
status = await self._run_in_executor(
|
||||
get_autotag_service().get_models_status
|
||||
)
|
||||
disk, _short, _size = await self._run_in_executor(
|
||||
functools.partial(
|
||||
disk_plan,
|
||||
self.db.model.db_path,
|
||||
backup=policy != "off",
|
||||
vacuum=False,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.exception("Tag output images could not be reviewed")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e)}, status=500
|
||||
)
|
||||
prompts = {row["prompt_id"]: row["prompt_tags"] for row in candidates}
|
||||
tagged = sum(1 for tags in prompts.values() if self._real_tags(tags))
|
||||
to_tag = len(prompts) - tagged if existing == "skip" else len(prompts)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"images": len(candidates),
|
||||
"prompts": len(prompts),
|
||||
"already_tagged": tagged,
|
||||
"to_tag": to_tag,
|
||||
"model": {
|
||||
"type": model_type,
|
||||
"label": config["label"],
|
||||
"size_label": size_label(config["size"]),
|
||||
"downloaded": bool(
|
||||
(status.get(model_type) or {}).get("downloaded")
|
||||
),
|
||||
},
|
||||
"existing": existing,
|
||||
"review_required": existing == "replace_after_review",
|
||||
"backup": {"policy": policy},
|
||||
"disk": disk,
|
||||
}
|
||||
)
|
||||
|
||||
async def start_autotag_job(self, request):
|
||||
"""``POST /jobs/autotag``: start Tag output images (202), or refuse.
|
||||
|
||||
In order: 400 body or options, 400 ``review-required``; 409
|
||||
``model-unavailable``; 507 ``insufficient-space``; 409 ``busy``
|
||||
(another records job, the legacy batch's hold included); 409
|
||||
``model-busy``. No await follows the busy check.
|
||||
"""
|
||||
from ..autotag import get_autotag_service
|
||||
from ..config import MAX_SUGGESTION_LIMIT, AutotagConfig
|
||||
from .annotation_suggest import OptionError, validate_options
|
||||
|
||||
data, error = await self._batch_body(
|
||||
request, ("model_type", "existing", "options"), BATCH_JOB_BODY_MAX
|
||||
)
|
||||
if error is not None:
|
||||
return error
|
||||
config, message = self._batch_choice(data)
|
||||
if message:
|
||||
return self._bad_request(message)
|
||||
model_type, existing, options = (
|
||||
data["model_type"],
|
||||
data["existing"],
|
||||
data["options"],
|
||||
)
|
||||
if isinstance(options, dict) and {"suggestion_limit", "use_gpu"} & set(options):
|
||||
return self._bad_request("options take no suggestion_limit or use_gpu")
|
||||
try:
|
||||
_model, options = validate_options(
|
||||
model_type,
|
||||
options,
|
||||
default_limit=AutotagConfig.suggestion_limit(),
|
||||
max_limit=MAX_SUGGESTION_LIMIT,
|
||||
)
|
||||
except OptionError as error:
|
||||
return self._bad_request(str(error))
|
||||
if existing == "replace_after_review":
|
||||
return web.json_response(
|
||||
{"success": False, "code": "review-required", "error": REVIEW_REQUIRED},
|
||||
status=400,
|
||||
)
|
||||
service = get_autotag_service()
|
||||
backup_on = self._safety_backup_policy() == "on"
|
||||
try:
|
||||
status = await self._run_in_executor(service.get_models_status)
|
||||
candidates = await self._run_in_executor(self.db.autotag_batch_candidates)
|
||||
_disk, short, _size = await self._run_in_executor(
|
||||
functools.partial(
|
||||
disk_plan, self.db.model.db_path, backup=backup_on, vacuum=False
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.exception("Tag output images could not start")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e)}, status=500
|
||||
)
|
||||
if not (status.get(model_type) or {}).get("downloaded"):
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "model-unavailable",
|
||||
"error": f"{config['label']} is not downloaded. "
|
||||
"Download it in System › Auto-tagging.",
|
||||
},
|
||||
status=409,
|
||||
)
|
||||
if backup_on and short is not None:
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "insufficient-space",
|
||||
"error": short_space_message(short),
|
||||
},
|
||||
status=507,
|
||||
)
|
||||
busy = self._jobs.blocker(AUTOTAG_BATCH_JOB_KIND, admission="records")
|
||||
if busy is not None:
|
||||
return busy_response(JobBusy(busy))
|
||||
owner = model_gate.new_owner(AUTOTAG_OWNER_KIND)
|
||||
if not model_gate.try_acquire(owner):
|
||||
return model_busy_response()
|
||||
opts = {
|
||||
"model_type": model_type,
|
||||
"existing": existing,
|
||||
"keep_in_memory": True,
|
||||
"gate_owner": owner,
|
||||
"general_threshold": options.get("general_threshold"),
|
||||
"character_threshold": options.get("character_threshold"),
|
||||
"prompt": options.get("instructions"),
|
||||
}
|
||||
try:
|
||||
job = self._jobs.start(
|
||||
AUTOTAG_BATCH_JOB_KIND,
|
||||
AUTOTAG_JOB_TITLE,
|
||||
admission="records",
|
||||
cancellable=True,
|
||||
scope={
|
||||
"model_type": model_type,
|
||||
"existing": existing,
|
||||
"images": len(candidates),
|
||||
"label": config["label"],
|
||||
},
|
||||
groups=AUTOTAG_GROUP_NAMES,
|
||||
runner=functools.partial(
|
||||
self._run_autotag_job, service, opts, backup_on
|
||||
),
|
||||
)
|
||||
except BaseException as refused:
|
||||
model_gate.release(owner)
|
||||
if isinstance(refused, JobBusy):
|
||||
return busy_response(refused)
|
||||
raise
|
||||
job.task.add_done_callback(lambda _task: model_gate.release(owner))
|
||||
return web.json_response({"success": True, "job": job.summary()}, status=202)
|
||||
|
||||
async def _run_autotag_job(self, service, opts, backup_on, job):
|
||||
ok, _backup = await self._safety_backup_step(job, self.db, backup_on)
|
||||
if not ok:
|
||||
return
|
||||
try:
|
||||
await self._tag_output_images(service, opts, job)
|
||||
except Exception:
|
||||
self.logger.exception("Tag output images failed")
|
||||
job.update(status_line=_INTERNAL_ERROR)
|
||||
job.log_line(f"Tagging failed · {_INTERNAL_ERROR}")
|
||||
job.finish("failed")
|
||||
finally:
|
||||
if backup_on:
|
||||
await self._prune_backup_pool(job, self.db)
|
||||
|
||||
async def _tag_output_images(self, service, opts, job):
|
||||
roots = _scan_roots(self._get_all_output_dirs())
|
||||
job.log_line(
|
||||
f"Tagging started · {job.scope['label']} · "
|
||||
+ ("skip prompts already tagged" if opts["existing"] == "skip" else "merge")
|
||||
)
|
||||
job.update(status_line="Loading the tagging model")
|
||||
|
||||
def on_item(index, total, group, image_path, detail):
|
||||
path, root, root_index = _scan_location(image_path or "", roots)
|
||||
job.update(
|
||||
processed=index + 1,
|
||||
total=total,
|
||||
progress=100 * (index + 1) // total,
|
||||
status_line=f"Tagging {index + 1} of {total}",
|
||||
)
|
||||
job.record(
|
||||
group,
|
||||
{
|
||||
"path": path,
|
||||
"root": root,
|
||||
"root_index": root_index,
|
||||
"detail": detail,
|
||||
},
|
||||
f"{path}: {detail}",
|
||||
)
|
||||
job.publish({"type": "progress", "progress": job.progress})
|
||||
|
||||
final = None
|
||||
async for event in self._autotag_batch_events(
|
||||
service,
|
||||
{**opts, "cancelled": lambda: job.cancelled, "on_item": on_item},
|
||||
):
|
||||
if event.get("type") in ("complete", "error"):
|
||||
final = event
|
||||
if final is not None and final["type"] == "error":
|
||||
job.update(status_line=final["message"])
|
||||
job.log_line(f"Tagging failed · {final['message']}")
|
||||
job.finish("failed")
|
||||
elif job.cancelled: # a cancel wins over a run that then completed
|
||||
job.update(status_line=BATCH_CANCELLED)
|
||||
job.log_line(BATCH_CANCELLED)
|
||||
job.finish("cancelled")
|
||||
else:
|
||||
done = f"Tagged {_plural(job.counters.get('tagged', 0), 'prompt')}"
|
||||
job.update(status_line=done, progress=100)
|
||||
job.log_line(done)
|
||||
job.finish("finished")
|
||||
|
||||
def _gated_runner(self, owner):
|
||||
"""Executor calls under the model gate *owner* holds (plain when None)."""
|
||||
if owner is None:
|
||||
@@ -1151,9 +606,10 @@ class AutotagRoutesMixin:
|
||||
general_threshold = data.get("general_threshold")
|
||||
character_threshold = data.get("character_threshold")
|
||||
try:
|
||||
general_threshold, character_threshold = _thresholds_or_saved(
|
||||
general_threshold, character_threshold
|
||||
)
|
||||
if general_threshold is not None:
|
||||
general_threshold = float(general_threshold)
|
||||
if character_threshold is not None:
|
||||
character_threshold = float(character_threshold)
|
||||
except (ValueError, TypeError):
|
||||
return self._bad_request(
|
||||
"general_threshold and character_threshold must be numeric"
|
||||
@@ -1336,16 +792,13 @@ class AutotagRoutesMixin:
|
||||
rel_path = image_path.relative_to(output_path)
|
||||
thumbnail_url = None
|
||||
if has_thumbnails:
|
||||
# Any preview the file owns, the current format first (#317)
|
||||
existing = thumbnail_settings.existing_thumbnail(
|
||||
output_path, rel_path
|
||||
thumb_rel = (
|
||||
f"thumbnails/{rel_path.with_suffix('').as_posix()}"
|
||||
f"_thumb{image_path.suffix}"
|
||||
)
|
||||
thumbnail_url = thumbnail_url_for(
|
||||
thumb_rel, output_path / thumb_rel
|
||||
)
|
||||
if existing is not None:
|
||||
thumb_rel = (
|
||||
thumbnail_settings.thumbnail_stem_rel(rel_path)
|
||||
+ existing.suffix
|
||||
)
|
||||
thumbnail_url = thumbnail_url_for(thumb_rel, existing)
|
||||
images.append(
|
||||
{
|
||||
"filename": image_path.name,
|
||||
|
||||
@@ -1,942 +0,0 @@
|
||||
"""Database backup and restore routes (#249 System S8).
|
||||
|
||||
Legacy routes, responses unchanged:
|
||||
|
||||
- ``GET /prompt_manager/backup`` downloads a consistent copy of the library.
|
||||
While the copy is made it holds the registry as ``backup`` ("Database
|
||||
backup", ``guarded``): it is refused while a restore, maintenance apply or
|
||||
other record-changing job runs. The copy is written in the library's
|
||||
folder and made one standalone file; the hold ends there, so a slow or
|
||||
stalled download never keeps other work waiting. A copy that fails its
|
||||
integrity check is a 500 with no bytes. One that fails only a semantic
|
||||
check (``verify_database_file``) is still sent: a user can always get a
|
||||
backup. The reason is logged and sent as ``X-PromptManager-Verify``
|
||||
(``ok`` otherwise), and the status records it as not verified. The file
|
||||
is streamed, then closed and removed.
|
||||
- ``POST /prompt_manager/restore`` (multipart ``database_file``) restores in
|
||||
one request through the same stage and publish as the restore job, holding
|
||||
the registry as ``restore`` (``exclusive``).
|
||||
|
||||
Reviewed restore (new routes answer errors as ``{success:false, code, error}``):
|
||||
|
||||
- ``POST /prompt_manager/restore/preview`` (multipart ``database_file``, the
|
||||
same disk-space cap) stages the upload in the library's folder
|
||||
(``database/restore_stage.py``) and answers what it holds beside the live
|
||||
library: ``{token, expires_at, summary, schema:{migrated, library_state},
|
||||
integrity:"ok", documents:{repaired, cleared, images}, live, size,
|
||||
safety_backup:{destination_label}}``. Nothing live is written or backed
|
||||
up. At most two previews are held (``too-large``, 429).
|
||||
- ``DELETE /prompt_manager/restore/preview/{token}`` discards a stage.
|
||||
- ``POST /prompt_manager/restore/apply {token}`` starts the restore job (202
|
||||
``{job}``): kind ``restore``, "Database restore", ``exclusive`` (it runs
|
||||
alone, like Move library), not cancellable, indeterminate. The runner waits
|
||||
for request leases to end, rechecks the live fingerprint (``plan-changed``),
|
||||
publishes the exact staged bytes, swaps ``self.db``, prunes old safety
|
||||
backups and logs what was restored. Once the copy replaced the library,
|
||||
the restore is done: the later steps only log their failures.
|
||||
- ``GET /prompt_manager/backup/status``: the safety backups on disk
|
||||
(``<db>.backup_*``, newest first, names only) and this session's last
|
||||
verified download (Q9: no sidecar file).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import datetime
|
||||
import os
|
||||
import sqlite3
|
||||
import tempfile
|
||||
from contextlib import closing
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from . import admin as _admin
|
||||
from .job_routes import busy_response
|
||||
from .jobs import RESTORE_JOB_KIND, JobBusy
|
||||
from .maintenance_plans import PlanExpired, PlanUsed
|
||||
from .restore_plans import RestorePlanStore
|
||||
|
||||
try:
|
||||
from ...database import restore_stage
|
||||
from ...database.operations import PromptDatabase
|
||||
from ...database.safety_backup import (
|
||||
SafetyBackupFailed,
|
||||
fsync_file,
|
||||
make_standalone,
|
||||
remove_quietly,
|
||||
safety_temp,
|
||||
)
|
||||
except ImportError:
|
||||
from database import restore_stage
|
||||
from database.operations import PromptDatabase
|
||||
from database.safety_backup import (
|
||||
SafetyBackupFailed,
|
||||
fsync_file,
|
||||
make_standalone,
|
||||
remove_quietly,
|
||||
safety_temp,
|
||||
)
|
||||
|
||||
BACKUP_JOB_KIND = "backup"
|
||||
BACKUP_HOLD_TITLE = "Database backup"
|
||||
LEGACY_RESTORE_TITLE = "Database restore (admin page)"
|
||||
RESTORE_JOB_TITLE = "Database restore"
|
||||
RESTORE_GROUPS = ("errors",)
|
||||
MAX_TOKEN_LENGTH = 128
|
||||
APPLY_BODY_MAX = 1024
|
||||
STREAM_CHUNK = 1024 * 1024
|
||||
MAX_LISTED_BACKUPS = 10
|
||||
SIDECARS = ("", "-wal", "-shm", "-journal")
|
||||
REQUESTS_BUSY = (
|
||||
"Other requests are still reading the library; nothing was changed. "
|
||||
"Try again in a moment."
|
||||
)
|
||||
# S9 I1: the review goes back to the page that started the restore
|
||||
# (``scope.review_kept``), which opens it again with Restore and Discard; no
|
||||
# other place can start a kept review, so the line says where it is.
|
||||
REVIEW_KEPT = (
|
||||
"Other requests are still reading the library; nothing was changed. "
|
||||
"The review was reopened in the page that started the restore: restore it "
|
||||
"again there, or discard it."
|
||||
)
|
||||
VERIFY_HEADER = "X-PromptManager-Verify"
|
||||
SQLITE_HEADER = b"SQLite format 3\x00"
|
||||
PLAN_CHANGED = (
|
||||
"The library changed since the review; nothing was changed. "
|
||||
"Review the restore again."
|
||||
)
|
||||
NOT_RESTORED = "The restore failed; the library was not changed."
|
||||
|
||||
|
||||
def _failure(message, status, code=None, **extra):
|
||||
return web.json_response(
|
||||
{"success": False, "code": code, "error": message, **extra}, status=status
|
||||
)
|
||||
|
||||
|
||||
def _iso(timestamp):
|
||||
return datetime.datetime.fromtimestamp(timestamp, datetime.timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _plural(count, word):
|
||||
return f"{count:,} {word}" + ("" if count == 1 else "s")
|
||||
|
||||
|
||||
def documents_line(documents):
|
||||
"""The repair report's sentence, or "" when nothing was damaged."""
|
||||
repaired, cleared = documents.get("repaired", 0), documents.get("cleared", 0)
|
||||
parts = []
|
||||
if repaired:
|
||||
parts.append(f"repaired {_plural(repaired, 'damaged workflow document')}")
|
||||
if cleared:
|
||||
images = documents.get("images", 0)
|
||||
parts.append(
|
||||
f"could not repair {_plural(cleared, 'workflow document')}: "
|
||||
f"{_plural(images, 'image')} {'shows' if images == 1 else 'show'} "
|
||||
"no workflow"
|
||||
)
|
||||
if not parts:
|
||||
return ""
|
||||
text = "; ".join(parts)
|
||||
return text[0].upper() + text[1:] + "."
|
||||
|
||||
|
||||
def restored_line(summary):
|
||||
return (
|
||||
f"Restored {_plural(summary['prompts'], 'prompt')}, "
|
||||
f"{_plural(summary['image_links'], 'image link')}, "
|
||||
f"{_plural(summary['pairs'], 'saved pair')} and "
|
||||
f"{_plural(summary['runs'], 'run')}"
|
||||
)
|
||||
|
||||
|
||||
def safety_backups(db_path, limit=MAX_LISTED_BACKUPS):
|
||||
"""``<db>.backup_*`` files beside the database, newest first (blocking)."""
|
||||
db_path = os.path.abspath(db_path)
|
||||
folder = os.path.dirname(db_path)
|
||||
prefix = os.path.basename(db_path) + ".backup_"
|
||||
found = []
|
||||
try:
|
||||
names = os.listdir(folder)
|
||||
except OSError:
|
||||
return []
|
||||
for name in names:
|
||||
if not name.startswith(prefix):
|
||||
continue
|
||||
try:
|
||||
stat = os.stat(os.path.join(folder, name))
|
||||
except OSError:
|
||||
continue
|
||||
if not os.path.isfile(os.path.join(folder, name)):
|
||||
continue
|
||||
found.append((stat.st_mtime, name, stat.st_size))
|
||||
found.sort(key=lambda item: (item[0], item[1]), reverse=True)
|
||||
return [
|
||||
{"name": name, "bytes": size, "modified_at": _iso(mtime)}
|
||||
for mtime, name, size in found[:limit]
|
||||
]
|
||||
|
||||
|
||||
def snapshot_corruption(path):
|
||||
"""Why the snapshot file is corrupt, or None (blocking).
|
||||
|
||||
Only physical damage refuses a download: not a SQLite file, or a failed
|
||||
``PRAGMA integrity_check``.
|
||||
"""
|
||||
try:
|
||||
with open(path, "rb") as handle:
|
||||
if handle.read(len(SQLITE_HEADER)) != SQLITE_HEADER:
|
||||
return "Not a SQLite database file"
|
||||
uri = Path(os.path.abspath(path)).as_uri() + "?mode=ro"
|
||||
with closing(sqlite3.connect(uri, uri=True)) as conn:
|
||||
row = conn.execute("PRAGMA integrity_check").fetchone()
|
||||
except (OSError, sqlite3.Error) as error:
|
||||
return f"Invalid SQLite database: {error}"
|
||||
verdict = row[0] if row else "no result"
|
||||
return None if verdict == "ok" else f"Integrity check failed: {verdict}"
|
||||
|
||||
|
||||
def write_backup_snapshot(model):
|
||||
"""A standalone copy of the library beside it: ``(path, warning)`` (blocking).
|
||||
|
||||
Raises SafetyBackupFailed with nothing left on disk when the copy cannot
|
||||
be written or is corrupt (``snapshot_corruption``). ``warning`` is the
|
||||
reason a semantic check (``verify_database_file``, workflow documents
|
||||
not read) refused the copy, or None. The caller removes the file.
|
||||
"""
|
||||
temp = safety_temp(model.db_path)
|
||||
try:
|
||||
if not model.backup_database(temp):
|
||||
raise SafetyBackupFailed("The backup could not be written")
|
||||
make_standalone(temp)
|
||||
corrupt = snapshot_corruption(temp)
|
||||
if corrupt is not None:
|
||||
raise SafetyBackupFailed("The backup is corrupt", corrupt)
|
||||
ok, reason = model.verify_database_file(temp, documents=False)
|
||||
fsync_file(temp)
|
||||
return temp, (None if ok else reason)
|
||||
except BaseException:
|
||||
remove_quietly(*(temp + suffix for suffix in SIDECARS))
|
||||
raise
|
||||
|
||||
|
||||
def _header_text(text):
|
||||
"""A reason as a one-line, latin-1 header value."""
|
||||
flat = " ".join(str(text).split())[:300]
|
||||
return flat.encode("latin-1", "replace").decode("latin-1")
|
||||
|
||||
|
||||
class BackupRestoreRoutesMixin:
|
||||
"""Mixin providing the backup and restore endpoints."""
|
||||
|
||||
def _register_backup_restore_routes(self, routes):
|
||||
@routes.post("/prompt_manager/restore/preview")
|
||||
async def restore_preview_route(request):
|
||||
return await self.preview_restore(request)
|
||||
|
||||
@routes.delete("/prompt_manager/restore/preview/{token}")
|
||||
async def restore_discard_route(request):
|
||||
return await self.discard_restore_preview(request)
|
||||
|
||||
@routes.post("/prompt_manager/restore/apply")
|
||||
async def restore_apply_route(request):
|
||||
return await self.apply_restore(request)
|
||||
|
||||
@routes.get("/prompt_manager/backup/status")
|
||||
async def backup_status_route(request):
|
||||
return await self.backup_status(request)
|
||||
|
||||
@property
|
||||
def _restore_plans(self):
|
||||
store = getattr(self, "_restore_plan_store", None)
|
||||
if store is None:
|
||||
store = self._restore_plan_store = RestorePlanStore()
|
||||
return store
|
||||
|
||||
# ── backup ────────────────────────────────────────────────────────
|
||||
|
||||
async def backup_database(self, request):
|
||||
"""Download a copy of the prompts database (WAL included)."""
|
||||
return await self._stream_backup(request)
|
||||
|
||||
async def _stream_backup(self, request):
|
||||
model = self.db.model
|
||||
if not os.path.exists(model.db_path):
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Database file not found"}, status=404
|
||||
)
|
||||
made = await self._run_held(
|
||||
BACKUP_JOB_KIND,
|
||||
BACKUP_HOLD_TITLE,
|
||||
"guarded",
|
||||
lambda: self._make_backup_snapshot(model),
|
||||
)
|
||||
if isinstance(made, web.StreamResponse):
|
||||
return made # refused, or the snapshot failed
|
||||
temp, warning = made
|
||||
return await self._send_snapshot(request, temp, warning)
|
||||
|
||||
async def _make_backup_snapshot(self, model):
|
||||
"""``(path, warning)``, or the error response (inside the hold)."""
|
||||
loop = asyncio.get_running_loop()
|
||||
future = loop.run_in_executor(None, write_backup_snapshot, model)
|
||||
try:
|
||||
temp, warning = await asyncio.shield(future)
|
||||
except asyncio.CancelledError:
|
||||
# The worker cannot be stopped: remove what it writes once it is
|
||||
# done (#335 re-review R4).
|
||||
future.add_done_callback(
|
||||
lambda done: loop.create_task(self._remove_finished_snapshot(done))
|
||||
)
|
||||
raise
|
||||
except SafetyBackupFailed as failed:
|
||||
self.logger.error(f"Backup failed: {failed} {failed.reason or ''}")
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Failed to create database backup"},
|
||||
status=500,
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Backup error: {e}", exc_info=True)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"error": f"Failed to backup database: {self._public_error(e)}",
|
||||
},
|
||||
status=500,
|
||||
)
|
||||
if warning is not None:
|
||||
self.logger.warning(f"Backup sent with a failed check: {warning}")
|
||||
return temp, warning
|
||||
|
||||
async def _remove_finished_snapshot(self, future):
|
||||
"""Remove the snapshot a cancelled download's worker finished."""
|
||||
if future.cancelled() or future.exception() is not None:
|
||||
return # nothing was left on disk
|
||||
temp, _warning = future.result()
|
||||
await self._run_in_executor(
|
||||
remove_quietly, *(temp + suffix for suffix in SIDECARS)
|
||||
)
|
||||
|
||||
async def _send_snapshot(self, request, temp, warning):
|
||||
"""Stream the snapshot (outside the hold), then close and remove it."""
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
handle = None
|
||||
try:
|
||||
size = os.path.getsize(temp)
|
||||
handle = open(temp, "rb")
|
||||
response = web.StreamResponse(
|
||||
headers={
|
||||
"Content-Type": "application/octet-stream",
|
||||
"Content-Disposition": (
|
||||
f'attachment; filename="prompts_backup_{timestamp}.db"'
|
||||
),
|
||||
"Content-Length": str(size),
|
||||
VERIFY_HEADER: "ok" if warning is None else _header_text(warning),
|
||||
}
|
||||
)
|
||||
await response.prepare(request)
|
||||
while True:
|
||||
chunk = await self._run_in_executor(handle.read, STREAM_CHUNK)
|
||||
if not chunk:
|
||||
break
|
||||
await response.write(chunk)
|
||||
await response.write_eof()
|
||||
self._last_backup_download = {
|
||||
"at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
||||
"bytes": size,
|
||||
"verified": warning is None,
|
||||
}
|
||||
return response
|
||||
finally:
|
||||
if handle is not None:
|
||||
handle.close() # before the unlink (Windows)
|
||||
await self._run_in_executor(
|
||||
remove_quietly, *(temp + suffix for suffix in SIDECARS)
|
||||
)
|
||||
|
||||
async def backup_status(self, request):
|
||||
try:
|
||||
backups = await self._run_in_executor(safety_backups, self.db.model.db_path)
|
||||
except Exception:
|
||||
self.logger.exception("Backup status failed")
|
||||
return _failure("Failed to read the backup status", 500)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"safety_backups": backups,
|
||||
"last_download": getattr(self, "_last_backup_download", None),
|
||||
},
|
||||
headers={"Cache-Control": "no-store"},
|
||||
)
|
||||
|
||||
# ── the upload ────────────────────────────────────────────────────
|
||||
|
||||
def _restore_refusal(self, size, ceiling, room, *, known):
|
||||
"""``400 too-large`` past the ceiling, else ``507 insufficient-space``."""
|
||||
from .maintenance import restore_space_message
|
||||
|
||||
if size > ceiling:
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "too-large",
|
||||
"error": (
|
||||
"File too large. Maximum size is "
|
||||
f"{ceiling // (1024 * 1024)}MB"
|
||||
),
|
||||
},
|
||||
status=400,
|
||||
)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "insufficient-space",
|
||||
"error": restore_space_message(size, room, known=known),
|
||||
},
|
||||
status=507,
|
||||
)
|
||||
|
||||
async def _receive_restore_upload(self, request, uploads):
|
||||
"""Stream ``database_file`` into the library's folder.
|
||||
|
||||
Returns ``(path, None)`` or ``(None, refusal)``. Each temp path is
|
||||
appended to ``uploads`` as soon as it exists, so the caller removes
|
||||
it (and its sidecars) whatever happens.
|
||||
"""
|
||||
from .maintenance import RESTORE_MAX_BYTES, restore_room
|
||||
|
||||
ceiling = getattr(self, "restore_max_bytes", RESTORE_MAX_BYTES)
|
||||
db_path = self.db.model.db_path
|
||||
# Free disk space decides how large a backup may be (#313): the
|
||||
# upload, its staged copy and the swap must all fit.
|
||||
room = await self._run_in_executor(restore_room, db_path)
|
||||
max_bytes = min(ceiling, room)
|
||||
length = request.content_length
|
||||
if length is not None and length > max_bytes:
|
||||
return None, self._restore_refusal(length, ceiling, room, known=True)
|
||||
reader = await request.multipart()
|
||||
field = await reader.next()
|
||||
if not field or field.name != "database_file":
|
||||
return None, web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"error": (
|
||||
"No database file uploaded. "
|
||||
"Expected field name: database_file"
|
||||
),
|
||||
},
|
||||
status=400,
|
||||
)
|
||||
fd, temp_path = tempfile.mkstemp(
|
||||
prefix=restore_stage.UPLOAD_PREFIX,
|
||||
suffix=".tmp",
|
||||
dir=os.path.dirname(os.path.abspath(db_path)),
|
||||
)
|
||||
uploads.append(temp_path)
|
||||
size = 0
|
||||
with os.fdopen(fd, "wb") as out:
|
||||
while True:
|
||||
chunk = await field.read_chunk()
|
||||
if not chunk:
|
||||
break
|
||||
size += len(chunk)
|
||||
if size > max_bytes:
|
||||
return None, self._restore_refusal(size, ceiling, room, known=False)
|
||||
out.write(chunk)
|
||||
if size == 0:
|
||||
return None, web.json_response(
|
||||
{"success": False, "error": "Uploaded file is empty"}, status=400
|
||||
)
|
||||
return temp_path, None
|
||||
|
||||
async def _upload_problem(self, path):
|
||||
"""None when the upload may be staged, else a readable reason."""
|
||||
ok, reason = await self._run_in_executor(
|
||||
self.db.model.verify_database_file, path, documents=False
|
||||
)
|
||||
if ok:
|
||||
reason = await self._run_in_executor(_admin._inspect_restore_upload, path)
|
||||
ok = reason is None
|
||||
return None if ok else reason
|
||||
|
||||
async def _remove_uploads(self, uploads):
|
||||
for path in uploads:
|
||||
await self._run_in_executor(
|
||||
remove_quietly, *(path + suffix for suffix in SIDECARS)
|
||||
)
|
||||
|
||||
async def _prune_after(self, db_path, *, failed=False):
|
||||
"""Prune old safety backups; a failure is only logged (never raises).
|
||||
|
||||
It runs after a restore replaced the library, or after one failed:
|
||||
either way the restore's outcome stands.
|
||||
"""
|
||||
try:
|
||||
removed = await self._run_in_executor(_admin._prune_safety_backups, db_path)
|
||||
except Exception as error:
|
||||
when = "failed" if failed else "completed"
|
||||
self.logger.warning(
|
||||
f"Could not prune safety backups after a {when} restore: "
|
||||
f"{type(error).__name__}"
|
||||
)
|
||||
return
|
||||
if removed:
|
||||
self.logger.info(f"Removed {len(removed)} old safety backups")
|
||||
|
||||
async def _publish(self, model, staged):
|
||||
"""Publish ``staged``; on a page-copy failure prune and re-raise."""
|
||||
try:
|
||||
return await self._run_in_executor(model.publish_staged, staged)
|
||||
except sqlite3.OperationalError as exc:
|
||||
if getattr(exc, "_prompt_manager_restore_backup_created", False) is True:
|
||||
await self._prune_after(model.db_path, failed=True)
|
||||
raise
|
||||
|
||||
async def _after_publish(self, db_path):
|
||||
"""Everything after the copy replaced the library: best effort, logged.
|
||||
|
||||
The restore is done once the copy succeeded, so nothing here fails it
|
||||
(#335 review I1). The library is adopted as ``self.db`` (the previous
|
||||
object, which points at the same file, stays when a new one cannot
|
||||
be opened), the generation is bumped (so the health report and
|
||||
reviewed plans of the old library are dropped), held restore
|
||||
previews are dropped and the previous object's connections closed.
|
||||
"""
|
||||
await self._prune_after(db_path)
|
||||
previous = self.db
|
||||
try:
|
||||
self.db = PromptDatabase(db_path)
|
||||
except Exception:
|
||||
self.logger.exception("Could not reopen the restored library")
|
||||
try:
|
||||
self._database_replaced()
|
||||
except Exception:
|
||||
self.logger.exception("Could not drop the old library's reports")
|
||||
try:
|
||||
await self._discard_held(self._restore_plans.pop_all())
|
||||
except Exception:
|
||||
self.logger.exception("Could not remove the old restore previews")
|
||||
if previous is not self.db:
|
||||
try:
|
||||
# Its connections would otherwise linger (a lock on Windows).
|
||||
await self._run_in_executor(previous.close_all)
|
||||
except Exception:
|
||||
self.logger.exception("Could not close the previous library")
|
||||
|
||||
async def _discard_held(self, held):
|
||||
"""Remove the stage folders of dropped store entries (off the loop)."""
|
||||
for entry in held:
|
||||
try:
|
||||
await self._run_in_executor(entry.staged.discard)
|
||||
except Exception:
|
||||
self.logger.exception("Could not remove a restore stage")
|
||||
|
||||
async def _discard_popped(self, token):
|
||||
"""Drop ``token`` and remove its stage, only if the store still held it.
|
||||
|
||||
A stage another apply's job has taken is not this request's to remove
|
||||
(#335 re-review R2).
|
||||
"""
|
||||
popped = self._restore_plans.pop(token)
|
||||
if popped is not None:
|
||||
await self._discard_held([popped])
|
||||
|
||||
async def _drop_expired_restores(self):
|
||||
"""Drop expired previews from the store (on the loop), then their files."""
|
||||
await self._discard_held(self._restore_plans.pop_expired())
|
||||
|
||||
def _schedule_restore_expiry(self):
|
||||
"""Remove each preview's stage when it expires, without a request.
|
||||
|
||||
A loop timer (``call_later``), so a server that stops with previews
|
||||
held leaves nothing pending; the startup sweep removes the folders.
|
||||
"""
|
||||
store = self._restore_plans
|
||||
expires = store.next_expiry()
|
||||
if expires is None:
|
||||
return
|
||||
handle = getattr(self, "_restore_expiry_handle", None)
|
||||
if handle is not None:
|
||||
handle.cancel()
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def fire():
|
||||
self._restore_expiry_handle = None
|
||||
task = loop.create_task(self._expire_restores())
|
||||
self._restore_expiry_task = task
|
||||
|
||||
delay = max(0.0, expires - store.now())
|
||||
self._restore_expiry_handle = loop.call_later(delay, fire)
|
||||
|
||||
async def _expire_restores(self):
|
||||
await self._drop_expired_restores()
|
||||
self._schedule_restore_expiry()
|
||||
|
||||
# ── legacy one-step restore ───────────────────────────────────────
|
||||
|
||||
async def restore_database(self, request):
|
||||
"""Restore the prompts database from an uploaded SQLite file.
|
||||
|
||||
The upload is streamed to a temp file in the library's folder. Free
|
||||
disk space decides how large it may be (``maintenance.restore_room``),
|
||||
under a 64 GiB ceiling (``restore_max_bytes``): a request whose
|
||||
Content-Length does not fit is refused before anything is written,
|
||||
and one without is refused as soon as its running size does not. The
|
||||
file is checked, staged, repaired and verified, then published over
|
||||
the live database, which is backed up first. The whole request holds
|
||||
the registry as ``restore`` (``exclusive``): it is refused while any
|
||||
job runs, and every job and request lease is refused while it runs.
|
||||
"""
|
||||
return await self._run_held(
|
||||
RESTORE_JOB_KIND,
|
||||
LEGACY_RESTORE_TITLE,
|
||||
"exclusive",
|
||||
lambda: self._restore_database(request),
|
||||
)
|
||||
|
||||
async def _restore_database(self, request):
|
||||
uploads = []
|
||||
try:
|
||||
temp_path, refused = await self._receive_restore_upload(request, uploads)
|
||||
if refused is not None:
|
||||
return refused
|
||||
reason = await self._upload_problem(temp_path)
|
||||
if reason is not None:
|
||||
return web.json_response(
|
||||
{"success": False, "error": reason}, status=400
|
||||
)
|
||||
if not await self._jobs.drain_leases():
|
||||
return web.json_response(
|
||||
{"success": False, "error": REQUESTS_BUSY}, status=409
|
||||
)
|
||||
model = self.db.model
|
||||
staged = await self._run_in_executor(model.stage_restore, temp_path)
|
||||
try:
|
||||
backup_path = await self._publish(model, staged)
|
||||
finally:
|
||||
await self._run_in_executor(staged.discard)
|
||||
# The library is replaced: from here on the restore succeeded.
|
||||
await self._after_publish(model.db_path)
|
||||
try:
|
||||
prompt_count = await self._run_in_executor(
|
||||
lambda: self.db.model.get_database_info().get("total_prompts", 0)
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("Could not count the restored prompts")
|
||||
prompt_count = staged.summary.get("prompts", 0)
|
||||
message = f"Database restored successfully. Found {prompt_count} prompts."
|
||||
body = {
|
||||
"success": True,
|
||||
"message": message,
|
||||
"prompt_count": prompt_count,
|
||||
"backup_created": (
|
||||
self._public_path(backup_path) if backup_path else None
|
||||
),
|
||||
}
|
||||
repaired = documents_line(staged.documents)
|
||||
if repaired:
|
||||
body["message"] = f"{message} {repaired}"
|
||||
body["workflow_documents"] = dict(staged.documents)
|
||||
return web.json_response(body)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Restore error: {e}", exc_info=True)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"error": f"Failed to restore database: {self._public_error(e)}",
|
||||
},
|
||||
status=500,
|
||||
)
|
||||
finally:
|
||||
await self._remove_uploads(uploads)
|
||||
|
||||
# ── reviewed restore ──────────────────────────────────────────────
|
||||
|
||||
def _live_fingerprint(self, db):
|
||||
return restore_stage.live_fingerprint(db.model.get_connection())
|
||||
|
||||
async def preview_restore(self, request):
|
||||
"""Stage an upload for review; nothing live is written (see module)."""
|
||||
await self._drop_expired_restores()
|
||||
if self._restore_plans.full():
|
||||
return _too_many_reviews()
|
||||
uploads = []
|
||||
try:
|
||||
db = self.db
|
||||
identity = self._database_identity(db)
|
||||
temp_path, refused = await self._receive_restore_upload(request, uploads)
|
||||
if refused is not None:
|
||||
return refused
|
||||
return await self._review_upload(db, identity, temp_path, uploads)
|
||||
except Exception:
|
||||
self.logger.exception("Restore preview failed")
|
||||
return _failure("Failed to review the restore", 500)
|
||||
finally:
|
||||
await self._remove_uploads(uploads)
|
||||
|
||||
async def _stage_for_review(self, db, temp_path, uploads):
|
||||
"""Stage ``temp_path`` in the executor; a cancel discards the result.
|
||||
|
||||
A cancelled handler (shutdown, a dropped client) cannot stop the
|
||||
worker thread: the stage it finishes is removed when it is done, and
|
||||
so is the upload it reads (``_orphan_stage_cleanup``).
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
future = loop.run_in_executor(None, db.model.stage_restore, temp_path)
|
||||
try:
|
||||
return await asyncio.shield(future)
|
||||
except asyncio.CancelledError:
|
||||
orphans = getattr(self, "_orphan_stages", None)
|
||||
if orphans is None:
|
||||
orphans = self._orphan_stages = set()
|
||||
entry = (future, tuple(uploads))
|
||||
orphans.add(entry)
|
||||
uploads.clear() # still being read: removed with the stage
|
||||
future.add_done_callback(
|
||||
lambda _f: loop.create_task(self._orphan_stage_cleanup())
|
||||
)
|
||||
raise
|
||||
|
||||
async def _orphan_stage_cleanup(self):
|
||||
"""Remove what cancelled previews staged, once their workers are done."""
|
||||
for entry in list(getattr(self, "_orphan_stages", ()) or ()):
|
||||
future, uploads = entry
|
||||
if not future.done():
|
||||
continue
|
||||
self._orphan_stages.discard(entry)
|
||||
if not future.cancelled() and future.exception() is None:
|
||||
await self._discard_held([_Orphan(future.result())])
|
||||
await self._remove_uploads(uploads)
|
||||
|
||||
async def _review_upload(self, db, identity, temp_path, uploads=None):
|
||||
"""Check, stage and compare an upload; the preview response.
|
||||
|
||||
Only the read of the live library takes a lease (#335 review M2):
|
||||
the upload and the staging never touch it, and a lease held for a
|
||||
long staging would make a restore's drain fail. A stage that is not
|
||||
handed to the store (a refusal, an error, a cancel) is removed.
|
||||
"""
|
||||
uploads = [] if uploads is None else uploads # the caller removes them
|
||||
reason = await self._upload_problem(temp_path)
|
||||
if reason is not None:
|
||||
return _failure(reason, 400, "invalid-backup")
|
||||
staged = None
|
||||
issued = False
|
||||
try:
|
||||
try:
|
||||
staged = await self._stage_for_review(db, temp_path, uploads)
|
||||
except ValueError as invalid:
|
||||
return _failure(str(invalid), 400, "invalid-backup")
|
||||
try:
|
||||
async with self._jobs.lease("Restore preview"):
|
||||
live, fingerprint = await self._run_in_executor(
|
||||
self._live_fingerprint, db
|
||||
)
|
||||
except JobBusy as busy:
|
||||
return busy_response(busy)
|
||||
store = self._restore_plans
|
||||
if identity != self._database_identity(self.db):
|
||||
return _failure(
|
||||
"The database was replaced while the restore was reviewed. "
|
||||
"Review again.",
|
||||
409,
|
||||
"storage-changed",
|
||||
)
|
||||
if store.full():
|
||||
return _too_many_reviews()
|
||||
token = store.issue(staged, fingerprint, identity, live)
|
||||
issued = True
|
||||
self._schedule_restore_expiry()
|
||||
finally:
|
||||
if staged is not None and not issued:
|
||||
await self._discard_held([_Orphan(staged)])
|
||||
held = store.entry(token)
|
||||
name = os.path.basename(db.model.db_path)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"token": token,
|
||||
"expires_at": _iso(held.expires_at),
|
||||
"summary": dict(staged.summary),
|
||||
"schema": dict(staged.schema),
|
||||
"integrity": "ok",
|
||||
"documents": dict(staged.documents),
|
||||
"live": dict(live),
|
||||
"size": staged.size,
|
||||
"safety_backup": {
|
||||
"destination_label": f"{name}.backup_<date and time>"
|
||||
},
|
||||
},
|
||||
headers={"Cache-Control": "no-store"},
|
||||
)
|
||||
|
||||
async def discard_restore_preview(self, request):
|
||||
token = request.match_info.get("token", "")
|
||||
if not 0 < len(token) <= MAX_TOKEN_LENGTH:
|
||||
return _failure("Unknown restore review.", 410, "plan-expired")
|
||||
held = self._restore_plans.pop(token)
|
||||
if held is None:
|
||||
await self._drop_expired_restores()
|
||||
return _failure("Unknown restore review.", 410, "plan-expired")
|
||||
await self._discard_held([held])
|
||||
return web.json_response({"success": True})
|
||||
|
||||
def _peek_restore(self, token):
|
||||
try:
|
||||
return self._restore_plans.peek(token), None
|
||||
except PlanUsed as used:
|
||||
job = self._jobs.get(used.job_id)
|
||||
return None, _failure(
|
||||
"This restore was already started.",
|
||||
409,
|
||||
"plan-used",
|
||||
job={"id": used.job_id, "title": job.title if job else None},
|
||||
)
|
||||
except PlanExpired:
|
||||
return None, _failure(
|
||||
"This review expired. Review the restore again.",
|
||||
410,
|
||||
"plan-expired",
|
||||
)
|
||||
|
||||
async def apply_restore(self, request):
|
||||
data, error = await self._read_json_body(request, max_bytes=APPLY_BODY_MAX)
|
||||
if error is not None:
|
||||
if error.status == 413:
|
||||
return _failure("Request body too large", 413, "too-large")
|
||||
return _failure(
|
||||
"Request body must be a JSON object", 400, "invalid-request"
|
||||
)
|
||||
token = data.get("token")
|
||||
if set(data) != {"token"} or not isinstance(token, str):
|
||||
return _failure("token must be the review's token", 400, "invalid-request")
|
||||
if not 0 < len(token) <= MAX_TOKEN_LENGTH:
|
||||
return _failure("token must be the review's token", 400, "invalid-request")
|
||||
await self._drop_expired_restores()
|
||||
held, refused = self._peek_restore(token)
|
||||
if refused is not None:
|
||||
await self._drop_expired_restores() # one peek found expired
|
||||
return refused
|
||||
db = self.db
|
||||
if held.db_identity != self._database_identity(db):
|
||||
await self._discard_popped(token)
|
||||
return _failure(
|
||||
"The database was replaced since the review. Review again.",
|
||||
409,
|
||||
"storage-changed",
|
||||
)
|
||||
try:
|
||||
async with self._jobs.lease("Restore review check"):
|
||||
_live, fingerprint = await self._run_in_executor(
|
||||
self._live_fingerprint, db
|
||||
)
|
||||
except JobBusy as busy:
|
||||
return busy_response(busy)
|
||||
if fingerprint != held.fingerprint:
|
||||
# It can never match again: free the stage for a new review.
|
||||
await self._discard_popped(token)
|
||||
return _failure(PLAN_CHANGED, 409, "plan-changed")
|
||||
# The token may have been used, expired or discarded meanwhile.
|
||||
held, refused = self._peek_restore(token)
|
||||
if refused is not None:
|
||||
await self._drop_expired_restores()
|
||||
return refused
|
||||
if db is not self.db:
|
||||
return _failure(
|
||||
"The database was replaced since the review. Review again.",
|
||||
409,
|
||||
"storage-changed",
|
||||
)
|
||||
try:
|
||||
job = self._jobs.start(
|
||||
RESTORE_JOB_KIND,
|
||||
RESTORE_JOB_TITLE,
|
||||
# It replaces the whole database: it runs alone, beside no
|
||||
# job of any class, like Move library.
|
||||
admission="exclusive",
|
||||
cancellable=False,
|
||||
scope={},
|
||||
groups=RESTORE_GROUPS,
|
||||
runner=lambda job: self._run_restore(job, db, held, token),
|
||||
)
|
||||
except JobBusy as busy:
|
||||
return busy_response(busy) # the token stays usable
|
||||
# Same loop turn as the peek and the start: the runner has not run.
|
||||
self._restore_plans.take(token, job.id)
|
||||
return web.json_response({"success": True, "job": job.summary()}, status=202)
|
||||
|
||||
async def _run_restore(self, job, db, held, token=None):
|
||||
from .maintenance import _fail_job
|
||||
|
||||
staged = held.staged
|
||||
keep_stage = False
|
||||
try:
|
||||
job.update(status_line="Waiting for other requests to finish")
|
||||
if not await self._jobs.drain_leases():
|
||||
# Nothing ran: the reviewed stage stays usable (#335 M2),
|
||||
# unless it expired meanwhile or its slot was reused.
|
||||
if token is not None and self._restore_plans.give_back(token, held):
|
||||
keep_stage = True
|
||||
self._schedule_restore_expiry()
|
||||
job.scope["review_kept"] = True
|
||||
return _fail_job(job, "busy", REVIEW_KEPT)
|
||||
return _fail_job(job, "busy", REQUESTS_BUSY)
|
||||
job.update(status_line="Checking the library")
|
||||
_live, fingerprint = await self._run_in_executor(self._live_fingerprint, db)
|
||||
if db is not self.db or fingerprint != held.fingerprint:
|
||||
return _fail_job(job, "plan-changed", PLAN_CHANGED)
|
||||
job.update(status_line="Restoring the database")
|
||||
job.log_line("Restoring the reviewed backup")
|
||||
model = db.model
|
||||
try:
|
||||
backup_path = await self._publish(model, staged)
|
||||
except restore_stage.StageChanged:
|
||||
return _fail_job(job, "plan-changed", PLAN_CHANGED)
|
||||
except SafetyBackupFailed:
|
||||
return _fail_job(
|
||||
job,
|
||||
"backup-failed",
|
||||
"The current library could not be backed up; nothing was "
|
||||
"changed.",
|
||||
)
|
||||
except sqlite3.OperationalError as error:
|
||||
kept = (
|
||||
" The safety backup was kept."
|
||||
if getattr(error, "_prompt_manager_restore_backup_created", False)
|
||||
else ""
|
||||
)
|
||||
self.logger.error(f"Restore failed: {error}")
|
||||
return _fail_job(job, "restore-failed", NOT_RESTORED + kept)
|
||||
except Exception:
|
||||
self.logger.exception("Restore job failed")
|
||||
return _fail_job(job, None, NOT_RESTORED)
|
||||
finally:
|
||||
if not keep_stage:
|
||||
await self._discard_held([_Orphan(staged)])
|
||||
# The copy replaced the library: the restore is done (#335 review I1).
|
||||
# Every later step only logs a failure.
|
||||
await self._after_publish(model.db_path)
|
||||
if backup_path:
|
||||
job.log_line(f"Safety backup {os.path.basename(backup_path)}")
|
||||
else:
|
||||
job.log_line("No safety backup: there was no library to keep")
|
||||
repaired = documents_line(staged.documents)
|
||||
if repaired:
|
||||
job.log_line(repaired)
|
||||
line = restored_line(staged.summary)
|
||||
job.log_line(line)
|
||||
job.update(progress=100, status_line=line)
|
||||
|
||||
|
||||
class _Orphan:
|
||||
"""A stage with no store entry, in the shape ``_discard_held`` takes."""
|
||||
|
||||
def __init__(self, staged):
|
||||
self.staged = staged
|
||||
|
||||
|
||||
def _too_many_reviews():
|
||||
return _failure(
|
||||
"Two restores are already being reviewed. Discard one first.",
|
||||
429,
|
||||
"too-large",
|
||||
)
|
||||
@@ -1,337 +0,0 @@
|
||||
"""One shared listing per output root, and the pure helpers behind it (#348).
|
||||
|
||||
The API keeps one ``RootListing`` per output root (``_root_listings`` in
|
||||
``py/api/__init__.py``). The record is read and written only on the event
|
||||
loop; the walk itself, the comparison with the previous listing and both
|
||||
folder sets run in the executor through ``walk_root``. Routes live in
|
||||
``py/api/images.py`` (``_listing``); this module has no aiohttp import so
|
||||
tests can exercise it without a server.
|
||||
|
||||
A settled output the image monitor reports is added to its root's stored
|
||||
listing without a walk (``output_settled``, ``insert_entry``), in the place
|
||||
a walk would list it. When that cannot be exact, the root is marked dirty
|
||||
(``mark_root_dirty``) instead; a dirty root gets one background walk, no
|
||||
sooner than ``rewalk_interval`` after its previous walk started (PR 2).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from bisect import bisect_left
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import PurePath
|
||||
from typing import Callable, NamedTuple
|
||||
|
||||
# A listing older than this (seconds) is served as it is, and one background
|
||||
# walk replaces it. Read as ``gallery_listing.REFRESH_AFTER`` so tests can
|
||||
# patch it.
|
||||
REFRESH_AFTER = 30
|
||||
|
||||
# The stale read is served at once; its refresh walk starts this many seconds
|
||||
# later, so the walk does not share the GIL with the page's own requests
|
||||
# (#348 live test: reads during the walk took 2.2 s instead of 0.07 s).
|
||||
REFRESH_DELAY = 1.0
|
||||
|
||||
# No background walk of a root starts sooner than rewalk_interval(last walk's
|
||||
# duration) after its previous walk started: FACTOR times that duration, at
|
||||
# least FLOOR and at most CAP seconds. Walks up to CAP / FACTOR (3 s) then
|
||||
# hold the GIL at most ~20% of the time; the cap keeps a new output listed
|
||||
# within ~20 s, at the cost of a higher share (back to back from 15 s) on
|
||||
# slower walks. Read through this module so tests can patch them.
|
||||
REWALK_FLOOR = 5
|
||||
REWALK_FACTOR = 5
|
||||
REWALK_CAP = 15
|
||||
|
||||
# A deferred walk is never armed sooner than this: no zero-delay loop.
|
||||
MIN_REWALK_DELAY = 0.5
|
||||
|
||||
# A reader with no listing gives up (500) after this many walks in a row were
|
||||
# discarded by invalidations made while they ran.
|
||||
MAX_DISCARDED_WALKS = 3
|
||||
|
||||
WALK_STOPPED = "The listing walk was stopped"
|
||||
|
||||
STORED = "stored"
|
||||
DISCARDED = "discarded"
|
||||
FAILED = "failed"
|
||||
|
||||
|
||||
class WalkOutcome(NamedTuple):
|
||||
"""What one walk ended with; a walk's flight future resolves to this."""
|
||||
|
||||
kind: str
|
||||
error: BaseException | None = None
|
||||
|
||||
@classmethod
|
||||
def stored(cls) -> "WalkOutcome":
|
||||
return cls(STORED)
|
||||
|
||||
@classmethod
|
||||
def discarded(cls) -> "WalkOutcome":
|
||||
return cls(DISCARDED)
|
||||
|
||||
@classmethod
|
||||
def failed(cls, error: BaseException) -> "WalkOutcome":
|
||||
return cls(FAILED, error)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RootListing:
|
||||
"""The shared listing of one output root.
|
||||
|
||||
A record is never deleted: ``drop()`` clears the listing and bumps
|
||||
``generation``, so a walk that started before a drop can never be stored
|
||||
after it.
|
||||
"""
|
||||
|
||||
# (path, mtime) newest first, as _scan_gallery_files_sync returns it.
|
||||
files: list | None = None
|
||||
folders: frozenset | None = None
|
||||
folders_with_ancestors: frozenset | None = None
|
||||
# Monotonic time the last walk finished, stored or failed.
|
||||
walked_at: float | None = None
|
||||
# Number of the last walk started; a running walk is always that one.
|
||||
walk_seq: int = 0
|
||||
# Number of the walk whose files are stored (Rescan waits on it).
|
||||
stored_seq: int = 0
|
||||
generation: int = 0
|
||||
# The running walk's future (resolves to a WalkOutcome), or None.
|
||||
flight: asyncio.Future | None = None
|
||||
# True while walks keep failing: one warning per streak.
|
||||
failing: bool = False
|
||||
# The one deferred background walk not started yet (loop.call_later), or
|
||||
# None: a stale refresh or a dirty root's rewalk, whichever is due first.
|
||||
rewalk_handle: asyncio.TimerHandle | None = None
|
||||
# normcase(realpath(root)), computed by the walk in the executor: monitor
|
||||
# paths are matched against it without file I/O on the loop.
|
||||
resolved_root: str | None = None
|
||||
# Monotonic start and duration of the last walk (rewalk_interval).
|
||||
last_started_at: float | None = None
|
||||
last_duration: float | None = None
|
||||
# A settled output was reported under this root; cleared when a walk starts.
|
||||
dirty: bool = False
|
||||
# Dropped because the output roots changed: no background walk until a
|
||||
# reader (or the warm-up) asks for this root again. The monitor keeps
|
||||
# watching a removed root until ComfyUI restarts.
|
||||
retired: bool = False
|
||||
# Relative paths of settled outputs waiting to be added to ``files``
|
||||
# without a walk, and the one task adding them (``output_settled``).
|
||||
pending_adds: list = field(default_factory=list)
|
||||
add_task: asyncio.Task | None = None
|
||||
# Settled outputs added to the stored listing since the last walk stored.
|
||||
patched: int = 0
|
||||
|
||||
def drop(self) -> None:
|
||||
self.files = None
|
||||
self.folders = None
|
||||
self.folders_with_ancestors = None
|
||||
self.walked_at = None
|
||||
self.pending_adds = []
|
||||
self.patched = 0
|
||||
self.generation += 1
|
||||
|
||||
|
||||
def folder_sets(root, files) -> tuple[frozenset, frozenset]:
|
||||
"""``(folders, folders_with_ancestors)`` of the files listed under *root*.
|
||||
|
||||
Folders keep the spelling ``/gallery/subfolders`` has always answered:
|
||||
``str(relative.parent)`` (native separators) for the folders holding
|
||||
files, plus every ancestor joined by ``/``.
|
||||
"""
|
||||
# From the path strings: relative_to per file was half of a large walk.
|
||||
sep, prefix = _root_prefix(root)
|
||||
folders = set()
|
||||
for path, _mtime in files:
|
||||
rel_dir = _folder_of(root, path, sep, prefix)
|
||||
if rel_dir:
|
||||
folders.add(rel_dir)
|
||||
with_ancestors = set(folders)
|
||||
for folder in folders:
|
||||
with_ancestors.update(_ancestors(folder))
|
||||
return frozenset(folders), frozenset(with_ancestors)
|
||||
|
||||
|
||||
def add_folder(root, path, folders, with_ancestors) -> tuple[frozenset, frozenset]:
|
||||
"""Both folder sets once *path* is listed too (the same sets when it adds none)."""
|
||||
sep, prefix = _root_prefix(root)
|
||||
rel_dir = _folder_of(root, path, sep, prefix)
|
||||
if not rel_dir or rel_dir in folders:
|
||||
return folders, with_ancestors
|
||||
return folders | {rel_dir}, with_ancestors | {rel_dir, *_ancestors(rel_dir)}
|
||||
|
||||
|
||||
def _root_prefix(root) -> tuple[str, str]:
|
||||
# The separator is the path flavour's own (tests use PureWindowsPath).
|
||||
sep = str(type(root)("a", "b"))[1]
|
||||
prefix = str(root)
|
||||
if not prefix.endswith(sep):
|
||||
prefix += sep
|
||||
return sep, prefix
|
||||
|
||||
|
||||
def _folder_of(root, path, sep, prefix) -> str:
|
||||
"""The folder ``/gallery/subfolders`` names for *path*; "" for the root."""
|
||||
text = str(path)
|
||||
if text.startswith(prefix):
|
||||
rel_dir = text[len(prefix) :].rpartition(sep)[0]
|
||||
else:
|
||||
rel_dir = str(path.relative_to(root).parent)
|
||||
return "" if rel_dir == "." else rel_dir
|
||||
|
||||
|
||||
def _ancestors(folder: str) -> list[str]:
|
||||
parts = folder.replace("\\", "/").split("/")
|
||||
return ["/".join(parts[:i]) for i in range(1, len(parts))]
|
||||
|
||||
|
||||
def relative_part(resolved_root: str, path, pathmod=os.path) -> str | None:
|
||||
"""The part of the resolved *path* below *resolved_root*, as a walk spells it.
|
||||
|
||||
Kept in *path*'s own case (``realpath`` gives a file's on-disk case on
|
||||
Windows) with the native separator. None when *path* is not below the
|
||||
root, or when the case folding changed the length so the slice would
|
||||
not be exact.
|
||||
"""
|
||||
sep = pathmod.sep
|
||||
text = str(path)
|
||||
if sep == "\\":
|
||||
text = _plain(text.replace("/", sep), pathmod)
|
||||
folded = pathmod.normcase(text)
|
||||
prefix = _plain(resolved_root, pathmod)
|
||||
if not prefix.endswith(sep):
|
||||
prefix += sep
|
||||
if len(folded) != len(text) or not folded.startswith(prefix):
|
||||
return None
|
||||
relative = text[len(prefix) :]
|
||||
if any(part in ("", ".", "..") for part in relative.split(sep)):
|
||||
return None
|
||||
return relative
|
||||
|
||||
|
||||
def insert_entry(files: list, entry: tuple, max_files: int) -> list | None:
|
||||
"""A new listing with *entry* ``(path, mtime)`` where a walk would list it.
|
||||
|
||||
*files* is newest first; an entry for the same path is replaced. Returns
|
||||
*files* itself when it already holds *entry*, and None when the place is
|
||||
not exact: another file has the same mtime (a walk orders those by its
|
||||
own visit order), or the listing is at the walk's *max_files* cap.
|
||||
Blocking (one pass over *files*): run it in the executor.
|
||||
"""
|
||||
path, mtime = entry
|
||||
text = str(path)
|
||||
kept = files
|
||||
for index, (listed, listed_mtime) in enumerate(files):
|
||||
if str(listed) == text:
|
||||
if listed_mtime == mtime:
|
||||
return files
|
||||
kept = files[:index] + files[index + 1 :]
|
||||
break
|
||||
else:
|
||||
if len(files) >= max_files:
|
||||
return None
|
||||
index = bisect_left(kept, -mtime, key=lambda item: -item[1])
|
||||
if index < len(kept) and kept[index][1] == mtime:
|
||||
return None
|
||||
return [*kept[:index], entry, *kept[index:]]
|
||||
|
||||
|
||||
# A listing that differs from the cached one by more entries than this is
|
||||
# rebuilt in full by the workspace caches (listing_delta).
|
||||
MAX_DELTA = 256
|
||||
|
||||
|
||||
def listing_delta(old: list, new: list, limit: int = MAX_DELTA):
|
||||
"""``(added, removed)`` entries from *old* to *new*, or None past *limit*.
|
||||
|
||||
A listing patched by ``insert_entry`` shares every unchanged entry object
|
||||
with the one it came from, so entries are told apart by identity: a walk
|
||||
builds new tuples, shares none and is rebuilt in full. *old* must still
|
||||
be alive (the
|
||||
caller's cache holds it), so no identity was reused.
|
||||
"""
|
||||
if abs(len(new) - len(old)) > limit:
|
||||
return None
|
||||
old_ids = set(map(id, old))
|
||||
added = []
|
||||
for entry in new:
|
||||
if id(entry) not in old_ids:
|
||||
added.append(entry)
|
||||
if len(added) > limit:
|
||||
return None
|
||||
if new and len(added) == len(new):
|
||||
return None # nothing shared: a walk
|
||||
if len(old) + len(added) == len(new):
|
||||
return added, []
|
||||
new_ids = set(map(id, new))
|
||||
return added, [entry for entry in old if id(entry) not in new_ids]
|
||||
|
||||
|
||||
def rewalk_interval(last_duration) -> float:
|
||||
"""Seconds from a walk's start to the next background walk of its root."""
|
||||
scaled = REWALK_FACTOR * (last_duration or 0.0)
|
||||
return min(REWALK_CAP, max(REWALK_FLOOR, scaled))
|
||||
|
||||
|
||||
def root_contains(resolved_root: str, path, pathmod=os.path) -> bool:
|
||||
"""True when the resolved *path* lies under *resolved_root* (no file I/O).
|
||||
|
||||
*resolved_root* is already ``normcase``d; *path* is folded the same way,
|
||||
so Windows matches case-insensitively and with either separator. The
|
||||
match ends at a separator: ``/out2/x.png`` is not under ``/out``.
|
||||
"""
|
||||
target = _plain(pathmod.normcase(str(path)), pathmod)
|
||||
prefix = _plain(resolved_root, pathmod)
|
||||
if not prefix.endswith(pathmod.sep):
|
||||
prefix += pathmod.sep
|
||||
return target.startswith(prefix)
|
||||
|
||||
|
||||
def _plain(path: str, pathmod) -> str:
|
||||
"""*path* (normcased or not) without a Windows extended-length prefix.
|
||||
|
||||
``realpath`` may keep ``\\\\?\\`` on a file past MAX_PATH while its shorter
|
||||
root resolves without it.
|
||||
"""
|
||||
if pathmod.sep != "\\":
|
||||
return path
|
||||
if path[:8].lower() == "\\\\?\\unc\\":
|
||||
return "\\\\" + path[8:]
|
||||
if path.startswith("\\\\?\\"):
|
||||
return path[4:]
|
||||
return path
|
||||
|
||||
|
||||
def resolve_root(root) -> str:
|
||||
"""``normcase(realpath(root))``, the spelling monitor paths are matched in.
|
||||
|
||||
File I/O (realpath follows links): call it in the executor.
|
||||
"""
|
||||
return os.path.normcase(os.path.realpath(str(root)))
|
||||
|
||||
|
||||
def same_listing(old, new) -> bool:
|
||||
"""True when two listings hold the same ``(path, mtime)`` pairs in order."""
|
||||
if old is None or new is None:
|
||||
return False
|
||||
return old is new or old == new
|
||||
|
||||
|
||||
def walk_root(
|
||||
scan_fn: Callable[[PurePath], list], root, previous_files
|
||||
) -> tuple[list, frozenset, frozenset, str]:
|
||||
"""The single executor job of a walk.
|
||||
|
||||
Returns ``(files, folders, with_ancestors, resolved_root)``.
|
||||
|
||||
Opens the root first, so a missing or unreadable root raises ``OSError``
|
||||
instead of listing nothing. Returns *previous_files* itself when nothing
|
||||
changed, which keeps the identity-keyed caches built on it warm.
|
||||
"""
|
||||
with os.scandir(root):
|
||||
pass
|
||||
files = scan_fn(root)
|
||||
if same_listing(previous_files, files):
|
||||
files = previous_files
|
||||
folders, with_ancestors = folder_sets(root, files)
|
||||
return files, folders, with_ancestors, resolve_root(root)
|
||||
+20
-56
@@ -20,7 +20,6 @@ from .history_current import (
|
||||
_OpenedCurrentMedia,
|
||||
_current_media_response,
|
||||
_prepare_current_media,
|
||||
accepts_jxl,
|
||||
_reauthorize_current_target,
|
||||
candidate_locator,
|
||||
resolve_current_candidate,
|
||||
@@ -143,18 +142,6 @@ def _query(request: web.Request) -> None:
|
||||
raise ValueError("Duplicate history query parameter")
|
||||
|
||||
|
||||
def _jxl_flag(request: web.Request) -> bool:
|
||||
"""The cover's ``jxl`` flag (#317): ``1`` when the workspace decoded a
|
||||
JPEG XL image (a browser may not list ``image/jxl`` in Accept), ``0`` or
|
||||
absent otherwise; anything else is a bad request."""
|
||||
value = request.query.get("jxl")
|
||||
if value is None:
|
||||
return False
|
||||
if value not in ("0", "1"):
|
||||
raise ValueError("Invalid jxl flag")
|
||||
return value == "1"
|
||||
|
||||
|
||||
def _page_limit(request: web.Request) -> int:
|
||||
limit = _decimal(request.query.get("limit", "40"), minimum=1)
|
||||
if limit > 100:
|
||||
@@ -747,10 +734,7 @@ class HistoryRoutesMixin:
|
||||
raise HistoryReadError(422, "response-limit") from None
|
||||
else:
|
||||
prepared = _prepare_current_media(
|
||||
resolved,
|
||||
scope,
|
||||
thumbnail=kind == "thumbnail",
|
||||
**media_options,
|
||||
resolved, scope, thumbnail=kind == "thumbnail"
|
||||
)
|
||||
_check_current(scope, self, fence, epoch, generation)
|
||||
return prepared
|
||||
@@ -777,45 +761,25 @@ class HistoryRoutesMixin:
|
||||
_check_current(scope, self, fence, epoch, generation)
|
||||
raise
|
||||
|
||||
async def execute():
|
||||
pending = asyncio.create_task(self._run_in_executor(project))
|
||||
try:
|
||||
return await asyncio.shield(pending)
|
||||
except asyncio.CancelledError:
|
||||
pending.add_done_callback(_close_prepared_result)
|
||||
raise
|
||||
except (
|
||||
HistoryReadError,
|
||||
HistoryEpochChanged,
|
||||
sqlite3.Error,
|
||||
OSError,
|
||||
ValueError,
|
||||
):
|
||||
raise
|
||||
except Exception as error:
|
||||
self.logger.warning(
|
||||
"Current history media unavailable (%s)", type(error).__name__
|
||||
)
|
||||
raise _current_unavailable("temporary_unavailable") from None
|
||||
|
||||
# A missing preview (#249) is noted on the first read, made through
|
||||
# the serve route's admission and single flight, then read again. A
|
||||
# video whose preview cannot be made is the placeholder frame.
|
||||
needed = []
|
||||
media_options = {}
|
||||
if kind == "thumbnail":
|
||||
media_options = {
|
||||
"ensure_preview": lambda *preview: needed.append(preview),
|
||||
"placeholder": self._video_placeholder_bytes(),
|
||||
"accepts_jxl": _jxl_flag(request) or accepts_jxl(request),
|
||||
}
|
||||
prepared = await execute()
|
||||
if needed:
|
||||
if not isinstance(prepared, bytes):
|
||||
prepared.close()
|
||||
await self._admit_preview(*needed[0])
|
||||
media_options["ensure_preview"] = None
|
||||
prepared = await execute()
|
||||
pending = asyncio.create_task(self._run_in_executor(project))
|
||||
try:
|
||||
prepared = await asyncio.shield(pending)
|
||||
except asyncio.CancelledError:
|
||||
pending.add_done_callback(_close_prepared_result)
|
||||
raise
|
||||
except (
|
||||
HistoryReadError,
|
||||
HistoryEpochChanged,
|
||||
sqlite3.Error,
|
||||
OSError,
|
||||
ValueError,
|
||||
):
|
||||
raise
|
||||
except Exception as error:
|
||||
self.logger.warning(
|
||||
"Current history media unavailable (%s)", type(error).__name__
|
||||
)
|
||||
raise _current_unavailable("temporary_unavailable") from None
|
||||
if isinstance(prepared, bytes):
|
||||
return _success_response(prepared)
|
||||
return await _current_media_response(request, prepared)
|
||||
|
||||
+23
-104
@@ -5,7 +5,6 @@ from dataclasses import dataclass
|
||||
import errno
|
||||
import hashlib
|
||||
import hmac
|
||||
import io
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
@@ -15,9 +14,7 @@ from typing import BinaryIO
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from . import thumbnail_settings
|
||||
from .images import _canonical, _public_path, SERVABLE_EXTENSIONS, VIDEO_EXTENSIONS
|
||||
from .images import existing_media_thumbnail, media_thumbnail_rel
|
||||
|
||||
try:
|
||||
from ...utils.execution_outputs import (
|
||||
@@ -241,19 +238,14 @@ def resolve_current_candidate(
|
||||
|
||||
|
||||
def _open_current_target(
|
||||
target: Path, namespace: Path, scope: CurrentOutputScope, preview: bool = False
|
||||
target: Path, namespace: Path, scope: CurrentOutputScope
|
||||
) -> BinaryIO:
|
||||
"""Open only the canonical target and verify authority against the actual handle.
|
||||
|
||||
A *preview* (a file under ``thumbnails/``) may also be JPEG XL (#317)."""
|
||||
"""Open only the canonical target and verify authority against the actual handle."""
|
||||
if _canonical(target) != target:
|
||||
raise CurrentLocationUnavailable("not_authorized", stale=True)
|
||||
if not _authorized(target, namespace, scope):
|
||||
raise CurrentLocationUnavailable("not_authorized")
|
||||
suffix = target.suffix.lower()
|
||||
if suffix not in SERVABLE_EXTENSIONS and not (
|
||||
preview and suffix in thumbnail_settings.THUMBNAIL_EXTENSIONS
|
||||
):
|
||||
if target.suffix.lower() not in SERVABLE_EXTENSIONS:
|
||||
raise CurrentLocationUnavailable("unsupported_media")
|
||||
flags = (
|
||||
os.O_RDONLY
|
||||
@@ -317,107 +309,34 @@ class _OpenedCurrentMedia:
|
||||
self.handle.close()
|
||||
|
||||
|
||||
def _preview_writable(root: Path, preview: Path) -> bool:
|
||||
"""A preview may be written: ``thumbnails`` and every link segment of the
|
||||
preview's folder resolve inside *root* (the serve route's rule, #249)."""
|
||||
folder = root / "thumbnails"
|
||||
return _canonical(folder).is_relative_to(_canonical(root)) and _canonical(
|
||||
preview
|
||||
).is_relative_to(_canonical(folder))
|
||||
|
||||
|
||||
def _is_jxl(path) -> bool:
|
||||
return Path(path).suffix.lower() == ".jxl"
|
||||
|
||||
|
||||
def _preview_without_jxl(root: Path, relative: Path) -> Path | None:
|
||||
"""The first existing preview *relative* owns that is not JPEG XL."""
|
||||
return next(
|
||||
(
|
||||
path
|
||||
for path in thumbnail_settings.owned_thumbnails(root, relative)
|
||||
if not _is_jxl(path)
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def accepts_jxl(request) -> bool:
|
||||
"""Whether the request's ``Accept`` lists ``image/jxl``."""
|
||||
accept = request.headers.get("Accept", "")
|
||||
return any(
|
||||
part.split(";", 1)[0].strip().lower() == "image/jxl"
|
||||
for part in accept.split(",")
|
||||
)
|
||||
|
||||
|
||||
def _placeholder_media(placeholder: bytes) -> "_OpenedCurrentMedia":
|
||||
return _OpenedCurrentMedia(io.BytesIO(placeholder), len(placeholder), "image/jpeg")
|
||||
|
||||
|
||||
def _cover_target(resolved, ensure_preview, placeholder, accepts_jxl):
|
||||
"""What the cover opens: a preview path, the placeholder media, or None
|
||||
for the original (see ``_prepare_current_media``)."""
|
||||
target = resolved.canonical_target
|
||||
relative = target.relative_to(resolved.gallery_root)
|
||||
is_video = target.suffix.lower() in VIDEO_EXTENSIONS
|
||||
# #317: any preview the file owns (the current format first); a missing
|
||||
# one is made under the current format's name.
|
||||
preview = existing_media_thumbnail(resolved.gallery_root, relative)
|
||||
if preview is not None and not accepts_jxl and _is_jxl(preview):
|
||||
preview = _preview_without_jxl(resolved.gallery_root, relative)
|
||||
if preview is None:
|
||||
if is_video and placeholder is not None:
|
||||
return _placeholder_media(placeholder)
|
||||
return None
|
||||
return _canonical(preview)
|
||||
if preview is not None:
|
||||
return _canonical(preview)
|
||||
preview = resolved.gallery_root / media_thumbnail_rel(relative)
|
||||
if ensure_preview is not None and _preview_writable(resolved.gallery_root, preview):
|
||||
ensure_preview(target, preview, is_video)
|
||||
if is_video and placeholder is not None:
|
||||
return _placeholder_media(placeholder)
|
||||
return None
|
||||
|
||||
|
||||
def _prepare_current_media(
|
||||
resolved: ResolvedCurrentCandidate,
|
||||
scope: CurrentOutputScope,
|
||||
*,
|
||||
thumbnail: bool,
|
||||
ensure_preview=None,
|
||||
placeholder: bytes | None = None,
|
||||
accepts_jxl: bool = True,
|
||||
resolved: ResolvedCurrentCandidate, scope: CurrentOutputScope, *, thumbnail: bool
|
||||
) -> _OpenedCurrentMedia:
|
||||
"""Open the current file, or its preview for ``thumbnail``.
|
||||
|
||||
A missing preview that may be written is handed to
|
||||
``ensure_preview(src, dst, is_video)`` (the route makes it and reads
|
||||
again, #249). Without one an image's original is opened and a video is
|
||||
the ``placeholder`` JPEG, never the video inside an ``<img>``.
|
||||
|
||||
The cover's URL names no format, so a JPEG XL preview goes only to a
|
||||
browser whose ``Accept`` lists ``image/jxl`` (*accepts_jxl*, review M3).
|
||||
Another one gets the file's preview in another format when there is one,
|
||||
else the original image or a video's placeholder: never a broken image.
|
||||
"""
|
||||
target = resolved.canonical_target
|
||||
is_preview = False
|
||||
if thumbnail:
|
||||
chosen = _cover_target(resolved, ensure_preview, placeholder, accepts_jxl)
|
||||
if isinstance(chosen, _OpenedCurrentMedia):
|
||||
return chosen
|
||||
if chosen is not None:
|
||||
target, is_preview = chosen, True
|
||||
handle = _open_current_target(target, resolved.namespace, scope, is_preview)
|
||||
relative = target.relative_to(resolved.gallery_root)
|
||||
extension = (
|
||||
".jpg"
|
||||
if target.suffix.lower() in VIDEO_EXTENSIONS
|
||||
else target.suffix.lower()
|
||||
)
|
||||
preview = (
|
||||
resolved.gallery_root
|
||||
/ "thumbnails"
|
||||
/ f"{relative.with_suffix('').as_posix()}_thumb{extension}"
|
||||
)
|
||||
try:
|
||||
preview.stat()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
else:
|
||||
target = _canonical(preview)
|
||||
handle = _open_current_target(target, resolved.namespace, scope)
|
||||
try:
|
||||
return _OpenedCurrentMedia(
|
||||
handle,
|
||||
os.fstat(handle.fileno()).st_size,
|
||||
thumbnail_settings.content_type(target)
|
||||
or mimetypes.guess_type(target.name)[0]
|
||||
or "application/octet-stream",
|
||||
mimetypes.guess_type(target.name)[0] or "application/octet-stream",
|
||||
)
|
||||
except BaseException:
|
||||
handle.close()
|
||||
|
||||
+359
-1248
File diff suppressed because it is too large
Load Diff
@@ -130,14 +130,6 @@ class JobRoutesMixin:
|
||||
async def start_image_tag_job_route(request):
|
||||
return await self.start_image_tag_job(request)
|
||||
|
||||
@routes.post("/prompt_manager/jobs/model-download")
|
||||
async def start_model_download_job_route(request):
|
||||
return await self.start_model_download_job(request)
|
||||
|
||||
@routes.post("/prompt_manager/jobs/autotag")
|
||||
async def start_autotag_job_route(request):
|
||||
return await self.start_autotag_job(request)
|
||||
|
||||
@routes.post("/prompt_manager/jobs/thumbnails")
|
||||
async def start_thumbnail_job_route(request):
|
||||
return await self.start_thumbnail_job(request)
|
||||
|
||||
-110
@@ -28,15 +28,6 @@ ones. One job of a kind runs at a time, and each job has an admission class
|
||||
A short legacy-page operation takes a ``hold``: a running, non-cancellable
|
||||
job that lasts as long as the request, so the Activity page lists it and a
|
||||
workspace job is refused while it runs (and the other way round).
|
||||
|
||||
Short request-scoped work that reads or writes the library beside any
|
||||
ordinary job (a Needs review decision, a maintenance or restore preview)
|
||||
takes a ``lease`` instead: it is not listed (so quick requests never push
|
||||
real jobs out of the finished history, and two of them never refuse each
|
||||
other), but it is refused while an ``exclusive`` job runs, and an
|
||||
``exclusive`` job that replaces the database (a restore) awaits
|
||||
``drain_leases`` before it switches the file, so no lease's connection is
|
||||
open across the switch (#249 S8).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -57,9 +48,6 @@ MAX_OUTCOME_ROWS = 200 # rows listed per outcome group; counters stay exact
|
||||
MAX_LOG_LINES = 2000
|
||||
LOG_CAP_LINE = "[Log capped at 2,000 lines]"
|
||||
MAX_FINISHED_JOBS = 20 # finished jobs this session keeps
|
||||
# The longest an exclusive job waits for running leases to end (a ceiling a
|
||||
# hung request reaches, not a pacing timer: leases end with their request).
|
||||
LEASE_DRAIN_SECONDS = 60
|
||||
|
||||
# The metadata scan's outcome groups, in counter-tile order of the spec's
|
||||
# mapping, and the per-file scan status each group collects. Files without
|
||||
@@ -87,8 +75,6 @@ SCAN_GROUPS = {
|
||||
# Record-changing: it writes image tag sets and new vocabulary names.
|
||||
IMAGE_TAG_JOB_KIND = "image-tags"
|
||||
IMAGE_TAG_GROUP_NAMES = ("tagged", "unchanged", "skipped", "errors")
|
||||
# Tag output images (#328) counts its prompts the same way.
|
||||
AUTOTAG_GROUP_NAMES = IMAGE_TAG_GROUP_NAMES
|
||||
|
||||
# Thumbnail generation (#249 System S5): previews written, already fresh, failed.
|
||||
THUMBNAIL_GROUP_NAMES = ("generated", "skipped", "errors")
|
||||
@@ -111,8 +97,6 @@ DUPLICATE_CLEANUP_JOB_KIND = "duplicate_cleanup"
|
||||
THUMBNAILS_JOB_KIND = "thumbnails"
|
||||
LORA_IMPORT_JOB_KIND = "lora_import"
|
||||
AUTOTAG_BATCH_JOB_KIND = "autotag"
|
||||
# A tagging model download (#328): read admission, not cancellable.
|
||||
MODEL_DOWNLOAD_JOB_KIND = "model_download"
|
||||
# Duplicate image files (#249 System S6, py/api/duplicate_files.py): the
|
||||
# find is a cancellable ``read`` job; the reviewed removal is ``records``.
|
||||
DUPLICATES_JOB_KIND = "duplicates"
|
||||
@@ -419,74 +403,6 @@ class JobRegistry:
|
||||
def __init__(self, logger_=None):
|
||||
self._logger = logger_ or logger
|
||||
self._jobs = {}
|
||||
self._leases = {} # title -> running count
|
||||
self._drain_waiters = []
|
||||
|
||||
# ── leases ────────────────────────────────────────────────────────
|
||||
|
||||
def lease(self, title):
|
||||
"""An async context manager for short request-scoped library work.
|
||||
|
||||
Entering raises JobBusy, naming the job, while an ``exclusive`` job
|
||||
runs; otherwise the lease is counted until the block exits. Leases
|
||||
are never listed and never refuse one another or any job.
|
||||
"""
|
||||
return _Lease(self, title)
|
||||
|
||||
def lease_blocker(self):
|
||||
"""The running ``exclusive`` job a new lease is refused for, or None."""
|
||||
for job in self.running():
|
||||
if job.admission == "exclusive":
|
||||
return job
|
||||
return None
|
||||
|
||||
@property
|
||||
def leases(self):
|
||||
"""Running leases by title (a copy)."""
|
||||
return dict(self._leases)
|
||||
|
||||
def _take_lease(self, title):
|
||||
blocker = self.lease_blocker()
|
||||
if blocker is not None:
|
||||
raise JobBusy(blocker)
|
||||
self._leases[title] = self._leases.get(title, 0) + 1
|
||||
|
||||
def _release_lease(self, title):
|
||||
count = self._leases.get(title, 0) - 1
|
||||
if count > 0:
|
||||
self._leases[title] = count
|
||||
else:
|
||||
self._leases.pop(title, None)
|
||||
if not self._leases:
|
||||
waiters, self._drain_waiters = self._drain_waiters, []
|
||||
for waiter in waiters:
|
||||
if not waiter.done():
|
||||
waiter.set_result(True)
|
||||
|
||||
async def drain_leases(self, timeout=None):
|
||||
"""Wait until no lease runs; False when ``timeout`` ended the wait.
|
||||
|
||||
For an ``exclusive`` job about to switch the database: once it runs
|
||||
no new lease is admitted, so the count only goes down.
|
||||
"""
|
||||
if not self._leases:
|
||||
return True
|
||||
waiter = asyncio.get_running_loop().create_future()
|
||||
self._drain_waiters.append(waiter)
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
asyncio.shield(waiter),
|
||||
LEASE_DRAIN_SECONDS if timeout is None else timeout,
|
||||
)
|
||||
return True
|
||||
except asyncio.TimeoutError:
|
||||
self._logger.warning(
|
||||
"Still running after the wait: " + ", ".join(sorted(self._leases))
|
||||
)
|
||||
return False
|
||||
finally:
|
||||
if waiter in self._drain_waiters:
|
||||
self._drain_waiters.remove(waiter)
|
||||
|
||||
def start(
|
||||
self,
|
||||
@@ -599,12 +515,6 @@ class JobRegistry:
|
||||
def get(self, job_id):
|
||||
return self._jobs.get(job_id)
|
||||
|
||||
def running(self):
|
||||
"""The running jobs, oldest first (a snapshot, safe off the event loop)."""
|
||||
jobs = [job for job in tuple(self._jobs.values()) if job.running]
|
||||
jobs.sort(key=lambda job: job._started_seq)
|
||||
return jobs
|
||||
|
||||
def active(self, kind):
|
||||
"""The running job of ``kind``, or None (safe off the event loop)."""
|
||||
for job in tuple(self._jobs.values()):
|
||||
@@ -649,23 +559,3 @@ class _Hold:
|
||||
if self._job.running:
|
||||
self._job.finish("finished" if exc_type is None else "failed")
|
||||
return False
|
||||
|
||||
|
||||
class _Lease:
|
||||
"""The async context manager ``JobRegistry.lease`` returns."""
|
||||
|
||||
def __init__(self, registry, title):
|
||||
self._registry = registry
|
||||
self._title = title
|
||||
self._held = False
|
||||
|
||||
async def __aenter__(self):
|
||||
self._registry._take_lease(self._title)
|
||||
self._held = True
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
if self._held:
|
||||
self._held = False
|
||||
self._registry._release_lease(self._title)
|
||||
return False
|
||||
|
||||
+8
-36
@@ -234,22 +234,9 @@ class LibraryRoutesMixin:
|
||||
)
|
||||
|
||||
def _library_media(
|
||||
self,
|
||||
records: list[dict],
|
||||
roots: list[Path],
|
||||
*,
|
||||
linked: bool = True,
|
||||
outputs: int | None = None,
|
||||
self, records: list[dict], roots: list[Path], *, linked: bool = True
|
||||
) -> list[dict]:
|
||||
"""Workspace descriptors for a prompt's linked media under ``roots``.
|
||||
|
||||
Roots from index ``outputs`` on are linked-only (LoRA media, #355):
|
||||
their media keeps its rooted identity, but gets no preview file (the
|
||||
card loads the original; a video shows its label), no favorite star
|
||||
and is marked ``linked_only`` (no Tag action).
|
||||
"""
|
||||
page = []
|
||||
linked_only = set()
|
||||
for record in records:
|
||||
try:
|
||||
raw_path = Path(record["image_path"])
|
||||
@@ -263,25 +250,10 @@ class LibraryRoutesMixin:
|
||||
canonical_root = _canonical(root)
|
||||
if path.is_relative_to(canonical_root):
|
||||
page.append((path, 0, canonical_root, index))
|
||||
if outputs is not None and index >= outputs:
|
||||
relative = path.relative_to(canonical_root).as_posix()
|
||||
linked_only.add((index, relative))
|
||||
break
|
||||
except (OSError, ValueError, RuntimeError):
|
||||
continue
|
||||
images = self._format_workspace_page(page, linked=linked)
|
||||
for image in images:
|
||||
if (image["root_index"], image["path"]) in linked_only:
|
||||
image["thumbnail_url"] = None
|
||||
image["thumbnail_ready"] = False
|
||||
image["linked_only"] = True
|
||||
image.pop("favorite", None)
|
||||
return images
|
||||
|
||||
def _library_roots(self) -> tuple[list[Path], int]:
|
||||
"""``(roots, number of output roots)`` for a prompt's linked media."""
|
||||
outputs = list(self._get_all_output_dirs())
|
||||
return self._media_roots(outputs), len(outputs)
|
||||
return self._format_workspace_page(page, linked=linked)
|
||||
|
||||
@pair_response
|
||||
async def get_library_options(self, request: web.Request) -> web.Response:
|
||||
@@ -322,13 +294,13 @@ class LibraryRoutesMixin:
|
||||
return result, records
|
||||
|
||||
result, records = await self._run_in_executor(read)
|
||||
roots = list(self._get_all_output_dirs())
|
||||
|
||||
def previews() -> None:
|
||||
roots, outputs = self._library_roots()
|
||||
for parent in result["items"]:
|
||||
# Previews belong to their parent; a per-parent link read is N+1.
|
||||
parent["previews"] = self._library_media(
|
||||
records[parent["id"]], roots, linked=False, outputs=outputs
|
||||
records[parent["id"]], roots, linked=False
|
||||
)
|
||||
|
||||
await self._run_in_executor(previews)
|
||||
@@ -370,9 +342,9 @@ class LibraryRoutesMixin:
|
||||
result, records = await self._run_in_executor(read)
|
||||
if result is None:
|
||||
return _missing()
|
||||
roots, outputs = await self._run_in_executor(self._library_roots)
|
||||
roots = list(self._get_all_output_dirs())
|
||||
result["positive"]["previews"] = await self._run_in_executor(
|
||||
lambda: self._library_media(records, roots, outputs=outputs)
|
||||
self._library_media, records, roots
|
||||
)
|
||||
return web.json_response({"success": True, **result})
|
||||
|
||||
@@ -386,7 +358,7 @@ class LibraryRoutesMixin:
|
||||
archive = _archive_scope(request)
|
||||
limit = _page_limit(request)
|
||||
cursor = request.query.get("cursor")
|
||||
roots, outputs = await self._run_in_executor(self._library_roots)
|
||||
roots = list(self._get_all_output_dirs())
|
||||
canonical_roots = [str(_canonical(root)) for root in roots]
|
||||
|
||||
def read() -> dict | None:
|
||||
@@ -419,7 +391,7 @@ class LibraryRoutesMixin:
|
||||
if result is None:
|
||||
return _missing()
|
||||
result["images"] = await self._run_in_executor(
|
||||
lambda: self._library_media(result["images"], roots, outputs=outputs)
|
||||
self._library_media, result["images"], roots
|
||||
)
|
||||
return web.json_response({"success": True, **result})
|
||||
|
||||
|
||||
@@ -1,420 +0,0 @@
|
||||
"""Library location routes (#249 S6b): where the library lives, Move, Keep.
|
||||
|
||||
- ``GET /prompt_manager/system/library`` → ``{success, location}``: the
|
||||
database's path (relative to ComfyUI, as the statistics show it), why it is
|
||||
there (``source``: ``default``, ``environment``, ``config`` or
|
||||
``node_folder``), whether it can move and where to (``target``), the
|
||||
"Keep it here" choice (``keep``), whether to offer the move (``offer``:
|
||||
movable and not kept), a ``problem`` code and the database's size. The
|
||||
System Overview carries the same object as ``location``.
|
||||
- ``POST /prompt_manager/system/library/move`` → 202 ``{job}``: the
|
||||
``library_move`` job, ``exclusive`` (it runs alone and every other job,
|
||||
legacy-page operations included, is refused while it runs), not
|
||||
cancellable. It copies, verifies and switches the live library
|
||||
(``database/library_move.py``) and ends with one truthful line; a failure
|
||||
leaves the library in the node folder, unchanged. Refused before any job
|
||||
starts: 400 ``not-movable``/``same-folder``/``target-invalid``, 409
|
||||
``target-has-library``/``target-conflict``/``busy``, 507
|
||||
``insufficient-space``.
|
||||
- ``POST /prompt_manager/system/library/keep {keep: bool}`` persists
|
||||
``library.keep_in_node_folder`` through the config store and answers the
|
||||
new location.
|
||||
|
||||
Errors are ``{success:false, code, error}`` and never name a path.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import os
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from .job_routes import busy_response
|
||||
from .jobs import JobBusy
|
||||
|
||||
try:
|
||||
from ...database import library_move as lm
|
||||
from ...database.operations import resolve_library_location
|
||||
from ...utils import library_location as ll
|
||||
except ImportError:
|
||||
from database import library_move as lm
|
||||
from database.operations import resolve_library_location
|
||||
from utils import library_location as ll
|
||||
|
||||
LIBRARY_MOVE_JOB_KIND = "library_move"
|
||||
LIBRARY_MOVE_TITLE = "Move library"
|
||||
LIBRARY_MOVE_GROUPS = ("files", "errors")
|
||||
KEEP_BODY_MAX = 1024
|
||||
|
||||
_REFUSAL_STATUS = {
|
||||
"not-movable": 400,
|
||||
"same-folder": 400,
|
||||
"target-invalid": 400,
|
||||
"target-has-library": 409,
|
||||
"target-conflict": 409,
|
||||
"insufficient-space": 507,
|
||||
}
|
||||
STILL_THERE = "The library is still in the node folder, unchanged."
|
||||
NOT_MOVED = "The library was not moved."
|
||||
OLD_LEFT = {
|
||||
"old-in-use": (
|
||||
"The old database is still open elsewhere, so it was left in the node "
|
||||
"folder; it is no longer used."
|
||||
),
|
||||
"old-locked": (
|
||||
"The old files could not be renamed, so they were left in the node "
|
||||
"folder; they are no longer used."
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _failure(message, status, code):
|
||||
return web.json_response(
|
||||
{"success": False, "code": code, "error": message}, status=status
|
||||
)
|
||||
|
||||
|
||||
def _plural(count, one, many=None):
|
||||
return f"{count:,} {one if count == 1 else (many or one + 's')}"
|
||||
|
||||
|
||||
def _same_file(first, second):
|
||||
try:
|
||||
return os.path.samefile(first, second)
|
||||
except OSError:
|
||||
return os.path.normcase(os.path.realpath(first)) == os.path.normcase(
|
||||
os.path.realpath(second)
|
||||
)
|
||||
|
||||
|
||||
CUSTOM_NODES = "custom_nodes"
|
||||
|
||||
|
||||
def _inside(path, folder):
|
||||
"""``path`` relative to ``folder`` by text (normalised spellings), or None."""
|
||||
child, parent = os.path.normcase(path), os.path.normcase(folder)
|
||||
if child == parent:
|
||||
return ""
|
||||
prefix = parent.rstrip(os.sep) + os.sep
|
||||
return path[len(prefix) :] if child.startswith(prefix) else None
|
||||
|
||||
|
||||
def _comfyui_custom_nodes():
|
||||
"""ComfyUI's custom_nodes folders as it lists them (unresolved), or []."""
|
||||
try:
|
||||
from ..config import _import_folder_paths
|
||||
|
||||
getter = getattr(_import_folder_paths(), "get_folder_paths", None)
|
||||
folders = getter(CUSTOM_NODES) if getter else []
|
||||
return [folder for folder in folders or [] if isinstance(folder, str)]
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def node_folder_public_path(path, extension_root, custom_nodes):
|
||||
"""``custom_nodes/<node folder>/...`` for a path in the node folder, or None.
|
||||
|
||||
The node folder is matched at its own unresolved place in a custom_nodes
|
||||
folder (one of ``custom_nodes``, or the parent of ``extension_root`` when
|
||||
that is named custom_nodes), so a node folder that is a symlink to a
|
||||
checkout outside ComfyUI still reads relative to ComfyUI. ``path`` may be
|
||||
spelled through the link or resolved. Never absolute, never ``..``:
|
||||
anything else is None and the caller uses ``_public_path``.
|
||||
"""
|
||||
if not path or not extension_root:
|
||||
return None
|
||||
try:
|
||||
root = os.path.normpath(os.path.abspath(extension_root))
|
||||
name = os.path.basename(root)
|
||||
folders = [os.path.normpath(os.path.abspath(f)) for f in custom_nodes]
|
||||
parent = os.path.dirname(root)
|
||||
if os.path.normcase(os.path.basename(parent)) == CUSTOM_NODES:
|
||||
folders.append(parent)
|
||||
spellings = [os.path.normpath(os.path.abspath(path))]
|
||||
spellings.append(os.path.realpath(spellings[0]))
|
||||
except (OSError, ValueError, TypeError):
|
||||
return None
|
||||
for folder in folders:
|
||||
node = os.path.join(folder, name)
|
||||
for spelling, node_spelling in zip(spellings, (node, os.path.realpath(node))):
|
||||
rest = _inside(spelling, node_spelling)
|
||||
if rest is None:
|
||||
continue
|
||||
parts = [os.path.basename(folder), name] + (
|
||||
rest.split(os.sep) if rest else []
|
||||
)
|
||||
if all(part and part not in (".", "..") for part in parts):
|
||||
return "/".join(parts)
|
||||
return None
|
||||
|
||||
|
||||
class LibraryLocationRoutesMixin:
|
||||
"""Mixin for the /prompt_manager/system/library endpoints."""
|
||||
|
||||
# Replaced in tests with a resolver over a temp tree.
|
||||
_library_resolver = staticmethod(resolve_library_location)
|
||||
|
||||
def _register_library_location_routes(self, routes):
|
||||
@routes.get("/prompt_manager/system/library")
|
||||
async def library_location_route(request):
|
||||
return await self.get_library_location(request)
|
||||
|
||||
@routes.post("/prompt_manager/system/library/move")
|
||||
async def library_move_route(request):
|
||||
return await self.start_library_move(request)
|
||||
|
||||
@routes.post("/prompt_manager/system/library/keep")
|
||||
async def library_keep_route(request):
|
||||
return await self.keep_library(request)
|
||||
|
||||
# ── location ──────────────────────────────────────────────────────
|
||||
|
||||
def _library_public_path(self, path):
|
||||
"""The library's path for display: through the node folder's place in
|
||||
custom_nodes when it is there (a symlinked node folder included), else
|
||||
``_public_path``. Only the library location uses this; sources and
|
||||
jobs keep ``_public_path``, which their requests round-trip."""
|
||||
shown = node_folder_public_path(
|
||||
path, self._extension_root, _comfyui_custom_nodes()
|
||||
)
|
||||
return shown or self._public_path(path)
|
||||
|
||||
def _live_location(self):
|
||||
"""The resolved location, unless the live database is another file.
|
||||
|
||||
A database the API was handed (tests, the browser fixture server) or
|
||||
one the resolver no longer names is reported as configured and is
|
||||
never moved.
|
||||
"""
|
||||
location = self._library_resolver()
|
||||
live = self.db.model.db_path
|
||||
if _same_file(location.db_path, live):
|
||||
return location
|
||||
return ll.LibraryLocation(
|
||||
ll.SOURCE_CONFIG,
|
||||
live,
|
||||
location.config_path,
|
||||
os.path.dirname(live),
|
||||
config_explicit=location.config_explicit,
|
||||
)
|
||||
|
||||
def _location_payload(self, location):
|
||||
"""The public view of ``location`` (blocking: stats the database)."""
|
||||
from ..config import PromptManagerConfig
|
||||
|
||||
keep = bool(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER)
|
||||
size = 0
|
||||
try:
|
||||
size = os.path.getsize(location.db_path)
|
||||
except OSError:
|
||||
pass
|
||||
return {
|
||||
"path": self._library_public_path(location.db_path),
|
||||
# Where Move library would put it, and whether config.json goes too.
|
||||
"target_source": location.target_source if location.movable else None,
|
||||
"config_moves": location.movable and not location.config_explicit,
|
||||
"source": location.source,
|
||||
"movable": location.movable,
|
||||
"keep": keep,
|
||||
"offer": location.movable and not keep,
|
||||
"target": (
|
||||
self._library_public_path(location.target_dir)
|
||||
if location.movable
|
||||
else None
|
||||
),
|
||||
"problem": location.problem or self._node_folder_risk(location),
|
||||
"size_bytes": size,
|
||||
}
|
||||
|
||||
# The node folder (tests point it at a temp tree).
|
||||
_extension_root = ll.EXTENSION_ROOT
|
||||
|
||||
def _node_folder_risk(self, location):
|
||||
"""``inside-node-folder`` for a library that stays in the node folder.
|
||||
|
||||
A configured database path, or a config.json the resolver still reads
|
||||
from the node folder, is deleted with the node by a reinstall, and
|
||||
it is never offered a move.
|
||||
"""
|
||||
if location.source == ll.SOURCE_NODE_FOLDER:
|
||||
return None
|
||||
paths = [location.db_path] + (
|
||||
[] if location.config_explicit else [location.config_path]
|
||||
)
|
||||
root = self._extension_root
|
||||
for path in paths:
|
||||
parent = os.path.dirname(os.path.abspath(path))
|
||||
if ll.same_folder(parent, root) or os.path.normcase(
|
||||
os.path.realpath(parent)
|
||||
).startswith(os.path.normcase(os.path.realpath(root)) + os.sep):
|
||||
return "inside-node-folder"
|
||||
return None
|
||||
|
||||
def _library_payload(self):
|
||||
return self._location_payload(self._live_location())
|
||||
|
||||
async def get_library_location(self, request):
|
||||
try:
|
||||
payload = await self._run_in_executor(self._library_payload)
|
||||
except Exception:
|
||||
self.logger.exception("Library location could not be read")
|
||||
return _failure(
|
||||
"The library location could not be read.", 500, "location-failed"
|
||||
)
|
||||
return web.json_response(
|
||||
{"success": True, "location": payload},
|
||||
headers={"Cache-Control": "no-store"},
|
||||
)
|
||||
|
||||
# ── move ──────────────────────────────────────────────────────────
|
||||
|
||||
def _plan_library_move(self):
|
||||
"""The checked plan (blocking); raises MoveFailed."""
|
||||
plan = lm.plan_move(self._live_location())
|
||||
if lm.space_problem(plan) is not None:
|
||||
raise lm.MoveFailed("insufficient-space")
|
||||
return plan
|
||||
|
||||
async def start_library_move(self, request):
|
||||
blocker = self._jobs.blocker(LIBRARY_MOVE_JOB_KIND, admission="exclusive")
|
||||
if blocker is not None:
|
||||
return busy_response(JobBusy(blocker))
|
||||
try:
|
||||
await self._run_in_executor(self._plan_library_move)
|
||||
except lm.MoveFailed as refused:
|
||||
status = _REFUSAL_STATUS.get(refused.code, 400)
|
||||
return _failure(refused.message, status, refused.code)
|
||||
except Exception:
|
||||
self.logger.exception("Move library could not be planned")
|
||||
return _failure(NOT_MOVED, 500, "move-failed")
|
||||
try:
|
||||
job = self._jobs.start(
|
||||
LIBRARY_MOVE_JOB_KIND,
|
||||
LIBRARY_MOVE_TITLE,
|
||||
admission="exclusive",
|
||||
cancellable=False,
|
||||
scope={},
|
||||
groups=LIBRARY_MOVE_GROUPS,
|
||||
runner=self._run_library_move,
|
||||
)
|
||||
except JobBusy as busy:
|
||||
return busy_response(busy)
|
||||
return web.json_response({"success": True, "job": job.summary()}, status=202)
|
||||
|
||||
def _after_library_switch(self):
|
||||
"""Forget what described the old file (best effort, after the switch).
|
||||
|
||||
Plans, the health report and the history read fence named the old
|
||||
file; the fence is closed before it is dropped.
|
||||
"""
|
||||
self._database_replaced()
|
||||
fence, self._history_read_fence = self._history_read_fence, None
|
||||
if fence is not None:
|
||||
fence.close()
|
||||
|
||||
def finish_interrupted_library_move(self):
|
||||
"""Startup: rename old files a crash left after a move (blocking)."""
|
||||
try:
|
||||
location = self._library_resolver()
|
||||
if _same_file(location.db_path, self.db.model.db_path):
|
||||
return lm.finish_interrupted_move(location)
|
||||
except Exception:
|
||||
self.logger.exception("Could not finish an interrupted library move")
|
||||
return []
|
||||
|
||||
async def _run_library_move(self, job):
|
||||
from ..config_store import _COMMIT_LOCK
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def progress(line):
|
||||
def show():
|
||||
job.log_line(line)
|
||||
job.update(status_line=line)
|
||||
|
||||
loop.call_soon_threadsafe(show)
|
||||
|
||||
job.update(status_line="Checking the new folder")
|
||||
try:
|
||||
plan = await self._run_in_executor(self._plan_library_move)
|
||||
target = self._library_public_path(plan.target_dir)
|
||||
job.log_line(f"Moving the library to {target}")
|
||||
result = await self._run_in_executor(
|
||||
functools.partial(
|
||||
lm.move_library,
|
||||
plan,
|
||||
after_switch=self._after_library_switch,
|
||||
config_lock=_COMMIT_LOCK,
|
||||
progress=progress,
|
||||
)
|
||||
)
|
||||
except lm.MoveFailed as failed:
|
||||
self.logger.error(f"Move library failed: {failed.code}")
|
||||
return self._fail_library_move(
|
||||
job, failed.code, f"{failed.message} {STILL_THERE}"
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("Move library failed")
|
||||
return self._fail_library_move(job, None, f"{NOT_MOVED} {STILL_THERE}")
|
||||
await asyncio.sleep(0) # the progress lines come first
|
||||
job.record_counted("files", len(result.files))
|
||||
line = (
|
||||
f"Moved the library to {target}: "
|
||||
f"{_plural(result.counts.get('prompts', 0), 'prompt')} and "
|
||||
f"{_plural(result.counts.get('generated_images', 0), 'image')} verified"
|
||||
)
|
||||
if result.retire_problem:
|
||||
line += ". " + OLD_LEFT.get(result.retire_problem, OLD_LEFT["old-locked"])
|
||||
elif result.renamed:
|
||||
line += (
|
||||
f". The old files were renamed {', '.join(result.renamed)}; "
|
||||
"nothing was deleted"
|
||||
)
|
||||
line += ". Logs move at the next restart."
|
||||
job.log_line(line)
|
||||
job.update(progress=100, status_line=line)
|
||||
|
||||
@staticmethod
|
||||
def _fail_library_move(job, code, message):
|
||||
from .maintenance import _fail_job
|
||||
|
||||
_fail_job(job, code, message)
|
||||
|
||||
# ── keep it here ─────────────────────────────────────────────────
|
||||
|
||||
async def keep_library(self, request):
|
||||
from ..config import PromptManagerConfig
|
||||
from ..config_store import get_config_store
|
||||
|
||||
try:
|
||||
if (request.content_length or 0) > KEEP_BODY_MAX:
|
||||
raise ValueError("too large")
|
||||
data = await request.json()
|
||||
except Exception:
|
||||
data = None
|
||||
if not isinstance(data, dict) or not isinstance(data.get("keep"), bool):
|
||||
return _failure("keep must be true or false.", 400, "invalid-request")
|
||||
keep = data["keep"]
|
||||
|
||||
def mutate(config):
|
||||
library = config.get("library")
|
||||
library = dict(library) if isinstance(library, dict) else {}
|
||||
library["keep_in_node_folder"] = keep
|
||||
config["library"] = library
|
||||
return config
|
||||
|
||||
def publish(_config):
|
||||
PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER = keep
|
||||
|
||||
try:
|
||||
await self._run_in_executor(
|
||||
functools.partial(
|
||||
get_config_store().commit,
|
||||
mutate,
|
||||
expected_revision=None,
|
||||
publish=publish,
|
||||
)
|
||||
)
|
||||
except Exception as exc:
|
||||
return self._config_commit_failure(exc)
|
||||
return await self.get_library_location(request)
|
||||
@@ -1,513 +0,0 @@
|
||||
"""The staged LoRA reference import as a reviewed Activity job (#221, S10).
|
||||
|
||||
``POST /prompt_manager/lora/import/preview`` (body ``{}``) reviews an import:
|
||||
it lists the LoraManager metadata files and parses them, and counts the
|
||||
current imports, those a saved pair uses and the example images an import
|
||||
would download. It writes nothing and requests nothing. It answers a
|
||||
``Plan`` (the maintenance plan's shape, ``py/api/maintenance.py``) plus
|
||||
``{lora_files, with_metadata, skipped, with_local_media, current_imports,
|
||||
blocked_by_pairs, planned_downloads, network:{planned_downloads, api_key},
|
||||
layout}`` (#355: ``layout`` says which LoraManager settings file, storage mode
|
||||
and library the import read, and ``error`` when they cannot be used). A plan
|
||||
that saved pairs block, or whose layout is unusable, carries no token
|
||||
(nothing can be applied).
|
||||
|
||||
``POST /prompt_manager/lora/import/apply {token}`` starts the job (202
|
||||
``{job}``): kind ``lora_import``, titled "LoRA reference import", admission
|
||||
``records``, cancellable while it downloads only. Refused before it starts:
|
||||
``plan-changed`` (409: the metadata files changed since the review),
|
||||
``plan-expired``, ``plan-used``, ``storage-changed``, ``busy``,
|
||||
``insufficient-space`` (as maintenance). The job stages the records
|
||||
(downloads into the cache first, each counted in ``downloaded`` or
|
||||
``download_errors``), takes the safety backup the policy promises, then
|
||||
replaces the collection in one transaction bound to the reviewed imports. Any
|
||||
failure leaves the previous imports exactly as they were, and the job says so.
|
||||
|
||||
The legacy ``POST /lora/scan`` runs the same stages (``stage_lora_import``
|
||||
and ``replace_lora_imports``) under its SSE event shape.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sqlite3
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from .. import lora_import, lora_utils
|
||||
from .inspection_service import SERVABLE_EXTENSIONS
|
||||
from .job_routes import busy_response
|
||||
from .jobs import LORA_IMPORT_JOB_KIND, JobBusy
|
||||
from .maintenance import (
|
||||
MAX_TOKEN_LENGTH,
|
||||
_failure,
|
||||
_fail_job,
|
||||
_invalid,
|
||||
_iso,
|
||||
disk_plan,
|
||||
short_space_message,
|
||||
)
|
||||
from .maintenance_plans import ScopeTooLarge
|
||||
|
||||
try:
|
||||
from ...database.category_import import EmptyReplacement
|
||||
from ...database.maintenance import PlanChanged
|
||||
from ...database.pairs import PairDependencyError
|
||||
except ImportError: # pragma: no cover - flat test imports
|
||||
from database.category_import import EmptyReplacement
|
||||
from database.maintenance import PlanChanged
|
||||
from database.pairs import PairDependencyError
|
||||
|
||||
LORA_OPERATION = "lora_import"
|
||||
# The review reads the library: a request lease, so it is refused while a
|
||||
# restore runs and a restore waits for it before switching the file (S8).
|
||||
LORA_PREVIEW_LEASE = "LoRA import review"
|
||||
CATEGORY = lora_import.LORA_CATEGORY
|
||||
SOURCE_CHANGED = (
|
||||
"The LoRA metadata files changed since the review; nothing was changed. "
|
||||
"Review the import again."
|
||||
)
|
||||
IMPORTS_CHANGED = (
|
||||
"The imported LoRA prompts changed since the review; the previous imports "
|
||||
"were kept. Review the import again."
|
||||
)
|
||||
BLOCKED = (
|
||||
"Saved pairs use some imported LoRA prompts, so the collection cannot be "
|
||||
"replaced; the previous imports were kept."
|
||||
)
|
||||
INTEGRATION_CHANGED = (
|
||||
"The LoraManager integration changed since the review; nothing was changed. "
|
||||
"Review the import again."
|
||||
)
|
||||
EMPTY = "No LoRA could be prepared for import; the previous imports were kept."
|
||||
FAILED = "The import failed; the previous LoRA imports were kept."
|
||||
CACHE_NOTE = "Images already downloaded stay in the image cache."
|
||||
CANCELLED = "Cancelled before anything was written; the previous imports were kept."
|
||||
DISABLED = "LoraManager integration is not enabled"
|
||||
NO_PATH = "LoraManager path not configured"
|
||||
|
||||
|
||||
def _noun(count, word, plural=None):
|
||||
return word if count == 1 else (plural or word + "s")
|
||||
|
||||
|
||||
def _count(count, word, plural=None):
|
||||
return f"{count:,} {_noun(count, word, plural)}"
|
||||
|
||||
|
||||
def _import_line(result, skipped):
|
||||
parts = [f"Imported {_count(result['imported'], 'LoRA')}"]
|
||||
if result["linked"] or result["duplicates"]:
|
||||
linked = result["linked"] + result["duplicates"]
|
||||
parts.append(f"{linked:,} linked to an existing prompt")
|
||||
if skipped:
|
||||
parts.append(f"{skipped:,} skipped")
|
||||
return " · ".join(parts)
|
||||
|
||||
|
||||
def _all_downloads_failed(errors):
|
||||
"""The warning when every planned download failed (#337)."""
|
||||
if errors == 1:
|
||||
return "Warning: the example image could not be downloaded"
|
||||
return f"Warning: none of the {errors:,} example images could be downloaded"
|
||||
|
||||
|
||||
def _with_downloads(line, downloaded, errors):
|
||||
"""``line`` with the failed downloads: counted when some failed, led by
|
||||
a warning when every planned download failed (#337). The job still
|
||||
finishes: the import itself succeeded and the downloads are best effort."""
|
||||
if not errors:
|
||||
return line
|
||||
if not downloaded:
|
||||
return f"{_all_downloads_failed(errors)} · {line}"
|
||||
failed = f"{_count(errors, 'example image')} could not be downloaded"
|
||||
return f"{line} · {failed}"
|
||||
|
||||
|
||||
class LoraImportRoutesMixin:
|
||||
"""Mixin providing the reviewed LoRA import (preview, apply, job)."""
|
||||
|
||||
def _register_lora_import_routes(self, routes):
|
||||
@routes.post("/prompt_manager/lora/import/preview")
|
||||
async def lora_import_preview_route(request):
|
||||
return await self.preview_lora_import(request)
|
||||
|
||||
@routes.post("/prompt_manager/lora/import/apply")
|
||||
async def lora_import_apply_route(request):
|
||||
return await self.apply_lora_import(request)
|
||||
|
||||
# ── shared stages ─────────────────────────────────────────────────
|
||||
|
||||
async def _lora_source(self):
|
||||
"""``(path, None)``: the configured, allowed LoraManager path, or
|
||||
``(None, (code, message))`` when the import cannot run."""
|
||||
from ..config import IntegrationConfig
|
||||
|
||||
if not IntegrationConfig.LORA_MANAGER_ENABLED:
|
||||
return None, ("lora-disabled", DISABLED)
|
||||
lm_path = IntegrationConfig.LORA_MANAGER_PATH
|
||||
if not lm_path:
|
||||
return None, ("lora-path", NO_PATH)
|
||||
allowed = await self._run_in_executor(
|
||||
lora_utils.resolve_lora_manager_path, lm_path
|
||||
)
|
||||
if allowed is None:
|
||||
from .lora_integration import _LORA_PATH_ERROR
|
||||
|
||||
return None, ("lora-path", _LORA_PATH_ERROR)
|
||||
return allowed, None
|
||||
|
||||
async def stage_lora_import(self, discovery, *, on_record=None, cancelled=None):
|
||||
"""Stage ``discovery``'s sidecars off the loop: records, skips and
|
||||
download counts."""
|
||||
from ..config import IntegrationConfig
|
||||
|
||||
cache_dir = await self._run_in_executor(lora_utils.get_lora_image_cache_dir)
|
||||
await self._run_in_executor(lora_utils.sweep_stale_downloads, cache_dir)
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def report(index, record, staged):
|
||||
if on_record is not None:
|
||||
loop.call_soon_threadsafe(on_record, index, record, staged)
|
||||
|
||||
return await self._run_in_executor(
|
||||
lambda: lora_import.stage(
|
||||
discovery,
|
||||
cache_dir,
|
||||
IntegrationConfig.CIVITAI_API_KEY,
|
||||
media_extensions=SERVABLE_EXTENSIONS,
|
||||
on_record=report,
|
||||
cancelled=cancelled or (lambda: False),
|
||||
logger=self.logger,
|
||||
)
|
||||
)
|
||||
|
||||
async def lora_files_unchanged(self, lm_path, expected):
|
||||
"""Whether LoraManager's sidecars for ``lm_path`` still fingerprint as
|
||||
``expected`` (found again, so an added file counts too)."""
|
||||
|
||||
def current():
|
||||
return lora_import.fingerprint(lora_import.discover(lm_path))
|
||||
|
||||
try:
|
||||
return await self._run_in_executor(current) == expected
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
async def replace_lora_imports(self, db, staged, *, expected_ids, allow_empty):
|
||||
"""Swap the collection for ``staged.records`` in one transaction."""
|
||||
return await self._run_in_executor(
|
||||
lambda: db.replace_category_prompts(
|
||||
CATEGORY,
|
||||
staged.records,
|
||||
expected_ids=expected_ids,
|
||||
allow_empty=allow_empty,
|
||||
)
|
||||
)
|
||||
|
||||
# ── preview ───────────────────────────────────────────────────────
|
||||
|
||||
async def preview_lora_import(self, request):
|
||||
try:
|
||||
async with self._jobs.lease(LORA_PREVIEW_LEASE):
|
||||
return await self._preview_lora_import(request)
|
||||
except JobBusy as busy:
|
||||
return busy_response(busy)
|
||||
|
||||
async def _preview_lora_import(self, request):
|
||||
data, error = await self._read_maintenance_body(request, ())
|
||||
if error is not None:
|
||||
return error
|
||||
lm_path, refused = await self._lora_source()
|
||||
if refused is not None:
|
||||
return _failure(refused[1], 400, refused[0])
|
||||
db = self.db
|
||||
identity = self._database_identity(db)
|
||||
try:
|
||||
cache_dir = await self._run_in_executor(
|
||||
lambda: lora_utils.get_lora_image_cache_dir(create=False)
|
||||
)
|
||||
found = await self._run_in_executor(
|
||||
lora_import.preview, lm_path, cache_dir, SERVABLE_EXTENSIONS
|
||||
)
|
||||
state = await self._run_in_executor(db.category_import_state, CATEGORY)
|
||||
except OSError as error:
|
||||
# A folder discovery cannot list refuses the import (it must not
|
||||
# replace the collection with part of it): say which one.
|
||||
self.logger.warning(
|
||||
f"LoRA import review: a folder could not be listed: {error}"
|
||||
)
|
||||
return _failure(
|
||||
f"A LoRA folder could not be listed ({self._public_error(error)}). "
|
||||
"Nothing was changed.",
|
||||
500,
|
||||
"folder-unreadable",
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("LoRA import preview failed")
|
||||
return _failure("The LoRA import could not be reviewed.", 500)
|
||||
if identity != self._database_identity(self.db):
|
||||
return _failure(
|
||||
"The database was replaced while it was reviewed. Review again.",
|
||||
409,
|
||||
"storage-changed",
|
||||
)
|
||||
return await self._lora_plan_response(db, identity, lm_path, found, state)
|
||||
|
||||
async def _lora_plan_response(self, db, identity, lm_path, found, state):
|
||||
from ..config import IntegrationConfig
|
||||
|
||||
current = len(state["ids"])
|
||||
blocked = state["paired"]
|
||||
policy = self._safety_backup_policy()
|
||||
try:
|
||||
disk, _short, size = await self._run_in_executor(
|
||||
disk_plan, db.model.db_path, backup=policy != "off", vacuum=False
|
||||
)
|
||||
except (OSError, sqlite3.Error):
|
||||
self.logger.exception("Could not measure the disk space for a plan")
|
||||
return _failure("The LoRA import could not be reviewed.", 500)
|
||||
token = None
|
||||
expires_at = None
|
||||
layout = found["layout"]
|
||||
# An unusable LoraManager layout found nothing it may replace with.
|
||||
if not blocked and layout["error"] is None:
|
||||
payload = {
|
||||
"path": lm_path,
|
||||
"fingerprint": found["fingerprint"],
|
||||
"expected_ids": state["ids"],
|
||||
"empty": found["with_metadata"] == 0,
|
||||
"backup": policy,
|
||||
}
|
||||
try:
|
||||
token = self._maintenance_plans.issue(
|
||||
LORA_OPERATION, found["fingerprint"], payload, identity
|
||||
)
|
||||
except ScopeTooLarge:
|
||||
return _failure(
|
||||
"This import is too large to review at once.", 413, "too-large"
|
||||
)
|
||||
expires_at = _iso(self._maintenance_plans.peek(token).expires_at)
|
||||
name = os.path.basename(db.model.db_path)
|
||||
downloads = found["planned_downloads"]
|
||||
reasons = (
|
||||
[f"{_count(blocked, 'imported prompt')} used by saved pairs"]
|
||||
if blocked
|
||||
else []
|
||||
)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"operation": LORA_OPERATION,
|
||||
"token": token,
|
||||
"expires_at": expires_at,
|
||||
"lora_files": found["lora_files"],
|
||||
"with_metadata": found["with_metadata"],
|
||||
"skipped": found["skipped"],
|
||||
"with_local_media": found["with_local_media"],
|
||||
"layout": layout,
|
||||
"current_imports": current,
|
||||
"blocked_by_pairs": blocked,
|
||||
"planned_downloads": downloads,
|
||||
"scope": {
|
||||
"count": found["with_metadata"],
|
||||
"label": (
|
||||
f"Replace {_count(current, 'imported prompt')} with "
|
||||
f"{_count(found['with_metadata'], 'LoRA')}"
|
||||
),
|
||||
},
|
||||
"protected": {"count": blocked, "reasons": reasons},
|
||||
"backup": {
|
||||
"policy": policy,
|
||||
"destination_label": f"{name}.backup_<date and time>",
|
||||
"expected_bytes": size if policy != "off" else None,
|
||||
},
|
||||
"media": {
|
||||
"deletes_files": False,
|
||||
"label": "No files are removed; image links are replaced",
|
||||
},
|
||||
"network": {
|
||||
"planned_downloads": downloads,
|
||||
"api_key": bool(IntegrationConfig.CIVITAI_API_KEY),
|
||||
},
|
||||
"details": [{"label": n, "detail": ""} for n in found["names"]],
|
||||
"disk": disk,
|
||||
},
|
||||
headers={"Cache-Control": "no-store"},
|
||||
)
|
||||
|
||||
# ── apply ─────────────────────────────────────────────────────────
|
||||
|
||||
async def apply_lora_import(self, request):
|
||||
data, error = await self._read_maintenance_body(request, ("token",))
|
||||
if error is not None:
|
||||
return error
|
||||
token = data.get("token")
|
||||
if not isinstance(token, str) or not 0 < len(token) <= MAX_TOKEN_LENGTH:
|
||||
return _invalid("token must be the plan's token")
|
||||
plan, refused = self._peek_plan(token)
|
||||
if refused is not None:
|
||||
return refused
|
||||
if plan.operation != LORA_OPERATION:
|
||||
return _failure(
|
||||
"This review expired. Review the plan again.", 410, "plan-expired"
|
||||
)
|
||||
store = self._maintenance_plans
|
||||
db = self.db
|
||||
if plan.db_identity != self._database_identity(db):
|
||||
store.discard(token)
|
||||
return _failure(
|
||||
"The database was replaced since the review. Review again.",
|
||||
409,
|
||||
"storage-changed",
|
||||
)
|
||||
lm_path = plan.payload["path"]
|
||||
# M5: the integration must still be on and at the reviewed install.
|
||||
current, _refused = await self._lora_source()
|
||||
if current != lm_path:
|
||||
store.discard(token)
|
||||
return _failure(INTEGRATION_CHANGED, 409, "plan-changed")
|
||||
try:
|
||||
discovery = await self._run_in_executor(lora_import.discover, lm_path)
|
||||
now = await self._run_in_executor(lora_import.fingerprint, discovery)
|
||||
backup = self._backup_promised(plan)
|
||||
_disk, short, _size = await self._run_in_executor(
|
||||
disk_plan, db.model.db_path, backup=backup, vacuum=False
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("LoRA import could not start")
|
||||
return _failure(
|
||||
"The import could not start. Nothing was changed.", 500, "apply-failed"
|
||||
)
|
||||
if now != plan.fingerprint:
|
||||
store.discard(token)
|
||||
return _failure(SOURCE_CHANGED, 409, "plan-changed")
|
||||
if short is not None:
|
||||
return _failure(short_space_message(short), 507, "insufficient-space")
|
||||
plan, refused = self._peek_plan(token)
|
||||
if refused is not None:
|
||||
return refused
|
||||
try:
|
||||
job = self._jobs.start(
|
||||
LORA_IMPORT_JOB_KIND,
|
||||
lora_import.JOB_TITLE,
|
||||
admission="records",
|
||||
cancellable=True,
|
||||
scope={
|
||||
"operation": LORA_OPERATION,
|
||||
"lora_files": len(discovery.sidecars),
|
||||
},
|
||||
groups=lora_import.GROUPS,
|
||||
runner=lambda job: self._run_lora_import(
|
||||
job, db, plan, discovery, backup
|
||||
),
|
||||
)
|
||||
except JobBusy as busy:
|
||||
return busy_response(busy)
|
||||
store.consume(token, LORA_OPERATION, job_id=job.id)
|
||||
return web.json_response({"success": True, "job": job.summary()}, status=202)
|
||||
|
||||
# ── the job ───────────────────────────────────────────────────────
|
||||
|
||||
async def _run_lora_import(self, job, db, plan, discovery, backup_on):
|
||||
total = len(discovery.sidecars)
|
||||
job.update(total=total, status_line="Preparing LoRA records")
|
||||
job.log_line(f"Found {_count(total, 'LoRA metadata file')}")
|
||||
|
||||
def on_record(index, record, staged):
|
||||
job.update(
|
||||
processed=index + 1,
|
||||
progress=int(90 * (index + 1) / max(total, 1)),
|
||||
status_line=f"Prepared {index + 1:,} of {total:,}",
|
||||
)
|
||||
job.publish({"type": "progress", "progress": job.progress})
|
||||
|
||||
try:
|
||||
staged = await self.stage_lora_import(
|
||||
discovery, on_record=on_record, cancelled=lambda: job.cancelled
|
||||
)
|
||||
except lora_import.ImportCancelled:
|
||||
job.log_line(f"{CANCELLED} {CACHE_NOTE}")
|
||||
job.update(status_line=CANCELLED)
|
||||
job.finish("cancelled")
|
||||
return
|
||||
except Exception:
|
||||
self.logger.exception("LoRA import staging failed")
|
||||
return _fail_job(job, "import-failed", f"{FAILED} {CACHE_NOTE}")
|
||||
self._count_downloads(job, staged)
|
||||
# M4: what staging skipped stays on the job whatever happens next.
|
||||
job.record_counted(
|
||||
"skipped",
|
||||
len(staged.skipped),
|
||||
[{"path": "", "detail": name} for name in staged.skipped],
|
||||
)
|
||||
# Past this point nothing stops halfway: the replace is one transaction.
|
||||
# Published, so a view that offered Cancel during the downloads drops it.
|
||||
job.cancellable = False
|
||||
job.publish({"type": "progress", "progress": job.progress})
|
||||
if job.cancelled:
|
||||
job.log_line(f"{CANCELLED} {CACHE_NOTE}")
|
||||
job.update(status_line=CANCELLED)
|
||||
job.finish("cancelled")
|
||||
return
|
||||
ok, _backup = await self._safety_backup_step(job, db, backup_on)
|
||||
if not ok:
|
||||
return
|
||||
# I1: the files are listed and fingerprinted again right before the
|
||||
# replace, so a file changed, removed or added while the downloads ran
|
||||
# (LoraManager refreshing its CivitAI info) aborts with nothing deleted.
|
||||
if not await self.lora_files_unchanged(plan.payload["path"], plan.fingerprint):
|
||||
return _fail_job(job, "plan-changed", f"{SOURCE_CHANGED} {CACHE_NOTE}")
|
||||
job.update(status_line="Replacing the LoRA collection", progress=95)
|
||||
await self._replace_in_job(job, db, plan, staged)
|
||||
|
||||
def _count_downloads(self, job, staged):
|
||||
downloaded = staged.stats.get("downloaded", 0)
|
||||
errors = staged.stats.get("download_errors", 0)
|
||||
job.record_counted("downloaded", downloaded)
|
||||
rows = []
|
||||
if errors and not downloaded:
|
||||
warning = _all_downloads_failed(errors)
|
||||
rows = [{"path": "", "detail": f"{warning}; the log has the details"}]
|
||||
job.log_line(warning)
|
||||
job.record_counted("download_errors", errors, rows)
|
||||
if downloaded or errors:
|
||||
job.log_line(
|
||||
f"Downloaded {_count(downloaded, 'example image')}"
|
||||
+ (f"; {errors:,} could not be downloaded" if errors else "")
|
||||
)
|
||||
|
||||
async def _replace_in_job(self, job, db, plan, staged):
|
||||
try:
|
||||
result = await self.replace_lora_imports(
|
||||
db,
|
||||
staged,
|
||||
expected_ids=plan.payload["expected_ids"],
|
||||
allow_empty=plan.payload["empty"],
|
||||
)
|
||||
except PlanChanged:
|
||||
return _fail_job(job, "plan-changed", f"{IMPORTS_CHANGED} {CACHE_NOTE}")
|
||||
except PairDependencyError:
|
||||
return _fail_job(job, "blocked-by-pairs", f"{BLOCKED} {CACHE_NOTE}")
|
||||
except EmptyReplacement:
|
||||
return _fail_job(job, "nothing-to-import", f"{EMPTY} {CACHE_NOTE}")
|
||||
except Exception:
|
||||
self.logger.exception("LoRA import failed")
|
||||
return _fail_job(job, "import-failed", f"{FAILED} {CACHE_NOTE}")
|
||||
rows = {"imported": [], "linked": []}
|
||||
for row in result["rows"]:
|
||||
group = "imported" if row["outcome"] == "imported" else "linked"
|
||||
rows[group].append({"path": "", "detail": row["name"]})
|
||||
job.record_counted("imported", result["imported"], rows["imported"])
|
||||
job.record_counted(
|
||||
"linked", result["linked"] + result["duplicates"], rows["linked"]
|
||||
)
|
||||
line = _import_line(result, len(staged.skipped))
|
||||
if result["removed"]:
|
||||
job.log_line(f"Replaced {_count(result['removed'], 'previous import')}")
|
||||
job.log_line(line)
|
||||
line = _with_downloads(
|
||||
line,
|
||||
staged.stats.get("downloaded", 0),
|
||||
staged.stats.get("download_errors", 0),
|
||||
)
|
||||
job.update(progress=100, processed=job.total, status_line=line)
|
||||
+136
-211
@@ -1,9 +1,7 @@
|
||||
"""LoraManager integration API routes for PromptManager."""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
@@ -39,30 +37,6 @@ def _lora_settings_mutation(values):
|
||||
return mutate
|
||||
|
||||
|
||||
# The admin page's own words for the refusals it can meet (it has no review).
|
||||
_FILES_CHANGED = (
|
||||
"The LoRA metadata files changed during the scan; the previous imports "
|
||||
"were kept. Run the scan again."
|
||||
)
|
||||
_IMPORTS_CHANGED = (
|
||||
"The imported LoRA prompts changed during the scan; the previous imports "
|
||||
"were kept. Run the scan again."
|
||||
)
|
||||
|
||||
|
||||
class LoraSettingsUnusable(Exception):
|
||||
"""LoraManager's settings cannot be used; nothing was imported (#355)."""
|
||||
|
||||
def __init__(self, reason):
|
||||
super().__init__(
|
||||
f"LoraManager settings: not usable: {reason}. Nothing was changed."
|
||||
)
|
||||
|
||||
|
||||
class LoraFilesChanged(Exception):
|
||||
"""A metadata file changed, appeared or went while the legacy scan staged."""
|
||||
|
||||
|
||||
def _lora_error_frame(message):
|
||||
"""An import stream's error event, as the admin page's modal reads it.
|
||||
|
||||
@@ -72,23 +46,6 @@ def _lora_error_frame(message):
|
||||
return {"type": "error", "status": message, "progress": 0, "message": message}
|
||||
|
||||
|
||||
def _closed(request):
|
||||
"""A check (safe off the loop) for whether ``request``'s client left."""
|
||||
|
||||
def closed():
|
||||
transport = request.transport
|
||||
return transport is None or transport.is_closing()
|
||||
|
||||
return closed
|
||||
|
||||
|
||||
def _stopped(error) -> bool:
|
||||
"""Whether a legacy scan ended because its page closed the connection."""
|
||||
from .. import lora_import
|
||||
|
||||
return isinstance(error, (ConnectionError, lora_import.ImportCancelled))
|
||||
|
||||
|
||||
class LoraIntegrationMixin:
|
||||
"""Mixin providing LoraManager detection, scanning, and trigger word endpoints."""
|
||||
|
||||
@@ -151,11 +108,7 @@ class LoraIntegrationMixin:
|
||||
"""Get current LoraManager integration status."""
|
||||
try:
|
||||
from ..config import IntegrationConfig
|
||||
from ..lora_utils import (
|
||||
detect_lora_manager,
|
||||
get_trigger_cache,
|
||||
resolve_lora_manager_path,
|
||||
)
|
||||
from ..lora_utils import detect_lora_manager, get_trigger_cache
|
||||
|
||||
config = IntegrationConfig.get_config()["lora_manager"]
|
||||
cache = get_trigger_cache()
|
||||
@@ -164,12 +117,6 @@ class LoraIntegrationMixin:
|
||||
detected_path = await self._run_in_executor(
|
||||
detect_lora_manager, config.get("path", "")
|
||||
)
|
||||
# A detected install reached through a link outside custom_nodes is
|
||||
# found but cannot be enabled (S10 review N5): say which.
|
||||
usable = bool(detected_path) and (
|
||||
await self._run_in_executor(resolve_lora_manager_path, detected_path)
|
||||
is not None
|
||||
)
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
@@ -179,7 +126,6 @@ class LoraIntegrationMixin:
|
||||
"trigger_words_enabled": config["trigger_words_enabled"],
|
||||
"has_civitai_api_key": bool(config.get("civitai_api_key")),
|
||||
"detected": detected_path is not None,
|
||||
"detected_usable": usable,
|
||||
"detected_path": (
|
||||
self._public_path(detected_path) if detected_path else ""
|
||||
),
|
||||
@@ -222,9 +168,6 @@ class LoraIntegrationMixin:
|
||||
async def lora_enable(self, request):
|
||||
"""Enable or disable LoraManager integration and save it to config.json.
|
||||
|
||||
``expected_revision`` (optional, GET /settings' ``revision``) refuses
|
||||
a save over a config.json changed meanwhile (``settings-changed``).
|
||||
|
||||
The values are committed through the config store first (#220) and
|
||||
published at runtime only once config.json holds them; a save that
|
||||
cannot be written reports ``persisted: false`` and changes nothing.
|
||||
@@ -237,11 +180,6 @@ class LoraIntegrationMixin:
|
||||
for flag in _LORA_FLAGS:
|
||||
if flag in data and not isinstance(data[flag], bool):
|
||||
return self._bad_lora_request(f"{flag} must be true or false")
|
||||
# The System Save row sends the revision it read (#220); the
|
||||
# legacy page sends none and merges into what is on disk.
|
||||
expected_revision = data.get("expected_revision")
|
||||
if expected_revision is not None and not isinstance(expected_revision, str):
|
||||
return self._bad_lora_request("expected_revision must be a string")
|
||||
enabled = data.get("enabled", False)
|
||||
path = data.get("path", "")
|
||||
trigger_words = data.get("trigger_words_enabled", False)
|
||||
@@ -288,7 +226,7 @@ class LoraIntegrationMixin:
|
||||
_saved, revision = await self._run_in_executor(
|
||||
get_config_store().commit,
|
||||
_lora_settings_mutation(values),
|
||||
expected_revision=expected_revision,
|
||||
expected_revision=None,
|
||||
publish=publish,
|
||||
)
|
||||
except (ConfigUnreadable, RevisionConflict, OSError) as exc:
|
||||
@@ -370,75 +308,43 @@ class LoraIntegrationMixin:
|
||||
) as job:
|
||||
await response.prepare(request)
|
||||
try:
|
||||
await self._lora_scan_stream(
|
||||
allowed, send_progress, disconnected=_closed(request)
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
# The server cancelled the handler when the page left.
|
||||
job.finish("cancelled")
|
||||
raise
|
||||
await self._lora_scan_stream(allowed, send_progress)
|
||||
except Exception as e:
|
||||
# A page that left stopped the scan; nothing failed (R2).
|
||||
job.finish("cancelled" if _stopped(e) else "failed")
|
||||
await self._send_scan_failure(e, send_progress)
|
||||
with contextlib.suppress(ConnectionError, RuntimeError):
|
||||
await response.write_eof()
|
||||
self.logger.exception("LoRA scan failed")
|
||||
job.finish("failed")
|
||||
await send_progress(_lora_error_frame(self._public_error(e)))
|
||||
|
||||
await response.write_eof()
|
||||
return response
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"LoRA scan failed: {e}")
|
||||
return self._error_response(e)
|
||||
|
||||
async def _send_scan_failure(self, error, send_progress):
|
||||
"""The error frame for a failed legacy scan; quiet when the page left.
|
||||
|
||||
Expected refusals (files or imports changed, saved pairs, nothing to
|
||||
import) are warnings in the admin page's own words; anything else is
|
||||
logged with its traceback.
|
||||
"""
|
||||
from .lora_import_routes import (
|
||||
EmptyReplacement,
|
||||
PairDependencyError,
|
||||
PlanChanged,
|
||||
async def _lora_scan_stream(self, lm_path, send_progress):
|
||||
"""Import every LoRA under ``lm_path``, reporting through ``send_progress``."""
|
||||
from ..config import IntegrationConfig
|
||||
from ..lora_utils import (
|
||||
download_civitai_images,
|
||||
find_lora_directories,
|
||||
get_example_prompt_from_metadata,
|
||||
get_lora_image_cache_dir,
|
||||
get_preview_images_from_metadata,
|
||||
get_trigger_words_from_metadata,
|
||||
get_model_name_from_metadata,
|
||||
read_lora_metadata,
|
||||
)
|
||||
|
||||
if _stopped(error):
|
||||
self.logger.info("LoRA scan stopped: the page closed its connection")
|
||||
return
|
||||
if isinstance(error, PlanChanged):
|
||||
message = _IMPORTS_CHANGED
|
||||
elif isinstance(error, (LoraFilesChanged, LoraSettingsUnusable)):
|
||||
message = str(error)
|
||||
else:
|
||||
message = self._public_error(error)
|
||||
if isinstance(
|
||||
error,
|
||||
(
|
||||
PlanChanged,
|
||||
LoraFilesChanged,
|
||||
LoraSettingsUnusable,
|
||||
EmptyReplacement,
|
||||
PairDependencyError,
|
||||
),
|
||||
):
|
||||
self.logger.warning(f"LoRA scan refused: {message}")
|
||||
else:
|
||||
self.logger.exception("LoRA scan failed")
|
||||
with contextlib.suppress(ConnectionError, RuntimeError):
|
||||
await send_progress(_lora_error_frame(message))
|
||||
await send_progress(
|
||||
{
|
||||
"type": "progress",
|
||||
"status": "Clearing previous lora-manager imports...",
|
||||
"progress": 0,
|
||||
}
|
||||
)
|
||||
|
||||
async def _lora_scan_stream(
|
||||
self, lm_path, send_progress, disconnected=lambda: False
|
||||
):
|
||||
"""Import every LoRA under ``lm_path``, reporting through ``send_progress``.
|
||||
|
||||
The staged path of the System import (#221, ``lora_import_routes``):
|
||||
discover, stage every record (downloads into the cache first), then
|
||||
replace the collection in one transaction. A failure at any stage
|
||||
raises with the previous imports untouched; an empty replacement of a
|
||||
non-empty collection is refused (it needs the reviewed import).
|
||||
"""
|
||||
from .. import lora_import
|
||||
# Clear previous imports so reimport is always clean
|
||||
await self._run_in_executor(self.db.delete_prompts_by_category, "lora-manager")
|
||||
|
||||
await send_progress(
|
||||
{
|
||||
@@ -447,11 +353,19 @@ class LoraIntegrationMixin:
|
||||
"progress": 2,
|
||||
}
|
||||
)
|
||||
discovery = await self._run_in_executor(lora_import.discover, lm_path)
|
||||
if discovery.layout.error is not None:
|
||||
raise LoraSettingsUnusable(discovery.layout.error)
|
||||
sidecars = discovery.sidecars
|
||||
total = len(sidecars)
|
||||
|
||||
lora_dirs = await self._run_in_executor(find_lora_directories, lm_path)
|
||||
|
||||
# Collect all metadata files
|
||||
meta_files = []
|
||||
for d in lora_dirs:
|
||||
dir_path = Path(d)
|
||||
meta_files.extend(dir_path.rglob("*.metadata.json"))
|
||||
|
||||
total = len(meta_files)
|
||||
imported = 0
|
||||
skipped = 0
|
||||
|
||||
await send_progress(
|
||||
{
|
||||
"type": "progress",
|
||||
@@ -461,105 +375,116 @@ class LoraIntegrationMixin:
|
||||
}
|
||||
)
|
||||
|
||||
reviewed = await self._run_in_executor(lora_import.fingerprint, discovery)
|
||||
events = asyncio.Queue()
|
||||
state = await self._run_in_executor(
|
||||
self.db.category_import_state, lora_import.LORA_CATEGORY
|
||||
)
|
||||
# Records whose prompt already exists outside the collection are
|
||||
# linked, not imported (the running counts; the end ones are exact).
|
||||
known = set()
|
||||
cache_dir = get_lora_image_cache_dir()
|
||||
|
||||
def on_record(index, record, _staged):
|
||||
events.put_nowait((index, record))
|
||||
for i, meta_file in enumerate(meta_files):
|
||||
metadata = await self._run_in_executor(read_lora_metadata, meta_file)
|
||||
if not metadata:
|
||||
skipped += 1
|
||||
continue
|
||||
|
||||
counts = {"imported": 0}
|
||||
model_name = get_model_name_from_metadata(metadata)
|
||||
trigger_words = get_trigger_words_from_metadata(metadata)
|
||||
|
||||
# Collect all images: local previews + downloaded civitai examples
|
||||
preview_paths = await self._run_in_executor(
|
||||
get_preview_images_from_metadata, metadata, meta_file
|
||||
)
|
||||
civitai_paths = await self._run_in_executor(
|
||||
download_civitai_images,
|
||||
metadata,
|
||||
meta_file,
|
||||
cache_dir,
|
||||
IntegrationConfig.CIVITAI_API_KEY,
|
||||
)
|
||||
|
||||
# Merge, local first, dedup
|
||||
seen = set(preview_paths)
|
||||
all_images = list(preview_paths)
|
||||
for cp in civitai_paths:
|
||||
if cp not in seen:
|
||||
all_images.append(cp)
|
||||
seen.add(cp)
|
||||
|
||||
# Build prompt text: prefer example prompt, then model name
|
||||
example_prompt = get_example_prompt_from_metadata(metadata)
|
||||
prompt_text = example_prompt or model_name
|
||||
|
||||
# Build tags
|
||||
tags = ["lora-manager", f"lora:{model_name}"]
|
||||
tags.extend(trigger_words)
|
||||
|
||||
# Save to database via existing mechanism
|
||||
try:
|
||||
import hashlib
|
||||
|
||||
prompt_hash = hashlib.sha256(
|
||||
prompt_text.strip().lower().encode("utf-8")
|
||||
).hexdigest()
|
||||
|
||||
async def report(index, record):
|
||||
if record is not None and record["hash"] not in known:
|
||||
known.add(record["hash"])
|
||||
existing = await self._run_in_executor(
|
||||
self.db.existing_prompt_hashes,
|
||||
[record["hash"]],
|
||||
lora_import.LORA_CATEGORY,
|
||||
self.db.get_prompt_by_hash, prompt_hash
|
||||
)
|
||||
if not existing:
|
||||
counts["imported"] += 1
|
||||
status = record["name"] if record else sidecars[index].metadata_path.name
|
||||
if record and record["images"]:
|
||||
status += f" ({len(record['images'])} images)"
|
||||
|
||||
if existing:
|
||||
# Link all images
|
||||
for pp in all_images:
|
||||
await self._run_in_executor(
|
||||
self.db.link_image_to_prompt,
|
||||
existing["id"],
|
||||
pp,
|
||||
)
|
||||
skipped += 1
|
||||
else:
|
||||
prompt_id = await self._run_in_executor(
|
||||
self.db.save_prompt,
|
||||
prompt_text,
|
||||
"lora-manager", # category
|
||||
tags,
|
||||
None, # rating
|
||||
None, # notes
|
||||
prompt_hash,
|
||||
)
|
||||
|
||||
if prompt_id:
|
||||
for pp in all_images:
|
||||
await self._run_in_executor(
|
||||
self.db.link_image_to_prompt,
|
||||
prompt_id,
|
||||
pp,
|
||||
)
|
||||
imported += 1
|
||||
else:
|
||||
skipped += 1
|
||||
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to import LoRA {model_name}: {e}")
|
||||
skipped += 1
|
||||
|
||||
# Progress update for every LoRA
|
||||
progress = int(5 + (90 * (i + 1) / max(total, 1)))
|
||||
img_count = len(all_images)
|
||||
status = f"{model_name}"
|
||||
if img_count:
|
||||
status += f" ({img_count} images)"
|
||||
await send_progress(
|
||||
{
|
||||
"type": "progress",
|
||||
"status": status,
|
||||
"progress": int(5 + (90 * (index + 1) / max(total, 1))),
|
||||
"processed": index + 1,
|
||||
"imported": counts["imported"],
|
||||
"skipped": index + 1 - counts["imported"],
|
||||
"progress": progress,
|
||||
"processed": i + 1,
|
||||
"imported": imported,
|
||||
"skipped": skipped,
|
||||
}
|
||||
)
|
||||
|
||||
# Each record's event is queued (through the loop) before staging
|
||||
# ends, so reporting until staging is done and then draining the
|
||||
# queue reports every record in order. I2: when the page goes away
|
||||
# (its connection closes, or a report cannot be sent) staging stops
|
||||
# before its next LoRA, and this returns only once it has stopped, so
|
||||
# the caller's records hold covers every download.
|
||||
stop = threading.Event()
|
||||
staging = asyncio.ensure_future(
|
||||
self.stage_lora_import(
|
||||
discovery,
|
||||
on_record=on_record,
|
||||
cancelled=lambda: stop.is_set() or disconnected(),
|
||||
)
|
||||
)
|
||||
try:
|
||||
while not staging.done():
|
||||
getter = asyncio.ensure_future(events.get())
|
||||
await asyncio.wait(
|
||||
{getter, staging}, return_when=asyncio.FIRST_COMPLETED
|
||||
)
|
||||
if getter.done():
|
||||
await report(*getter.result())
|
||||
else:
|
||||
getter.cancel()
|
||||
while not events.empty():
|
||||
await report(*events.get_nowait())
|
||||
finally:
|
||||
stop.set()
|
||||
# Staging must have stopped before the records hold is released,
|
||||
# even when the handler is cancelled again while it waits (R3).
|
||||
interrupted = False
|
||||
while not staging.done():
|
||||
try:
|
||||
await asyncio.shield(staging)
|
||||
except asyncio.CancelledError:
|
||||
interrupted = True
|
||||
except Exception:
|
||||
pass
|
||||
if interrupted:
|
||||
raise asyncio.CancelledError
|
||||
staged = await staging
|
||||
if not await self.lora_files_unchanged(lm_path, reviewed):
|
||||
raise LoraFilesChanged(_FILES_CHANGED)
|
||||
|
||||
await send_progress(
|
||||
{
|
||||
"type": "progress",
|
||||
"status": "Replacing previous lora-manager imports...",
|
||||
"progress": 96,
|
||||
}
|
||||
)
|
||||
result = await self.replace_lora_imports(
|
||||
self.db, staged, expected_ids=state["ids"], allow_empty=False
|
||||
)
|
||||
await send_progress(
|
||||
{
|
||||
"type": "complete",
|
||||
"progress": 100,
|
||||
"total": total,
|
||||
"imported": result["imported"],
|
||||
"skipped": total - result["imported"],
|
||||
"imported": imported,
|
||||
"skipped": skipped,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
+169
-88
@@ -37,8 +37,7 @@ runner runs); a refused start keeps it. The runner:
|
||||
|
||||
The job log names the backup by its file name only, and the job says "No
|
||||
safety backup" when none was taken. Apply errors: ``invalid-request`` (400),
|
||||
``apply-failed`` (500: nothing started), ``plan-expired`` (410: unknown,
|
||||
expired or forgotten after a restart),
|
||||
``plan-expired`` (410: unknown, expired or forgotten after a restart),
|
||||
``plan-used`` (409, with the job that applied it), ``storage-changed`` (409:
|
||||
a restore replaced the database since the review) and ``busy`` (409).
|
||||
"""
|
||||
@@ -48,6 +47,8 @@ import json
|
||||
import os
|
||||
import shutil
|
||||
import sqlite3
|
||||
import tempfile
|
||||
from contextlib import closing
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
@@ -57,32 +58,16 @@ from .jobs import MAINTENANCE_JOB_KIND, JobBusy
|
||||
from .maintenance_plans import PlanExpired, PlanUsed, ScopeTooLarge
|
||||
|
||||
try:
|
||||
from ...database import safety_backup as _backups
|
||||
from ...database.safety_backup import ( # noqa: F401 - re-exported
|
||||
SAFETY_TEMP_PREFIX,
|
||||
SafetyBackupFailed,
|
||||
sweep_safety_temps,
|
||||
)
|
||||
from ...database.workflow_storage import PRE_MIGRATION_LABEL, pre_migration_backups
|
||||
from ...database.health import readonly_uri
|
||||
from ...database.maintenance import MAINTENANCE_OPERATIONS, PlanChanged
|
||||
except ImportError:
|
||||
from database import safety_backup as _backups
|
||||
from database.safety_backup import ( # noqa: F401 - re-exported
|
||||
SAFETY_TEMP_PREFIX,
|
||||
SafetyBackupFailed,
|
||||
sweep_safety_temps,
|
||||
)
|
||||
from database.workflow_storage import PRE_MIGRATION_LABEL, pre_migration_backups
|
||||
from database.health import readonly_uri
|
||||
from database.maintenance import MAINTENANCE_OPERATIONS, PlanChanged
|
||||
|
||||
PREVIEW_BODY_MAX = 4096
|
||||
MAINTENANCE_PREVIEW_LEASE = "Maintenance preview"
|
||||
# Duplicate-file previews carry the reviewed keepers (py/api/duplicate_files.py).
|
||||
DUPLICATE_PREVIEW_BODY_MAX = 1024 * 1024
|
||||
DUPLICATE_FILES = "duplicate_files"
|
||||
WORKFLOW_STORAGE = "workflow_storage" # applied by py/api/workflow_storage.py
|
||||
PREVIEW_OPERATIONS = MAINTENANCE_OPERATIONS + (DUPLICATE_FILES,)
|
||||
# Room left over beyond what a step writes, so a disk is never filled to 0.
|
||||
DISK_MARGIN_BYTES = 16 * 2**20
|
||||
@@ -92,7 +77,6 @@ MAINTENANCE_TITLES = {
|
||||
"optimize": "Optimize database",
|
||||
"unused_components": "Unused prompt components",
|
||||
"missing_links": "Missing-file links",
|
||||
"workflow_storage": "Store each workflow once",
|
||||
}
|
||||
_OPERATION_ERROR = "operation must be one of: " + ", ".join(PREVIEW_OPERATIONS)
|
||||
MAINTENANCE_GROUPS = ("changed", "protected", "errors")
|
||||
@@ -105,16 +89,16 @@ PLAN_CHANGED = (
|
||||
"Review the plan again."
|
||||
)
|
||||
APPLY_FAILED = "Maintenance failed; nothing was changed."
|
||||
NOTHING_TO_DO = "Every image record already stores its workflow once; nothing to do."
|
||||
NO_BACKUP = "No safety backup"
|
||||
_NAME_ATTEMPTS = 1000
|
||||
|
||||
|
||||
def backup_failed_message(error):
|
||||
"""The job line for a failed backup, naming a damaged workflow document."""
|
||||
reason = getattr(error, "reason", None) or ""
|
||||
if reason.startswith("Workflow document "):
|
||||
return f"{BACKUP_FAILED} {reason}"
|
||||
return BACKUP_FAILED
|
||||
class SafetyBackupFailed(Exception):
|
||||
"""The safety backup could not be written, verified or put in place."""
|
||||
|
||||
|
||||
def _now():
|
||||
return datetime.datetime.now()
|
||||
|
||||
|
||||
def _free_bytes(path):
|
||||
@@ -176,7 +160,7 @@ def _megabytes(size):
|
||||
return f"{size / 2**20:,.1f} MB"
|
||||
|
||||
|
||||
def disk_plan(db_path, *, backup, vacuum, vacuum_bytes=None):
|
||||
def disk_plan(db_path, *, backup, vacuum):
|
||||
"""Whether the disk has room for an apply's writes (blocking).
|
||||
|
||||
``backup``: the safety backup, one database-sized file beside the
|
||||
@@ -184,9 +168,7 @@ def disk_plan(db_path, *, backup, vacuum, vacuum_bytes=None):
|
||||
beside it and builds it in a temporary copy in SQLite's temporary folder
|
||||
(``sqlite_temp_dir``; SQLite's documentation: up to twice the database
|
||||
in free space). Needs
|
||||
on one device add up. ``vacuum_bytes``, when given, is what VACUUM
|
||||
copies (the live pages, for Reclaim space) instead of the whole file.
|
||||
Returns ``(public, short)``: ``public`` is
|
||||
on one device add up. Returns ``(public, short)``: ``public`` is
|
||||
``{free_bytes, required_bytes, sufficient}`` for the database folder,
|
||||
``short`` None or ``(place, what, required, free)`` of the first place
|
||||
that is short, and ``size`` the database's size (0 when nothing is
|
||||
@@ -199,10 +181,8 @@ def disk_plan(db_path, *, backup, vacuum, vacuum_bytes=None):
|
||||
if backup:
|
||||
needs.append(("database folder", folder, size, "the safety backup"))
|
||||
if vacuum:
|
||||
copy = size if vacuum_bytes is None else vacuum_bytes
|
||||
what = "Optimize" if vacuum_bytes is None else "Reclaim space"
|
||||
needs.append(("database folder", folder, copy, what))
|
||||
needs.append(("temporary folder", sqlite_temp_dir(), copy, what))
|
||||
needs.append(("database folder", folder, size, "Optimize"))
|
||||
needs.append(("temporary folder", sqlite_temp_dir(), size, "Optimize"))
|
||||
devices = {}
|
||||
for place, path, needed, what in needs:
|
||||
entry = devices.setdefault(_device(path), {"place": place, "path": path})
|
||||
@@ -239,18 +219,29 @@ def size_label(size):
|
||||
|
||||
|
||||
def restore_room(db_path):
|
||||
"""The largest restore upload the disk has room for (blocking).
|
||||
"""The largest restore upload the disks have room for (blocking).
|
||||
|
||||
Everything a restore writes is in the library's folder (#318): the
|
||||
upload, its staged copy (``database/restore_stage.py``) and the WAL
|
||||
growth while the stage is copied over the live database (3x the upload),
|
||||
plus the safety backup of the current database (1x its size), with
|
||||
``DISK_MARGIN_BYTES`` to spare.
|
||||
The upload is written to a temp file and staged once more in the
|
||||
temporary folder (2x), then copied over the live database beside it,
|
||||
whose current contents are first kept as a safety backup (1x the upload
|
||||
for the WAL, plus the current database). Needs on one device add up,
|
||||
each with ``DISK_MARGIN_BYTES``.
|
||||
"""
|
||||
db_path = os.path.abspath(db_path)
|
||||
folder = os.path.dirname(db_path)
|
||||
current = database_bytes(db_path) if os.path.exists(db_path) else 0
|
||||
room = (_free_bytes(folder) - current - DISK_MARGIN_BYTES) // 3
|
||||
devices = {}
|
||||
for path, per_upload, fixed in (
|
||||
(tempfile.gettempdir(), 2, 0),
|
||||
(folder, 1, current),
|
||||
):
|
||||
entry = devices.setdefault(_device(path), {"path": path, "per": 0, "fixed": 0})
|
||||
entry["per"] += per_upload
|
||||
entry["fixed"] += fixed
|
||||
room = min(
|
||||
(_free_bytes(e["path"]) - e["fixed"] - DISK_MARGIN_BYTES) // e["per"]
|
||||
for e in devices.values()
|
||||
)
|
||||
return max(0, room)
|
||||
|
||||
|
||||
@@ -290,13 +281,139 @@ def _changed_line(operation, count):
|
||||
return f"Removed {count:,} {noun}"
|
||||
|
||||
|
||||
def take_safety_backup(model, label="backup"):
|
||||
def _remove_quietly(*paths):
|
||||
for path in paths:
|
||||
try:
|
||||
if path and os.path.exists(path):
|
||||
os.remove(path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _reserve_backup_name(db_path):
|
||||
"""Create and return an unused ``<db>.backup_<stamp>[_n]`` (O_EXCL)."""
|
||||
for candidate in _backup_names(db_path):
|
||||
try:
|
||||
fd = os.open(candidate, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
|
||||
except FileExistsError:
|
||||
continue
|
||||
os.close(fd)
|
||||
return candidate
|
||||
raise SafetyBackupFailed("No free backup name")
|
||||
|
||||
|
||||
def _backup_names(db_path):
|
||||
base = f"{db_path}.backup_{_now().strftime('%Y%m%d_%H%M%S_%f')}"
|
||||
for attempt in range(_NAME_ATTEMPTS):
|
||||
yield base if attempt == 0 else f"{base}_{attempt}"
|
||||
|
||||
|
||||
def _place_backup(temp, db_path):
|
||||
"""Give the verified ``temp`` its unique backup name; return that path.
|
||||
|
||||
``os.link`` is atomic and refuses an existing name, so no empty or half
|
||||
file ever carries a backup name, even across a crash. A filesystem
|
||||
without hard links falls back to reserving the name (``O_EXCL``) and
|
||||
``os.replace``. The temp file is removed by the caller.
|
||||
"""
|
||||
for candidate in _backup_names(db_path):
|
||||
try:
|
||||
os.link(temp, candidate)
|
||||
except FileExistsError:
|
||||
continue
|
||||
except (OSError, NotImplementedError):
|
||||
break # no hard links here
|
||||
else:
|
||||
return candidate
|
||||
else:
|
||||
raise SafetyBackupFailed("No free backup name")
|
||||
final = _reserve_backup_name(db_path)
|
||||
try:
|
||||
os.replace(temp, final)
|
||||
except OSError:
|
||||
_remove_quietly(final)
|
||||
raise
|
||||
return final
|
||||
|
||||
|
||||
SAFETY_TEMP_PREFIX = ".pm-safety-"
|
||||
|
||||
|
||||
def sweep_safety_temps(db_path):
|
||||
"""Remove the temp copies an interrupted safety backup left (blocking).
|
||||
|
||||
Called once at startup, before any job can take a backup. Returns the
|
||||
names removed.
|
||||
"""
|
||||
folder = os.path.dirname(os.path.abspath(db_path))
|
||||
removed = []
|
||||
try:
|
||||
names = os.listdir(folder)
|
||||
except OSError:
|
||||
return removed
|
||||
for name in names:
|
||||
if not (name.startswith(SAFETY_TEMP_PREFIX) and ".tmp" in name):
|
||||
continue
|
||||
try:
|
||||
os.remove(os.path.join(folder, name))
|
||||
removed.append(name)
|
||||
except OSError:
|
||||
continue
|
||||
return removed
|
||||
|
||||
|
||||
def _standalone(path):
|
||||
"""Make the copy one self-contained file (rollback journal, not WAL).
|
||||
|
||||
The backup API copies the live database's WAL flag; a WAL file opened
|
||||
read-only later (a verify, a restore preview) would leave ``-wal`` and
|
||||
``-shm`` sidecars beside it.
|
||||
"""
|
||||
try:
|
||||
with closing(sqlite3.connect(path)) as conn:
|
||||
conn.execute("PRAGMA journal_mode=DELETE").fetchone()
|
||||
except sqlite3.Error as exc:
|
||||
raise SafetyBackupFailed("The backup could not be finished") from exc
|
||||
|
||||
|
||||
def _fsync(path):
|
||||
with open(path, "rb+") as handle:
|
||||
os.fsync(handle.fileno())
|
||||
|
||||
|
||||
def take_safety_backup(model):
|
||||
"""Write and verify a safety backup; return its file name (blocking).
|
||||
|
||||
See ``database/safety_backup.py``: a backup file on disk is always a
|
||||
verified one. Raises SafetyBackupFailed, leaving no file.
|
||||
The copy is made in a temp file beside the database and only renamed
|
||||
onto its final, unique name once it verified, so a backup file on disk
|
||||
is always a verified one. Raises SafetyBackupFailed, leaving no file.
|
||||
"""
|
||||
return os.path.basename(_backups.take_safety_backup(model, label))
|
||||
db_path = os.path.abspath(model.db_path)
|
||||
fd, temp = tempfile.mkstemp(
|
||||
dir=os.path.dirname(db_path), prefix=SAFETY_TEMP_PREFIX, suffix=".tmp"
|
||||
)
|
||||
os.close(fd)
|
||||
final = None
|
||||
try:
|
||||
if not model.backup_database(temp):
|
||||
raise SafetyBackupFailed("The backup could not be written")
|
||||
_standalone(temp)
|
||||
ok, _reason = model.verify_database_file(temp)
|
||||
if not ok:
|
||||
raise SafetyBackupFailed("The backup did not verify")
|
||||
_fsync(temp)
|
||||
if os.name != "nt":
|
||||
os.chmod(temp, 0o600)
|
||||
final = _place_backup(temp, db_path)
|
||||
return os.path.basename(final)
|
||||
except OSError as exc:
|
||||
_remove_quietly(final)
|
||||
raise SafetyBackupFailed("The backup could not be put in place") from exc
|
||||
except SafetyBackupFailed:
|
||||
_remove_quietly(final)
|
||||
raise
|
||||
finally:
|
||||
_remove_quietly(temp, temp + "-journal", temp + "-wal", temp + "-shm")
|
||||
|
||||
|
||||
def _failure(message, status, code=None, **extra):
|
||||
@@ -359,19 +476,6 @@ class MaintenanceRoutesMixin:
|
||||
return data, None
|
||||
|
||||
async def preview_maintenance(self, request):
|
||||
"""Review one operation under a registry lease (#249 S8).
|
||||
|
||||
The preview's read-only connections must not be open while a restore
|
||||
switches the database: it is refused (409 ``busy``) while an
|
||||
``exclusive`` job runs, and a restore waits for running previews.
|
||||
"""
|
||||
try:
|
||||
async with self._jobs.lease(MAINTENANCE_PREVIEW_LEASE):
|
||||
return await self._preview_maintenance(request)
|
||||
except JobBusy as busy:
|
||||
return busy_response(busy)
|
||||
|
||||
async def _preview_maintenance(self, request):
|
||||
data, error = await self._read_maintenance_body(
|
||||
request,
|
||||
("operation", "options"),
|
||||
@@ -415,7 +519,6 @@ class MaintenanceRoutesMixin:
|
||||
"targets": draft.targets,
|
||||
"roots": roots,
|
||||
"protected": draft.protected["count"],
|
||||
"rows": draft.scope["count"],
|
||||
},
|
||||
)
|
||||
|
||||
@@ -435,18 +538,11 @@ class MaintenanceRoutesMixin:
|
||||
"storage-changed",
|
||||
)
|
||||
policy = self._safety_backup_policy()
|
||||
destination = "backup"
|
||||
if operation == WORKFLOW_STORAGE:
|
||||
# Always backed up first, whatever the policy (the way back);
|
||||
# one kept pre-migration backup ever, pooled ones after it.
|
||||
policy = "required"
|
||||
kept = await self._run_in_executor(pre_migration_backups, db.model.db_path)
|
||||
destination = "backup" if kept else PRE_MIGRATION_LABEL
|
||||
try:
|
||||
disk, _short, size = await self._run_in_executor(
|
||||
disk_plan,
|
||||
db.model.db_path,
|
||||
backup=policy != "off",
|
||||
backup=policy == "on",
|
||||
vacuum=operation == "optimize",
|
||||
)
|
||||
except (OSError, sqlite3.Error):
|
||||
@@ -471,8 +567,8 @@ class MaintenanceRoutesMixin:
|
||||
**public,
|
||||
"backup": {
|
||||
"policy": policy,
|
||||
"destination_label": f"{name}.{destination}_<date and time>",
|
||||
"expected_bytes": size if policy != "off" else None,
|
||||
"destination_label": f"{name}.backup_<date and time>",
|
||||
"expected_bytes": size if policy == "on" else None,
|
||||
},
|
||||
"disk": disk,
|
||||
},
|
||||
@@ -500,9 +596,6 @@ class MaintenanceRoutesMixin:
|
||||
409,
|
||||
"storage-changed",
|
||||
)
|
||||
if plan.operation == WORKFLOW_STORAGE and not plan.payload.get("rows"):
|
||||
store.discard(token) # no job, no backup: nothing to rewrite
|
||||
return _failure(NOTHING_TO_DO, 409, "nothing-to-do")
|
||||
backup = self._backup_promised(plan)
|
||||
try:
|
||||
_disk, short, _size = await self._run_in_executor(
|
||||
@@ -513,9 +606,7 @@ class MaintenanceRoutesMixin:
|
||||
)
|
||||
except (OSError, sqlite3.Error):
|
||||
self.logger.exception("Could not measure the disk space for an apply")
|
||||
# Coded: the workspace tells the server's own refusal (nothing
|
||||
# started) from a proxy's 5xx (the job may have started).
|
||||
return _failure(APPLY_FAILED, 500, "apply-failed")
|
||||
return _failure(APPLY_FAILED, 500)
|
||||
if short is not None:
|
||||
# Refused before any job starts: nothing fails halfway for space.
|
||||
return _failure(short_space_message(short), 507, "insufficient-space")
|
||||
@@ -527,8 +618,6 @@ class MaintenanceRoutesMixin:
|
||||
try:
|
||||
if plan.operation == DUPLICATE_FILES:
|
||||
job = self._start_duplicate_cleanup(db, plan, backup)
|
||||
elif plan.operation == WORKFLOW_STORAGE:
|
||||
job = self._start_workflow_storage(db, backup)
|
||||
else:
|
||||
job = self._jobs.start(
|
||||
MAINTENANCE_JOB_KIND,
|
||||
@@ -574,8 +663,6 @@ class MaintenanceRoutesMixin:
|
||||
Decided once, when the apply starts (the disk check measured it), and
|
||||
handed to the job.
|
||||
"""
|
||||
if plan.operation == WORKFLOW_STORAGE:
|
||||
return True # never without its pre-migration backup
|
||||
return plan.payload["backup"] == "on" or self._safety_backup_policy() == "on"
|
||||
|
||||
async def _safety_backup_step(self, job, db, backup_on):
|
||||
@@ -592,9 +679,9 @@ class MaintenanceRoutesMixin:
|
||||
job.log_line("Taking a safety backup of the database")
|
||||
try:
|
||||
backup = await self._run_in_executor(take_safety_backup, db.model)
|
||||
except SafetyBackupFailed as error:
|
||||
except SafetyBackupFailed:
|
||||
self.logger.error("Safety backup failed", exc_info=True)
|
||||
_fail_job(job, "backup-failed", backup_failed_message(error))
|
||||
_fail_job(job, "backup-failed", BACKUP_FAILED)
|
||||
return False, None
|
||||
job.log_line(f"Safety backup {backup} verified")
|
||||
return True, backup
|
||||
@@ -679,13 +766,7 @@ def backup_note(backup):
|
||||
|
||||
|
||||
def _fail_job(job, code, message):
|
||||
"""End ``job`` failed with ``message`` (one error row) and its code.
|
||||
|
||||
The code is also the job's ``scope["failure"]``, so Activity can offer
|
||||
Review again for a plan the library overtook.
|
||||
"""
|
||||
if code:
|
||||
job.scope["failure"] = code
|
||||
"""End ``job`` failed with ``message`` (one error row) and its code."""
|
||||
job.record("errors", {"path": "", "detail": message}, f"! {message}")
|
||||
job.update(status_line=message)
|
||||
job.publish({"type": "error", "code": code, "message": message})
|
||||
|
||||
@@ -1,146 +0,0 @@
|
||||
"""Reviewed restores: staged uploads held under single-use tokens (#249 S8).
|
||||
|
||||
A restore preview stages the upload (``database/restore_stage.py``) and
|
||||
keeps it here under a ``secrets.token_urlsafe(32)`` token with the live
|
||||
library's fingerprint and the database identity it was reviewed against.
|
||||
``take`` hands the stage to the restore job exactly once; a second apply of
|
||||
the token names that job (``PlanUsed``). An expired, discarded or unknown
|
||||
token raises ``PlanExpired``; its stage folder is removed. A restart forgets
|
||||
every token (the startup sweep removes the folders).
|
||||
|
||||
At most ``max_stages`` stages are held. Bytes are bounded by free disk space
|
||||
(the upload is refused when the disk has no room, and held stages use that
|
||||
disk), not by a fixed budget. A stage being applied belongs to its job and is
|
||||
never evicted or counted.
|
||||
|
||||
The store is only read and changed on the event loop (it has no lock). Its
|
||||
entries own stage folders, whose removal is blocking: every method that
|
||||
drops an entry returns it (or keeps it in ``take_doomed``), and the caller
|
||||
removes the folder in the executor (``HeldRestore.staged.discard``).
|
||||
"""
|
||||
|
||||
import secrets
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
from .maintenance_plans import PlanExpired, PlanUsed
|
||||
|
||||
STAGE_TTL_SECONDS = 900
|
||||
MAX_STAGES = 2
|
||||
_USED_KEPT = 32 # applied tokens remembered for plan-used answers
|
||||
|
||||
|
||||
class TooManyStages(Exception):
|
||||
"""Every stage slot is held by a reviewed restore."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class HeldRestore:
|
||||
staged: Any # database.restore_stage.StagedRestore
|
||||
fingerprint: str
|
||||
db_identity: Any
|
||||
live: Dict[str, int]
|
||||
expires_at: float
|
||||
|
||||
|
||||
class RestorePlanStore:
|
||||
def __init__(
|
||||
self,
|
||||
clock: Optional[Callable[[], float]] = None,
|
||||
ttl: float = STAGE_TTL_SECONDS,
|
||||
max_stages: int = MAX_STAGES,
|
||||
):
|
||||
self._ttl = ttl
|
||||
self._max = max_stages
|
||||
self._clock = clock or time.time
|
||||
self._held: Dict[str, HeldRestore] = {}
|
||||
self._used: Dict[str, str] = {}
|
||||
self._doomed: List[HeldRestore] = []
|
||||
|
||||
@property
|
||||
def ttl(self) -> float:
|
||||
return self._ttl
|
||||
|
||||
def now(self) -> float:
|
||||
return self._clock()
|
||||
|
||||
def __len__(self):
|
||||
return len(self._held)
|
||||
|
||||
def full(self) -> bool:
|
||||
return len(self._held) >= self._max
|
||||
|
||||
def issue(self, staged, fingerprint, db_identity, live) -> str:
|
||||
"""Hold ``staged`` and return its token; TooManyStages when full."""
|
||||
if self.full():
|
||||
raise TooManyStages()
|
||||
token = secrets.token_urlsafe(32)
|
||||
self._held[token] = HeldRestore(
|
||||
staged=staged,
|
||||
fingerprint=fingerprint,
|
||||
db_identity=db_identity,
|
||||
live=dict(live),
|
||||
expires_at=self._clock() + self._ttl,
|
||||
)
|
||||
return token
|
||||
|
||||
def peek(self, token) -> HeldRestore:
|
||||
if token in self._used:
|
||||
raise PlanUsed(self._used[token])
|
||||
held = self._held.get(token)
|
||||
if held is None:
|
||||
raise PlanExpired()
|
||||
if held.expires_at <= self._clock():
|
||||
self._doomed.append(self._held.pop(token))
|
||||
raise PlanExpired()
|
||||
return held
|
||||
|
||||
def take(self, token, job_id) -> HeldRestore:
|
||||
"""Hand the held stage to the job applying it (single use)."""
|
||||
held = self.peek(token)
|
||||
del self._held[token]
|
||||
self._used[token] = job_id
|
||||
while len(self._used) > _USED_KEPT:
|
||||
self._used.pop(next(iter(self._used)))
|
||||
return held
|
||||
|
||||
def give_back(self, token, held) -> bool:
|
||||
"""Hold a taken stage again under its token (its job did not apply it).
|
||||
|
||||
False, holding nothing, when the review expired meanwhile or its slot
|
||||
was reused (the store is full): the caller then discards the stage.
|
||||
"""
|
||||
if held.expires_at <= self._clock() or self.full():
|
||||
return False
|
||||
self._used.pop(token, None)
|
||||
self._held[token] = held
|
||||
return True
|
||||
|
||||
def entry(self, token) -> Optional[HeldRestore]:
|
||||
"""The held entry, expired or not (for the response that issued it)."""
|
||||
return self._held.get(token)
|
||||
|
||||
def pop(self, token) -> Optional[HeldRestore]:
|
||||
"""Drop the token; its entry (to discard) or None when it was not held."""
|
||||
return self._held.pop(token, None)
|
||||
|
||||
def pop_expired(self) -> List[HeldRestore]:
|
||||
"""Drop every expired entry; them, with any ``peek`` found expired."""
|
||||
now = self._clock()
|
||||
expired = [t for t, held in self._held.items() if held.expires_at <= now]
|
||||
return [self._held.pop(token) for token in expired] + self.take_doomed()
|
||||
|
||||
def pop_all(self) -> List[HeldRestore]:
|
||||
"""Drop every entry (the library was replaced); them, to discard."""
|
||||
held = list(self._held.values()) + self.take_doomed()
|
||||
self._held.clear()
|
||||
return held
|
||||
|
||||
def take_doomed(self) -> List[HeldRestore]:
|
||||
doomed, self._doomed = self._doomed, []
|
||||
return doomed
|
||||
|
||||
def next_expiry(self) -> Optional[float]:
|
||||
"""When the soonest held entry expires (clock time), or None."""
|
||||
return min((held.expires_at for held in self._held.values()), default=None)
|
||||
+1
-19
@@ -31,8 +31,6 @@ from aiohttp import web
|
||||
from .annotations import _HEX64, _TOKEN, _UUID, _capture_error, _error
|
||||
from .images import media_thumbnail_url
|
||||
from .inspection_service import VIDEO_EXTENSIONS
|
||||
from .job_routes import busy_response
|
||||
from .jobs import JobBusy
|
||||
from .prompts import bad_request
|
||||
|
||||
try:
|
||||
@@ -241,9 +239,6 @@ def _side_label(titles: dict, side, empty: str) -> str:
|
||||
return display_title(titles.get(stored, ()), text) or empty
|
||||
|
||||
|
||||
REVIEW_DECISION_LEASE = "Needs review decision"
|
||||
|
||||
|
||||
class ReviewRoutesMixin:
|
||||
"""The Needs review queue over E1's schema and E2's recorded items."""
|
||||
|
||||
@@ -451,20 +446,7 @@ class ReviewRoutesMixin:
|
||||
# -- decision --------------------------------------------------------------
|
||||
|
||||
async def decide_review_item(self, request: web.Request) -> web.Response:
|
||||
"""POST one explicit decision on one pending item.
|
||||
|
||||
The decision runs under a registry lease (#249 S8): it is refused
|
||||
(409 ``busy``) while a restore or another ``exclusive`` job runs, and
|
||||
a restore waits for decisions already running before it switches the
|
||||
database. Decisions never refuse one another or a scan.
|
||||
"""
|
||||
try:
|
||||
async with self._jobs.lease(REVIEW_DECISION_LEASE):
|
||||
return await self._decide_review_item(request)
|
||||
except JobBusy as busy:
|
||||
return busy_response(busy)
|
||||
|
||||
async def _decide_review_item(self, request: web.Request) -> web.Response:
|
||||
"""POST one explicit decision on one pending item."""
|
||||
data, error = await self._read_json_body(request, max_bytes=MAX_BODY_BYTES)
|
||||
if error is not None:
|
||||
return error
|
||||
|
||||
+2
-158
@@ -13,7 +13,6 @@ checks one path and writes nothing.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import functools
|
||||
import os
|
||||
|
||||
from aiohttp import web
|
||||
@@ -24,11 +23,9 @@ from ..config_store import (
|
||||
get_config_store,
|
||||
)
|
||||
from . import sources
|
||||
from . import thumbnail_encode
|
||||
from . import thumbnail_settings
|
||||
from .admin import _current_image_monitor
|
||||
from .job_routes import busy_response
|
||||
from .jobs import THUMBNAILS_JOB_KIND, JobBusy
|
||||
from .jobs import JobBusy
|
||||
|
||||
try:
|
||||
from ...utils.validators import validate_result_timeout, validate_worker_threads
|
||||
@@ -44,26 +41,9 @@ _SETTINGS_FILE_KEYS = {
|
||||
"WORKER_THREADS": ("performance", "worker_threads"),
|
||||
"MONITORING_ENABLED": ("gallery", "monitoring", "enabled"),
|
||||
"SAFETY_BACKUP": ("maintenance", "safety_backup"),
|
||||
"THUMBNAIL_FORMAT": ("thumbnails", "format"),
|
||||
"THUMBNAIL_QUALITY": ("thumbnails", "quality"),
|
||||
"THUMBNAIL_EDGE": ("thumbnails", "edge"),
|
||||
# The one tagging configuration (#328), published to AutotagConfig.
|
||||
"AUTOTAG_MODEL_TYPE": ("autotag", "defaults", "model_type"),
|
||||
"AUTOTAG_EXISTING": ("autotag", "defaults", "existing"),
|
||||
"AUTOTAG_GENERAL_THRESHOLD": ("autotag", "defaults", "general_threshold"),
|
||||
"AUTOTAG_CHARACTER_THRESHOLD": ("autotag", "defaults", "character_threshold"),
|
||||
"AUTOTAG_TEMPLATE": ("autotag", "defaults", "template"),
|
||||
"AUTOTAG_SUGGESTION_LIMIT": ("autotag", "defaults", "suggestion_limit"),
|
||||
}
|
||||
# The thumbnail settings a running Convert outdated job fixes (#317).
|
||||
_THUMBNAIL_SETTINGS = ("THUMBNAIL_FORMAT", "THUMBNAIL_QUALITY", "THUMBNAIL_EDGE")
|
||||
# Attributes that live on GalleryConfig; the others on PromptManagerConfig.
|
||||
_GALLERY_SETTINGS = frozenset({"MONITORING_ENABLED"})
|
||||
# Attributes that live on AutotagConfig, named without the AUTOTAG_ prefix.
|
||||
_AUTOTAG_SETTINGS = frozenset(
|
||||
name for name in _SETTINGS_FILE_KEYS if name.startswith("AUTOTAG_")
|
||||
)
|
||||
_AUTOTAG_PREFIX = "AUTOTAG_"
|
||||
SOURCE_CHANGE_JOB_KIND = "source_change"
|
||||
SOURCE_CHANGE_JOB_TITLE = "Source change"
|
||||
|
||||
@@ -163,63 +143,6 @@ def _set_nested(config, sections, key, value):
|
||||
config[sections[0]] = merged
|
||||
|
||||
|
||||
def _validate_autotag_settings(data):
|
||||
"""The ``autotag_*`` settings as ``{AUTOTAG_ATTR: value}``.
|
||||
|
||||
The first bad value raises ValueError with a message that names the
|
||||
setting by its label in System › Auto-tagging.
|
||||
"""
|
||||
from ..config import (
|
||||
EXISTING_POLICIES,
|
||||
MAX_SUGGESTION_LIMIT,
|
||||
MODEL_TYPES,
|
||||
validate_template,
|
||||
validate_threshold,
|
||||
)
|
||||
|
||||
updates = {}
|
||||
if "autotag_model_type" in data:
|
||||
value = data["autotag_model_type"]
|
||||
if not isinstance(value, str) or value not in MODEL_TYPES:
|
||||
raise ValueError(
|
||||
"Tagging model must be one of " + ", ".join(sorted(MODEL_TYPES)) + "."
|
||||
)
|
||||
updates["AUTOTAG_MODEL_TYPE"] = value
|
||||
if "autotag_existing" in data:
|
||||
value = data["autotag_existing"]
|
||||
if not isinstance(value, str) or value not in EXISTING_POLICIES:
|
||||
raise ValueError(
|
||||
"Existing tags must be skip, merge or replace_after_review."
|
||||
)
|
||||
updates["AUTOTAG_EXISTING"] = value
|
||||
for key, label in (
|
||||
("autotag_general_threshold", "General-tag confidence"),
|
||||
("autotag_character_threshold", "Character-tag confidence"),
|
||||
):
|
||||
if key in data:
|
||||
try:
|
||||
value = validate_threshold(label, data[key])
|
||||
except ValueError as error:
|
||||
raise ValueError(f"{error}.") from None
|
||||
updates[key.upper()] = value
|
||||
if "autotag_template" in data:
|
||||
try:
|
||||
updates["AUTOTAG_TEMPLATE"] = validate_template(
|
||||
data["autotag_template"], "LLM tag template"
|
||||
)
|
||||
except ValueError as error:
|
||||
raise ValueError(f"{error}.") from None
|
||||
if "autotag_suggestion_limit" in data:
|
||||
value = data["autotag_suggestion_limit"]
|
||||
if type(value) is not int or not 1 <= value <= MAX_SUGGESTION_LIMIT:
|
||||
raise ValueError(
|
||||
"Suggestions per image must be a whole number from 1 to "
|
||||
f"{MAX_SUGGESTION_LIMIT}."
|
||||
)
|
||||
updates["AUTOTAG_SUGGESTION_LIMIT"] = value
|
||||
return updates
|
||||
|
||||
|
||||
async def _read_json(request):
|
||||
"""Parse a JSON object body through the package's capped reader."""
|
||||
from . import _read_json_body
|
||||
@@ -246,22 +169,19 @@ class SettingsRoutesMixin:
|
||||
async def get_settings(self, request):
|
||||
"""Get current settings."""
|
||||
try:
|
||||
from ..config import AutotagConfig, PromptManagerConfig, GalleryConfig
|
||||
from ..config import PromptManagerConfig, GalleryConfig
|
||||
|
||||
root_paths = [
|
||||
self._public_path(d) for d in GalleryConfig.MONITORING_DIRECTORIES
|
||||
]
|
||||
revision = await self._run_in_executor(self._config_revision)
|
||||
source_status, monitor = await self._run_in_executor(self._source_status)
|
||||
# #317: which thumbnail formats can be written (a cached probe).
|
||||
encoders = await self._run_in_executor(thumbnail_encode.encoder_status)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
**revision,
|
||||
"sources": source_status,
|
||||
"monitor": monitor,
|
||||
"thumbnail_encoders": encoders,
|
||||
"settings": {
|
||||
"monitoring_enabled": GalleryConfig.MONITORING_ENABLED,
|
||||
"result_timeout": PromptManagerConfig.RESULT_TIMEOUT,
|
||||
@@ -269,17 +189,6 @@ class SettingsRoutesMixin:
|
||||
"infinite_scroll": PromptManagerConfig.INFINITE_SCROLL,
|
||||
"worker_threads": PromptManagerConfig.WORKER_THREADS,
|
||||
"safety_backup": PromptManagerConfig.SAFETY_BACKUP,
|
||||
"thumbnail_format": PromptManagerConfig.THUMBNAIL_FORMAT,
|
||||
"thumbnail_quality": PromptManagerConfig.THUMBNAIL_QUALITY,
|
||||
"thumbnail_edge": PromptManagerConfig.THUMBNAIL_EDGE,
|
||||
"autotag_model_type": AutotagConfig.MODEL_TYPE,
|
||||
"autotag_existing": AutotagConfig.EXISTING,
|
||||
"autotag_general_threshold": AutotagConfig.GENERAL_THRESHOLD,
|
||||
"autotag_character_threshold": (
|
||||
AutotagConfig.CHARACTER_THRESHOLD
|
||||
),
|
||||
"autotag_template": AutotagConfig.TEMPLATE,
|
||||
"autotag_suggestion_limit": AutotagConfig.SUGGESTION_LIMIT,
|
||||
"cpu_count": PromptManagerConfig.max_worker_threads(),
|
||||
"gallery_root_paths": root_paths,
|
||||
"gallery_root_path": root_paths[0] if root_paths else "",
|
||||
@@ -386,26 +295,6 @@ class SettingsRoutesMixin:
|
||||
return web.json_response(
|
||||
{"success": False, "error": str(ve)}, status=400
|
||||
)
|
||||
if updates.get("THUMBNAIL_FORMAT") == "jxl":
|
||||
# #317: JPEG XL only with an encoder (a cached probe).
|
||||
encoder = await self._run_in_executor(
|
||||
functools.partial(
|
||||
thumbnail_encode.detect_jxl_encoder, recheck_missing=True
|
||||
)
|
||||
)
|
||||
if encoder is None:
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "jxl-unavailable",
|
||||
"error": thumbnail_encode.JXL_UNAVAILABLE,
|
||||
},
|
||||
status=400,
|
||||
)
|
||||
|
||||
refused = self._thumbnail_change_refused(updates)
|
||||
if refused is not None:
|
||||
return refused
|
||||
|
||||
# Off the loop: an accepted root is resolved (filesystem access).
|
||||
new_roots, error = await self._run_in_executor(
|
||||
@@ -466,33 +355,6 @@ class SettingsRoutesMixin:
|
||||
status=500,
|
||||
)
|
||||
|
||||
def _thumbnail_change_refused(self, updates):
|
||||
"""409 ``busy`` for a thumbnail format, quality or size change while
|
||||
Convert outdated runs (review I4): the job converts to the settings
|
||||
it started with, and a change under it would make it remove previews
|
||||
in the newly chosen format. Unchanged values pass."""
|
||||
from ..config import PromptManagerConfig
|
||||
|
||||
changed = any(
|
||||
key in updates and updates[key] != getattr(PromptManagerConfig, key)
|
||||
for key in _THUMBNAIL_SETTINGS
|
||||
)
|
||||
job = self._jobs.active(THUMBNAILS_JOB_KIND) if changed else None
|
||||
if job is None or job.scope.get("mode") != "convert":
|
||||
return None
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "busy",
|
||||
"error": (
|
||||
f"{job.title} is running. Change the thumbnail settings "
|
||||
"when it finishes, or cancel it first."
|
||||
),
|
||||
"running": {"id": job.id, "title": job.title},
|
||||
},
|
||||
status=409,
|
||||
)
|
||||
|
||||
def _publish_settings(self, updates, new_roots):
|
||||
"""Install committed settings at runtime (inside the config commit).
|
||||
|
||||
@@ -502,12 +364,7 @@ class SettingsRoutesMixin:
|
||||
"""
|
||||
from ..config import GalleryConfig, PromptManagerConfig
|
||||
|
||||
from ..config import AutotagConfig
|
||||
|
||||
for name, value in updates.items():
|
||||
if name in _AUTOTAG_SETTINGS:
|
||||
setattr(AutotagConfig, name[len(_AUTOTAG_PREFIX) :], value)
|
||||
continue
|
||||
owner = GalleryConfig if name in _GALLERY_SETTINGS else PromptManagerConfig
|
||||
setattr(owner, name, value)
|
||||
restart_fields = []
|
||||
@@ -634,19 +491,6 @@ class SettingsRoutesMixin:
|
||||
updates["WORKER_THREADS"] = validate_worker_threads(
|
||||
data["worker_threads"], PromptManagerConfig.max_worker_threads()
|
||||
)
|
||||
if "thumbnail_format" in data:
|
||||
updates["THUMBNAIL_FORMAT"] = thumbnail_settings.validate_format(
|
||||
data["thumbnail_format"]
|
||||
)
|
||||
if "thumbnail_quality" in data:
|
||||
updates["THUMBNAIL_QUALITY"] = thumbnail_settings.validate_quality(
|
||||
data["thumbnail_quality"]
|
||||
)
|
||||
if "thumbnail_edge" in data:
|
||||
updates["THUMBNAIL_EDGE"] = thumbnail_settings.validate_edge(
|
||||
data["thumbnail_edge"]
|
||||
)
|
||||
updates.update(_validate_autotag_settings(data))
|
||||
return updates
|
||||
|
||||
def _parse_gallery_roots(self, data):
|
||||
|
||||
+1
-25
@@ -13,9 +13,7 @@
|
||||
- ``describe_sources`` (blocking: it resolves paths) reports the configured
|
||||
sources with their validity, the effective sources with their listing
|
||||
index and watch state, and the image monitor's state. Every path in it is
|
||||
a public path (``_public_path``), never an absolute one; ``in_comfyui``
|
||||
says whether an effective source's public path is relative to the ComfyUI
|
||||
base (``in_comfyui_base``), so the page can read it as ``ComfyUI/<path>``.
|
||||
a public path (``_public_path``), never an absolute one.
|
||||
|
||||
A source's watch state is ``watching`` when the monitor's observer is alive,
|
||||
scheduled that directory and monitoring is on; ``stopping`` when it still
|
||||
@@ -211,27 +209,6 @@ def watch_state(path, *, present, alive, scheduled_keys, enabled):
|
||||
return "restart_pending"
|
||||
|
||||
|
||||
def in_comfyui_base(path):
|
||||
"""Whether ``_public_path(path)`` is relative to the ComfyUI base directory.
|
||||
|
||||
The base is the first anchor ``_public_path`` tries, so this agrees with
|
||||
it: the same canonical (realpath + normcase) spelling on both sides.
|
||||
"""
|
||||
from pathlib import Path
|
||||
|
||||
from ..config import _comfyui_base_directory
|
||||
|
||||
base = _comfyui_base_directory()
|
||||
if not path or not base:
|
||||
return False
|
||||
try:
|
||||
canonical = Path(os.path.normcase(os.path.realpath(str(path))))
|
||||
anchor = Path(os.path.normcase(os.path.realpath(base)))
|
||||
except (OSError, ValueError):
|
||||
return False
|
||||
return canonical.is_relative_to(anchor)
|
||||
|
||||
|
||||
def describe_sources(*, configured, output_dirs, auto_detected, monitor, enabled):
|
||||
"""``(sources, monitor)`` for GET /settings (blocking; public paths only).
|
||||
|
||||
@@ -256,7 +233,6 @@ def describe_sources(*, configured, output_dirs, auto_detected, monitor, enabled
|
||||
scheduled_keys=scheduled_keys,
|
||||
enabled=enabled,
|
||||
),
|
||||
"in_comfyui": in_comfyui_base(path),
|
||||
}
|
||||
for index, path in enumerate(output_dirs)
|
||||
]
|
||||
|
||||
+1
-4
@@ -168,10 +168,9 @@ class SystemRoutesMixin:
|
||||
from ..config import GalleryConfig
|
||||
|
||||
try:
|
||||
(output_dirs, auto_detected), counts, location = await _gather_settled(
|
||||
(output_dirs, auto_detected), counts = await _gather_settled(
|
||||
self._run_in_executor(self._output_dirs_with_origin),
|
||||
self._run_in_executor(self.db.library_counts),
|
||||
self._run_in_executor(self._library_payload),
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("System overview failed")
|
||||
@@ -195,8 +194,6 @@ class SystemRoutesMixin:
|
||||
"effective": len(output_dirs),
|
||||
"auto_detected": auto_detected,
|
||||
},
|
||||
# Where the library lives and why (py/api/library_location.py).
|
||||
"location": location,
|
||||
},
|
||||
headers={"Cache-Control": "no-store"},
|
||||
)
|
||||
|
||||
@@ -1,865 +0,0 @@
|
||||
"""Automatic thumbnails (#249): previews are a cache the server keeps itself.
|
||||
|
||||
The approved mock has no repeated global thumbnail prompt
|
||||
(``docs/concepts/workspace/README.md``), so the workspace never asks. Every
|
||||
listing gives an image or video a thumbnail URL (``media_thumbnail``), and the
|
||||
previews behind those URLs are made here, in three ways:
|
||||
|
||||
- **On first view.** ``GET /prompt_manager/images/serve/thumbnails/<rel>_thumb
|
||||
<ext>`` for a preview that does not exist yet is admitted on the event loop
|
||||
(``_PreviewGate``): at most ``ON_DEMAND_LIMIT`` previews are written at once,
|
||||
on this API's own small pool (never the default executor), one flight runs
|
||||
per file, and the requests waiting are capped (``MAX_WAITING``); past the
|
||||
cap a request gets its fallback at once. Inside the admission the preview is
|
||||
mapped back to its media file under the same root (the request text is
|
||||
checked first: no traversal, drive, UNC or empty segment; then a regular
|
||||
file inside the root, never through a link), made and served. When it
|
||||
cannot be made an image's original is served and a video's placeholder
|
||||
frame, both ``Cache-Control: no-store``: a video is never sent to an
|
||||
``<img>``. A file that cannot be decoded is remembered for this session
|
||||
(keyed by its path, size and mtime) and counted for the System Overview
|
||||
(``failed`` in ``/images/thumbnails/status``); I/O errors and a preview
|
||||
another program holds open (Windows) are not remembered. The recorded
|
||||
history's thumbnail route goes through the same admission and flight key.
|
||||
- **On capture.** The image monitor hands every new output it settles to
|
||||
``queue_capture_thumbnail``, which queues its preview for one background
|
||||
worker thread (``AutoThumbnailer``). Each wake takes at most Worker threads
|
||||
(``WORKER_THREADS``) files and writes them on that many threads; enqueuing
|
||||
is O(1), so the monitor and ComfyUI's queue never wait. A file already
|
||||
queued is not queued again, the queue is bounded (``MAX_QUEUED``) and a
|
||||
preview left out is made on its first view instead.
|
||||
- **After a scan.** A scan that added files checks the missing previews; from
|
||||
``SCAN_THUMBNAIL_JOB_THRESHOLD`` up it starts the thumbnail Activity job
|
||||
(missing only, visible and cancellable) instead of writing them silently,
|
||||
and fewer go through the worker.
|
||||
|
||||
Every preview is written by ``_generate_one`` with the server's thumbnail
|
||||
settings (format, quality, size limit; ``thumbnail_settings``, #317), into a
|
||||
temp sibling verified and renamed over the old one, so a reader never sees a
|
||||
partial file. A request for an older preview name (``_thumb.png``) is answered
|
||||
with the preview in the current format.
|
||||
|
||||
Holds: automatic previews only write files under ``thumbnails/`` and never
|
||||
touch the database, like the thumbnail job, so they run beside a scan, image
|
||||
tagging and the Needs review recording (S5 M7). They are not written while a
|
||||
job runs that replaces the database or removes media and previews: a restore,
|
||||
a duplicate removal, a sources change, the gallery's thumbnail clear or an
|
||||
``exclusive`` job (``_auto_thumbnail_blocker``). The holds are checked when a
|
||||
write starts (after its slot) and again just before its rename.
|
||||
|
||||
Shutdown: ``close_auto_thumbnails`` stops the worker at a file boundary,
|
||||
kills a running ffmpeg, waits a bounded time (and logs when the wait ran
|
||||
out) and shuts the view pool down; after the app's cleanup
|
||||
(``final=True``) nothing starts a worker again. Importing this module again
|
||||
(a package reload) stops the workers the previous import left.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import collections
|
||||
import functools
|
||||
import io
|
||||
import os
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from . import thumbnail_settings
|
||||
from .images import (
|
||||
MEDIA_EXTENSIONS,
|
||||
THUMBNAIL_CLEAR_TITLE,
|
||||
VIDEO_EXTENSIONS,
|
||||
_has_traversal,
|
||||
_is_within,
|
||||
_thumbnail_is_fresh,
|
||||
_video_placeholder_image,
|
||||
_worker_threads,
|
||||
media_thumbnail_rel,
|
||||
)
|
||||
from .jobs import (
|
||||
DUPLICATE_CLEANUP_JOB_KIND,
|
||||
RESTORE_JOB_KIND,
|
||||
THUMBNAILS_JOB_KIND,
|
||||
JobBusy,
|
||||
)
|
||||
from .settings_routes import SOURCE_CHANGE_JOB_KIND
|
||||
from .thumbnail_encode import shutdown_encoders, writer_settings
|
||||
from .thumbnail_jobs import ThumbnailOptions
|
||||
|
||||
try:
|
||||
from ...utils.parallel import map_parallel
|
||||
except ImportError: # tests import the package top-level
|
||||
from utils.parallel import map_parallel
|
||||
|
||||
# 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)
|
||||
PLACEHOLDER_QUALITY = 85
|
||||
# Source extensions in the order a shared preview name looks for them (R2).
|
||||
_SOURCE_ORDER = thumbnail_settings.IMAGE_SOURCE_ORDER + tuple(sorted(VIDEO_EXTENSIONS))
|
||||
# Previews views may write at once (the size of the view pool).
|
||||
ON_DEMAND_LIMIT = 3
|
||||
# View requests that may wait on the gate; one more gets its fallback at once.
|
||||
MAX_WAITING = 64
|
||||
# How long a pool thread waits for the capture worker's flight of the same file.
|
||||
FLIGHT_WAIT_SECONDS = 20
|
||||
# How long close() waits for the worker's batch in hand (ffmpeg's own limit).
|
||||
CLOSE_WAIT_SECONDS = 30
|
||||
# Files the capture worker holds; more are made on their first view.
|
||||
MAX_QUEUED = 10_000
|
||||
# Undecodable files remembered per session (oldest forgotten first).
|
||||
MAX_FAILURES_REMEMBERED = 4096
|
||||
# Missing previews after a scan from which the visible thumbnail job starts:
|
||||
# a fresh install or a newly added source. Fewer go through the worker.
|
||||
SCAN_THUMBNAIL_JOB_THRESHOLD = 200
|
||||
# Jobs beside which no automatic preview is written (see the module docstring).
|
||||
AUTO_REFUSED_KINDS = frozenset(
|
||||
{RESTORE_JOB_KIND, DUPLICATE_CLEANUP_JOB_KIND, SOURCE_CHANGE_JOB_KIND}
|
||||
)
|
||||
WORKER_NAME = "PromptManagerThumbnailWorker"
|
||||
POOL_PREFIX = "pm-preview"
|
||||
# Identifies this import's workers: a reload stops the ones it did not start.
|
||||
_MODULE_TOKEN = object()
|
||||
|
||||
READY, FAILED, BUSY, DEFERRED = "ready", "failed", "busy", "deferred"
|
||||
# What the gate answers a request past MAX_WAITING.
|
||||
OVER = object()
|
||||
|
||||
|
||||
def _fingerprint(src):
|
||||
"""``(path, size, mtime_ns)`` of *src*, or None when it cannot be read."""
|
||||
try:
|
||||
stat = os.stat(src)
|
||||
except OSError:
|
||||
return None
|
||||
return (os.path.normcase(str(src)), stat.st_size, stat.st_mtime_ns)
|
||||
|
||||
|
||||
def flight_key(root, rel_parts):
|
||||
"""The single-flight key of ``root/rel_parts``: its folder resolved, case
|
||||
folded. The serve route, the history route and the worker share it."""
|
||||
path = Path(root).joinpath(*rel_parts)
|
||||
return os.path.normcase(os.path.join(os.path.realpath(path.parent), path.name))
|
||||
|
||||
|
||||
def _dst_key(dst):
|
||||
return flight_key(Path(dst).parent, (Path(dst).name,))
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def video_placeholder_jpeg():
|
||||
"""The placeholder frame as JPEG bytes (a video whose preview is missing)."""
|
||||
buffer = io.BytesIO()
|
||||
_video_placeholder_image(PLACEHOLDER_SIZE).save(
|
||||
buffer, "JPEG", quality=PLACEHOLDER_QUALITY, optimize=True
|
||||
)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def _is_preview_request(filepath):
|
||||
"""Text check only: ``thumbnails/<...>/<name>_thumb<lower-case ext>``.
|
||||
|
||||
Refuses, as text and before any path is joined, what a joined path could
|
||||
take as an anchor: a backslash, a drive or stream colon, and empty, ``.``
|
||||
or ``..`` segments (so a leading ``/`` and ``//server/share`` too).
|
||||
"""
|
||||
if "\\" in filepath or ":" in filepath:
|
||||
return False
|
||||
parts = filepath.split("/")
|
||||
if len(parts) < 2 or parts[0] != "thumbnails":
|
||||
return False
|
||||
if any(part in ("", ".", "..") for part in parts) or "thumbnails" in parts[1:-1]:
|
||||
return False
|
||||
stem, extension = os.path.splitext(parts[-1])
|
||||
return (
|
||||
stem.endswith("_thumb")
|
||||
and len(stem) > len("_thumb")
|
||||
and extension == extension.lower()
|
||||
and extension in thumbnail_settings.THUMBNAIL_EXTENSIONS
|
||||
)
|
||||
|
||||
|
||||
def _regular_file(path):
|
||||
"""True for a regular file that is not a link (one lstat)."""
|
||||
try:
|
||||
return not path.is_symlink() and path.is_file()
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
def _candidates(parts):
|
||||
"""``(folder parts, base, source extensions)`` of a checked preview
|
||||
request: every media extension that owns the requested preview name
|
||||
(``thumbnail_settings.owns``), images first (ruling R2)."""
|
||||
stem, extension = os.path.splitext(parts[-1])
|
||||
wanted = [
|
||||
source
|
||||
for source in _SOURCE_ORDER
|
||||
if source in MEDIA_EXTENSIONS and thumbnail_settings.owns(extension, source)
|
||||
]
|
||||
return parts[1:-1], stem[: -len("_thumb")], wanted
|
||||
|
||||
|
||||
def _current_preview(root, parts):
|
||||
"""The preview *parts* names, in the format written now (#317): a request
|
||||
for an older name (``_thumb.png`` before a conversion) is answered with
|
||||
the current file."""
|
||||
_folder, base, _wanted = _candidates(parts)
|
||||
name = base + thumbnail_settings.THUMB_SUFFIX + writer_settings().extension
|
||||
return Path(root).joinpath(*parts[:-1], name)
|
||||
|
||||
|
||||
def _quick_source(root, parts):
|
||||
"""``(src, is_video)`` by exact name only (no folder listing), or None."""
|
||||
root = Path(root)
|
||||
folder_parts, base, wanted = _candidates(parts)
|
||||
folder = root.joinpath(*folder_parts)
|
||||
for extension in wanted:
|
||||
for name in (base + extension, base + extension.upper()):
|
||||
path = folder / name
|
||||
if _regular_file(path) and _is_within(path, root):
|
||||
return path, extension in VIDEO_EXTENSIONS
|
||||
return None
|
||||
|
||||
|
||||
def _thumbnail_source(root, parts):
|
||||
"""``(src, dst, is_video)`` for the preview *parts* under *root*, or None.
|
||||
|
||||
*parts* passed ``_is_preview_request``. The source has the preview's name
|
||||
without ``_thumb`` and an extension that owns the requested one (any case:
|
||||
exact names first, then one listing of the folder; ruling R4). It must be
|
||||
a regular file inside *root*. ``dst`` is the preview in the format written
|
||||
now (#317), whatever extension was requested; it is None when the
|
||||
preview's folder resolves outside ``root/thumbnails`` (a link): nothing is
|
||||
written there.
|
||||
"""
|
||||
root = Path(root)
|
||||
folder_parts, base, wanted = _candidates(parts)
|
||||
quick = _quick_source(root, parts)
|
||||
if quick is not None:
|
||||
found = quick[0]
|
||||
else:
|
||||
found = _scan_folder(root.joinpath(*folder_parts), base, set(wanted))
|
||||
if found is None or not _is_within(found, root):
|
||||
return None
|
||||
rel = Path(*folder_parts, found.name)
|
||||
dst = root / media_thumbnail_rel(rel)
|
||||
if dst != _current_preview(root, parts):
|
||||
return None
|
||||
if not _is_within(root / "thumbnails", root) or not _is_within(
|
||||
dst, root / "thumbnails"
|
||||
):
|
||||
dst = None
|
||||
return found, dst, found.suffix.lower() in VIDEO_EXTENSIONS
|
||||
|
||||
|
||||
def _scan_folder(folder, base, extensions):
|
||||
"""A mixed-case extension (``.Png``): one listing of *folder*, or None."""
|
||||
try:
|
||||
with os.scandir(folder) as entries:
|
||||
for entry in entries:
|
||||
stem, extension = os.path.splitext(entry.name)
|
||||
if stem == base and extension.lower() in extensions:
|
||||
path = Path(entry.path)
|
||||
if _regular_file(path):
|
||||
return path
|
||||
except OSError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _stop_stale_workers():
|
||||
"""Stop capture workers a previous import of this module started."""
|
||||
for thread in threading.enumerate():
|
||||
token = getattr(thread, "pm_preview_token", None)
|
||||
close = getattr(thread, "pm_preview_close", None)
|
||||
if token is not None and token is not _MODULE_TOKEN and close is not None:
|
||||
close()
|
||||
|
||||
|
||||
class AutoThumbnailer:
|
||||
"""Single-flight preview writes and the capture worker (thread-safe).
|
||||
|
||||
``generate(src, dst, is_video, *, stop, guard, track)`` writes one
|
||||
preview and returns ``_generate_one``'s result dict; ``may_write()``
|
||||
says whether automatic previews may be written now. Views call
|
||||
``ensure`` on the view pool, the monitor calls ``enqueue`` from its
|
||||
thread.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
generate,
|
||||
may_write,
|
||||
logger,
|
||||
*,
|
||||
workers=_worker_threads,
|
||||
max_queued=MAX_QUEUED,
|
||||
start_worker=True,
|
||||
):
|
||||
self._generate = generate
|
||||
self._may_write = may_write
|
||||
self._logger = logger
|
||||
self._workers = workers
|
||||
self._max_queued = max_queued
|
||||
self._start_worker = start_worker
|
||||
self._changed = threading.Condition()
|
||||
self._inflight = {}
|
||||
self._failures = collections.OrderedDict()
|
||||
self._queue = collections.deque()
|
||||
self._queued = set()
|
||||
self._processes = []
|
||||
self._worker = None
|
||||
# True from a worker's start until it leaves (under the lock), so an
|
||||
# idle worker can exit and the next enqueue starts a fresh one.
|
||||
self._worker_running = False
|
||||
self._closed = False
|
||||
self._busy = False
|
||||
self._batch = 0
|
||||
|
||||
@property
|
||||
def failure_count(self):
|
||||
with self._changed:
|
||||
return len(self._failures)
|
||||
|
||||
@property
|
||||
def closed(self):
|
||||
return self._closed
|
||||
|
||||
def _allowed(self):
|
||||
return not self._closed and self._may_write()
|
||||
|
||||
def _track(self, process):
|
||||
"""Keep an ffmpeg process so close() can kill it."""
|
||||
with self._changed:
|
||||
self._processes = [p for p in self._processes if p.poll() is None]
|
||||
self._processes.append(process)
|
||||
closed = self._closed
|
||||
if closed:
|
||||
process.kill()
|
||||
|
||||
# ── one preview ───────────────────────────────────────────────────
|
||||
|
||||
def ensure(self, src, dst, is_video):
|
||||
"""Make *dst* for *src* unless it is fresh; returns the outcome.
|
||||
|
||||
``ready`` (the preview is there), ``failed`` (the file cannot be
|
||||
decoded, remembered), ``busy`` (an I/O error, the file is gone, or
|
||||
another flight did not finish) or ``deferred`` (a hold refuses
|
||||
automatic writes, checked when the write starts and before its
|
||||
rename). Blocking: call it off the event loop.
|
||||
"""
|
||||
if _thumbnail_is_fresh(src, dst):
|
||||
return READY
|
||||
fingerprint = _fingerprint(src)
|
||||
if fingerprint is None:
|
||||
return BUSY
|
||||
key = _dst_key(dst)
|
||||
with self._changed:
|
||||
if fingerprint in self._failures:
|
||||
return FAILED
|
||||
flight = self._inflight.get(key)
|
||||
if flight is None:
|
||||
self._inflight[key] = done = threading.Event()
|
||||
if flight is not None:
|
||||
flight.wait(FLIGHT_WAIT_SECONDS)
|
||||
if _thumbnail_is_fresh(src, dst):
|
||||
return READY
|
||||
with self._changed:
|
||||
return FAILED if fingerprint in self._failures else BUSY
|
||||
try:
|
||||
return self._lead(src, dst, is_video, fingerprint)
|
||||
finally:
|
||||
with self._changed:
|
||||
self._inflight.pop(key, None)
|
||||
self._changed.notify_all()
|
||||
done.set()
|
||||
|
||||
def _lead(self, src, dst, is_video, fingerprint):
|
||||
if not self._allowed():
|
||||
return DEFERRED
|
||||
if _thumbnail_is_fresh(src, dst):
|
||||
return READY
|
||||
result = self._generate(
|
||||
src,
|
||||
dst,
|
||||
is_video,
|
||||
stop=lambda: self._closed,
|
||||
guard=self._allowed,
|
||||
track=self._track,
|
||||
)
|
||||
if result.get("action") in ("generated", "skipped"):
|
||||
return READY
|
||||
if result.get("held"):
|
||||
return DEFERRED
|
||||
if not result.get("decode"):
|
||||
return BUSY # an I/O error says nothing about the file
|
||||
error = result.get("error") or "failed"
|
||||
with self._changed:
|
||||
self._failures[fingerprint] = error
|
||||
while len(self._failures) > MAX_FAILURES_REMEMBERED:
|
||||
self._failures.popitem(last=False)
|
||||
self._logger.info(
|
||||
f"No preview for {os.path.basename(str(src))} ({error}); "
|
||||
"the original is shown"
|
||||
)
|
||||
return FAILED
|
||||
|
||||
# ── the capture worker ────────────────────────────────────────────
|
||||
|
||||
def enqueue(self, src, dst, is_video):
|
||||
"""Queue *dst*'s preview for the worker; False when already queued,
|
||||
full or closed. Never blocks on a write."""
|
||||
key = _dst_key(dst)
|
||||
with self._changed:
|
||||
if self._closed or key in self._queued:
|
||||
return False
|
||||
if len(self._queue) >= self._max_queued:
|
||||
return False
|
||||
self._queued.add(key)
|
||||
self._queue.append((key, Path(src), Path(dst), bool(is_video)))
|
||||
self._busy = True
|
||||
if self._start_worker and not self._worker_running:
|
||||
self._start_worker_locked()
|
||||
self._changed.notify_all()
|
||||
return True
|
||||
|
||||
def _start_worker_locked(self):
|
||||
"""Start the worker (lock held). A start that fails leaves the file
|
||||
queued and the next enqueue tries again."""
|
||||
worker = threading.Thread(target=self._run, name=WORKER_NAME, daemon=True)
|
||||
worker.pm_preview_token = _MODULE_TOKEN
|
||||
worker.pm_preview_close = self.close
|
||||
self._worker_running = True
|
||||
try:
|
||||
worker.start()
|
||||
except RuntimeError:
|
||||
self._worker_running = False
|
||||
self._busy = False
|
||||
self._logger.warning(
|
||||
"The automatic preview worker could not start; the next output "
|
||||
"tries again",
|
||||
exc_info=True,
|
||||
)
|
||||
return
|
||||
self._worker = worker
|
||||
|
||||
def _run(self):
|
||||
try:
|
||||
self._drain()
|
||||
except Exception:
|
||||
self._logger.warning("The automatic preview worker stopped", exc_info=True)
|
||||
finally:
|
||||
# Whatever ended the loop, the next enqueue starts a fresh worker.
|
||||
with self._changed:
|
||||
self._worker_running = False
|
||||
self._busy = False
|
||||
self._batch = 0
|
||||
self._changed.notify_all()
|
||||
|
||||
def _drain(self):
|
||||
while True:
|
||||
with self._changed:
|
||||
if self._closed or not self._queue:
|
||||
# Idle or closed: leave, so no thread outlives its work.
|
||||
self._busy = False
|
||||
self._worker_running = False
|
||||
self._changed.notify_all()
|
||||
return
|
||||
take = min(len(self._queue), max(1, int(self._workers() or 1)))
|
||||
batch = [self._queue.popleft() for _ in range(take)]
|
||||
self._batch = len(batch)
|
||||
self._changed.notify_all()
|
||||
try:
|
||||
map_parallel(self._write_queued, batch, len(batch))
|
||||
finally:
|
||||
with self._changed:
|
||||
self._batch = 0
|
||||
for key, *_rest in batch:
|
||||
self._queued.discard(key)
|
||||
|
||||
def _write_queued(self, item):
|
||||
_key, src, dst, is_video = item
|
||||
try:
|
||||
self.ensure(src, dst, is_video)
|
||||
except Exception: # one file never stops the worker
|
||||
self._logger.warning(
|
||||
f"Automatic preview failed for {src.name}", exc_info=True
|
||||
)
|
||||
|
||||
def close(self):
|
||||
"""Stop the worker at its next file boundary (queued files are
|
||||
dropped), kill a running ffmpeg and wait up to CLOSE_WAIT_SECONDS;
|
||||
a worker still running after that is logged."""
|
||||
with self._changed:
|
||||
self._closed = True
|
||||
self._queue.clear()
|
||||
self._queued.clear()
|
||||
processes, self._processes = self._processes, []
|
||||
self._changed.notify_all()
|
||||
worker = self._worker
|
||||
for process in processes:
|
||||
try:
|
||||
process.kill()
|
||||
except OSError:
|
||||
pass
|
||||
if worker is None or worker is threading.current_thread():
|
||||
return
|
||||
worker.join(CLOSE_WAIT_SECONDS)
|
||||
if worker.is_alive():
|
||||
self._logger.warning(
|
||||
"The automatic preview worker is still running after "
|
||||
f"{CLOSE_WAIT_SECONDS} s; it stops at its next file"
|
||||
)
|
||||
|
||||
# ── waits (tests and shutdown; never timers) ──────────────────────
|
||||
|
||||
def _wait_for(self, predicate, timeout):
|
||||
with self._changed:
|
||||
return self._changed.wait_for(predicate, timeout)
|
||||
|
||||
def wait_idle(self, timeout=None):
|
||||
"""True once the worker has nothing queued or in hand."""
|
||||
return self._wait_for(lambda: not self._busy, timeout)
|
||||
|
||||
def wait_batch(self, size, timeout=None):
|
||||
"""True once the worker holds a batch of *size* files."""
|
||||
return self._wait_for(lambda: self._batch == size, timeout)
|
||||
|
||||
|
||||
class _PreviewGate:
|
||||
"""Admission of view requests on the event loop (#249 I2).
|
||||
|
||||
Waiters are coroutines: at most ``limit`` previews are written at once
|
||||
on the view pool, one flight runs per key, and past ``MAX_WAITING``
|
||||
waiting requests ``run`` answers ``OVER`` at once.
|
||||
"""
|
||||
|
||||
def __init__(self, limit, logger):
|
||||
self.loop = asyncio.get_running_loop()
|
||||
self._slots = asyncio.Semaphore(limit)
|
||||
self._flights = {}
|
||||
self._tasks = set()
|
||||
self._logger = logger
|
||||
self._tick = asyncio.Event()
|
||||
self.waiting = 0
|
||||
|
||||
def _changed(self):
|
||||
tick, self._tick = self._tick, asyncio.Event()
|
||||
tick.set()
|
||||
|
||||
async def settled(self, predicate):
|
||||
"""Return once ``predicate()`` holds (re-checked at each change)."""
|
||||
while not predicate():
|
||||
await self._tick.wait()
|
||||
|
||||
async def run(self, key, work, pool):
|
||||
"""``work()``'s answer for *key* (shared by every request of the
|
||||
flight), None when it raised, or OVER past the cap."""
|
||||
if self.waiting >= MAX_WAITING:
|
||||
return OVER
|
||||
flight = self._flights.get(key)
|
||||
if flight is None:
|
||||
flight = self.loop.create_future()
|
||||
self._flights[key] = flight
|
||||
task = self.loop.create_task(self._lead(key, flight, work, pool))
|
||||
self._tasks.add(task)
|
||||
task.add_done_callback(self._tasks.discard)
|
||||
self.waiting += 1
|
||||
self._changed()
|
||||
try:
|
||||
return await asyncio.shield(flight)
|
||||
finally:
|
||||
self.waiting -= 1
|
||||
self._changed()
|
||||
|
||||
async def _lead(self, key, flight, work, pool):
|
||||
result = None
|
||||
try:
|
||||
async with self._slots:
|
||||
result = await self.loop.run_in_executor(pool, work)
|
||||
except Exception:
|
||||
self._logger.warning("Automatic preview failed", exc_info=True)
|
||||
finally:
|
||||
self._flights.pop(key, None)
|
||||
if not flight.done():
|
||||
flight.set_result(result)
|
||||
|
||||
|
||||
class AutoThumbnailMixin:
|
||||
"""Automatic previews on the API: views, the monitor's queue, scans."""
|
||||
|
||||
_auto_thumbnailer = None
|
||||
_auto_thumbnails_shut = False
|
||||
_preview_gate_obj = None
|
||||
_preview_pool_obj = None
|
||||
_auto_thumbnailer_lock = threading.Lock()
|
||||
|
||||
def _auto_thumbnails(self):
|
||||
"""This API's AutoThumbnailer, created on first use (no worker after
|
||||
the app's cleanup)."""
|
||||
with AutoThumbnailMixin._auto_thumbnailer_lock:
|
||||
if self._auto_thumbnailer is None:
|
||||
self._auto_thumbnailer = AutoThumbnailer(
|
||||
self._write_auto_thumbnail,
|
||||
lambda: self._auto_thumbnail_blocker() is None,
|
||||
self.logger,
|
||||
start_worker=not self._auto_thumbnails_shut,
|
||||
)
|
||||
return self._auto_thumbnailer
|
||||
|
||||
def close_auto_thumbnails(self, final=False):
|
||||
"""Stop the capture worker and the view pool, and forget them.
|
||||
|
||||
Bounded (see AutoThumbnailer.close). ``final`` (the app's cleanup)
|
||||
also keeps a late monitor event from starting a worker again; a
|
||||
test reset passes False and the next use starts fresh ones.
|
||||
"""
|
||||
with AutoThumbnailMixin._auto_thumbnailer_lock:
|
||||
if final:
|
||||
self._auto_thumbnails_shut = True
|
||||
thumbnailer, self._auto_thumbnailer = self._auto_thumbnailer, None
|
||||
pool, self._preview_pool_obj = self._preview_pool_obj, None
|
||||
self._preview_gate_obj = None
|
||||
if thumbnailer is not None:
|
||||
thumbnailer.close()
|
||||
if pool is not None:
|
||||
pool.shutdown(wait=False, cancel_futures=True)
|
||||
if final:
|
||||
shutdown_encoders() # kill JPEG XL / ffmpeg encoders still running
|
||||
|
||||
def _auto_thumbnail_failures(self):
|
||||
thumbnailer = self._auto_thumbnailer
|
||||
return 0 if thumbnailer is None else thumbnailer.failure_count
|
||||
|
||||
def _write_auto_thumbnail(self, src, dst, is_video, *, stop, guard, track):
|
||||
return self._generate_one(
|
||||
src,
|
||||
dst,
|
||||
writer_settings(),
|
||||
is_video,
|
||||
stop=stop,
|
||||
guard=guard,
|
||||
track=track,
|
||||
)
|
||||
|
||||
def _auto_thumbnail_blocker(self):
|
||||
"""The running job automatic previews wait for, or None (any thread)."""
|
||||
for job in self._jobs.running():
|
||||
if job.admission == "exclusive" or job.kind in AUTO_REFUSED_KINDS:
|
||||
return job
|
||||
if job.kind == THUMBNAILS_JOB_KIND and job.title == THUMBNAIL_CLEAR_TITLE:
|
||||
return job
|
||||
return None
|
||||
|
||||
# ── the view admission ────────────────────────────────────────────
|
||||
|
||||
def _preview_gate(self):
|
||||
"""The gate of the running loop (a new loop gets a new one)."""
|
||||
gate = self._preview_gate_obj
|
||||
if gate is None or gate.loop is not asyncio.get_running_loop():
|
||||
gate = self._preview_gate_obj = _PreviewGate(ON_DEMAND_LIMIT, self.logger)
|
||||
return gate
|
||||
|
||||
def _preview_pool(self):
|
||||
with AutoThumbnailMixin._auto_thumbnailer_lock:
|
||||
if self._preview_pool_obj is None:
|
||||
self._preview_pool_obj = ThreadPoolExecutor(
|
||||
ON_DEMAND_LIMIT, thread_name_prefix=POOL_PREFIX
|
||||
)
|
||||
return self._preview_pool_obj
|
||||
|
||||
def _preview_pool_threads(self):
|
||||
pool = self._preview_pool_obj
|
||||
return [] if pool is None else list(getattr(pool, "_threads", ()))
|
||||
|
||||
def _video_placeholder_bytes(self):
|
||||
return video_placeholder_jpeg()
|
||||
|
||||
async def _admit_preview(self, src, dst, is_video):
|
||||
"""Make one preview through the view admission (the history route).
|
||||
|
||||
Shares the serve route's flight key; every flight answers the
|
||||
``ensure`` outcome (or None, or OVER), whichever route started it.
|
||||
"""
|
||||
return await self._preview_gate().run(
|
||||
flight_key(Path(dst).parent, (Path(dst).name,)),
|
||||
functools.partial(self._ensure_view, src, dst, is_video),
|
||||
self._preview_pool(),
|
||||
)
|
||||
|
||||
def _ensure_view(self, src, dst, is_video):
|
||||
"""``ensure``, reporting a preview removed under us as busy."""
|
||||
outcome = self._auto_thumbnails().ensure(src, dst, is_video)
|
||||
if outcome == READY and not Path(dst).is_file():
|
||||
return BUSY # removed under us: the fallback this once
|
||||
return outcome
|
||||
|
||||
# ── on first view ─────────────────────────────────────────────────
|
||||
|
||||
def _preview_on_demand(self, root, parts):
|
||||
"""The ``ensure`` outcome for the preview *parts* under *root*
|
||||
(``deferred`` when its folder is linked out of the root), or None
|
||||
when no media file previews as *parts*."""
|
||||
located = _thumbnail_source(root, parts)
|
||||
if located is None:
|
||||
return None
|
||||
src, dst, is_video = located
|
||||
if dst is None:
|
||||
return DEFERRED
|
||||
return self._ensure_view(src, dst, is_video)
|
||||
|
||||
def _preview_response(self, outcome, src, dst, is_video):
|
||||
if outcome == READY and dst is not None:
|
||||
return self._file_response(dst)
|
||||
if is_video:
|
||||
return web.Response(
|
||||
body=video_placeholder_jpeg(),
|
||||
content_type="image/jpeg",
|
||||
headers={"Cache-Control": "no-store"},
|
||||
)
|
||||
response = web.FileResponse(src)
|
||||
response.headers["Cache-Control"] = "no-store"
|
||||
return response
|
||||
|
||||
async def _serve_missing_thumbnail(self, output_dirs, filepath):
|
||||
"""The made preview, or its fallback; None when *filepath* names no
|
||||
preview of a media file in *output_dirs*.
|
||||
|
||||
The flight answers only the outcome (a history flight may be the one
|
||||
joined); the file it previews is derived here: the preview path is
|
||||
the request's own, and a fallback looks its source up again on the
|
||||
view pool (bounded by it, never the default executor).
|
||||
"""
|
||||
if _has_traversal(filepath) or not _is_preview_request(filepath):
|
||||
return None
|
||||
parts = tuple(filepath.split("/"))
|
||||
gate = self._preview_gate()
|
||||
pool = self._preview_pool()
|
||||
for root in output_dirs:
|
||||
root = Path(root)
|
||||
outcome = await gate.run(
|
||||
flight_key(root, parts),
|
||||
functools.partial(self._preview_on_demand, root, parts),
|
||||
pool,
|
||||
)
|
||||
if outcome is None:
|
||||
continue
|
||||
if outcome is OVER:
|
||||
# Past the cap: exact names only, no listing, no write.
|
||||
quick = _quick_source(root, parts)
|
||||
if quick is None:
|
||||
continue
|
||||
return self._preview_response(BUSY, quick[0], None, quick[1])
|
||||
dst = _current_preview(root, parts)
|
||||
if outcome == READY:
|
||||
return self._file_response(dst)
|
||||
located = await gate.loop.run_in_executor(
|
||||
pool, _thumbnail_source, root, parts
|
||||
)
|
||||
if located is None:
|
||||
continue
|
||||
src, _dst, is_video = located
|
||||
return self._preview_response(outcome, src, None, is_video)
|
||||
return None
|
||||
|
||||
# ── on capture ────────────────────────────────────────────────────
|
||||
|
||||
def queue_capture_thumbnail(self, image_path):
|
||||
"""Queue the preview of a new output (the monitor's thread).
|
||||
|
||||
Returns ``queued``, ``fresh`` (its preview is current) or ``skipped``
|
||||
(not a media file under a current root, a link, a preview itself,
|
||||
already queued, or after the app's cleanup). Never writes on the
|
||||
caller's thread.
|
||||
"""
|
||||
if self._auto_thumbnails_shut:
|
||||
return "skipped"
|
||||
located = self._capture_target(image_path)
|
||||
if located is None:
|
||||
return "skipped"
|
||||
src, dst, is_video = located
|
||||
if _thumbnail_is_fresh(src, dst):
|
||||
return "fresh"
|
||||
if self._auto_thumbnails().enqueue(src, dst, is_video):
|
||||
return "queued"
|
||||
return "skipped"
|
||||
|
||||
def _capture_target(self, image_path):
|
||||
path = Path(os.path.abspath(str(image_path)))
|
||||
if path.suffix.lower() not in MEDIA_EXTENSIONS or not _regular_file(path):
|
||||
return None
|
||||
folded = os.path.normcase(str(path))
|
||||
for root in self._get_all_output_dirs():
|
||||
root = Path(os.path.abspath(str(root)))
|
||||
base = os.path.normcase(str(root)).rstrip("\\/") + os.sep
|
||||
if not folded.startswith(base) or not _is_within(path, root):
|
||||
continue
|
||||
rel = Path(os.path.relpath(path, root))
|
||||
if "thumbnails" in rel.parts[:-1]:
|
||||
return None
|
||||
dst = root / media_thumbnail_rel(rel)
|
||||
if not _is_within(root / "thumbnails", root) or not _is_within(
|
||||
dst, root / "thumbnails"
|
||||
):
|
||||
return None
|
||||
return path, dst, path.suffix.lower() in VIDEO_EXTENSIONS
|
||||
return None
|
||||
|
||||
# ── after a scan ──────────────────────────────────────────────────
|
||||
|
||||
async def _thumbnails_after_scan(self, job, output_dirs=None):
|
||||
"""Hand a finished scan's missing previews on; returns what it did.
|
||||
|
||||
None (the scan added nothing), ``job`` (the thumbnail job started for
|
||||
``SCAN_THUMBNAIL_JOB_THRESHOLD`` or more missing previews), ``busy``
|
||||
(a job refused it; the previews are made on view), ``queued`` (fewer,
|
||||
through the worker) or ``none`` (nothing missing). It counts the
|
||||
sources the scan walked (``output_dirs``, narrowed to the job's
|
||||
``roots``; every source when None). The server log says which.
|
||||
"""
|
||||
if job.counters.get("added", 0) <= 0:
|
||||
return None
|
||||
if output_dirs is None:
|
||||
output_dirs = await self._run_in_executor(self._get_all_output_dirs)
|
||||
roots = job.scope.get("roots")
|
||||
if roots:
|
||||
output_dirs = [output_dirs[i] for i in roots if 0 <= i < len(output_dirs)]
|
||||
plan = await self._run_in_executor(self._plan_gallery_thumbnails, output_dirs)
|
||||
missing = len(plan["targets"]) + plan["remaining"]
|
||||
if missing >= SCAN_THUMBNAIL_JOB_THRESHOLD:
|
||||
options = ThumbnailOptions("missing", writer_settings())
|
||||
try:
|
||||
self._start_thumbnail_job(options)
|
||||
except JobBusy as busy:
|
||||
self.logger.info(
|
||||
f"{missing:,} previews missing after the scan; the thumbnail "
|
||||
f"job did not start while {busy.running.title} runs"
|
||||
)
|
||||
return "busy"
|
||||
self.logger.info(
|
||||
f"{missing:,} previews missing after the scan; thumbnail job started"
|
||||
)
|
||||
return "job"
|
||||
if not missing:
|
||||
return "none"
|
||||
thumbnailer = self._auto_thumbnails()
|
||||
for src, dst, is_video in plan["targets"]:
|
||||
thumbnailer.enqueue(src, dst, is_video)
|
||||
self.logger.info(f"{missing:,} previews missing after the scan; queued")
|
||||
return "queued"
|
||||
|
||||
async def _hand_off_scan_thumbnails(self, job, output_dirs=None):
|
||||
"""``_thumbnails_after_scan`` for the scan runner: never raises."""
|
||||
try:
|
||||
return await self._thumbnails_after_scan(job, output_dirs)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
self.logger.warning("Previews after the scan failed", exc_info=True)
|
||||
return None
|
||||
|
||||
|
||||
# A package reload re-runs this module: stop the workers the old one left.
|
||||
_stop_stale_workers()
|
||||
@@ -1,404 +0,0 @@
|
||||
"""Convert outdated thumbnails (#317): an explicit, cancellable Activity job.
|
||||
|
||||
``POST /prompt_manager/jobs/thumbnails {mode: "convert"}`` (and
|
||||
``GET /prompt_manager/images/thumbnails/plan?mode=convert`` for its count)
|
||||
starts "Convert outdated thumbnails": the thumbnail job's kind, ``guarded``
|
||||
admission (never beside a record-changing job such as a scan, a restore or
|
||||
a duplicate removal; one thumbnail job at a time), its worker threads and
|
||||
its outcome groups. Nothing converts on its own: changing the format in
|
||||
Appearance only makes previews "outdated".
|
||||
|
||||
A preview is outdated when the file owns one in another format than the one
|
||||
written now (``writer_settings``: a legacy ``_thumb.png``, or a format
|
||||
switched away from). For each such file, at a file boundary (Cancel keeps
|
||||
what was converted):
|
||||
|
||||
1. the new preview is written from the SOURCE media in the configured
|
||||
format (``_generate_one``: atomic, verified to decode; a video whose frame
|
||||
cannot be decoded fails rather than getting the placeholder), unless a
|
||||
current preview at least as new as the source is already there and
|
||||
decodes;
|
||||
2. only then are that file's old-format previews removed: regular files (no
|
||||
links) in its ``thumbnails/`` folder, owned by that source
|
||||
(``thumbnail_settings.owns``), with the source still present. One already
|
||||
gone counts as removed.
|
||||
|
||||
Files sharing a stem in one folder (``x.png`` and ``x.mp4``) share one
|
||||
preview (ruling R2), so they are converted once, as a group: the preview is
|
||||
made from the first source in ``thumbnail_settings.source_rank`` order
|
||||
(``.png``, ``.jpg``, ``.jpeg``, ``.webp``, ``.gif``, then videos), the old
|
||||
previews of every source in the group are removed, and each removed or
|
||||
written file's bytes count once.
|
||||
|
||||
A failure keeps the old preview. Previews whose source is gone (orphans) are
|
||||
never planned, so never touched. The format the job converts to is fixed
|
||||
when it starts; Appearance refuses a thumbnail settings change while it runs
|
||||
(409), and should the format still change underneath, the job stops at the
|
||||
next file without removing anything in another format. The final status line
|
||||
reads "N converted · F failed · <bytes before> → <bytes after>".
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from . import thumbnail_settings
|
||||
from .thumbnail_encode import ThumbnailVerifyError, verify_thumbnail, writer_settings
|
||||
from .admin import _plural, _scan_location, _scan_roots, _scan_status_line
|
||||
from .images import _thumbnail_is_fresh, _worker_threads, _thumbnail_batch_size
|
||||
from .prompts import safe_error_message
|
||||
|
||||
try:
|
||||
from ...utils.parallel import map_parallel
|
||||
except ImportError: # tests import the package top-level
|
||||
from utils.parallel import map_parallel
|
||||
|
||||
CONVERT_JOB_TITLE = "Convert outdated thumbnails"
|
||||
NOTHING_OUTDATED = "No outdated previews to convert"
|
||||
CANCELLED_BEFORE_CONVERTING = "Cancelled before any preview was converted"
|
||||
SOURCE_GONE = "The file is gone; its old preview was kept"
|
||||
OLD_NOT_REMOVED = "Converted; the old preview could not be removed"
|
||||
FORMAT_CHANGED = (
|
||||
"Stopped: the thumbnail format changed while converting · "
|
||||
"{converted} converted. Run Convert outdated again."
|
||||
)
|
||||
_UNITS = ("B", "KB", "MB", "GB", "TB")
|
||||
|
||||
|
||||
def readable_bytes(count):
|
||||
"""``1.5 KB``-style size (whole bytes below 1 KB)."""
|
||||
size = float(count)
|
||||
for unit in _UNITS:
|
||||
if size < 1024 or unit == _UNITS[-1]:
|
||||
return f"{int(size)} B" if unit == "B" else f"{size:.1f} {unit}"
|
||||
size /= 1024
|
||||
return f"{size:.1f} {_UNITS[-1]}" # pragma: no cover - loop always returns
|
||||
|
||||
|
||||
def _old_previews(src, dst, settings):
|
||||
"""``[(path, size)]``: *src*'s previews beside *dst* in another format."""
|
||||
stem = dst.name[: -len(dst.suffix)] if dst.suffix else dst.name
|
||||
old = []
|
||||
for path in thumbnail_settings.owned_in_folder(
|
||||
dst.parent, stem, src.suffix, settings.extension
|
||||
):
|
||||
if path.suffix == settings.extension:
|
||||
continue
|
||||
try:
|
||||
old.append((path, path.stat().st_size))
|
||||
except OSError:
|
||||
continue
|
||||
return old
|
||||
|
||||
|
||||
def _group_targets(targets):
|
||||
"""``[[(src, dst, is_video)]]``: targets sharing a preview name (one
|
||||
folder, one stem) in one group, its first source first (R2)."""
|
||||
groups = {}
|
||||
for target in targets:
|
||||
dst = Path(target[1])
|
||||
key = os.path.normcase(str(dst.with_suffix("")))
|
||||
groups.setdefault(key, []).append(target)
|
||||
ordered = []
|
||||
for group in groups.values():
|
||||
group.sort(
|
||||
key=lambda t: (thumbnail_settings.source_rank(Path(t[0]).suffix), str(t[0]))
|
||||
)
|
||||
ordered.append(group)
|
||||
return ordered
|
||||
|
||||
|
||||
def _same_file(a, b):
|
||||
"""Whether *a* and *b* name one file (re-review N1: on a case-insensitive
|
||||
disk ``x_thumb.JPG`` is the ``x_thumb.jpg`` just written). Without inode
|
||||
numbers (st_ino 0 on some SMB/FAT mounts) the resolved, normcased paths
|
||||
are compared."""
|
||||
try:
|
||||
first, second = os.stat(a), os.stat(b)
|
||||
except OSError:
|
||||
return False
|
||||
if first.st_ino and second.st_ino:
|
||||
return (first.st_dev, first.st_ino) == (second.st_dev, second.st_ino)
|
||||
return os.path.normcase(os.path.realpath(a)) == os.path.normcase(
|
||||
os.path.realpath(b)
|
||||
)
|
||||
|
||||
|
||||
def _format_changed(settings):
|
||||
"""Whether previews are written in another format than *settings* now."""
|
||||
return writer_settings().extension != settings.extension
|
||||
|
||||
|
||||
def _finished_line(totals, remaining):
|
||||
line = (
|
||||
f"{totals['converted']:,} converted · {totals['failed']:,} failed · "
|
||||
f"{readable_bytes(totals['before'])} → {readable_bytes(totals['after'])}"
|
||||
)
|
||||
if remaining:
|
||||
line += f" · {remaining:,} more on the next run"
|
||||
return line
|
||||
|
||||
|
||||
class ThumbnailConvertMixin:
|
||||
"""The Convert outdated job (started through ``_start_thumbnail_job``)."""
|
||||
|
||||
async def _run_convert_job(self, job, options):
|
||||
try:
|
||||
state = await self._convert_job_body(job, options)
|
||||
except Exception:
|
||||
self.logger.exception("Thumbnail conversion failed")
|
||||
job.update(status_line="Thumbnail conversion stopped on an internal error.")
|
||||
job.log_line("Conversion failed · internal error")
|
||||
state = "failed"
|
||||
job.finish(state)
|
||||
|
||||
async def _convert_job_body(self, job, options):
|
||||
output_dirs = await self._run_in_executor(self._get_all_output_dirs)
|
||||
if not output_dirs:
|
||||
message = "ComfyUI output directory not found"
|
||||
job.update(status_line=message)
|
||||
job.log_line(f"Conversion failed · {message}")
|
||||
return "failed"
|
||||
folders = _plural(len(output_dirs), "output folder")
|
||||
job.scope["output_folders"] = len(output_dirs)
|
||||
job.update(status_line=f"Looking for outdated previews in {folders}")
|
||||
job.log_line(
|
||||
f"Conversion started · {folders} · to {options.settings.describe()}"
|
||||
)
|
||||
plan = await self._plan_conversion(job, options, output_dirs)
|
||||
if job.cancelled:
|
||||
return self._end_conversion(job, 0, CANCELLED_BEFORE_CONVERTING)
|
||||
if not plan["targets"]:
|
||||
job.update(progress=100, status_line=NOTHING_OUTDATED)
|
||||
job.log_line(f"Conversion finished · {NOTHING_OUTDATED.lower()}")
|
||||
return "finished"
|
||||
totals, left = await self._convert_in_batches(job, options, plan, output_dirs)
|
||||
if totals["stopped"]:
|
||||
job.update(status_line=FORMAT_CHANGED.format(converted=totals["converted"]))
|
||||
job.log_line(
|
||||
"Conversion stopped · the thumbnail format changed · "
|
||||
f"{_plural(totals['converted'], 'preview')} converted"
|
||||
)
|
||||
return "cancelled"
|
||||
if job.cancelled:
|
||||
return self._end_conversion(
|
||||
job,
|
||||
totals["converted"],
|
||||
f"Cancelled · {_plural(totals['converted'], 'preview')} converted",
|
||||
)
|
||||
line = _finished_line(totals, left)
|
||||
job.update(progress=100, status_line=line)
|
||||
job.log_line(f"Conversion finished · {line}")
|
||||
self.logger.info(f"Thumbnail conversion: {line}")
|
||||
return "finished"
|
||||
|
||||
@staticmethod
|
||||
def _end_conversion(job, converted, status_line):
|
||||
job.update(status_line=status_line)
|
||||
job.log_line(
|
||||
f"Conversion cancelled · {_plural(converted, 'preview')} converted"
|
||||
)
|
||||
return "cancelled"
|
||||
|
||||
async def _plan_conversion(self, job, options, output_dirs, exclude=None):
|
||||
"""One batch's plan (at most MAX_THUMBNAIL_FILES targets), off-loop."""
|
||||
return await self._run_in_executor(
|
||||
self._plan_gallery_thumbnails,
|
||||
output_dirs,
|
||||
"convert",
|
||||
stop=lambda: job.cancelled,
|
||||
settings=options.settings,
|
||||
exclude=exclude,
|
||||
)
|
||||
|
||||
async def _convert_in_batches(self, job, options, plan, output_dirs):
|
||||
"""Convert ``plan``, then plan again after each batch until nothing
|
||||
is left, the job is cancelled or stopped, or a batch converted
|
||||
nothing (files that fail are never planned again). Returns
|
||||
``(totals across batches, files left)``."""
|
||||
totals = {"converted": 0, "failed": 0, "before": 0, "after": 0}
|
||||
totals["stopped"] = False
|
||||
run = {
|
||||
"index": 0,
|
||||
"batches": 0,
|
||||
"total": len(plan["targets"]) + plan["remaining"],
|
||||
}
|
||||
run["start"] = time.monotonic()
|
||||
job.update(
|
||||
total=run["total"], status_line=_scan_status_line(0, run["total"], 0, 0)
|
||||
)
|
||||
attempted = set()
|
||||
while True:
|
||||
before = totals["converted"]
|
||||
await self._convert_targets(
|
||||
job, options, plan["targets"], output_dirs, totals, run
|
||||
)
|
||||
attempted.update(src for src, _dst, _video in plan["targets"])
|
||||
stop = job.cancelled or totals["stopped"]
|
||||
if stop or not plan["remaining"] or totals["converted"] == before:
|
||||
return totals, plan["remaining"]
|
||||
plan = await self._plan_conversion(
|
||||
job, options, output_dirs, exclude=attempted
|
||||
)
|
||||
if job.cancelled:
|
||||
return totals, 0
|
||||
run["total"] = run["index"] + len(plan["targets"]) + plan["remaining"]
|
||||
job.update(total=run["total"])
|
||||
if not plan["targets"]:
|
||||
return totals, plan["remaining"]
|
||||
|
||||
async def _convert_targets(self, job, options, targets, output_dirs, totals, run):
|
||||
"""Convert one batch's *targets* in worker batches of preview groups,
|
||||
recording each file against the job's overall ``run`` counters."""
|
||||
roots = _scan_roots(output_dirs)
|
||||
groups = _group_targets(targets)
|
||||
batch_size = _thumbnail_batch_size(_worker_threads())
|
||||
for batch_start in range(0, len(groups), batch_size):
|
||||
if job.cancelled or totals["stopped"]:
|
||||
break
|
||||
batch = groups[batch_start : batch_start + batch_size]
|
||||
results = await self._run_in_executor(
|
||||
self._convert_batch, batch, options.settings, lambda: job.cancelled
|
||||
)
|
||||
for group, group_results in zip(batch, results):
|
||||
for (src, _dst, _video), result in zip(group, group_results):
|
||||
if result["action"] == "stopped":
|
||||
totals["stopped"] = True
|
||||
if result["action"] in ("cancelled", "stopped"):
|
||||
continue
|
||||
run["index"] += 1
|
||||
index, total = run["index"], max(run["total"], run["index"])
|
||||
self._record_conversion(job, src, result, roots, totals)
|
||||
job.update(processed=index, progress=index * 100 // total)
|
||||
run["batches"] += 1
|
||||
job.update(
|
||||
status_line=_scan_status_line(
|
||||
run["index"],
|
||||
run["total"],
|
||||
run["batches"],
|
||||
time.monotonic() - run["start"],
|
||||
)
|
||||
)
|
||||
await asyncio.sleep(0) # let other requests run between batches
|
||||
|
||||
@staticmethod
|
||||
def _record_conversion(job, src, result, roots, totals):
|
||||
path, label, index = _scan_location(src, roots)
|
||||
where = f"{label}/{path}" if label else path
|
||||
row = {"path": path, "root": label, "root_index": index}
|
||||
if result["action"] == "converted":
|
||||
totals["converted"] += 1
|
||||
totals["before"] += result["before"]
|
||||
totals["after"] += result["after"]
|
||||
detail = result["detail"]
|
||||
job.record("generated", {**row, "detail": detail}, f"+ {where} · {detail}")
|
||||
return
|
||||
totals["failed"] += 1
|
||||
detail = result.get("error") or "failed"
|
||||
job.record("errors", {**row, "detail": detail}, f"! {where} · {detail}")
|
||||
|
||||
def _convert_batch(self, groups, settings, stop):
|
||||
"""Convert preview ``groups`` on the configured worker threads
|
||||
(blocking); one result list per group."""
|
||||
convert = functools.partial(self._convert_group, settings=settings, stop=stop)
|
||||
return map_parallel(convert, groups, _worker_threads())
|
||||
|
||||
def _convert_group(self, group, settings, stop):
|
||||
"""Convert the preview a group of same-stem files shares, once.
|
||||
|
||||
The first source writes it; the others are recorded with the same
|
||||
outcome and no bytes of their own."""
|
||||
src, dst, is_video = group[0]
|
||||
siblings = [Path(target[0]) for target in group[1:]]
|
||||
result = self._convert_one(
|
||||
src, dst, is_video, settings, stop=stop, siblings=siblings
|
||||
)
|
||||
if result["action"] != "converted":
|
||||
return [result] * len(group)
|
||||
shared = {**result, "before": 0, "after": 0}
|
||||
shared["detail"] = f"{result['detail']} (shares {Path(src).name}'s preview)"
|
||||
return [result] + [shared] * len(siblings)
|
||||
|
||||
def _convert_one(self, src, dst, is_video, settings, stop=None, siblings=()):
|
||||
"""Convert one file's preview (blocking; never raises).
|
||||
|
||||
``converted`` with the bytes removed and written, ``error`` (the old
|
||||
preview kept), ``cancelled`` (a file boundary) or ``stopped`` (the
|
||||
format changed underneath, review I4). *siblings* share the preview
|
||||
(R2): their old previews go too.
|
||||
"""
|
||||
if stop is not None and stop():
|
||||
return {"action": "cancelled"}
|
||||
if _format_changed(settings):
|
||||
return {"action": "stopped"}
|
||||
dst = Path(dst)
|
||||
try:
|
||||
old = self._old_group_previews([Path(src), *siblings], dst, settings)
|
||||
if not self._current_preview_usable(src, dst, settings):
|
||||
made = self._generate_one(
|
||||
src, dst, settings, is_video, rebuild=True, placeholder=False
|
||||
)
|
||||
if made["action"] != "generated":
|
||||
return {"action": "error", "error": made.get("error")}
|
||||
dst = Path(made["path"])
|
||||
if dst.suffix != settings.extension or not dst.is_file():
|
||||
return {"action": "error", "error": "The new preview was not written"}
|
||||
if not os.path.isfile(src):
|
||||
return {"action": "error", "error": SOURCE_GONE}
|
||||
if _format_changed(settings):
|
||||
return {"action": "stopped"}
|
||||
old = [(path, size) for path, size in old if not _same_file(path, dst)]
|
||||
removed, kept = self._remove_old_previews(old)
|
||||
if kept:
|
||||
return {"action": "error", "error": OLD_NOT_REMOVED}
|
||||
label = thumbnail_settings.FORMAT_LABELS[settings.format]
|
||||
return {
|
||||
"action": "converted",
|
||||
"before": removed,
|
||||
"after": dst.stat().st_size,
|
||||
"detail": f"Converted to {label}",
|
||||
}
|
||||
except Exception as exc:
|
||||
self.logger.warning(f"Conversion failed for {src.name}", exc_info=True)
|
||||
return {"action": "error", "error": safe_error_message(exc)}
|
||||
|
||||
@staticmethod
|
||||
def _old_group_previews(sources, dst, settings):
|
||||
"""``[(path, size)]`` of every source's old previews, each path once."""
|
||||
seen, old = set(), []
|
||||
for source in sources:
|
||||
for path, size in _old_previews(source, dst, settings):
|
||||
key = os.path.normcase(str(path))
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
old.append((path, size))
|
||||
return old
|
||||
|
||||
@staticmethod
|
||||
def _current_preview_usable(src, dst, settings):
|
||||
"""A current preview at least as new as *src* that decodes (M4)."""
|
||||
if not _thumbnail_is_fresh(src, dst):
|
||||
return False
|
||||
try:
|
||||
verify_thumbnail(dst, settings.format)
|
||||
except ThumbnailVerifyError:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _remove_old_previews(self, old):
|
||||
"""``(bytes removed, [kept paths])``; a file already gone counts as
|
||||
removed, with no bytes (another run or a sibling removed it)."""
|
||||
removed, kept = 0, []
|
||||
for path, size in old:
|
||||
try:
|
||||
os.remove(path)
|
||||
except FileNotFoundError:
|
||||
continue
|
||||
except OSError as exc:
|
||||
self.logger.warning(f"Could not remove old preview {path}: {exc}")
|
||||
kept.append(path)
|
||||
continue
|
||||
removed += size
|
||||
return removed, kept
|
||||
@@ -1,524 +0,0 @@
|
||||
"""Thumbnail encoders (#317): every preview file is written here.
|
||||
|
||||
``write_thumbnail(frame, dst, settings)`` encodes a prepared frame
|
||||
(``thumbnail_frame``: the first frame, RGB, fitted inside the edge) in the
|
||||
configured format into a hidden temp sibling, closes it, verifies that it
|
||||
decodes, and renames it over ``dst`` (``os.replace``), so a reader never sees
|
||||
a partial file and a failure keeps the old preview. ``guard()`` is asked just
|
||||
before the rename (automatic previews' holds, #249).
|
||||
|
||||
Formats: WebP (Pillow, ``method=4``) and JPEG (Pillow, ``optimize``) always
|
||||
work. JPEG XL needs no new dependency; the first encoder found is used:
|
||||
|
||||
1. Pillow, natively or through the ``pillow_jxl`` plugin (imported when
|
||||
installed: ``pip install pillow-jxl-plugin``), given the quality itself;
|
||||
2. ffmpeg built with ``libjxl`` (``-distance``), fed a PNG on stdin;
|
||||
3. ``cjxl`` from libjxl (``-d``), fed a temp PNG that is closed first.
|
||||
|
||||
Quality maps to a JPEG XL distance with libjxl's own
|
||||
``JxlEncoderDistanceFromQuality``: ``0.1 + (100 - q) * 0.09`` from 30 up
|
||||
(q80 → d1.9, q85 → d1.45, q90 → d1.0), the quadratic below 30, and q100 →
|
||||
d0.1 (a preview is never lossless). Pillow's plugin applies the same mapping
|
||||
inside libjxl. Detection is cached (``detect_jxl_encoder(refresh=True)``
|
||||
probes again); "none found" is probed again on the settings read
|
||||
(``encoder_status``) and after MISSING_ENCODER_TTL, so one slow probe never
|
||||
turns JPEG XL off until a restart. An encoder that fails because its binary
|
||||
or module is gone makes that write detect again: another encoder found is
|
||||
used for that file, none at all writes it as WebP (logged once) and
|
||||
refreshes what Appearance reads. Any other failure (one bad file) raises
|
||||
without a probe.
|
||||
|
||||
Encoder processes (and ffmpeg's video frames, images.py) go through
|
||||
``run_encoder``: tracked, bounded by a timeout, and killed by
|
||||
``shutdown_encoders()`` at the app's cleanup, after which none starts.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import io
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess # nosec B404 - fixed argv lists, never a shell
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import replace
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from . import thumbnail_settings
|
||||
|
||||
try:
|
||||
from ...utils.logging_config import get_logger
|
||||
except ImportError: # tests import the package top-level
|
||||
from utils.logging_config import get_logger
|
||||
|
||||
logger = get_logger("prompt_manager.thumbnails")
|
||||
|
||||
PILLOW_JXL_FORMAT = "JXL"
|
||||
JXL_UNAVAILABLE = (
|
||||
"No JPEG XL encoder found. Install pillow-jxl-plugin "
|
||||
"(pip install pillow-jxl-plugin), ffmpeg built with libjxl, or cjxl from "
|
||||
"libjxl, then restart ComfyUI."
|
||||
)
|
||||
# Seconds an encoder process may run (ffmpeg's video frame shares it).
|
||||
ENCODE_TIMEOUT = 30
|
||||
PROBE_TIMEOUT = 10
|
||||
JXL_EFFORT = 7
|
||||
WEBP_METHOD = 4
|
||||
JXL_SIGNATURES = (b"\xff\x0a", b"\x00\x00\x00\x0cJXL \r\n\x87\n")
|
||||
_NO_WINDOW = getattr(subprocess, "CREATE_NO_WINDOW", 0)
|
||||
|
||||
|
||||
class PreviewWriteHeld(Exception):
|
||||
"""A preview write stopped before its rename: ``guard()`` said no (#249)."""
|
||||
|
||||
|
||||
# The result error of a write held before its rename (never remembered).
|
||||
PREVIEW_HELD = "Preview writes are held while another job runs"
|
||||
|
||||
|
||||
class ThumbnailVerifyError(RuntimeError):
|
||||
"""The encoded preview did not decode; it was never renamed in."""
|
||||
|
||||
|
||||
class EncoderFailed(RuntimeError):
|
||||
"""An encoder process failed or wrote nothing."""
|
||||
|
||||
|
||||
class EncoderClosed(RuntimeError):
|
||||
"""The app is shutting down: no encoder process runs any more."""
|
||||
|
||||
|
||||
# ── temp files ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def temp_sibling(dst, tag="part"):
|
||||
"""Hidden temp path next to *dst* that keeps its extension."""
|
||||
dst = Path(dst)
|
||||
token = f"{os.getpid()}-{threading.get_ident()}"
|
||||
return dst.with_name(f".{dst.stem}.{tag}-{token}{dst.suffix}")
|
||||
|
||||
|
||||
def replace_unless_held(tmp, dst, guard=None):
|
||||
"""Rename *tmp* over *dst*, unless ``guard()`` is false (PreviewWriteHeld)."""
|
||||
if guard is not None and not guard():
|
||||
raise PreviewWriteHeld(PREVIEW_HELD)
|
||||
os.replace(tmp, dst)
|
||||
|
||||
|
||||
# ``temp_sibling``'s names: ``.<stem>.<tag>-<pid>-<thread id><ext>``.
|
||||
_TEMP_NAME = re.compile(r"^\..+\.(?:part|src)-(\d+)-\d+\.[A-Za-z0-9]+$")
|
||||
# A temp at least this old (seconds) outlived its write: no encode runs longer.
|
||||
STRAY_TEMP_AGE = 4 * ENCODE_TIMEOUT
|
||||
# When this process started (this module's import): a temp carrying its PID
|
||||
# that is older was left by an earlier run with the same PID (containers).
|
||||
PROCESS_STARTED = time.time()
|
||||
|
||||
|
||||
def sweep_stray_temps(folder, min_age=None):
|
||||
"""Remove the temp files a crashed write left under *folder* (a preview
|
||||
folder, walked without following links); returns ``(count, bytes)``.
|
||||
|
||||
Only ``temp_sibling`` names are touched, never a link. With *min_age*
|
||||
(seconds, the startup sweep) a temp younger than that, or one this
|
||||
process wrote since it started (it may still be writing it), is left;
|
||||
``min_age=0`` (the clear, which holds every preview write) removes them
|
||||
all.
|
||||
"""
|
||||
min_age = STRAY_TEMP_AGE if min_age is None else min_age
|
||||
now = time.time()
|
||||
own = str(os.getpid())
|
||||
count = size = 0
|
||||
for root, _dirs, files in os.walk(folder):
|
||||
for name in files:
|
||||
match = _TEMP_NAME.match(name)
|
||||
if match is None:
|
||||
continue
|
||||
path = Path(root) / name
|
||||
try:
|
||||
stat = path.lstat()
|
||||
if path.is_symlink() or not path.is_file():
|
||||
continue
|
||||
if min_age and (
|
||||
now - stat.st_mtime < min_age
|
||||
or (match.group(1) == own and stat.st_mtime >= PROCESS_STARTED)
|
||||
):
|
||||
continue
|
||||
path.unlink()
|
||||
except OSError as exc:
|
||||
logger.debug(f"Could not remove stray temp {name}: {exc}")
|
||||
continue
|
||||
count += 1
|
||||
size += stat.st_size
|
||||
return count, size
|
||||
|
||||
|
||||
def _unlink_quietly(path):
|
||||
try:
|
||||
Path(path).unlink()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
except OSError as exc:
|
||||
logger.debug(f"Could not remove temp file {Path(path).name}: {exc}")
|
||||
|
||||
|
||||
# ── encoder processes ────────────────────────────────────────────────
|
||||
|
||||
_process_lock = threading.Lock()
|
||||
_processes = set()
|
||||
_closed = False
|
||||
|
||||
|
||||
def run_encoder(cmd, *, input=None, timeout=ENCODE_TIMEOUT, track=None):
|
||||
"""Run *cmd* (a fixed argv) with *input* on stdin; returns a
|
||||
CompletedProcess with bytes output.
|
||||
|
||||
The process is tracked so ``shutdown_encoders`` can kill it; past
|
||||
*timeout* it is killed and TimeoutExpired raised. After the shutdown
|
||||
nothing starts (EncoderClosed). ``track(process)`` sees it too.
|
||||
"""
|
||||
with _process_lock:
|
||||
if _closed:
|
||||
raise EncoderClosed("Encoders are shut down")
|
||||
process = subprocess.Popen( # nosec B603 - fixed argv, no shell
|
||||
cmd,
|
||||
stdin=subprocess.PIPE if input is not None else subprocess.DEVNULL,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
creationflags=_NO_WINDOW,
|
||||
)
|
||||
with _process_lock:
|
||||
_processes.add(process)
|
||||
closed = _closed
|
||||
try:
|
||||
if closed:
|
||||
process.kill()
|
||||
if track is not None:
|
||||
track(process)
|
||||
try:
|
||||
out, err = process.communicate(input=input, timeout=timeout)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
process.communicate()
|
||||
raise
|
||||
finally:
|
||||
with _process_lock:
|
||||
_processes.discard(process)
|
||||
closed = _closed
|
||||
if closed:
|
||||
raise EncoderClosed("Encoder stopped by the shutdown")
|
||||
return subprocess.CompletedProcess(cmd, process.returncode, out, err)
|
||||
|
||||
|
||||
def shutdown_encoders():
|
||||
"""Kill every running encoder process; none starts afterwards."""
|
||||
global _closed
|
||||
with _process_lock:
|
||||
_closed = True
|
||||
running = list(_processes)
|
||||
for process in running:
|
||||
try:
|
||||
process.kill()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def reset_encoders():
|
||||
"""Allow encoders again and forget the detection (tests, a reload)."""
|
||||
global _closed, _detected, _fallback_logged
|
||||
with _process_lock:
|
||||
_closed = False
|
||||
with _detect_lock:
|
||||
_detected = _UNSET
|
||||
_fallback_logged = False
|
||||
|
||||
|
||||
# ── JPEG XL detection ────────────────────────────────────────────────
|
||||
|
||||
_UNSET = object()
|
||||
_detect_lock = threading.Lock()
|
||||
_detected = _UNSET
|
||||
_detected_at = 0.0
|
||||
# Seconds after which "no JPEG XL encoder" is probed again.
|
||||
MISSING_ENCODER_TTL = 60
|
||||
_fallback_logged = False
|
||||
|
||||
|
||||
def _pillow_jxl_format():
|
||||
"""Pillow's JPEG XL format id when it can save one, else None."""
|
||||
try:
|
||||
native = Image.registered_extensions().get(".jxl")
|
||||
except Exception:
|
||||
native = None
|
||||
if native and native in Image.SAVE:
|
||||
return native
|
||||
return PILLOW_JXL_FORMAT if PILLOW_JXL_FORMAT in Image.SAVE else None
|
||||
|
||||
|
||||
def _pillow_can_save_jxl():
|
||||
if _pillow_jxl_format():
|
||||
return True
|
||||
try:
|
||||
importlib.import_module("pillow_jxl")
|
||||
except Exception:
|
||||
return False
|
||||
return _pillow_jxl_format() is not None
|
||||
|
||||
|
||||
def _pillow_decodes_jxl():
|
||||
try:
|
||||
native = Image.registered_extensions().get(".jxl")
|
||||
except Exception:
|
||||
native = None
|
||||
return bool(native and native in Image.OPEN) or PILLOW_JXL_FORMAT in Image.OPEN
|
||||
|
||||
|
||||
def _ffmpeg_has_libjxl():
|
||||
"""Whether the ffmpeg on PATH lists the ``libjxl`` encoder."""
|
||||
if not shutil.which("ffmpeg"):
|
||||
return False
|
||||
try:
|
||||
done = run_encoder(
|
||||
["ffmpeg", "-hide_banner", "-encoders"], timeout=PROBE_TIMEOUT
|
||||
)
|
||||
except (OSError, subprocess.SubprocessError, EncoderClosed):
|
||||
return False
|
||||
listing = done.stdout
|
||||
if isinstance(listing, bytes):
|
||||
listing = listing.decode("utf-8", "replace")
|
||||
return any(
|
||||
len(line.split()) > 1 and line.split()[1] == "libjxl"
|
||||
for line in listing.splitlines()
|
||||
)
|
||||
|
||||
|
||||
def _detect():
|
||||
if _pillow_can_save_jxl():
|
||||
return "pillow"
|
||||
if _ffmpeg_has_libjxl():
|
||||
return "ffmpeg"
|
||||
if shutil.which("cjxl"):
|
||||
return "cjxl"
|
||||
return None
|
||||
|
||||
|
||||
def detect_jxl_encoder(refresh=False, recheck_missing=False):
|
||||
"""``pillow``, ``ffmpeg``, ``cjxl`` or None (cached; blocking on a probe).
|
||||
|
||||
A cached None is probed again with *recheck_missing* or once it is older
|
||||
than MISSING_ENCODER_TTL."""
|
||||
global _detected, _detected_at
|
||||
with _detect_lock:
|
||||
now = time.monotonic()
|
||||
stale_missing = _detected is None and (
|
||||
recheck_missing or now - _detected_at >= MISSING_ENCODER_TTL
|
||||
)
|
||||
if refresh or _detected is _UNSET or stale_missing:
|
||||
_detected = _detect()
|
||||
_detected_at = now
|
||||
return _detected
|
||||
|
||||
|
||||
def _encoder_gone(encoder):
|
||||
"""Whether *encoder*'s binary or module is no longer there."""
|
||||
if encoder == "pillow":
|
||||
return not _pillow_can_save_jxl()
|
||||
return shutil.which(encoder) is None
|
||||
|
||||
|
||||
def encoder_status():
|
||||
"""Which formats can be written now (``GET /prompt_manager/settings``;
|
||||
a "none found" is probed again here)."""
|
||||
encoder = detect_jxl_encoder(recheck_missing=True)
|
||||
return {
|
||||
"webp": {"available": True},
|
||||
"jpeg": {"available": True},
|
||||
"jxl": {
|
||||
"available": encoder is not None,
|
||||
"encoder": encoder,
|
||||
"reason": None if encoder is not None else JXL_UNAVAILABLE,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def jxl_distance(quality):
|
||||
"""libjxl's distance for a 1–100 quality (q100 gives 0.1, never lossless)."""
|
||||
if quality >= 100:
|
||||
return 0.1
|
||||
if quality >= 30:
|
||||
return 0.1 + (100 - quality) * 0.09
|
||||
return 53.0 / 3000.0 * quality * quality - 23.0 / 20.0 * quality + 25.0
|
||||
|
||||
|
||||
# ── frames and writers ───────────────────────────────────────────────
|
||||
|
||||
|
||||
def thumbnail_frame(img, edge):
|
||||
"""The first frame of *img* as RGB, fitted inside ``edge`` x ``edge``."""
|
||||
try:
|
||||
img.seek(0)
|
||||
except (EOFError, AttributeError, ValueError):
|
||||
pass
|
||||
try:
|
||||
img.draft("RGB", (edge, edge)) # JPEG: decode at a reduced scale
|
||||
except (AttributeError, ValueError, OSError):
|
||||
pass
|
||||
frame = img.convert("RGB") if img.mode != "RGB" else img.copy()
|
||||
frame.thumbnail((edge, edge), Image.Resampling.LANCZOS)
|
||||
return frame
|
||||
|
||||
|
||||
def _png_bytes(img):
|
||||
buffer = io.BytesIO()
|
||||
img.save(buffer, "PNG", compress_level=1)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def _write_webp(img, tmp, settings, track):
|
||||
img.save(tmp, "WEBP", quality=settings.quality, method=WEBP_METHOD)
|
||||
|
||||
|
||||
def _write_jpeg(img, tmp, settings, track):
|
||||
img.save(tmp, "JPEG", quality=settings.quality, optimize=True)
|
||||
|
||||
|
||||
def _check(done, tmp, name):
|
||||
if done.returncode != 0 or not Path(tmp).is_file():
|
||||
detail = (done.stderr or b"")[-300:]
|
||||
if isinstance(detail, bytes):
|
||||
detail = detail.decode("utf-8", "replace")
|
||||
raise EncoderFailed(f"{name} could not encode the preview: {detail.strip()}")
|
||||
|
||||
|
||||
class _EncoderGone(RuntimeError):
|
||||
"""No JPEG XL encoder is left: this file is written as WebP."""
|
||||
|
||||
|
||||
def _write_jxl(img, tmp, settings, track):
|
||||
"""Encode with the detected encoder. One that failed because it is gone
|
||||
makes the write detect again once (another one is used, none raises
|
||||
_EncoderGone); any other failure raises without a probe."""
|
||||
encoder = detect_jxl_encoder()
|
||||
try:
|
||||
_encode_jxl(encoder, img, tmp, settings, track)
|
||||
return
|
||||
except (OSError, EncoderFailed) as exc:
|
||||
if encoder is not None and not _encoder_gone(encoder):
|
||||
raise
|
||||
again = detect_jxl_encoder(refresh=True)
|
||||
if again is None:
|
||||
raise _EncoderGone(str(exc)) from exc
|
||||
if again == encoder:
|
||||
raise
|
||||
logger.info(f"JPEG XL encoder {encoder} failed; using {again}")
|
||||
_unlink_quietly(tmp)
|
||||
_encode_jxl(again, img, tmp, settings, track)
|
||||
|
||||
|
||||
def _encode_jxl(encoder, img, tmp, settings, track):
|
||||
distance = f"{jxl_distance(settings.quality):.2f}"
|
||||
if encoder == "pillow":
|
||||
img.save(
|
||||
tmp,
|
||||
_pillow_jxl_format() or PILLOW_JXL_FORMAT,
|
||||
quality=min(settings.quality, 99),
|
||||
effort=JXL_EFFORT,
|
||||
num_threads=1,
|
||||
)
|
||||
return
|
||||
if encoder == "ffmpeg":
|
||||
cmd = [
|
||||
"ffmpeg", "-v", "error", "-f", "png_pipe", "-i", "-",
|
||||
"-frames:v", "1", "-c:v", "libjxl", "-distance", distance,
|
||||
"-effort", str(JXL_EFFORT), "-f", "image2", "-update", "1",
|
||||
"-y", str(tmp),
|
||||
] # fmt: skip
|
||||
_check(run_encoder(cmd, input=_png_bytes(img), track=track), tmp, "ffmpeg")
|
||||
return
|
||||
if encoder == "cjxl":
|
||||
source = temp_sibling(Path(tmp).with_suffix(".png"), "src")
|
||||
try:
|
||||
with open(source, "wb") as handle:
|
||||
handle.write(_png_bytes(img))
|
||||
cmd = ["cjxl", str(source), str(tmp), "-d", distance, "-e", str(JXL_EFFORT)]
|
||||
_check(run_encoder(cmd, track=track), tmp, "cjxl")
|
||||
finally:
|
||||
_unlink_quietly(source)
|
||||
return
|
||||
raise EncoderFailed(JXL_UNAVAILABLE)
|
||||
|
||||
|
||||
_WRITERS = {"webp": _write_webp, "jpeg": _write_jpeg, "jxl": _write_jxl}
|
||||
|
||||
|
||||
def verify_thumbnail(path, fmt):
|
||||
"""Raise ThumbnailVerifyError unless *path* decodes as a preview.
|
||||
|
||||
A JPEG XL file Pillow cannot read (no plugin) is checked by signature.
|
||||
"""
|
||||
try:
|
||||
if fmt == "jxl" and not _pillow_decodes_jxl():
|
||||
with open(path, "rb") as handle:
|
||||
head = handle.read(12)
|
||||
if not head.startswith(JXL_SIGNATURES):
|
||||
raise ThumbnailVerifyError("The JPEG XL preview has no signature")
|
||||
return
|
||||
with Image.open(path) as written:
|
||||
written.load()
|
||||
except ThumbnailVerifyError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise ThumbnailVerifyError(f"The new preview does not decode: {exc}") from exc
|
||||
|
||||
|
||||
def _jxl_fallback():
|
||||
"""Log once that JPEG XL previews are being written as WebP."""
|
||||
global _fallback_logged
|
||||
with _detect_lock:
|
||||
first, _fallback_logged = not _fallback_logged, True
|
||||
if first:
|
||||
logger.warning(
|
||||
"JPEG XL is the thumbnail format but no JPEG XL encoder is available "
|
||||
"any more; previews are written as WebP until one is installed"
|
||||
)
|
||||
|
||||
|
||||
def writer_settings(settings=None):
|
||||
"""The settings a preview is written with now: the configured ones, but
|
||||
WebP while JPEG XL has no encoder (logged once). Its extension names the
|
||||
file, so "outdated" means a preview this server would not write now."""
|
||||
settings = settings or thumbnail_settings.current()
|
||||
if settings.format == "jxl" and detect_jxl_encoder() is None:
|
||||
_jxl_fallback()
|
||||
return replace(settings, format="webp")
|
||||
return settings
|
||||
|
||||
|
||||
def write_thumbnail(frame, dst, settings, *, guard=None, track=None):
|
||||
"""Encode *frame* into *dst* in ``settings``' format; returns the path
|
||||
written (``.webp`` when JPEG XL fell back). Atomic and verified; raises
|
||||
PreviewWriteHeld, ThumbnailVerifyError, EncoderFailed, EncoderClosed or
|
||||
the OSError of the write."""
|
||||
dst = Path(dst)
|
||||
fmt = settings.format
|
||||
if fmt == "jxl" and detect_jxl_encoder() is None:
|
||||
_jxl_fallback()
|
||||
fmt, dst = "webp", dst.with_suffix(".webp")
|
||||
try:
|
||||
return _write_verified(frame, dst, fmt, settings, guard, track)
|
||||
except _EncoderGone:
|
||||
_jxl_fallback()
|
||||
dst = dst.with_suffix(".webp")
|
||||
return _write_verified(frame, dst, "webp", settings, guard, track)
|
||||
|
||||
|
||||
def _write_verified(frame, dst, fmt, settings, guard, track):
|
||||
tmp = temp_sibling(dst)
|
||||
try:
|
||||
_WRITERS[fmt](frame, tmp, settings, track)
|
||||
verify_thumbnail(tmp, fmt)
|
||||
replace_unless_held(tmp, dst, guard)
|
||||
finally:
|
||||
_unlink_quietly(tmp)
|
||||
return dst
|
||||
+119
-143
@@ -1,32 +1,29 @@
|
||||
"""Thumbnail generation as an Activity job (#249 System S5, Task 12; #317).
|
||||
"""Thumbnail generation as an Activity job (#249 System S5, Task 12).
|
||||
|
||||
- ``GET /prompt_manager/images/thumbnails/plan?mode=`` counts what a run
|
||||
would do over every gallery root and writes nothing: ``{success, targets,
|
||||
fresh, outdated, remaining, roots, mode, format, quality, size}``.
|
||||
- ``POST /prompt_manager/jobs/thumbnails {mode}`` starts the ``thumbnails``
|
||||
job (202 ``{job}``). Its admission class is ``guarded``: it never runs
|
||||
beside a record-changing job, and one thumbnail job runs at a time. It
|
||||
stops at a file boundary when cancelled, keeping the previews it wrote.
|
||||
Its outcome groups are ``generated``, ``skipped`` (already fresh) and
|
||||
``errors``.
|
||||
- ``GET /prompt_manager/images/thumbnails/plan?mode=&size=&quality=`` counts
|
||||
what a run would do over every gallery root and writes nothing:
|
||||
``{success, targets, fresh, remaining, roots, mode, size, quality}``.
|
||||
- ``POST /prompt_manager/jobs/thumbnails {mode, size, quality}`` starts the
|
||||
``thumbnails`` job (202 ``{job}``). Its admission class is ``guarded``: it
|
||||
never runs beside a record-changing job, and one thumbnail job runs at a
|
||||
time. It stops at a file boundary when cancelled, keeping the previews it
|
||||
wrote. Its outcome groups are ``generated``, ``skipped`` (already fresh)
|
||||
and ``errors``.
|
||||
|
||||
``mode`` is ``missing`` (previews that are absent or older than their file),
|
||||
``rebuild`` (every preview again, fresh ones too) or ``convert`` (Convert
|
||||
outdated: previews in another format, see ``thumbnail_convert.py``). Format, quality and
|
||||
size limit are the server's thumbnail settings (``thumbnail_settings``,
|
||||
#317), read once when the job starts and shown in its scope; a ``size`` or
|
||||
``quality`` an older client sends is ignored. Every preview is written next
|
||||
to the old one, verified and renamed over it, so nothing is cleared first and
|
||||
a failure keeps the old preview; originals are only read.
|
||||
``mode`` is ``missing`` (previews that are absent or older than their file)
|
||||
or ``rebuild`` (every preview again, fresh ones too). ``size`` is the
|
||||
preview's longest edge, 128 to 2048 pixels. ``quality`` picks the encoder
|
||||
settings (``THUMBNAIL_ENCODINGS``). Every preview is written next to the old
|
||||
one and renamed over it, so nothing is cleared first and a failure keeps the
|
||||
old preview; originals are only read.
|
||||
|
||||
The legacy gallery stream ``POST /images/generate-thumbnails/progress``
|
||||
starts this job (missing previews, the server's settings; its ``quality``
|
||||
preset is ignored, #317 R3) or attaches to the running one. Its first frame,
|
||||
``job {id, started}``, names that job and whether this stream started it (the
|
||||
gallery cancels by that id, and only a job it started); then come the job's
|
||||
legacy ``status``/``start``/``progress``/``file_error``/``complete`` frames. A
|
||||
client that goes away no longer stops the work: the job outlives the request,
|
||||
as a scan does.
|
||||
starts this job (its quality preset still sizes the previews) or attaches to
|
||||
the running one. Its first frame, ``job {id, started}``, names that job and
|
||||
whether this stream started it (the gallery cancels by that id, and only a
|
||||
job it started); then come the job's legacy ``status``/``start``/
|
||||
``progress``/``file_error``/``complete`` frames. A client that goes away no
|
||||
longer stops the work: the job outlives the request, as a scan does.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -39,6 +36,7 @@ from aiohttp import web
|
||||
|
||||
from .admin import _plural, _scan_location, _scan_roots, _scan_status_line
|
||||
from .images import (
|
||||
THUMBNAIL_ENCODINGS,
|
||||
_BodyTooLarge,
|
||||
_new_thumbnail_stats,
|
||||
_read_json_body,
|
||||
@@ -50,13 +48,24 @@ from .images import (
|
||||
)
|
||||
from .job_routes import busy_response
|
||||
from .jobs import THUMBNAIL_GROUP_NAMES, THUMBNAILS_JOB_KIND, JobBusy
|
||||
from .thumbnail_convert import CONVERT_JOB_TITLE
|
||||
from .thumbnail_encode import writer_settings
|
||||
from .thumbnail_settings import ThumbnailSettings
|
||||
|
||||
THUMBNAIL_JOB_TITLE = "Thumbnail generation"
|
||||
THUMBNAIL_MODES = ("missing", "rebuild", "convert")
|
||||
_OPTIONS_ERROR = 'mode must be "missing", "rebuild" or "convert"'
|
||||
THUMBNAIL_MODES = ("missing", "rebuild")
|
||||
THUMBNAIL_QUALITIES = tuple(THUMBNAIL_ENCODINGS)
|
||||
MIN_THUMBNAIL_EDGE = 128
|
||||
MAX_THUMBNAIL_EDGE = 2048
|
||||
# The concept's "Preview size limit" default.
|
||||
DEFAULT_THUMBNAIL_EDGE = 512
|
||||
DEFAULT_THUMBNAIL_OPTIONS = {
|
||||
"mode": "missing",
|
||||
"size": DEFAULT_THUMBNAIL_EDGE,
|
||||
"quality": "medium",
|
||||
}
|
||||
_OPTIONS_ERROR = (
|
||||
'mode must be "missing" or "rebuild", quality "low", "medium" or "high", '
|
||||
f"and size a whole number of pixels from {MIN_THUMBNAIL_EDGE} "
|
||||
f"to {MAX_THUMBNAIL_EDGE}"
|
||||
)
|
||||
CANCELLED_ERROR = "Thumbnail generation was cancelled"
|
||||
CANCELLED_BEFORE_WRITING = "Cancelled before any preview was written"
|
||||
NOTHING_TO_PREVIEW = "No images or videos to preview"
|
||||
@@ -71,35 +80,56 @@ _LOG_PREFIXES = {"generated": "+", "errors": "!"}
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ThumbnailOptions:
|
||||
"""One run's choices: the mode and the thumbnail settings it writes with."""
|
||||
"""One run's choices: the mode, the preview's longest edge and quality."""
|
||||
|
||||
mode: str
|
||||
settings: ThumbnailSettings
|
||||
edge: int
|
||||
quality: str
|
||||
|
||||
@property
|
||||
def size(self):
|
||||
return (self.edge, self.edge)
|
||||
|
||||
@property
|
||||
def encoding(self):
|
||||
return THUMBNAIL_ENCODINGS[self.quality]
|
||||
|
||||
@property
|
||||
def rebuild(self):
|
||||
return self.mode == "rebuild"
|
||||
|
||||
def scope(self):
|
||||
return {
|
||||
"mode": self.mode,
|
||||
"format": self.settings.format,
|
||||
"quality": self.settings.quality,
|
||||
"size": self.settings.edge,
|
||||
}
|
||||
return {"mode": self.mode, "size": self.edge, "quality": self.quality}
|
||||
|
||||
|
||||
def _edge(value, *, from_query):
|
||||
"""The preview edge as an int, or None when it is not one in range."""
|
||||
if from_query and isinstance(value, str):
|
||||
if not value.isascii() or not value.isdigit() or len(value) > 5:
|
||||
return None
|
||||
value = int(value)
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
return None
|
||||
return value if MIN_THUMBNAIL_EDGE <= value <= MAX_THUMBNAIL_EDGE else None
|
||||
|
||||
|
||||
def thumbnail_options(raw, *, from_query=False):
|
||||
"""``(ThumbnailOptions, None)`` from a body or query, or ``(None, message)``.
|
||||
|
||||
An absent mode is ``missing``; the settings are the server's now.
|
||||
Absent fields take ``DEFAULT_THUMBNAIL_OPTIONS``. Query values are
|
||||
strings, so a query ``size`` may be decimal digits; a JSON ``size`` must
|
||||
be an int.
|
||||
"""
|
||||
if not isinstance(raw, dict) and not hasattr(raw, "get"):
|
||||
return None, _OPTIONS_ERROR
|
||||
mode = raw.get("mode", "missing")
|
||||
if mode not in THUMBNAIL_MODES:
|
||||
mode = raw.get("mode", DEFAULT_THUMBNAIL_OPTIONS["mode"])
|
||||
quality = raw.get("quality", DEFAULT_THUMBNAIL_OPTIONS["quality"])
|
||||
edge = _edge(raw.get("size", DEFAULT_THUMBNAIL_EDGE), from_query=from_query)
|
||||
if mode not in THUMBNAIL_MODES or quality not in THUMBNAIL_QUALITIES:
|
||||
return None, _OPTIONS_ERROR
|
||||
return ThumbnailOptions(mode, writer_settings()), None
|
||||
if edge is None:
|
||||
return None, _OPTIONS_ERROR
|
||||
return ThumbnailOptions(mode, edge, quality), None
|
||||
|
||||
|
||||
def _invalid(message):
|
||||
@@ -138,10 +168,7 @@ class ThumbnailJobsMixin:
|
||||
try:
|
||||
output_dirs = await self._run_in_executor(self._get_all_output_dirs)
|
||||
plan = await self._run_in_executor(
|
||||
self._plan_gallery_thumbnails,
|
||||
output_dirs,
|
||||
options.mode,
|
||||
settings=options.settings,
|
||||
self._plan_gallery_thumbnails, output_dirs, rebuild=options.rebuild
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("Thumbnail plan failed")
|
||||
@@ -158,7 +185,6 @@ class ThumbnailJobsMixin:
|
||||
"success": True,
|
||||
"targets": len(plan["targets"]),
|
||||
"fresh": plan["fresh"],
|
||||
"outdated": plan["outdated"],
|
||||
"remaining": plan["remaining"],
|
||||
"roots": [self._public_path(root) for root in output_dirs],
|
||||
**options.scope(),
|
||||
@@ -189,20 +215,14 @@ class ThumbnailJobsMixin:
|
||||
return web.json_response({"success": True, "job": job.summary()}, status=202)
|
||||
|
||||
def _start_thumbnail_job(self, options):
|
||||
"""Register the thumbnail job synchronously; raises JobBusy.
|
||||
|
||||
``convert`` runs Convert outdated (``thumbnail_convert.py``) under the
|
||||
same kind, admission, cancel and groups.
|
||||
"""
|
||||
convert = options.mode == "convert"
|
||||
runner = self._run_convert_job if convert else self._run_thumbnail_job
|
||||
"""Register the thumbnail job synchronously; raises JobBusy."""
|
||||
return self._jobs.start(
|
||||
THUMBNAILS_JOB_KIND,
|
||||
CONVERT_JOB_TITLE if convert else THUMBNAIL_JOB_TITLE,
|
||||
THUMBNAIL_JOB_TITLE,
|
||||
admission="guarded",
|
||||
cancellable=True,
|
||||
scope=options.scope(),
|
||||
runner=functools.partial(runner, options=options),
|
||||
runner=functools.partial(self._run_thumbnail_job, options=options),
|
||||
groups=THUMBNAIL_GROUP_NAMES,
|
||||
)
|
||||
|
||||
@@ -245,7 +265,7 @@ class ThumbnailJobsMixin:
|
||||
job.update(status_line=f"Looking for media files in {folders}")
|
||||
job.log_line(
|
||||
f"Thumbnails started · {folders} · {_MODE_LOG[options.mode]} · "
|
||||
f"{options.settings.describe()}"
|
||||
f"{options.edge} px · {options.quality}"
|
||||
)
|
||||
job.publish(
|
||||
_frame(
|
||||
@@ -253,83 +273,41 @@ class ThumbnailJobsMixin:
|
||||
{"phase": "scanning", "message": "Scanning for images and videos..."},
|
||||
)
|
||||
)
|
||||
plan = await self._plan_thumbnail_batch(job, options, output_dirs)
|
||||
plan = await self._run_in_executor(
|
||||
self._plan_gallery_thumbnails,
|
||||
output_dirs,
|
||||
rebuild=options.rebuild,
|
||||
stop=lambda: job.cancelled,
|
||||
)
|
||||
if job.cancelled: # M2: stopped during the walk, before any start frame
|
||||
return self._end_cancelled(job, 0, CANCELLED_BEFORE_WRITING)
|
||||
roots = _scan_roots(output_dirs)
|
||||
self._record_fresh(job, plan, roots)
|
||||
stats, start, left, attempted = await self._generate_in_batches(
|
||||
job, options, plan, roots, output_dirs
|
||||
)
|
||||
stats, start = await self._generate_thumbnail_targets(job, options, plan, roots)
|
||||
self.invalidate_gallery_cache()
|
||||
generated = _plural(stats["generated"], "preview")
|
||||
if job.cancelled:
|
||||
return self._end_cancelled(
|
||||
job, stats["generated"], f"Cancelled · {generated} written"
|
||||
)
|
||||
overall = {"targets": attempted, "skipped": plan["skipped"], "remaining": left}
|
||||
payload = _thumbnail_complete_payload(
|
||||
len(attempted) + plan["skipped"] + left,
|
||||
len(plan["targets"]) + plan["skipped"] + plan["remaining"],
|
||||
stats,
|
||||
start,
|
||||
remaining=left,
|
||||
processed=len(attempted),
|
||||
remaining=plan["remaining"],
|
||||
processed=len(plan["targets"]),
|
||||
)
|
||||
job.update(progress=100, status_line=_finished_line(overall, stats))
|
||||
if left:
|
||||
job.log_line(f"{left:,} more previews are left for the next run")
|
||||
job.update(progress=100, status_line=_finished_line(plan, stats))
|
||||
if plan["remaining"]:
|
||||
job.log_line(
|
||||
f"{plan['remaining']:,} more previews are left for the next run"
|
||||
)
|
||||
errors = _plural(len(stats["errors"]), "error")
|
||||
job.log_line(f"Thumbnails finished · {generated} written · {errors}")
|
||||
self.logger.info(payload["message"])
|
||||
job.publish(_frame("complete", payload))
|
||||
return "finished"
|
||||
|
||||
async def _plan_thumbnail_batch(self, job, options, output_dirs, exclude=None):
|
||||
"""One batch's plan (at most MAX_THUMBNAIL_FILES targets), off-loop."""
|
||||
return await self._run_in_executor(
|
||||
self._plan_gallery_thumbnails,
|
||||
output_dirs,
|
||||
options.mode,
|
||||
stop=lambda: job.cancelled,
|
||||
settings=options.settings,
|
||||
exclude=exclude,
|
||||
)
|
||||
|
||||
async def _generate_in_batches(self, job, options, plan, roots, output_dirs):
|
||||
"""Work through ``plan``, then plan again after each batch until
|
||||
nothing is left, the job is cancelled, or a batch made no progress
|
||||
(every file failed: they are not planned again, and the rest waits
|
||||
for the next run). Files worked on are never planned twice. Returns
|
||||
``(stats, start, files left, files worked on)``; progress counts
|
||||
against the overall total."""
|
||||
run = {
|
||||
"index": 0,
|
||||
"batches": 0,
|
||||
"total": len(plan["targets"]) + plan["remaining"],
|
||||
}
|
||||
stats = _new_thumbnail_stats()
|
||||
stats["skipped"] = plan["skipped"]
|
||||
start = time.monotonic()
|
||||
self._publish_thumbnail_start(job, plan, run["total"])
|
||||
attempted = set()
|
||||
while True:
|
||||
before = stats["generated"] + stats["skipped"]
|
||||
await self._generate_thumbnail_targets(
|
||||
job, options, plan, roots, stats, run, start
|
||||
)
|
||||
attempted.update(src for src, _dst, _video in plan["targets"])
|
||||
progressed = stats["generated"] + stats["skipped"] > before
|
||||
if job.cancelled or not plan["remaining"] or not progressed:
|
||||
return stats, start, plan["remaining"], attempted
|
||||
plan = await self._plan_thumbnail_batch(
|
||||
job, options, output_dirs, exclude=attempted
|
||||
)
|
||||
if job.cancelled:
|
||||
return stats, start, 0, attempted
|
||||
run["total"] = run["index"] + len(plan["targets"]) + plan["remaining"]
|
||||
job.update(total=run["total"])
|
||||
if not plan["targets"]:
|
||||
return stats, start, plan["remaining"], attempted
|
||||
|
||||
@staticmethod
|
||||
def _end_cancelled(job, generated, status_line):
|
||||
"""Close a cancelled job: its status line, log line and error frame."""
|
||||
@@ -357,10 +335,10 @@ class ThumbnailJobsMixin:
|
||||
if plan["skipped"]:
|
||||
job.log_line(f"= {plan['skipped']:,} previews already up to date")
|
||||
|
||||
@staticmethod
|
||||
def _publish_thumbnail_start(job, plan, total):
|
||||
"""The legacy stream's start frame, with the job's overall total."""
|
||||
async def _generate_thumbnail_targets(self, job, options, plan, roots):
|
||||
"""Generate ``plan``'s targets in worker batches, publishing each file."""
|
||||
targets = plan["targets"]
|
||||
total = len(targets)
|
||||
video_count = sum(1 for _, _, is_video in targets if is_video)
|
||||
message = (
|
||||
f"Found {total - video_count} images and {video_count} videos to process"
|
||||
@@ -376,7 +354,7 @@ class ThumbnailJobsMixin:
|
||||
{
|
||||
"total_images": total,
|
||||
"skipped": plan["skipped"],
|
||||
"remaining": 0,
|
||||
"remaining": plan["remaining"],
|
||||
"phase": "processing",
|
||||
"image_count": total - video_count,
|
||||
"video_count": video_count,
|
||||
@@ -384,30 +362,28 @@ class ThumbnailJobsMixin:
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
async def _generate_thumbnail_targets(
|
||||
self, job, options, plan, roots, stats, run, start
|
||||
):
|
||||
"""Generate one batch (``plan``'s targets) in worker batches,
|
||||
publishing each file against the job's overall ``run`` counters."""
|
||||
targets = plan["targets"]
|
||||
stats = _new_thumbnail_stats()
|
||||
stats["skipped"] = plan["skipped"]
|
||||
start = time.monotonic()
|
||||
batch_size = _thumbnail_batch_size(_worker_threads())
|
||||
generate = {
|
||||
"encoding": options.encoding,
|
||||
"rebuild": options.rebuild,
|
||||
"stop": lambda: job.cancelled,
|
||||
}
|
||||
for batch_start in range(0, len(targets), batch_size):
|
||||
index = 0
|
||||
batches = 0
|
||||
for batch_start in range(0, total, batch_size):
|
||||
if job.cancelled:
|
||||
break
|
||||
batch = targets[batch_start : batch_start + batch_size]
|
||||
results = await self._run_in_executor(
|
||||
self._generate_batch, batch, options.settings, **generate
|
||||
self._generate_batch, batch, options.size, **generate
|
||||
)
|
||||
for (src, _dst, _video), result in zip(batch, results):
|
||||
if result["action"] == "cancelled":
|
||||
continue
|
||||
run["index"] += 1
|
||||
index, total = run["index"], max(run["total"], run["index"])
|
||||
index += 1
|
||||
_record_thumbnail_result(stats, result)
|
||||
self._record_thumbnail_outcome(job, src, result, roots)
|
||||
job.update(processed=index, progress=index * 100 // total)
|
||||
@@ -424,13 +400,14 @@ class ThumbnailJobsMixin:
|
||||
_thumbnail_progress_payload(index, total, stats, result, start),
|
||||
)
|
||||
)
|
||||
run["batches"] += 1
|
||||
batches += 1
|
||||
job.update(
|
||||
status_line=_scan_status_line(
|
||||
run["index"], run["total"], run["batches"], time.monotonic() - start
|
||||
index, total, batches, time.monotonic() - start
|
||||
)
|
||||
)
|
||||
await asyncio.sleep(0) # let other requests run between batches
|
||||
return stats, start
|
||||
|
||||
@staticmethod
|
||||
def _record_thumbnail_outcome(job, src, result, roots):
|
||||
@@ -453,13 +430,14 @@ class ThumbnailJobsMixin:
|
||||
async def generate_thumbnails_with_progress(self, request):
|
||||
"""The gallery's thumbnail stream (POST, server-sent events).
|
||||
|
||||
Starts the thumbnail job (missing previews with the server's
|
||||
thumbnail settings; the gallery's ``quality`` preset is ignored, R3),
|
||||
or attaches to the running thumbnail job, and streams its frames. A
|
||||
Starts the thumbnail job with the gallery's quality preset as the
|
||||
preview size (missing previews, medium encoding, as before), or
|
||||
attaches to the running thumbnail job, and streams its frames. A
|
||||
legacy hold (the blocking generate or clear) or a record-changing
|
||||
job refuses it with the stream's error frame. The job goes on when
|
||||
the client goes away.
|
||||
"""
|
||||
thumbnail_size = self._thumbnail_size(request.query.get("quality", "medium"))
|
||||
response = web.StreamResponse(
|
||||
status=200,
|
||||
headers={
|
||||
@@ -471,7 +449,7 @@ class ThumbnailJobsMixin:
|
||||
await response.prepare(request)
|
||||
# M3: read the registry after the last await, so a job that started
|
||||
# while the headers went out is attached to, not refused.
|
||||
job, started, busy = self._legacy_stream_job()
|
||||
job, started, busy = self._legacy_stream_job(thumbnail_size)
|
||||
if busy is not None:
|
||||
message = f"{busy.title} already running"
|
||||
await self._emit_sse(
|
||||
@@ -498,21 +476,19 @@ class ThumbnailJobsMixin:
|
||||
await events.aclose()
|
||||
return response
|
||||
|
||||
def _legacy_stream_job(self):
|
||||
def _legacy_stream_job(self, thumbnail_size):
|
||||
"""``(job, started, busy)`` for the legacy stream, synchronously.
|
||||
|
||||
Attaches to the running cancellable thumbnail job, or starts one for
|
||||
the missing previews with the server's settings, or names the job in
|
||||
the way.
|
||||
Attaches to the running cancellable thumbnail job, or starts one with
|
||||
the gallery's preset as the size (missing previews, medium encoding,
|
||||
as the gallery always wrote), or names the job in the way.
|
||||
"""
|
||||
running = self._jobs.active(THUMBNAILS_JOB_KIND)
|
||||
if running is not None and running.scope.get("mode") == "convert":
|
||||
return None, False, running # no legacy frames to follow (#317)
|
||||
if running is not None and running.cancellable:
|
||||
return running, False, None
|
||||
busy = self._jobs.blocker(THUMBNAILS_JOB_KIND, admission="guarded")
|
||||
if busy is not None:
|
||||
return None, False, busy
|
||||
# blocker() said free and nothing awaited since, so the start is admitted.
|
||||
options = ThumbnailOptions("missing", writer_settings())
|
||||
options = ThumbnailOptions("missing", thumbnail_size[0], "medium")
|
||||
return self._start_thumbnail_job(options), True, None
|
||||
|
||||
@@ -1,228 +0,0 @@
|
||||
"""Thumbnail settings and names (#317): the one place every preview writer reads.
|
||||
|
||||
Settings (System › Appearance, ``POST /prompt_manager/settings``, config.json
|
||||
``thumbnails``): ``thumbnail_format`` (``webp`` default, ``jxl``, ``jpeg``),
|
||||
``thumbnail_quality`` (1–100, default 80) and ``thumbnail_edge`` (the longest
|
||||
edge, 128–2048 px, default 512). ``current()`` is the snapshot a writer takes
|
||||
once per preview; automatic previews, the Activity job, the history route and
|
||||
the legacy gallery routes all go through it.
|
||||
|
||||
Names: a media file ``<rel>/<stem>.<ext>`` previews as
|
||||
``thumbnails/<rel>/<stem>_thumb<extension of the format>`` (``.webp``,
|
||||
``.jxl``, ``.jpg``), images and videos alike. Two files with one stem in one
|
||||
folder share that name (ruling R2). Older releases kept the source's own
|
||||
extension (``.png``, ``.jpeg``, ``.gif``, ``.webp``, ``.jpg``; ``.jpg`` for a
|
||||
video): those are read, served and converted, never written. A preview with a
|
||||
format extension (``.webp``, ``.jxl``, ``.jpg``) belongs to every same-stem
|
||||
source; a legacy-only one (``.png``, ``.jpeg``, ``.gif``) only to a source
|
||||
with that extension (``owns``, ruling R4). Lookups take the first existing
|
||||
preview the source owns, preferring the configured format.
|
||||
"""
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
FORMATS = ("webp", "jxl", "jpeg")
|
||||
FORMAT_EXTENSIONS = {"webp": ".webp", "jxl": ".jxl", "jpeg": ".jpg"}
|
||||
FORMAT_LABELS = {"webp": "WebP", "jxl": "JPEG XL", "jpeg": "JPEG"}
|
||||
DEFAULT_FORMAT = "webp"
|
||||
DEFAULT_QUALITY = 80
|
||||
DEFAULT_EDGE = 512
|
||||
MIN_QUALITY, MAX_QUALITY = 1, 100
|
||||
MIN_EDGE, MAX_EDGE = 128, 2048
|
||||
|
||||
# Extensions a format writes, then the ones only older releases wrote.
|
||||
WRITTEN_EXTENSIONS = (".webp", ".jxl", ".jpg")
|
||||
LEGACY_EXTENSIONS = (".png", ".jpeg", ".gif")
|
||||
THUMBNAIL_EXTENSIONS = WRITTEN_EXTENSIONS + LEGACY_EXTENSIONS
|
||||
CONTENT_TYPES = {
|
||||
".webp": "image/webp",
|
||||
".jxl": "image/jxl",
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".png": "image/png",
|
||||
".gif": "image/gif",
|
||||
}
|
||||
THUMB_SUFFIX = "_thumb"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ThumbnailSettings:
|
||||
"""One preview's settings: the format, the quality and the longest edge."""
|
||||
|
||||
format: str
|
||||
quality: int
|
||||
edge: int
|
||||
|
||||
@property
|
||||
def extension(self):
|
||||
return FORMAT_EXTENSIONS[self.format]
|
||||
|
||||
@property
|
||||
def size(self):
|
||||
return (self.edge, self.edge)
|
||||
|
||||
def describe(self):
|
||||
"""``WebP · quality 80 · 512 px`` (log lines and the job's scope)."""
|
||||
return f"{FORMAT_LABELS[self.format]} · quality {self.quality} · {self.edge} px"
|
||||
|
||||
|
||||
def current():
|
||||
"""The configured settings now (read once per preview or per run)."""
|
||||
from ..config import PromptManagerConfig
|
||||
|
||||
fmt = PromptManagerConfig.THUMBNAIL_FORMAT
|
||||
return ThumbnailSettings(
|
||||
fmt if fmt in FORMATS else DEFAULT_FORMAT,
|
||||
int(PromptManagerConfig.THUMBNAIL_QUALITY),
|
||||
int(PromptManagerConfig.THUMBNAIL_EDGE),
|
||||
)
|
||||
|
||||
|
||||
def _whole(value, low, high, message):
|
||||
if type(value) is not int or not low <= value <= high:
|
||||
raise ValueError(message)
|
||||
return value
|
||||
|
||||
|
||||
def validate_format(value):
|
||||
if not isinstance(value, str) or value not in FORMATS:
|
||||
raise ValueError('thumbnail_format must be "webp", "jxl" or "jpeg"')
|
||||
return value
|
||||
|
||||
|
||||
def validate_quality(value):
|
||||
return _whole(
|
||||
value,
|
||||
MIN_QUALITY,
|
||||
MAX_QUALITY,
|
||||
f"thumbnail_quality must be a whole number from {MIN_QUALITY} to {MAX_QUALITY}",
|
||||
)
|
||||
|
||||
|
||||
def validate_edge(value):
|
||||
return _whole(
|
||||
value,
|
||||
MIN_EDGE,
|
||||
MAX_EDGE,
|
||||
"thumbnail_edge must be a whole number of pixels "
|
||||
f"from {MIN_EDGE} to {MAX_EDGE}",
|
||||
)
|
||||
|
||||
|
||||
# ── names ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def thumbnail_stem_rel(rel_path):
|
||||
"""``thumbnails/<rel without extension>_thumb`` (POSIX)."""
|
||||
return f"thumbnails/{Path(rel_path).with_suffix('').as_posix()}{THUMB_SUFFIX}"
|
||||
|
||||
|
||||
def thumbnail_rel(rel_path, extension=None):
|
||||
"""The preview of *rel_path* in *extension* (the configured one when None)."""
|
||||
return thumbnail_stem_rel(rel_path) + (extension or current().extension)
|
||||
|
||||
|
||||
# Ruling R2: files sharing a stem share one preview, made from the first
|
||||
# source in this order: these images, then videos (by extension name).
|
||||
IMAGE_SOURCE_ORDER = (".png", ".jpg", ".jpeg", ".webp", ".gif")
|
||||
|
||||
|
||||
def source_rank(source_extension):
|
||||
"""Sort key of a source extension in R2's order (images first)."""
|
||||
extension = source_extension.lower()
|
||||
if extension in IMAGE_SOURCE_ORDER:
|
||||
return (0, IMAGE_SOURCE_ORDER.index(extension), extension)
|
||||
return (1, 0, extension)
|
||||
|
||||
|
||||
def owns(thumb_extension, source_extension):
|
||||
"""Whether a ``_thumb<thumb_extension>`` preview can belong to a source
|
||||
with *source_extension* (ruling R4)."""
|
||||
thumb_extension = thumb_extension.lower()
|
||||
if thumb_extension in WRITTEN_EXTENSIONS:
|
||||
return True
|
||||
if thumb_extension in LEGACY_EXTENSIONS:
|
||||
return thumb_extension == source_extension.lower()
|
||||
return False
|
||||
|
||||
|
||||
def _legacy_extension(source_extension):
|
||||
from .inspection_service import VIDEO_EXTENSIONS
|
||||
|
||||
source_extension = source_extension.lower()
|
||||
return ".jpg" if source_extension in VIDEO_EXTENSIONS else source_extension
|
||||
|
||||
|
||||
def candidate_extensions(source_extension, preferred=None):
|
||||
"""The preview extensions a source owns, in lookup order: the configured
|
||||
one, the name an older release gave it, then the other written ones."""
|
||||
preferred = preferred or current().extension
|
||||
ordered = [preferred, _legacy_extension(source_extension), *WRITTEN_EXTENSIONS]
|
||||
seen = []
|
||||
for extension in ordered:
|
||||
if extension not in seen and owns(extension, source_extension):
|
||||
seen.append(extension)
|
||||
return tuple(seen)
|
||||
|
||||
|
||||
def owned_extensions(source_extension, preferred=None):
|
||||
"""``candidate_extensions`` plus the other spellings an older release
|
||||
may have given a legacy preview (review M1): the source's own
|
||||
(``X_thumb.JPG`` for ``X.JPG``) and the upper-case legacy one."""
|
||||
extensions = list(candidate_extensions(source_extension, preferred))
|
||||
legacy = _legacy_extension(source_extension)
|
||||
for spelling in (source_extension, legacy.upper()):
|
||||
if spelling not in extensions and owns(spelling, source_extension):
|
||||
extensions.append(spelling)
|
||||
return tuple(extensions)
|
||||
|
||||
|
||||
def _is_preview_file(path):
|
||||
try:
|
||||
return not path.is_symlink() and path.is_file()
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
def owned_in_folder(folder, stem, source_extension, preferred=None):
|
||||
"""Yield the existing previews ``<folder>/<stem><ext>`` a source with
|
||||
*source_extension* owns, in lookup order (regular files, no links; one
|
||||
file found under two spellings on a case-insensitive disk once)."""
|
||||
seen = set()
|
||||
for extension in owned_extensions(source_extension, preferred):
|
||||
path = Path(folder) / (stem + extension)
|
||||
if not _is_preview_file(path):
|
||||
continue
|
||||
try:
|
||||
stat = path.stat()
|
||||
identity = (stat.st_dev, stat.st_ino) if stat.st_ino else str(path)
|
||||
except OSError:
|
||||
continue
|
||||
if identity in seen:
|
||||
continue # the same file under another case (Windows, macOS)
|
||||
seen.add(identity)
|
||||
yield path
|
||||
|
||||
|
||||
def _owned(root, rel_path, preferred):
|
||||
rel_path = Path(rel_path)
|
||||
stem = Path(root) / thumbnail_stem_rel(rel_path)
|
||||
return owned_in_folder(stem.parent, stem.name, rel_path.suffix, preferred)
|
||||
|
||||
|
||||
def owned_thumbnails(root, rel_path, preferred=None):
|
||||
"""Every existing preview of *rel_path* under *root*, in lookup order."""
|
||||
return list(_owned(root, rel_path, preferred))
|
||||
|
||||
|
||||
def existing_thumbnail(root, rel_path, preferred=None):
|
||||
"""The preview a listing shows for *rel_path*: the first owned one that
|
||||
exists (the configured format first), or None."""
|
||||
return next(_owned(root, rel_path, preferred), None)
|
||||
|
||||
|
||||
def content_type(path):
|
||||
"""The Content-Type of a preview file by its extension, or None."""
|
||||
return CONTENT_TYPES.get(os.path.splitext(str(path))[1].lower())
|
||||
@@ -1,277 +0,0 @@
|
||||
"""Store each workflow once: the migration and Reclaim space jobs (#249 S6).
|
||||
|
||||
Reviewed like every maintenance operation (``POST /prompt_manager/
|
||||
maintenance/preview {operation:"workflow_storage"}``, a ``Plan`` with the
|
||||
records to rewrite, the backup and the disk check) and applied with its token
|
||||
(``POST /prompt_manager/maintenance/apply``). The apply starts:
|
||||
|
||||
1. **Store each workflow once**: kind ``maintenance``, admission ``records``,
|
||||
cancellable, groups ``changed``/``protected``/``errors``. It always takes
|
||||
a verified pre-migration backup, whatever the safety backup policy, named
|
||||
``<db>.pre-workflow-storage_<stamp>`` outside the pruned ``<db>.backup_*``
|
||||
pool and kept until the user removes it (``backup-failed`` ends the job
|
||||
with nothing rewritten). It then rewrites old-form rows batch by batch
|
||||
(``database/workflow_storage.py``): one transaction per batch of at most
|
||||
``BATCH_SECONDS``, then ``BATCH_PAUSE`` with the write lock free, progress
|
||||
after each, Cancel honoured between batches. A row that cannot be
|
||||
rewritten stays inline and is listed in ``errors``; any other failure
|
||||
rolls its batch back whole and fails the job. Every record still reads as
|
||||
before, and the next apply continues from the first old-form row.
|
||||
2. **Reclaim space**, once every record was looked at: kind ``maintenance``,
|
||||
admission ``exclusive`` (it runs alone; VACUUM holds the write lock), not
|
||||
cancellable. It checks the disk for VACUUM's copies of the live pages
|
||||
(``disk_plan``), then sweeps, runs VACUUM and truncates the WAL. When
|
||||
another job holds the registry, it does not start and the finished
|
||||
migration says so and names Optimize database.
|
||||
|
||||
Only the database changes: no image or video file is opened.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import sqlite3
|
||||
|
||||
from .jobs import MAINTENANCE_JOB_KIND, JobBusy
|
||||
from .maintenance import (
|
||||
MAINTENANCE_GROUPS,
|
||||
PRE_MIGRATION_LABEL,
|
||||
SafetyBackupFailed,
|
||||
backup_failed_message,
|
||||
_fail_job,
|
||||
_megabytes,
|
||||
disk_plan,
|
||||
short_space_message,
|
||||
take_safety_backup,
|
||||
)
|
||||
|
||||
try:
|
||||
from ...database import workflow_storage
|
||||
from ...database.health import readonly_uri
|
||||
except ImportError:
|
||||
from database import workflow_storage
|
||||
from database.health import readonly_uri
|
||||
|
||||
OPERATION = workflow_storage.OPERATION
|
||||
TITLE = "Store each workflow once"
|
||||
RECLAIM_TITLE = "Reclaim space"
|
||||
REWRITE_FAILED = (
|
||||
"Storing workflows once stopped on an error. The batch in progress was "
|
||||
"rolled back; every record reads as before. Run it again to continue."
|
||||
)
|
||||
RECLAIM_FAILED = "Reclaim space failed; the database is unchanged."
|
||||
ROW_KEPT = "could not be rewritten; it keeps its inline text and reads as before"
|
||||
|
||||
|
||||
NOTHING_LEFT = "Every image record already stores its workflow once"
|
||||
|
||||
|
||||
def _plural(count, one, many=None):
|
||||
return f"{count:,} {one if count == 1 else (many or one + 's')}"
|
||||
|
||||
|
||||
def live_bytes(db_path):
|
||||
"""The pages VACUUM copies: page_count minus the free list (blocking)."""
|
||||
conn = sqlite3.connect(readonly_uri(db_path), uri=True)
|
||||
try:
|
||||
page_size = conn.execute("PRAGMA page_size").fetchone()[0]
|
||||
pages = conn.execute("PRAGMA page_count").fetchone()[0]
|
||||
free = conn.execute("PRAGMA freelist_count").fetchone()[0]
|
||||
return page_size * (pages - free)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
class WorkflowStorageRoutesMixin:
|
||||
"""Mixin for the maintenance routes: the workflow storage jobs."""
|
||||
|
||||
def _start_workflow_storage(self, db, backup):
|
||||
"""Start the migration job (raises JobBusy); the token is the caller's."""
|
||||
return self._jobs.start(
|
||||
MAINTENANCE_JOB_KIND,
|
||||
TITLE,
|
||||
admission="records",
|
||||
cancellable=True,
|
||||
scope={"operation": OPERATION},
|
||||
groups=MAINTENANCE_GROUPS,
|
||||
runner=lambda job: self._run_workflow_storage(job, db, backup),
|
||||
)
|
||||
|
||||
async def _run_workflow_storage(self, job, db, _backup_on):
|
||||
# Always backed up first, whatever the policy: the way back. The first
|
||||
# apply keeps a pre-migration backup; later ones (after a Cancel, a
|
||||
# rerun) take a pooled safety backup, so only one copy is never pruned.
|
||||
counts = await self._run_in_executor(db.workflow_storage_status)
|
||||
if not counts["rows"]:
|
||||
job.log_line(NOTHING_LEFT)
|
||||
job.update(progress=100, status_line=NOTHING_LEFT)
|
||||
return
|
||||
kept = await self._run_in_executor(
|
||||
workflow_storage.pre_migration_backups, db.model.db_path
|
||||
)
|
||||
label = "backup" if kept else PRE_MIGRATION_LABEL
|
||||
backup = await self._backup_before_rewrite(job, db, label, kept)
|
||||
if backup is None:
|
||||
return
|
||||
try:
|
||||
done = await self._rewrite_workflows(job, db, backup)
|
||||
finally:
|
||||
if kept:
|
||||
await self._prune_backup_pool(job, db)
|
||||
if done:
|
||||
job.finish("finished")
|
||||
self._start_reclaim_space(job, db)
|
||||
|
||||
async def _backup_before_rewrite(self, job, db, label, kept):
|
||||
"""The verified backup's status note, or None (the job failed).
|
||||
|
||||
``label`` ``pre-workflow-storage``: kept until the user removes it.
|
||||
``backup``: a pooled safety backup, while the earliest kept
|
||||
pre-migration backup (``kept[0]``) stays. A failure ends the job
|
||||
``backup-failed``, nothing changed.
|
||||
"""
|
||||
noun = "safety backup" if kept else "pre-migration backup"
|
||||
job.update(status_line=f"Taking the {noun}")
|
||||
job.log_line(f"Taking the {noun} of the database")
|
||||
try:
|
||||
name = await self._run_in_executor(take_safety_backup, db.model, label)
|
||||
except SafetyBackupFailed as error:
|
||||
self.logger.error("Backup before the rewrite failed", exc_info=True)
|
||||
_fail_job(job, "backup-failed", backup_failed_message(error))
|
||||
return None
|
||||
if kept:
|
||||
job.log_line(
|
||||
f"Safety backup {name} verified; the pre-migration backup "
|
||||
f"{kept[0]} is kept"
|
||||
)
|
||||
return f"Safety backup {name}"
|
||||
job.log_line(
|
||||
f"Pre-migration backup {name} verified; it is kept until you remove it"
|
||||
)
|
||||
return f"Pre-migration backup {name}"
|
||||
|
||||
async def _rewrite_workflows(self, job, db, backup):
|
||||
"""Rewrite every old-form row; True when every row was looked at."""
|
||||
counts = await self._run_in_executor(db.workflow_storage_status)
|
||||
job.update(total=counts["rows"], processed=0)
|
||||
job.update(status_line="Rewriting image records")
|
||||
job.log_line(f"{_plural(counts['rows'], 'image record')} to rewrite")
|
||||
rows = protected = errors = 0
|
||||
last_id = 0
|
||||
while True:
|
||||
if job.cancelled:
|
||||
return self._stopped(job, rows, backup)
|
||||
try:
|
||||
batch = await self._run_in_executor(
|
||||
db.migrate_workflow_batch,
|
||||
last_id,
|
||||
workflow_storage.BATCH_ROWS,
|
||||
workflow_storage.BATCH_BYTES,
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("Storing workflows once failed")
|
||||
_fail_job(job, None, REWRITE_FAILED)
|
||||
return False
|
||||
if batch["done"]:
|
||||
break
|
||||
last_id = batch["last_id"]
|
||||
rows += batch["rows"]
|
||||
protected += batch["protected"]
|
||||
errors += len(batch["errors"])
|
||||
job.record_counted("changed", batch["rows"])
|
||||
job.record_counted("protected", batch["protected"])
|
||||
for image_id in batch["errors"]:
|
||||
job.record(
|
||||
"errors",
|
||||
{"path": "", "detail": f"Image record {image_id}: {ROW_KEPT}"},
|
||||
f"! Image record {image_id}: {ROW_KEPT}",
|
||||
)
|
||||
total = max(job.total, rows)
|
||||
job.update(
|
||||
processed=rows,
|
||||
total=total,
|
||||
progress=100 * rows // total if total else 100,
|
||||
status_line=f"Rewrote {_plural(rows, 'image record')}",
|
||||
)
|
||||
job.publish({"type": "progress", "processed": rows, "total": total})
|
||||
# Leave the write lock free between batches (monitor, run counts).
|
||||
await asyncio.sleep(workflow_storage.BATCH_PAUSE)
|
||||
removed = await self._run_in_executor(db.sweep_workflow_documents)
|
||||
line = f"Rewrote {_plural(rows, 'image record')}"
|
||||
if protected:
|
||||
line += f"; {_plural(protected, 'stored value')} not text, kept as is"
|
||||
if errors:
|
||||
kept = "keeps its" if errors == 1 else "keep their"
|
||||
line += (
|
||||
f"; {_plural(errors, 'image record')} could not be rewritten and "
|
||||
f"{kept} inline text"
|
||||
)
|
||||
job.log_line(line)
|
||||
if removed:
|
||||
job.log_line(f"Removed {_plural(removed, 'unused workflow document')}")
|
||||
job.update(progress=100, status_line=f"{line} · {backup}")
|
||||
return True
|
||||
|
||||
def _stopped(self, job, rows, backup):
|
||||
line = (
|
||||
f"Stopped after {_plural(rows, 'image record')}; "
|
||||
"run it again to continue"
|
||||
)
|
||||
job.log_line(line)
|
||||
job.update(status_line=f"{line} · {backup}")
|
||||
job.finish("cancelled")
|
||||
return False
|
||||
|
||||
def _start_reclaim_space(self, finished, db):
|
||||
"""Start Reclaim space after the migration ``finished``, or say why not."""
|
||||
try:
|
||||
self._jobs.start(
|
||||
MAINTENANCE_JOB_KIND,
|
||||
RECLAIM_TITLE,
|
||||
admission="exclusive",
|
||||
cancellable=False,
|
||||
scope={"operation": OPERATION},
|
||||
groups=MAINTENANCE_GROUPS,
|
||||
runner=lambda job: self._run_reclaim_space(job, db),
|
||||
)
|
||||
except JobBusy as busy:
|
||||
line = (
|
||||
f"Reclaim space did not start: {busy.running.title} is running. "
|
||||
"Run Optimize database to give the space back."
|
||||
)
|
||||
finished.log_line(line)
|
||||
finished.update(status_line=f"{finished.status_line} · {line}")
|
||||
return
|
||||
finished.log_line("Reclaim space started")
|
||||
|
||||
async def _run_reclaim_space(self, job, db):
|
||||
job.update(status_line="Checking free disk space")
|
||||
try:
|
||||
live = await self._run_in_executor(live_bytes, db.model.db_path)
|
||||
_disk, short, _size = await self._run_in_executor(
|
||||
disk_plan,
|
||||
db.model.db_path,
|
||||
backup=False,
|
||||
vacuum=True,
|
||||
vacuum_bytes=live,
|
||||
)
|
||||
except (OSError, sqlite3.Error):
|
||||
self.logger.exception("Could not measure the disk space for Reclaim space")
|
||||
return _fail_job(job, None, RECLAIM_FAILED)
|
||||
if short is not None:
|
||||
return _fail_job(job, "insufficient-space", short_space_message(short))
|
||||
job.update(status_line="Reclaiming space (the database is locked)")
|
||||
job.log_line("Rebuilding the database file (VACUUM)")
|
||||
try:
|
||||
result = await self._run_in_executor(db.reclaim_space)
|
||||
except Exception:
|
||||
self.logger.exception("Reclaim space failed")
|
||||
return _fail_job(job, None, RECLAIM_FAILED)
|
||||
saved = max(0, result["before_bytes"] - result["after_bytes"])
|
||||
line = (
|
||||
f"Reclaimed {_megabytes(saved)} · the database is now "
|
||||
f"{_megabytes(result['after_bytes'])}"
|
||||
)
|
||||
if result["checkpoint_busy"]:
|
||||
line += " (the rest after the next restart)"
|
||||
job.record_counted("changed", 1)
|
||||
job.log_line(line)
|
||||
job.update(progress=100, status_line=line)
|
||||
+28
-30
@@ -25,23 +25,6 @@ except ImportError:
|
||||
from database.annotations import MAX_IMAGE_TAGS
|
||||
from database.vocabulary import validate_vocabulary_tag_name
|
||||
|
||||
# The tag template and thresholds live with the saved tagging defaults
|
||||
# (#328), so config.py never imports PIL.
|
||||
try:
|
||||
from .tagging_defaults import DEFAULT_TAG_TEMPLATE
|
||||
except ImportError:
|
||||
from py.tagging_defaults import DEFAULT_TAG_TEMPLATE
|
||||
|
||||
|
||||
def _autotag_config():
|
||||
"""The saved defaults; imported late, as config needs ComfyUI's server."""
|
||||
try:
|
||||
from .config import AutotagConfig
|
||||
except ImportError:
|
||||
from py.config import AutotagConfig
|
||||
return AutotagConfig
|
||||
|
||||
|
||||
# Try to import logging from utils, fallback to standard logging
|
||||
try:
|
||||
from ..utils.logging_config import get_logger
|
||||
@@ -62,12 +45,10 @@ except ImportError:
|
||||
return logger
|
||||
|
||||
|
||||
# Model configurations - same as standalone_tagger.py. ``label`` is the short
|
||||
# name System, Activity and a 409 model-busy show (#328).
|
||||
# Model configurations - same as standalone_tagger.py
|
||||
MODELS = {
|
||||
"gguf": {
|
||||
"name": "GGUF (Recommended)",
|
||||
"label": "GGUF JoyCaption",
|
||||
"description": "Quantized model, ~4.5GB, faster loading",
|
||||
"size": "~4.5GB",
|
||||
"repo": "mradermacher/llama-joycaption-beta-one-hf-llava-GGUF",
|
||||
@@ -78,7 +59,6 @@ MODELS = {
|
||||
},
|
||||
"hf": {
|
||||
"name": "HuggingFace",
|
||||
"label": "HuggingFace JoyCaption",
|
||||
"description": "Full precision model, ~16GB, higher quality",
|
||||
"size": "~16GB",
|
||||
"repo": "fancyfeast/llama-joycaption-beta-one-hf-llava",
|
||||
@@ -86,7 +66,6 @@ MODELS = {
|
||||
},
|
||||
"wd14-swinv2": {
|
||||
"name": "WD14 SwinV2",
|
||||
"label": "WD14 SwinV2",
|
||||
"description": "Fast classifier tagger, ~400MB, no prompt needed",
|
||||
"size": "~400MB",
|
||||
"repo": "SmilingWolf/wd-swinv2-tagger-v3",
|
||||
@@ -96,7 +75,6 @@ MODELS = {
|
||||
},
|
||||
"wd14-vit": {
|
||||
"name": "WD14 ViT",
|
||||
"label": "WD14 ViT",
|
||||
"description": "Fast classifier tagger, ~400MB, ViT architecture",
|
||||
"size": "~400MB",
|
||||
"repo": "SmilingWolf/wd-vit-tagger-v3",
|
||||
@@ -107,7 +85,17 @@ MODELS = {
|
||||
}
|
||||
|
||||
# Default prompts
|
||||
DEFAULT_PROMPT = DEFAULT_TAG_TEMPLATE
|
||||
DEFAULT_PROMPT = (
|
||||
"Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces).\n"
|
||||
"Include: subject_count (1girl/1boy/etc), gender, ethnicity (asian, caucasian, "
|
||||
"etc), age_approx (teen, young_woman), body_attributes, face_features, "
|
||||
"hairstyle, appearance, clothing, accessories, pose, expression, action, "
|
||||
"camera_angle, camera_shot, scene_composition (wide_shot, long_shot, "
|
||||
"dutch_angle), environment, background_elements, lighting, mood.\n"
|
||||
"Prioritize visual accuracy.\n"
|
||||
"No sentences, no commentary, no captions. Only tags. Keep length short but "
|
||||
"descriptive."
|
||||
)
|
||||
|
||||
# Upper bound of one suggestion request: the image tag set cap.
|
||||
MAX_SUGGESTION_LIMIT = MAX_IMAGE_TAGS
|
||||
@@ -191,6 +179,10 @@ class AutoTagService:
|
||||
self._current_model_type: Optional[str] = None
|
||||
self._custom_prompt: str = DEFAULT_PROMPT
|
||||
|
||||
# WD14 threshold defaults
|
||||
self._wd14_general_threshold: float = 0.35
|
||||
self._wd14_character_threshold: float = 0.85
|
||||
|
||||
@property
|
||||
def models_config(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Get the models configuration dictionary."""
|
||||
@@ -213,13 +205,19 @@ class AutoTagService:
|
||||
|
||||
@property
|
||||
def wd14_general_threshold(self) -> float:
|
||||
"""The saved General-tag confidence (``autotag.defaults``)."""
|
||||
return _autotag_config().GENERAL_THRESHOLD
|
||||
return self._wd14_general_threshold
|
||||
|
||||
@wd14_general_threshold.setter
|
||||
def wd14_general_threshold(self, value: float):
|
||||
self._wd14_general_threshold = max(0.0, min(1.0, value))
|
||||
|
||||
@property
|
||||
def wd14_character_threshold(self) -> float:
|
||||
"""The saved Character-tag confidence (``autotag.defaults``)."""
|
||||
return _autotag_config().CHARACTER_THRESHOLD
|
||||
return self._wd14_character_threshold
|
||||
|
||||
@wd14_character_threshold.setter
|
||||
def wd14_character_threshold(self, value: float):
|
||||
self._wd14_character_threshold = max(0.0, min(1.0, value))
|
||||
|
||||
def get_models_status(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Get availability status for all model types.
|
||||
@@ -862,12 +860,12 @@ class AutoTagService:
|
||||
gen_thresh = (
|
||||
general_threshold
|
||||
if general_threshold is not None
|
||||
else _autotag_config().GENERAL_THRESHOLD
|
||||
else self._wd14_general_threshold
|
||||
)
|
||||
char_thresh = (
|
||||
character_threshold
|
||||
if character_threshold is not None
|
||||
else _autotag_config().CHARACTER_THRESHOLD
|
||||
else self._wd14_character_threshold
|
||||
)
|
||||
return self._generate_wd14(
|
||||
session, tags_list, image, gen_thresh, char_thresh
|
||||
|
||||
+28
-146
@@ -36,14 +36,12 @@ from typing import Any, Dict, List, Optional, Tuple # noqa: E402
|
||||
# Import logging system
|
||||
try:
|
||||
from ..utils.logging_config import get_logger
|
||||
from ..utils import library_location
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, current_dir)
|
||||
from utils.logging_config import get_logger
|
||||
from utils import library_location
|
||||
|
||||
# Initialize logger for config operations
|
||||
config_logger = get_logger("prompt_manager.config")
|
||||
@@ -202,6 +200,7 @@ class GalleryConfig:
|
||||
AUTO_CLEANUP_MISSING_FILES (bool): Automatically remove missing file records
|
||||
MAX_IMAGE_AGE_DAYS (int): Maximum age in days before cleaning up images
|
||||
IMAGES_PER_PAGE (int): Number of images to display per page in web UI
|
||||
THUMBNAIL_SIZE (int): Size in pixels for generated thumbnails
|
||||
ENABLE_SEARCH (bool): Enable search functionality in web interface
|
||||
ENABLE_METADATA_VIEW (bool): Enable metadata viewing for images
|
||||
MAX_CONCURRENT_PROCESSING (int): Maximum concurrent image processing tasks
|
||||
@@ -226,6 +225,7 @@ class GalleryConfig:
|
||||
|
||||
# Web interface settings
|
||||
IMAGES_PER_PAGE = 20
|
||||
THUMBNAIL_SIZE = 256
|
||||
ENABLE_SEARCH = True
|
||||
ENABLE_METADATA_VIEW = True
|
||||
|
||||
@@ -265,6 +265,7 @@ class GalleryConfig:
|
||||
},
|
||||
"web_interface": {
|
||||
"images_per_page": cls.IMAGES_PER_PAGE,
|
||||
"thumbnail_size": cls.THUMBNAIL_SIZE,
|
||||
"enable_search": cls.ENABLE_SEARCH,
|
||||
"enable_metadata_view": cls.ENABLE_METADATA_VIEW,
|
||||
},
|
||||
@@ -327,6 +328,8 @@ class GalleryConfig:
|
||||
web_interface = new_config.get("web_interface", {})
|
||||
if "images_per_page" in web_interface:
|
||||
cls.IMAGES_PER_PAGE = web_interface["images_per_page"]
|
||||
if "thumbnail_size" in web_interface:
|
||||
cls.THUMBNAIL_SIZE = web_interface["thumbnail_size"]
|
||||
if "enable_search" in web_interface:
|
||||
cls.ENABLE_SEARCH = web_interface["enable_search"]
|
||||
if "enable_metadata_view" in web_interface:
|
||||
@@ -563,90 +566,18 @@ class IntegrationConfig:
|
||||
DEFAULT_SUGGESTION_LIMIT = 10
|
||||
MAX_SUGGESTION_LIMIT = 100 # pinned to MAX_IMAGE_TAGS by tests/test_config.py
|
||||
|
||||
# The tagging engines, options, template and validators (server-free, #328).
|
||||
try:
|
||||
from .tagging_defaults import ( # noqa: F401 - re-exported
|
||||
DEFAULT_TAG_TEMPLATE,
|
||||
EXISTING_POLICIES,
|
||||
LLM_MODEL_TYPES,
|
||||
MAX_INSTRUCTION_BYTES,
|
||||
MAX_INSTRUCTION_CHARS,
|
||||
MODEL_TYPES,
|
||||
WD14_MODEL_TYPES,
|
||||
validate_template,
|
||||
validate_threshold,
|
||||
)
|
||||
except ImportError:
|
||||
from py.tagging_defaults import ( # noqa: F401 - re-exported
|
||||
DEFAULT_TAG_TEMPLATE,
|
||||
EXISTING_POLICIES,
|
||||
LLM_MODEL_TYPES,
|
||||
MAX_INSTRUCTION_BYTES,
|
||||
MAX_INSTRUCTION_CHARS,
|
||||
MODEL_TYPES,
|
||||
WD14_MODEL_TYPES,
|
||||
validate_template,
|
||||
validate_threshold,
|
||||
)
|
||||
|
||||
|
||||
def _valid_model_type(value: Any) -> str:
|
||||
if not isinstance(value, str) or value not in MODEL_TYPES:
|
||||
raise ValueError("must be one of " + ", ".join(sorted(MODEL_TYPES)))
|
||||
return value
|
||||
|
||||
|
||||
def _valid_existing(value: Any) -> str:
|
||||
if not isinstance(value, str) or value not in EXISTING_POLICIES:
|
||||
raise ValueError("must be one of " + ", ".join(EXISTING_POLICIES))
|
||||
return value
|
||||
|
||||
|
||||
def _valid_suggestion_limit(value: Any) -> int:
|
||||
if type(value) is not int or not 1 <= value <= MAX_SUGGESTION_LIMIT:
|
||||
raise ValueError(f"must be an integer from 1 to {MAX_SUGGESTION_LIMIT}")
|
||||
return value
|
||||
|
||||
|
||||
class AutotagConfig:
|
||||
"""Persisted tagging defaults under ``autotag.defaults`` in config.json.
|
||||
|
||||
One tagging configuration (#328): the model, what a batch does with
|
||||
prompts already tagged, the two WD14 thresholds, the language-model
|
||||
template and how many suggestions Generate returns. Each key is checked
|
||||
alone on load; a hand-edited value that fails its check falls back to
|
||||
its default with one warning rather than failing the load.
|
||||
``suggestion_limit`` is how many suggestions Generate returns (the top N
|
||||
by score), an integer in ``1..MAX_SUGGESTION_LIMIT``. A hand-edited value
|
||||
outside that range, or of another type, falls back to the default with a
|
||||
warning rather than failing the load.
|
||||
"""
|
||||
|
||||
MODEL_TYPE = "wd14-swinv2"
|
||||
EXISTING = "skip"
|
||||
GENERAL_THRESHOLD = 0.35
|
||||
CHARACTER_THRESHOLD = 0.85
|
||||
TEMPLATE = DEFAULT_TAG_TEMPLATE
|
||||
SUGGESTION_LIMIT = DEFAULT_SUGGESTION_LIMIT
|
||||
|
||||
# config.json key -> (attribute, validator, default)
|
||||
_KEYS = {
|
||||
"model_type": ("MODEL_TYPE", _valid_model_type, "wd14-swinv2"),
|
||||
"existing": ("EXISTING", _valid_existing, "skip"),
|
||||
"general_threshold": (
|
||||
"GENERAL_THRESHOLD",
|
||||
lambda v: validate_threshold("value", v),
|
||||
0.35,
|
||||
),
|
||||
"character_threshold": (
|
||||
"CHARACTER_THRESHOLD",
|
||||
lambda v: validate_threshold("value", v),
|
||||
0.85,
|
||||
),
|
||||
"template": ("TEMPLATE", validate_template, DEFAULT_TAG_TEMPLATE),
|
||||
"suggestion_limit": (
|
||||
"SUGGESTION_LIMIT",
|
||||
_valid_suggestion_limit,
|
||||
DEFAULT_SUGGESTION_LIMIT,
|
||||
),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def suggestion_limit(cls) -> int:
|
||||
return cls.SUGGESTION_LIMIT
|
||||
@@ -659,14 +590,7 @@ class AutotagConfig:
|
||||
|
||||
@classmethod
|
||||
def get_config(cls) -> Dict[str, Any]:
|
||||
"""The saved defaults; the template only when it is not the built-in
|
||||
one, so a later built-in template still reaches who never set one."""
|
||||
defaults = {
|
||||
key: getattr(cls, attr) for key, (attr, _v, _d) in cls._KEYS.items()
|
||||
}
|
||||
if defaults["template"] == DEFAULT_TAG_TEMPLATE:
|
||||
del defaults["template"]
|
||||
return {"defaults": defaults}
|
||||
return {"defaults": {"suggestion_limit": cls.SUGGESTION_LIMIT}}
|
||||
|
||||
@classmethod
|
||||
def update_config(cls, new_config: Any) -> None:
|
||||
@@ -677,18 +601,17 @@ class AutotagConfig:
|
||||
if not isinstance(defaults, dict):
|
||||
config_logger.warning("Ignoring autotag.defaults that is not an object")
|
||||
return
|
||||
for key, (attr, validate, default) in cls._KEYS.items():
|
||||
if key not in defaults:
|
||||
continue
|
||||
try:
|
||||
setattr(cls, attr, validate(defaults[key]))
|
||||
except ValueError as error:
|
||||
config_logger.warning(
|
||||
f"autotag.defaults.{key} {error}; using {default!r}"
|
||||
if key != "template"
|
||||
else f"autotag.defaults.template {error}; using the built-in one"
|
||||
)
|
||||
setattr(cls, attr, default)
|
||||
if "suggestion_limit" not in defaults:
|
||||
return
|
||||
value = defaults["suggestion_limit"]
|
||||
if cls.valid_suggestion_limit(value):
|
||||
cls.SUGGESTION_LIMIT = value
|
||||
return
|
||||
config_logger.warning(
|
||||
"autotag.defaults.suggestion_limit must be an integer from 1 to "
|
||||
f"{MAX_SUGGESTION_LIMIT}; using {DEFAULT_SUGGESTION_LIMIT}"
|
||||
)
|
||||
cls.SUGGESTION_LIMIT = DEFAULT_SUGGESTION_LIMIT
|
||||
|
||||
|
||||
class PromptManagerConfig:
|
||||
@@ -740,19 +663,6 @@ class PromptManagerConfig:
|
||||
# database first (System › Backup & export › Recovery policy, #249 S6).
|
||||
SAFETY_BACKUP = True
|
||||
|
||||
# "Keep it here" (#249 S6b): the user chose to keep a node-folder library
|
||||
# where it is, so the workspace stops offering to move it.
|
||||
KEEP_LIBRARY_IN_NODE_FOLDER = False
|
||||
|
||||
# Thumbnails (#317): every preview writer reads these through
|
||||
# py/api/thumbnail_settings.py (System › Appearance). The legacy
|
||||
# ``gallery.web_interface.thumbnail_size`` was never read and is ignored.
|
||||
THUMBNAIL_FORMATS = ("webp", "jxl", "jpeg")
|
||||
THUMBNAIL_FORMAT = "webp"
|
||||
THUMBNAIL_QUALITY = 80
|
||||
THUMBNAIL_EDGE = 512
|
||||
THUMBNAIL_EDGE_RANGE = (128, 2048)
|
||||
|
||||
@classmethod
|
||||
def max_worker_threads(cls) -> int:
|
||||
"""Largest accepted worker-thread count: the detected core count."""
|
||||
@@ -808,12 +718,6 @@ class PromptManagerConfig:
|
||||
"worker_threads": cls.WORKER_THREADS,
|
||||
},
|
||||
"maintenance": {"safety_backup": cls.SAFETY_BACKUP},
|
||||
"library": {"keep_in_node_folder": cls.KEEP_LIBRARY_IN_NODE_FOLDER},
|
||||
"thumbnails": {
|
||||
"format": cls.THUMBNAIL_FORMAT,
|
||||
"quality": cls.THUMBNAIL_QUALITY,
|
||||
"edge": cls.THUMBNAIL_EDGE,
|
||||
},
|
||||
"gallery": GalleryConfig.get_config(),
|
||||
"integrations": IntegrationConfig.get_config(),
|
||||
"autotag": AutotagConfig.get_config(),
|
||||
@@ -823,12 +727,14 @@ class PromptManagerConfig:
|
||||
def get_config_path(cls) -> str:
|
||||
"""Path of the persisted config.json.
|
||||
|
||||
Honours ``PROMPT_MANAGER_CONFIG_PATH`` when set; otherwise it lives in
|
||||
the library's data folder (``utils.library_location``): the node
|
||||
folder for a library that was not moved, else ComfyUI's
|
||||
``user/__promptmanager``.
|
||||
Honours ``PROMPT_MANAGER_CONFIG_PATH`` when set; otherwise the file
|
||||
lives at the repository root next to ``pyproject.toml``.
|
||||
"""
|
||||
return library_location.config_location()[0]
|
||||
override = os.environ.get(CONFIG_PATH_ENV, "").strip()
|
||||
if override:
|
||||
return override
|
||||
repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
return os.path.join(repo_root, "config.json")
|
||||
|
||||
@classmethod
|
||||
def load_from_file(cls, config_path: Optional[str] = None):
|
||||
@@ -953,14 +859,6 @@ class PromptManagerConfig:
|
||||
):
|
||||
cls.SAFETY_BACKUP = maintenance["safety_backup"]
|
||||
|
||||
library = new_config.get("library", {})
|
||||
if isinstance(library, dict) and isinstance(
|
||||
library.get("keep_in_node_folder"), bool
|
||||
):
|
||||
cls.KEEP_LIBRARY_IN_NODE_FOLDER = library["keep_in_node_folder"]
|
||||
|
||||
cls._update_thumbnails(new_config.get("thumbnails"))
|
||||
|
||||
# Update gallery config
|
||||
if "gallery" in new_config:
|
||||
GalleryConfig.update_config(new_config["gallery"])
|
||||
@@ -972,22 +870,6 @@ class PromptManagerConfig:
|
||||
if "autotag" in new_config:
|
||||
AutotagConfig.update_config(new_config["autotag"])
|
||||
|
||||
@classmethod
|
||||
def _update_thumbnails(cls, section: Any) -> None:
|
||||
"""Apply a ``thumbnails`` section; a hand-edited bad value is ignored."""
|
||||
if not isinstance(section, dict):
|
||||
return
|
||||
fmt = section.get("format")
|
||||
if fmt in cls.THUMBNAIL_FORMATS:
|
||||
cls.THUMBNAIL_FORMAT = fmt
|
||||
quality = section.get("quality")
|
||||
if type(quality) is int and 1 <= quality <= 100:
|
||||
cls.THUMBNAIL_QUALITY = quality
|
||||
edge = section.get("edge")
|
||||
low, high = cls.THUMBNAIL_EDGE_RANGE
|
||||
if type(edge) is int and low <= edge <= high:
|
||||
cls.THUMBNAIL_EDGE = edge
|
||||
|
||||
|
||||
# Load configuration on import
|
||||
try:
|
||||
|
||||
@@ -1,223 +0,0 @@
|
||||
"""The staged LoRA reference import (#221, #249 System S10, #355).
|
||||
|
||||
The import reads LoraManager's ``*.metadata.json`` sidecars and turns each
|
||||
LoRA into one prompt of the ``lora-manager`` category, tagged with its name and
|
||||
trigger words and linked to its images. It runs in three stages, and only the
|
||||
last writes to the database:
|
||||
|
||||
1. **discover**: read LoraManager's settings the way LoraManager does
|
||||
(``py/lora_manager_layout.py``: the settings file from the environment
|
||||
folder, the portable plugin file or the user config folder; the LoRA
|
||||
roots, the storage mode, the sidecar and example-images roots) and find
|
||||
the sidecars (``discover``). In ``alongside`` mode a sidecar beside its
|
||||
model is kept when ``<base>.safetensors`` exists; in ``centralized`` mode
|
||||
every sidecar under the sidecar root is kept when its ``file_path`` names
|
||||
a model inside a LoRA root. ``fingerprint`` covers the kept sidecars (path,
|
||||
size and modification time) and the layout; ``preview`` adds the counts a
|
||||
review shows. Nothing is written and nothing is requested.
|
||||
2. **stage**: read each sidecar and build its record. Its LoraManager preview
|
||||
and example images are linked where they are (``local_media``); only a
|
||||
LoRA without any local image downloads its CivitAI example images into the
|
||||
image cache first (best effort, counted). A cancelled or failed stage
|
||||
leaves the database untouched; downloaded images stay in the cache.
|
||||
3. **replace**: the sidecars are found and fingerprinted once more (a file
|
||||
changed, removed or added while staging aborts the import), then
|
||||
``PromptDatabase.replace_category_prompts`` swaps the old collection for
|
||||
the staged records in one transaction (or changes nothing).
|
||||
|
||||
Both the System job (``POST /lora/import/apply``) and the legacy ``POST
|
||||
/lora/scan`` run these stages (``py/api/lora_integration.py``).
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import AbstractSet, Callable, Dict, List, Optional
|
||||
|
||||
from . import lora_utils
|
||||
from .lora_manager_layout import (
|
||||
Discovery,
|
||||
discover_sidecars,
|
||||
layout_summary,
|
||||
local_media,
|
||||
)
|
||||
|
||||
try:
|
||||
from ..utils.hashing import generate_prompt_hash
|
||||
except ImportError: # pragma: no cover - flat test imports
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
LORA_CATEGORY = "lora-manager"
|
||||
JOB_TITLE = "LoRA reference import"
|
||||
# The job's outcome groups, in its counter-tile order.
|
||||
GROUPS = ("imported", "linked", "skipped", "downloaded", "download_errors", "errors")
|
||||
|
||||
|
||||
class ImportCancelled(Exception):
|
||||
"""The import was cancelled while staging; nothing was written."""
|
||||
|
||||
|
||||
def discover(lora_manager_path: Optional[str]) -> Discovery:
|
||||
"""The sidecars of every LoRA whose model exists inside a LoRA root.
|
||||
|
||||
``lora_manager_path`` is resolved as the integration resolves it (an
|
||||
install under ``custom_nodes``; empty means auto-detection). Raises
|
||||
whatever listing a folder raised: a discovery that cannot see every
|
||||
folder must not replace the collection with part of it.
|
||||
"""
|
||||
return discover_sidecars(lora_utils.lora_manager_layout(lora_manager_path))
|
||||
|
||||
|
||||
def fingerprint(discovery: Discovery) -> str:
|
||||
"""What a review saw: the layout, then each kept sidecar's path, size and
|
||||
modification time (a LoraManager save that changes unrelated settings
|
||||
leaves it as it was)."""
|
||||
layout = discovery.layout
|
||||
digest = hashlib.sha256()
|
||||
for value in (
|
||||
layout.settings_path,
|
||||
layout.source,
|
||||
layout.error,
|
||||
layout.mode,
|
||||
"\0".join(layout.lora_roots),
|
||||
layout.sidecar_root,
|
||||
layout.example_images_root,
|
||||
layout.library,
|
||||
):
|
||||
digest.update(f"{value}\n".encode("utf-8", "surrogatepass"))
|
||||
for name in discovery.unreadable:
|
||||
digest.update(f"unreadable\0{name}\n".encode("utf-8", "surrogatepass"))
|
||||
for sidecar in discovery.sidecars:
|
||||
path = sidecar.metadata_path
|
||||
try:
|
||||
stat = path.stat()
|
||||
entry = f"{path}\0{stat.st_size}\0{stat.st_mtime_ns}\n"
|
||||
except OSError:
|
||||
entry = f"{path}\0missing\n"
|
||||
digest.update(entry.encode("utf-8", "surrogatepass"))
|
||||
return "sha256:" + digest.hexdigest()
|
||||
|
||||
|
||||
def preview(
|
||||
lora_manager_path: Optional[str],
|
||||
cache_dir: Path,
|
||||
media_extensions: AbstractSet[str],
|
||||
) -> Dict:
|
||||
"""The review's counts (blocking; reads only, requests nothing).
|
||||
|
||||
``{lora_files, with_metadata, skipped, with_local_media,
|
||||
planned_downloads, fingerprint, names, layout}``: the metadata files
|
||||
seen, the kept LoRAs whose sidecar parses, the rest, the LoRAs with a
|
||||
local image, the CivitAI images an import would download for the LoRAs
|
||||
without one (not yet in the cache), the first LoRA names and what the
|
||||
layout says (``layout_summary``).
|
||||
"""
|
||||
discovery = discover(lora_manager_path)
|
||||
layout = discovery.layout
|
||||
with_metadata = 0
|
||||
with_local_media = 0
|
||||
planned = 0
|
||||
names = []
|
||||
for sidecar in discovery.sidecars:
|
||||
metadata = lora_utils.read_lora_metadata(sidecar.metadata_path)
|
||||
if not metadata:
|
||||
continue
|
||||
with_metadata += 1
|
||||
if local_media(sidecar, metadata, layout, media_extensions):
|
||||
with_local_media += 1
|
||||
else:
|
||||
planned += lora_utils.planned_civitai_downloads(
|
||||
metadata, sidecar.metadata_path, cache_dir
|
||||
)
|
||||
if len(names) < 50:
|
||||
names.append(lora_utils.get_model_name_from_metadata(metadata))
|
||||
return {
|
||||
"lora_files": discovery.scanned,
|
||||
"with_metadata": with_metadata,
|
||||
"skipped": discovery.scanned - with_metadata,
|
||||
"with_local_media": with_local_media,
|
||||
"planned_downloads": planned,
|
||||
"fingerprint": fingerprint(discovery),
|
||||
"names": names,
|
||||
"layout": layout_summary(layout, Path.home()),
|
||||
}
|
||||
|
||||
|
||||
def build_record(metadata: Dict, path: Path, images: List[str]) -> Dict:
|
||||
"""The prompt a LoRA becomes: its example prompt (else its name), tags, images."""
|
||||
model_name = lora_utils.get_model_name_from_metadata(metadata)
|
||||
text = lora_utils.get_example_prompt_from_metadata(metadata) or model_name
|
||||
tags = ["lora-manager", f"lora:{model_name}"]
|
||||
tags.extend(lora_utils.get_trigger_words_from_metadata(metadata))
|
||||
return {
|
||||
"name": model_name,
|
||||
"text": text,
|
||||
"hash": generate_prompt_hash(text),
|
||||
"tags": list(dict.fromkeys(tags)),
|
||||
"images": images,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class StagedImport:
|
||||
"""The records an import will write, and what staging counted."""
|
||||
|
||||
total: int
|
||||
records: List[Dict] = field(default_factory=list)
|
||||
skipped: List[str] = field(default_factory=list)
|
||||
stats: Dict[str, int] = field(
|
||||
default_factory=lambda: {"downloaded": 0, "download_errors": 0}
|
||||
)
|
||||
|
||||
|
||||
def stage(
|
||||
discovery: Discovery,
|
||||
cache_dir: Path,
|
||||
api_key: str = "",
|
||||
*,
|
||||
media_extensions: AbstractSet[str],
|
||||
on_record: Optional[Callable[[int, Dict, "StagedImport"], None]] = None,
|
||||
cancelled: Callable[[], bool] = lambda: False,
|
||||
logger=None,
|
||||
) -> StagedImport:
|
||||
"""Read each sidecar and build its record with its images.
|
||||
|
||||
Blocking. Local images (``local_media``) come first; a LoRA without any
|
||||
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,
|
||||
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():
|
||||
raise ImportCancelled()
|
||||
record = None
|
||||
path = sidecar.metadata_path
|
||||
metadata = lora_utils.read_lora_metadata(path)
|
||||
if not metadata:
|
||||
staged.skipped.append(path.name)
|
||||
else:
|
||||
try:
|
||||
local = local_media(
|
||||
sidecar, metadata, discovery.layout, media_extensions
|
||||
)
|
||||
remote = []
|
||||
if not local:
|
||||
remote = lora_utils.download_civitai_images(
|
||||
metadata, path, cache_dir, api_key, stats=staged.stats
|
||||
)
|
||||
images = list(dict.fromkeys([*local, *remote]))
|
||||
record = build_record(metadata, path, images)
|
||||
staged.records.append(record)
|
||||
except Exception as error: # one LoRA never stops the import
|
||||
name = lora_utils.get_model_name_from_metadata(metadata)
|
||||
if logger is not None:
|
||||
logger.warning(f"Failed to prepare LoRA {name}: {error}")
|
||||
staged.skipped.append(name)
|
||||
if on_record is not None:
|
||||
on_record(index, record, staged)
|
||||
return staged
|
||||
@@ -1,602 +0,0 @@
|
||||
"""Where LoraManager keeps its files (#355).
|
||||
|
||||
LoraManager (``custom_nodes/comfyui-lora-manager``) stores one
|
||||
``<name>.metadata.json`` sidecar per model, a preview image or video beside
|
||||
it, and each model's downloaded CivitAI example images. Where it stores them
|
||||
is a user setting, and its settings file itself can live in three places.
|
||||
This module repeats LoraManager's own rules, reading its files only (never
|
||||
importing it), so the LoRA import, the trigger-word cache and the image
|
||||
server share one definition:
|
||||
|
||||
- ``load_layout`` finds the settings file (environment folder, portable
|
||||
plugin file, or the user config folder, in LoraManager's order), keeps only
|
||||
the seven keys the import needs, and resolves the LoRA roots, the sidecar
|
||||
root (centralized storage), the example-images root and the active
|
||||
LoraManager library. Every folder from the settings passes the folder
|
||||
guard (exists, not a filesystem root, not the home folder itself).
|
||||
- ``discover_sidecars`` finds every sidecar whose model file exists inside a
|
||||
LoRA root: beside the models in ``alongside`` mode, under the sidecar root
|
||||
(matched to its model by ``file_path``) in ``centralized`` mode.
|
||||
- ``local_media`` lists a LoRA's preview and example images that the image
|
||||
server may serve.
|
||||
|
||||
Standard library plus ``platformdirs`` only (LoraManager depends on it), so
|
||||
it is testable outside ComfyUI.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from dataclasses import dataclass, replace
|
||||
from pathlib import Path
|
||||
from typing import (
|
||||
AbstractSet,
|
||||
Callable,
|
||||
Dict,
|
||||
Iterator,
|
||||
List,
|
||||
Mapping,
|
||||
Optional,
|
||||
Sequence,
|
||||
Tuple,
|
||||
)
|
||||
|
||||
# py/utils/settings_paths.py: SETTINGS_DIR_ENV, _LM_PORTABLE_ENV, APP_NAME.
|
||||
SETTINGS_DIR_ENV = "LORA_MANAGER_SETTINGS_DIR"
|
||||
PORTABLE_ENV = "LORA_MANAGER_PORTABLE"
|
||||
APP_NAME = "ComfyUI-LoRA-Manager"
|
||||
# py/utils/constants.py PREVIEW_EXTENSIONS, in LoraManager's order.
|
||||
PREVIEW_EXTENSIONS = (
|
||||
".webp",
|
||||
".preview.webp",
|
||||
".preview.png",
|
||||
".preview.jpeg",
|
||||
".preview.jpg",
|
||||
".preview.mp4",
|
||||
".png",
|
||||
".jpeg",
|
||||
".jpg",
|
||||
".mp4",
|
||||
".gif",
|
||||
".webm",
|
||||
".avif",
|
||||
".jxl",
|
||||
)
|
||||
# LoraManager's LoRA scanner accepts .safetensors only (py/services/lora_scanner.py).
|
||||
MODEL_SUFFIX = ".safetensors"
|
||||
# py/utils/sidecar_paths.py METADATA_SUFFIX.
|
||||
METADATA_SUFFIX = ".metadata.json"
|
||||
MAX_SETTINGS_BYTES = 5 * 1024 * 1024
|
||||
# The only settings the import reads; every other key (the file can hold API
|
||||
# keys) is dropped right after parsing.
|
||||
SETTINGS_KEYS = frozenset(
|
||||
{
|
||||
"sidecar_storage_mode",
|
||||
"sidecar_storage_path",
|
||||
"example_images_path",
|
||||
"folder_paths",
|
||||
"extra_folder_paths",
|
||||
"libraries",
|
||||
"active_library",
|
||||
}
|
||||
)
|
||||
# py/utils/sidecar_paths.py STORAGE_MODE_ALONGSIDE / STORAGE_MODE_CENTRALIZED.
|
||||
MODE_ALONGSIDE = "alongside"
|
||||
MODE_CENTRALIZED = "centralized"
|
||||
# utils/video_metadata.py VIDEO_EXTENSIONS (not imported: that package needs
|
||||
# the node's relative imports).
|
||||
VIDEO_SUFFIXES = frozenset({".mp4", ".mov", ".m4v", ".webm", ".mkv", ".avi", ".wmv"})
|
||||
|
||||
SOURCE_ENVIRONMENT = "environment"
|
||||
SOURCE_PORTABLE = "portable"
|
||||
SOURCE_USER_CONFIG = "user config"
|
||||
|
||||
NO_INSTALL = "no LoraManager install found"
|
||||
UNREADABLE = "the settings file could not be read"
|
||||
NOT_JSON = "the settings file is not valid JSON"
|
||||
NOT_OBJECT = "the settings file is not a JSON object"
|
||||
TOO_LARGE = "the settings file is larger than 5 MB"
|
||||
NO_PLATFORMDIRS = "platformdirs is not available"
|
||||
NO_SIDECAR_ROOT = "the sidecar folder does not exist"
|
||||
BAD_SIDECAR_ROOT = "the sidecar folder is a filesystem root or the home folder"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LoraManagerLayout:
|
||||
"""What LoraManager's settings say about where its files are."""
|
||||
|
||||
settings_path: Optional[Path]
|
||||
source: Optional[str]
|
||||
error: Optional[str]
|
||||
mode: str
|
||||
lora_roots: Tuple[str, ...]
|
||||
sidecar_root: Optional[str]
|
||||
example_images_root: Optional[str]
|
||||
library: str
|
||||
media_roots: Tuple[str, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Sidecar:
|
||||
"""One LoRA's metadata file and the model file it describes."""
|
||||
|
||||
metadata_path: Path
|
||||
model_path: Path
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Discovery:
|
||||
"""The sidecars kept, how many metadata files were seen, and those unread."""
|
||||
|
||||
layout: LoraManagerLayout
|
||||
sidecars: List[Sidecar]
|
||||
scanned: int
|
||||
unreadable: List[str]
|
||||
|
||||
|
||||
class _SettingsError(Exception):
|
||||
"""The settings file cannot be used, for the reason it carries."""
|
||||
|
||||
|
||||
# ── paths ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _canonical(path: str, pathmod=os.path) -> str:
|
||||
# py/config.py _canonical_path (not imported: py.config imports ComfyUI's
|
||||
# ``server`` at import time).
|
||||
return pathmod.normcase(pathmod.realpath(path))
|
||||
|
||||
|
||||
def _is_filesystem_root(path: str, pathmod=os.path) -> bool:
|
||||
# py/config.py _is_filesystem_root: '/', 'C:\\' and UNC share roots.
|
||||
return pathmod.dirname(path) == path
|
||||
|
||||
|
||||
def _guard_folder(
|
||||
path: str,
|
||||
home: str,
|
||||
*,
|
||||
pathmod=os.path,
|
||||
isdir: Optional[Callable[[str], bool]] = None,
|
||||
) -> Optional[str]:
|
||||
"""The canonical form of ``path`` when it may be read and served, else None.
|
||||
|
||||
Kept only when it is a folder, not a filesystem root, and neither the
|
||||
home folder nor one of its ancestors (``/home``, ``C:\\Users``: a
|
||||
security review ruled that a folder holding the home folder must never
|
||||
be read or served).
|
||||
"""
|
||||
isdir = isdir or os.path.isdir
|
||||
try:
|
||||
canonical = _canonical(path, pathmod)
|
||||
if not isdir(canonical) or _is_filesystem_root(canonical, pathmod):
|
||||
return None
|
||||
if _inside(home, (canonical,), pathmod):
|
||||
return None
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
return canonical
|
||||
|
||||
|
||||
def _inside(path, roots: Sequence[str], pathmod=os.path) -> bool:
|
||||
"""True when ``path`` (canonical) is one of ``roots`` or below one."""
|
||||
try:
|
||||
candidate = _canonical(str(path), pathmod)
|
||||
except (OSError, ValueError):
|
||||
return False
|
||||
for root in roots:
|
||||
prefix = root if root.endswith(pathmod.sep) else root + pathmod.sep
|
||||
if candidate == root or candidate.startswith(prefix):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _absolute(value) -> Optional[str]:
|
||||
"""A settings path, stripped, ``~``-expanded and absolute; None when unset."""
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
return None
|
||||
return os.path.abspath(os.path.expanduser(value.strip()))
|
||||
|
||||
|
||||
# ── the settings file ───────────────────────────────────────────────
|
||||
|
||||
|
||||
def _platform_user_config_dir() -> str:
|
||||
from platformdirs import user_config_dir
|
||||
|
||||
return user_config_dir(APP_NAME, appauthor=False)
|
||||
|
||||
|
||||
def _read_bytes(path: Path) -> bytes:
|
||||
with open(path, "rb") as handle:
|
||||
return handle.read(MAX_SETTINGS_BYTES + 1)
|
||||
|
||||
|
||||
def _read_settings(path: Path) -> Optional[Dict]:
|
||||
"""The parsed object in ``path``; None when it does not exist."""
|
||||
try:
|
||||
raw = _read_bytes(path)
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
except OSError:
|
||||
raise _SettingsError(UNREADABLE) from None
|
||||
if len(raw) > MAX_SETTINGS_BYTES:
|
||||
raise _SettingsError(TOO_LARGE)
|
||||
try:
|
||||
data = json.loads(raw.decode("utf-8"))
|
||||
except (UnicodeDecodeError, ValueError):
|
||||
raise _SettingsError(NOT_JSON) from None
|
||||
if not isinstance(data, dict):
|
||||
raise _SettingsError(NOT_OBJECT)
|
||||
return data
|
||||
|
||||
|
||||
def _settings_subset(data: Mapping) -> Dict:
|
||||
"""Only the keys the import reads."""
|
||||
return {key: data[key] for key in SETTINGS_KEYS if key in data}
|
||||
|
||||
|
||||
def _portable_flag(environ: Mapping[str, str], plugin_settings: Path) -> bool:
|
||||
"""LoraManager's portable mode (settings_paths._should_use_portable_settings)."""
|
||||
forced = environ.get(PORTABLE_ENV)
|
||||
if forced == "1":
|
||||
return True
|
||||
if forced == "0":
|
||||
return False
|
||||
try:
|
||||
data = _read_settings(plugin_settings)
|
||||
except _SettingsError:
|
||||
return False
|
||||
return bool(data) and data.get("use_portable_settings") is True
|
||||
|
||||
|
||||
def _settings_location(
|
||||
lora_manager_path: str,
|
||||
environ: Mapping[str, str],
|
||||
user_config_dir: Callable[[], str],
|
||||
) -> Tuple[Path, str]:
|
||||
"""``(settings file, source)`` in LoraManager's lookup order."""
|
||||
override = environ.get(SETTINGS_DIR_ENV)
|
||||
if override:
|
||||
folder = os.path.abspath(os.path.expanduser(override))
|
||||
return Path(folder) / "settings.json", SOURCE_ENVIRONMENT
|
||||
plugin_settings = Path(lora_manager_path) / "settings.json"
|
||||
if _portable_flag(environ, plugin_settings):
|
||||
return plugin_settings, SOURCE_PORTABLE
|
||||
try:
|
||||
folder = user_config_dir()
|
||||
except ImportError:
|
||||
raise _SettingsError(NO_PLATFORMDIRS) from None
|
||||
return Path(folder) / "settings.json", SOURCE_USER_CONFIG
|
||||
|
||||
|
||||
# ── libraries and folders ───────────────────────────────────────────
|
||||
|
||||
|
||||
def _sanitize_library_name(name: str) -> str:
|
||||
# py/utils/example_images_paths.py sanitize_library_name.
|
||||
safe = re.sub(r"[^A-Za-z0-9_.-]", "_", name or "")
|
||||
return safe or "default"
|
||||
|
||||
|
||||
def _active_library(settings: Mapping) -> Tuple[str, Optional[Dict], int]:
|
||||
"""``(name, library dict or None, number of libraries)``."""
|
||||
libraries = settings.get("libraries")
|
||||
if not isinstance(libraries, dict) or not libraries:
|
||||
return "default", None, 0
|
||||
active = settings.get("active_library")
|
||||
if not isinstance(active, str) or active not in libraries:
|
||||
active = next(iter(libraries))
|
||||
library = libraries[active]
|
||||
return str(active), library if isinstance(library, dict) else None, len(libraries)
|
||||
|
||||
|
||||
def _lora_list(settings: Mapping, library: Optional[Dict], key: str) -> List[str]:
|
||||
"""The ``loras`` folders of ``key`` (the library's, else the top level)."""
|
||||
source = library if library is not None and key in library else settings
|
||||
section = source.get(key)
|
||||
if not isinstance(section, dict):
|
||||
return []
|
||||
value = section.get("loras")
|
||||
if isinstance(value, str):
|
||||
value = [value]
|
||||
if not isinstance(value, list):
|
||||
return []
|
||||
return [entry for entry in value if isinstance(entry, str) and entry.strip()]
|
||||
|
||||
|
||||
def _guarded(paths, home: str) -> Tuple[str, ...]:
|
||||
kept = []
|
||||
for path in paths:
|
||||
absolute = _absolute(path)
|
||||
guarded = _guard_folder(absolute, home) if absolute else None
|
||||
if guarded and guarded not in kept:
|
||||
kept.append(guarded)
|
||||
return tuple(kept)
|
||||
|
||||
|
||||
def _example_images_root(
|
||||
settings: Mapping, library: str, libraries: int, home: str
|
||||
) -> Optional[str]:
|
||||
"""py/utils/example_images_paths.py get_library_root, when the folder exists."""
|
||||
configured = _absolute(settings.get("example_images_path"))
|
||||
root = _guard_folder(configured, home) if configured else None
|
||||
if root is None or libraries <= 1:
|
||||
return root
|
||||
# Guarded (canonical) like its root: a library name keeps its case, and
|
||||
# the media check compares normcased paths.
|
||||
return _guard_folder(os.path.join(root, _sanitize_library_name(library)), home)
|
||||
|
||||
|
||||
def _media_roots(*groups) -> Tuple[str, ...]:
|
||||
roots = []
|
||||
for group in groups:
|
||||
for root in group:
|
||||
if root and root not in roots:
|
||||
roots.append(root)
|
||||
return tuple(roots)
|
||||
|
||||
|
||||
def load_layout(
|
||||
lora_manager_path: Optional[str],
|
||||
*,
|
||||
default_roots: Callable[[], List[str]],
|
||||
environ: Mapping[str, str] = os.environ,
|
||||
user_config_dir: Callable[[], str] = _platform_user_config_dir,
|
||||
) -> LoraManagerLayout:
|
||||
"""Read LoraManager's settings the way LoraManager does (read only).
|
||||
|
||||
``lora_manager_path`` is the resolved install (``None`` when none was
|
||||
found). An unusable layout carries ``error``, alongside mode and the
|
||||
default LoRA roots, so serving and trigger words keep working.
|
||||
"""
|
||||
home = str(Path.home())
|
||||
defaults = _guarded(default_roots(), home)
|
||||
fallback = LoraManagerLayout(
|
||||
settings_path=None,
|
||||
source=None,
|
||||
error=NO_INSTALL,
|
||||
mode=MODE_ALONGSIDE,
|
||||
lora_roots=defaults,
|
||||
sidecar_root=None,
|
||||
example_images_root=None,
|
||||
library="default",
|
||||
media_roots=defaults,
|
||||
)
|
||||
if not lora_manager_path:
|
||||
return fallback
|
||||
try:
|
||||
settings_path, source = _settings_location(
|
||||
lora_manager_path, environ, user_config_dir
|
||||
)
|
||||
except _SettingsError as error:
|
||||
return replace(fallback, error=str(error))
|
||||
fallback = replace(fallback, settings_path=settings_path, source=source)
|
||||
try:
|
||||
raw = _read_settings(settings_path)
|
||||
except _SettingsError as error:
|
||||
return replace(fallback, error=str(error))
|
||||
settings = _settings_subset(raw or {})
|
||||
return _layout_from_settings(settings, settings_path, source, defaults, home)
|
||||
|
||||
|
||||
def _layout_from_settings(settings, settings_path, source, defaults, home):
|
||||
library, entry, libraries = _active_library(settings)
|
||||
lora_roots = _media_roots(
|
||||
defaults,
|
||||
_guarded(_lora_list(settings, entry, "folder_paths"), home),
|
||||
_guarded(_lora_list(settings, entry, "extra_folder_paths"), home),
|
||||
)
|
||||
mode = settings.get("sidecar_storage_mode")
|
||||
mode = MODE_CENTRALIZED if mode == MODE_CENTRALIZED else MODE_ALONGSIDE
|
||||
example_root = _example_images_root(settings, library, libraries, home)
|
||||
layout = LoraManagerLayout(
|
||||
settings_path=settings_path,
|
||||
source=source,
|
||||
error=None,
|
||||
mode=mode,
|
||||
lora_roots=lora_roots,
|
||||
sidecar_root=None,
|
||||
example_images_root=example_root,
|
||||
library=library,
|
||||
media_roots=_media_roots(lora_roots, [example_root]),
|
||||
)
|
||||
if mode != MODE_CENTRALIZED:
|
||||
return layout
|
||||
configured = _absolute(settings.get("sidecar_storage_path"))
|
||||
configured = configured or os.path.join(str(settings_path.parent), "sidecars")
|
||||
sidecar_root = _guard_folder(configured, home)
|
||||
if sidecar_root is None:
|
||||
reason = NO_SIDECAR_ROOT if not os.path.isdir(configured) else BAD_SIDECAR_ROOT
|
||||
return replace(
|
||||
layout,
|
||||
error=reason,
|
||||
lora_roots=defaults,
|
||||
example_images_root=None,
|
||||
media_roots=defaults,
|
||||
)
|
||||
return replace(
|
||||
layout,
|
||||
sidecar_root=sidecar_root,
|
||||
media_roots=_media_roots(lora_roots, [sidecar_root], [example_root]),
|
||||
)
|
||||
|
||||
|
||||
# ── sidecars ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _metadata_files(folder: str) -> Iterator[Path]:
|
||||
"""Every ``*.metadata.json`` under ``folder``; an unlistable folder raises.
|
||||
|
||||
Symlinked folders are followed, each real folder once (no loops).
|
||||
"""
|
||||
|
||||
def fail(error):
|
||||
raise error
|
||||
|
||||
visited = set()
|
||||
for current, dirs, files in os.walk(folder, onerror=fail, followlinks=True):
|
||||
real = _canonical(current)
|
||||
if real in visited:
|
||||
dirs[:] = []
|
||||
continue
|
||||
visited.add(real)
|
||||
dirs.sort()
|
||||
for name in sorted(files):
|
||||
if name.endswith(METADATA_SUFFIX):
|
||||
yield Path(current) / name
|
||||
|
||||
|
||||
def _read_sidecar(path: Path) -> Optional[Dict]:
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as handle:
|
||||
data = json.load(handle)
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
return data if isinstance(data, dict) else None
|
||||
|
||||
|
||||
def _model_for(path: Path, layout: LoraManagerLayout, unreadable: List[str]):
|
||||
"""The model file a sidecar describes, when it exists inside a LoRA root."""
|
||||
if layout.mode == MODE_ALONGSIDE:
|
||||
base = path.name[: -len(METADATA_SUFFIX)]
|
||||
model = path.parent / f"{base}{MODEL_SUFFIX}"
|
||||
return model if model.is_file() else None
|
||||
metadata = _read_sidecar(path)
|
||||
if metadata is None:
|
||||
unreadable.append(path.name)
|
||||
return None
|
||||
file_path = metadata.get("file_path")
|
||||
if not isinstance(file_path, str) or not file_path.strip():
|
||||
return None
|
||||
model = Path(file_path)
|
||||
if not model.is_file() or not _inside(model, layout.lora_roots):
|
||||
return None
|
||||
return model
|
||||
|
||||
|
||||
def discover_sidecars(layout: LoraManagerLayout) -> Discovery:
|
||||
"""Every sidecar whose model exists inside a LoRA root, sorted, once each.
|
||||
|
||||
Raises ``OSError`` when a folder cannot be listed: a discovery that
|
||||
cannot see every folder must not replace the collection with part of it.
|
||||
"""
|
||||
if layout.error is not None:
|
||||
return Discovery(layout, [], 0, [])
|
||||
if layout.mode == MODE_CENTRALIZED:
|
||||
folders = [layout.sidecar_root]
|
||||
else:
|
||||
folders = list(layout.lora_roots)
|
||||
seen = set()
|
||||
scanned = 0
|
||||
unreadable: List[str] = []
|
||||
sidecars = []
|
||||
for folder in folders:
|
||||
for path in _metadata_files(folder):
|
||||
resolved = Path(os.path.realpath(path))
|
||||
key = os.path.normcase(str(resolved))
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
scanned += 1
|
||||
model = _model_for(resolved, layout, unreadable)
|
||||
if model is not None:
|
||||
sidecars.append(Sidecar(resolved, Path(os.path.realpath(model))))
|
||||
sidecars.sort(key=lambda sidecar: str(sidecar.metadata_path))
|
||||
return Discovery(layout, sidecars, scanned, unreadable)
|
||||
|
||||
|
||||
# ── images ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _natural_key(name: str):
|
||||
parts = re.split(r"(\d+)", name)
|
||||
return [int(part) if part.isdigit() else part.lower() for part in parts], name
|
||||
|
||||
|
||||
def _example_files(folder: str) -> List[str]:
|
||||
try:
|
||||
with os.scandir(folder) as entries:
|
||||
names = [entry.name for entry in entries if entry.is_file()]
|
||||
except OSError:
|
||||
return []
|
||||
return [os.path.join(folder, name) for name in sorted(names, key=_natural_key)]
|
||||
|
||||
|
||||
def _preview_candidates(sidecar: Sidecar, metadata: Dict) -> Iterator[str]:
|
||||
"""The paths a preview is looked for at, in LoraManager's order."""
|
||||
url = metadata.get("preview_url")
|
||||
if isinstance(url, str) and url.strip() and os.path.isabs(url.strip()):
|
||||
yield url.strip()
|
||||
base = sidecar.metadata_path.name[: -len(METADATA_SUFFIX)]
|
||||
for suffix in PREVIEW_EXTENSIONS:
|
||||
yield str(sidecar.metadata_path.parent / f"{base}{suffix}")
|
||||
|
||||
|
||||
def local_media(
|
||||
sidecar: Sidecar,
|
||||
metadata: Dict,
|
||||
layout: LoraManagerLayout,
|
||||
media_extensions: AbstractSet[str],
|
||||
) -> List[str]:
|
||||
"""The LoRA's preview, then its example images: absolute, resolved paths.
|
||||
|
||||
A file counts only when its extension is in ``media_extensions`` (what the
|
||||
image server serves) and it lies inside ``layout.media_roots``.
|
||||
"""
|
||||
|
||||
def usable(path: str) -> bool:
|
||||
return (
|
||||
os.path.splitext(path)[1].lower() in media_extensions
|
||||
and os.path.isfile(path)
|
||||
and _inside(path, layout.media_roots)
|
||||
)
|
||||
|
||||
found: List[str] = []
|
||||
for path in _preview_candidates(sidecar, metadata):
|
||||
if usable(path):
|
||||
found.append(path)
|
||||
break
|
||||
sha256 = metadata.get("sha256")
|
||||
if layout.example_images_root and isinstance(sha256, str) and sha256.strip():
|
||||
folder = os.path.join(layout.example_images_root, sha256.strip().lower())
|
||||
found.extend(path for path in _example_files(folder) if usable(path))
|
||||
media: List[str] = []
|
||||
for path in found:
|
||||
resolved = os.path.realpath(path)
|
||||
if resolved not in media:
|
||||
media.append(resolved)
|
||||
return _cover_first(media)
|
||||
|
||||
|
||||
def _cover_first(media: List[str]) -> List[str]:
|
||||
"""A still image in front of a leading video (no thumbnails for these yet)."""
|
||||
if not media or os.path.splitext(media[0])[1].lower() not in VIDEO_SUFFIXES:
|
||||
return media
|
||||
for index, path in enumerate(media):
|
||||
if os.path.splitext(path)[1].lower() not in VIDEO_SUFFIXES:
|
||||
return [path] + media[:index] + media[index + 1 :]
|
||||
return media
|
||||
|
||||
|
||||
# ── the review ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def layout_summary(layout: LoraManagerLayout, home: Path) -> Dict:
|
||||
"""What the review shows about the layout (no settings values)."""
|
||||
settings_file = None
|
||||
if layout.settings_path is not None:
|
||||
settings_file = str(layout.settings_path)
|
||||
try:
|
||||
relative = Path(settings_file).relative_to(home)
|
||||
settings_file = os.path.join("~", str(relative))
|
||||
except ValueError:
|
||||
pass
|
||||
return {
|
||||
"source": layout.source,
|
||||
"settings_file": settings_file,
|
||||
"error": layout.error,
|
||||
"mode": layout.mode,
|
||||
"library": layout.library,
|
||||
"example_images": layout.example_images_root is not None,
|
||||
}
|
||||
+162
-425
@@ -8,17 +8,14 @@ return empty results rather than raising.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import http.client
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
from urllib.parse import urljoin, urlsplit
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
try:
|
||||
from ..utils.logging_config import get_logger
|
||||
@@ -28,16 +25,6 @@ except ImportError:
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
from utils.logging_config import get_logger
|
||||
|
||||
from dataclasses import replace
|
||||
|
||||
from .lora_manager_layout import (
|
||||
MODE_ALONGSIDE,
|
||||
METADATA_SUFFIX,
|
||||
LoraManagerLayout,
|
||||
discover_sidecars,
|
||||
load_layout,
|
||||
)
|
||||
|
||||
logger = get_logger("prompt_manager.lora_utils")
|
||||
|
||||
# ── LoraManager detection ────────────────────────────────────────────
|
||||
@@ -95,46 +82,23 @@ def detect_lora_manager(custom_path: str = "") -> Optional[str]:
|
||||
if p.is_dir() and _looks_like_lora_manager(p):
|
||||
return str(p.resolve())
|
||||
|
||||
# 2. Auto-detect in every custom_nodes folder (case-insensitive scan):
|
||||
# ComfyUI's own first, then the extra ones it registered (#144).
|
||||
folders = []
|
||||
# 2. Auto-detect via custom_nodes (case-insensitive scan)
|
||||
root = find_comfyui_root()
|
||||
if root:
|
||||
folders.append(root / "custom_nodes")
|
||||
folders.extend(custom_nodes_directories())
|
||||
seen = set()
|
||||
for custom_nodes in folders:
|
||||
key = os.path.normcase(os.path.abspath(str(custom_nodes)))
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
found = _lora_manager_in(Path(custom_nodes))
|
||||
if found:
|
||||
return found
|
||||
custom_nodes = root / "custom_nodes"
|
||||
if custom_nodes.is_dir():
|
||||
for entry in custom_nodes.iterdir():
|
||||
if (
|
||||
entry.is_dir()
|
||||
and "lora" in entry.name.lower()
|
||||
and "manager" in entry.name.lower()
|
||||
and _looks_like_lora_manager(entry)
|
||||
):
|
||||
return str(entry.resolve())
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _lora_manager_in(custom_nodes: Path) -> Optional[str]:
|
||||
"""The LoraManager install directly under ``custom_nodes``, or None."""
|
||||
try:
|
||||
entries = sorted(custom_nodes.iterdir()) if custom_nodes.is_dir() else []
|
||||
except OSError:
|
||||
return None
|
||||
for entry in entries:
|
||||
try:
|
||||
if (
|
||||
entry.is_dir()
|
||||
and "lora" in entry.name.lower()
|
||||
and "manager" in entry.name.lower()
|
||||
and _looks_like_lora_manager(entry)
|
||||
):
|
||||
return str(entry.resolve())
|
||||
except OSError:
|
||||
continue
|
||||
return None
|
||||
|
||||
|
||||
def _looks_like_lora_manager(path: Path) -> bool:
|
||||
"""Heuristic: does this directory look like a LoraManager install?"""
|
||||
# Must have __init__.py (ComfyUI extension) or README.md
|
||||
@@ -245,22 +209,23 @@ def resolve_lora_manager_path(path) -> Optional[str]:
|
||||
# ── Metadata reading ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def default_lora_roots() -> List[str]:
|
||||
"""The LoRA folders ComfyUI knows, without LoraManager's settings.
|
||||
def find_lora_directories(lora_manager_path: str) -> List[str]:
|
||||
"""Find directories that contain LoRA models (with .metadata.json files).
|
||||
|
||||
ComfyUI's ``models/loras``, the ``loras`` entries of
|
||||
``extra_model_paths.yaml`` and, inside ComfyUI, every folder
|
||||
``folder_paths`` registers for ``loras``. ``lora_manager_layout`` adds
|
||||
the folders LoraManager's own settings name.
|
||||
Searches: ComfyUI models/loras, extra_model_paths.yaml lora dirs,
|
||||
and the LoraManager extension dir itself.
|
||||
"""
|
||||
dirs = set()
|
||||
lm_path = Path(lora_manager_path)
|
||||
|
||||
root = find_comfyui_root()
|
||||
if root:
|
||||
# Default models/loras
|
||||
models_loras = root / "models" / "loras"
|
||||
if models_loras.is_dir():
|
||||
dirs.add(str(models_loras.resolve()))
|
||||
|
||||
# Extra model paths from ComfyUI config
|
||||
for extra_dir in _get_extra_lora_paths(root):
|
||||
if extra_dir.is_dir():
|
||||
dirs.add(str(extra_dir.resolve()))
|
||||
@@ -277,28 +242,13 @@ def default_lora_roots() -> List[str]:
|
||||
# folder_paths unavailable — not running inside ComfyUI runtime
|
||||
pass
|
||||
|
||||
# Check for any .metadata.json in the LoraManager dir tree
|
||||
for meta in lm_path.rglob("*.metadata.json"):
|
||||
dirs.add(str(meta.parent.resolve()))
|
||||
|
||||
return sorted(dirs)
|
||||
|
||||
|
||||
def lora_manager_install(lora_manager_path: Optional[str]) -> Optional[str]:
|
||||
"""The LoraManager install to read: the configured path when it is an
|
||||
allowed install under ``custom_nodes``, auto-detection when it is empty."""
|
||||
if isinstance(lora_manager_path, str) and lora_manager_path.strip():
|
||||
return resolve_lora_manager_path(lora_manager_path)
|
||||
detected = detect_lora_manager()
|
||||
return resolve_lora_manager_path(detected) if detected else None
|
||||
|
||||
|
||||
def layout_for_install(install: Optional[str]) -> LoraManagerLayout:
|
||||
"""The layout of an already resolved install (None: none was found)."""
|
||||
return load_layout(install, default_roots=default_lora_roots)
|
||||
|
||||
|
||||
def lora_manager_layout(lora_manager_path: Optional[str]) -> LoraManagerLayout:
|
||||
"""Where LoraManager keeps its files (``py/lora_manager_layout.py``)."""
|
||||
return layout_for_install(lora_manager_install(lora_manager_path))
|
||||
|
||||
|
||||
def _get_extra_lora_paths(comfyui_root: Path) -> List[Path]:
|
||||
"""Parse extra_model_paths.yaml for additional LoRA directories."""
|
||||
results = []
|
||||
@@ -404,42 +354,61 @@ def get_model_name_from_metadata(metadata: Dict) -> str:
|
||||
return name
|
||||
|
||||
|
||||
def get_preview_images_from_metadata(metadata: Dict, metadata_path: Path) -> List[str]:
|
||||
"""Find all local preview/example image paths for a LoRA.
|
||||
|
||||
Returns:
|
||||
List of absolute path strings to image files.
|
||||
"""
|
||||
results = []
|
||||
lora_dir = metadata_path.parent
|
||||
file_name = metadata.get("file_name", "")
|
||||
if not file_name:
|
||||
stem = metadata_path.name.replace(".metadata.json", "")
|
||||
file_name = stem
|
||||
|
||||
base_name = Path(file_name).stem
|
||||
|
||||
# Check standard preview naming conventions
|
||||
for ext in (
|
||||
".png",
|
||||
".jpg",
|
||||
".jpeg",
|
||||
".webp",
|
||||
".preview.png",
|
||||
".preview.jpg",
|
||||
".preview.jpeg",
|
||||
):
|
||||
candidate = lora_dir / f"{base_name}{ext}"
|
||||
if candidate.exists():
|
||||
results.append(str(candidate.resolve()))
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def get_preview_image_from_metadata(
|
||||
metadata: Dict, metadata_path: Path
|
||||
) -> Optional[str]:
|
||||
"""Find the first preview image path for a LoRA (backward compat)."""
|
||||
images = get_preview_images_from_metadata(metadata, metadata_path)
|
||||
return images[0] if images else None
|
||||
|
||||
|
||||
_THUMB_MAX_SIZE = 512
|
||||
|
||||
# CivitAI's own domains (#340: civitai.red and civitai.green serve the same
|
||||
# content as civitai.com).
|
||||
CIVITAI_DOMAINS = ("civitai.com", "civitai.red", "civitai.green")
|
||||
# Only these hosts may ever receive the CivitAI API key, and only over HTTPS.
|
||||
CIVITAI_HOSTS = frozenset(
|
||||
f"{prefix}{domain}"
|
||||
for domain in CIVITAI_DOMAINS
|
||||
for prefix in ("", "www.", "api.", "image.")
|
||||
{"civitai.com", "www.civitai.com", "api.civitai.com", "image.civitai.com"}
|
||||
)
|
||||
MAX_CIVITAI_DOWNLOAD_BYTES = 50 * 1024 * 1024
|
||||
_DOWNLOAD_CHUNK_BYTES = 64 * 1024
|
||||
|
||||
|
||||
# A redirect is followed (#337: image.civitai.com now answers with a redirect
|
||||
# to blobs-b2.civitai.com) only over HTTPS to one of CIVITAI_DOMAINS or a
|
||||
# subdomain of one, at most this many hops, and never with the API key.
|
||||
MAX_CIVITAI_REDIRECTS = 3
|
||||
_REDIRECT_CODES = frozenset({301, 302, 303, 307, 308})
|
||||
_USER_AGENT = "ComfyUI-PromptManager/1.0"
|
||||
|
||||
|
||||
class _DownloadFailed(Exception):
|
||||
"""A download failed for the short, loggable ``reason`` it carries."""
|
||||
|
||||
def __init__(self, reason: str):
|
||||
super().__init__(reason)
|
||||
self.reason = reason
|
||||
|
||||
|
||||
class _RefuseRedirects(urllib.request.HTTPRedirectHandler):
|
||||
"""Never follow a redirect inside urllib: it would re-send the Bearer
|
||||
header. A 3xx surfaces as ``HTTPError``; ``_download_one`` decides."""
|
||||
"""Never follow a redirect: urllib would re-send the Bearer header to it."""
|
||||
|
||||
def redirect_request(self, req, fp, code, msg, headers, newurl):
|
||||
logger.warning(f"Refusing redirect from {req.full_url} to {newurl}")
|
||||
return None
|
||||
|
||||
|
||||
@@ -456,84 +425,6 @@ def _open_url(req: urllib.request.Request, timeout: float = 10):
|
||||
return _opener.open(req, timeout=timeout)
|
||||
|
||||
|
||||
def _loggable(url: str) -> str:
|
||||
"""``url`` without its query, fragment or credentials (signed URLs)."""
|
||||
try:
|
||||
parts = urlsplit(url)
|
||||
host = parts.hostname or ""
|
||||
except ValueError:
|
||||
return "<unparsable URL>"
|
||||
return f"{parts.scheme}://{host}{parts.path}"
|
||||
|
||||
|
||||
def _on_civitai_domain(hostname: str) -> bool:
|
||||
"""True when ``hostname`` is one of CIVITAI_DOMAINS or a subdomain of one."""
|
||||
return any(
|
||||
hostname == domain or hostname.endswith(f".{domain}")
|
||||
for domain in CIVITAI_DOMAINS
|
||||
)
|
||||
|
||||
|
||||
def is_civitai_redirect_target(url: str) -> bool:
|
||||
"""True when a redirect to ``url`` may be followed: HTTPS on one of
|
||||
CIVITAI_DOMAINS or a subdomain of one, matched on the parsed host name,
|
||||
default port and no credentials in the URL."""
|
||||
try:
|
||||
parts = urlsplit(url)
|
||||
port = parts.port
|
||||
except ValueError:
|
||||
return False
|
||||
hostname = (parts.hostname or "").lower()
|
||||
return (
|
||||
parts.scheme == "https"
|
||||
and _on_civitai_domain(hostname)
|
||||
and parts.username is None
|
||||
and parts.password is None
|
||||
and port in (None, 443)
|
||||
)
|
||||
|
||||
|
||||
def _redirect_target(url: str, error) -> Optional[str]:
|
||||
"""The URL a 3xx ``error`` for ``url`` points to, or None when absent."""
|
||||
if error.code not in _REDIRECT_CODES:
|
||||
return None
|
||||
location = (error.headers.get("Location") or "").strip()
|
||||
if not location:
|
||||
return None
|
||||
return urljoin(url, location)
|
||||
|
||||
|
||||
def _open_following_civitai_redirects(url: str, api_key: str):
|
||||
"""Open ``url``, following only allowed CivitAI redirects (keyless).
|
||||
|
||||
The first request carries the key; every redirected one carries the
|
||||
User-Agent alone. Returns the open response; raises ``_DownloadFailed``
|
||||
when a redirect is refused or the hop limit is reached. Other HTTP
|
||||
errors propagate.
|
||||
"""
|
||||
headers = {"User-Agent": _USER_AGENT}
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
current = url
|
||||
for hop in range(MAX_CIVITAI_REDIRECTS + 1):
|
||||
req = urllib.request.Request(current, headers=headers)
|
||||
try:
|
||||
return _open_url(req, timeout=10)
|
||||
except urllib.error.HTTPError as error:
|
||||
with error:
|
||||
target = _redirect_target(current, error)
|
||||
if target is None:
|
||||
raise
|
||||
if hop == MAX_CIVITAI_REDIRECTS:
|
||||
raise _DownloadFailed(f"more than {MAX_CIVITAI_REDIRECTS} redirects")
|
||||
if not is_civitai_redirect_target(target):
|
||||
raise _DownloadFailed(f"redirect refused to {_loggable(target)}")
|
||||
logger.debug(f"Following redirect to {_loggable(target)}")
|
||||
current = target
|
||||
headers = {"User-Agent": _USER_AGENT} # never the key across a redirect
|
||||
raise AssertionError("unreachable") # pragma: no cover - loop ends above
|
||||
|
||||
|
||||
def is_civitai_url(url: str) -> bool:
|
||||
"""True when ``url`` is an HTTPS URL on an allow-listed CivitAI host."""
|
||||
try:
|
||||
@@ -574,160 +465,73 @@ def _read_capped(resp, cap: int) -> Optional[bytes]:
|
||||
return b"".join(chunks)
|
||||
|
||||
|
||||
def _network_reason(error: BaseException) -> str:
|
||||
"""A short reason for a failed request: a timeout or a network error."""
|
||||
cause = getattr(error, "reason", error) # URLError wraps the socket error
|
||||
if isinstance(cause, TimeoutError):
|
||||
return "timed out"
|
||||
return f"network error: {cause}"
|
||||
|
||||
|
||||
def _fetch_civitai_body(url: str, api_key: str) -> bytes:
|
||||
"""The body behind ``url`` (redirects as allowed, size capped), or
|
||||
``_DownloadFailed`` with the reason it could not be fetched."""
|
||||
try:
|
||||
with _open_following_civitai_redirects(url, api_key) as resp:
|
||||
raw = _read_capped(resp, MAX_CIVITAI_DOWNLOAD_BYTES)
|
||||
except urllib.error.HTTPError as error:
|
||||
error.close()
|
||||
raise _DownloadFailed(f"HTTP {error.code}") from None
|
||||
except (urllib.error.URLError, OSError, http.client.HTTPException) as error:
|
||||
raise _DownloadFailed(_network_reason(error)) from None
|
||||
if raw is None:
|
||||
cap = MAX_CIVITAI_DOWNLOAD_BYTES
|
||||
size = f"{cap // (1024 * 1024)} MB" if cap >= 1024 * 1024 else f"{cap} byte"
|
||||
raise _DownloadFailed(f"larger than the {size} cap")
|
||||
return raw
|
||||
|
||||
|
||||
def _thumbnail(raw: bytes):
|
||||
"""``raw`` decoded and shrunk to an RGB thumbnail, or ``_DownloadFailed``."""
|
||||
from io import BytesIO
|
||||
|
||||
from PIL import Image
|
||||
|
||||
try:
|
||||
with Image.open(BytesIO(raw)) as img:
|
||||
img.thumbnail((_THUMB_MAX_SIZE, _THUMB_MAX_SIZE), Image.LANCZOS)
|
||||
return img.convert("RGB")
|
||||
except Exception as error: # PIL raises several types for bad data
|
||||
raise _DownloadFailed(
|
||||
f"image could not be decoded ({type(error).__name__})"
|
||||
) from None
|
||||
|
||||
|
||||
def _download_one(url: str, local_path: Path, api_key: str) -> Optional[str]:
|
||||
"""Download a single CivitAI image, resize to thumbnail, save as JPEG.
|
||||
|
||||
Refuses any URL that is not HTTPS on an allow-listed CivitAI host so the
|
||||
API key is never sent elsewhere, follows a redirect only as
|
||||
``_open_following_civitai_redirects`` allows (without the key), and
|
||||
refuses bodies larger than MAX_CIVITAI_DOWNLOAD_BYTES. A failure logs
|
||||
one warning with a short reason and the URL without query or
|
||||
credentials (#340), and returns None.
|
||||
API key is never sent elsewhere, and refuses bodies larger than
|
||||
MAX_CIVITAI_DOWNLOAD_BYTES.
|
||||
"""
|
||||
if not is_civitai_url(url):
|
||||
logger.warning(f"Refusing non-CivitAI download URL: {url}")
|
||||
return None
|
||||
try:
|
||||
if not is_civitai_url(url):
|
||||
raise _DownloadFailed("not an allowed CivitAI URL")
|
||||
image = _thumbnail(_fetch_civitai_body(url, api_key))
|
||||
try:
|
||||
_save_atomically(image, local_path)
|
||||
except OSError as error:
|
||||
raise _DownloadFailed(
|
||||
f"image could not be saved ({error.strerror or type(error).__name__})"
|
||||
) from None
|
||||
except _DownloadFailed as failure:
|
||||
reason = failure.reason
|
||||
except Exception as error: # best effort: an import never fails on one image
|
||||
reason = f"unexpected {type(error).__name__}"
|
||||
else:
|
||||
headers = {"User-Agent": "ComfyUI-PromptManager/1.0"}
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
req = urllib.request.Request(url, headers=headers)
|
||||
with _open_url(req, timeout=10) as resp:
|
||||
raw = _read_capped(resp, MAX_CIVITAI_DOWNLOAD_BYTES)
|
||||
if raw is None:
|
||||
logger.warning(f"Refusing oversized download: {url}")
|
||||
return None
|
||||
|
||||
# Resize to thumbnail to save disk space
|
||||
from io import BytesIO
|
||||
|
||||
from PIL import Image
|
||||
|
||||
img = Image.open(BytesIO(raw))
|
||||
img.thumbnail((_THUMB_MAX_SIZE, _THUMB_MAX_SIZE), Image.LANCZOS)
|
||||
img = img.convert("RGB")
|
||||
img.save(str(local_path), "JPEG", quality=85)
|
||||
|
||||
return str(local_path.resolve())
|
||||
logger.warning(f"CivitAI image download failed: {_loggable(url)}: {reason}")
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to download {url}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
_PART_PREFIX = ".pm-"
|
||||
_PART_SUFFIX = ".part"
|
||||
# A temp download older than this was left by a killed process, not a live one
|
||||
# (each download has a 10 s timeout).
|
||||
STALE_PART_SECONDS = 3600
|
||||
def download_civitai_images(
|
||||
metadata: Dict, metadata_path: Path, cache_dir: Path, api_key: str = ""
|
||||
) -> List[str]:
|
||||
"""Download civitai example images to a local cache directory.
|
||||
|
||||
Uses thumbnail URLs (512px) instead of full-size originals, and
|
||||
downloads in parallel (up to 8 concurrent) for speed.
|
||||
|
||||
def _save_atomically(image, local_path: Path) -> None:
|
||||
"""Write ``image`` as JPEG so ``local_path`` only ever names a whole file.
|
||||
Args:
|
||||
api_key: CivitAI API key for authenticated downloads (NSFW content).
|
||||
|
||||
The image goes to a ``.pm-*.part`` file in the same folder, which is
|
||||
closed before ``os.replace`` puts it in place (Windows cannot replace an
|
||||
open file); a failure removes the temp file. A reader (another import's
|
||||
cache check) never sees a half-written image. The temp file is opened
|
||||
like ``img.save`` opens a file, so the image keeps the umask's mode.
|
||||
Returns:
|
||||
List of absolute paths to downloaded image files.
|
||||
"""
|
||||
import secrets
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
while True:
|
||||
temp = local_path.parent / (
|
||||
f"{_PART_PREFIX}{secrets.token_hex(8)}{_PART_SUFFIX}"
|
||||
)
|
||||
try:
|
||||
handle = open(temp, "xb")
|
||||
except FileExistsError:
|
||||
continue
|
||||
break
|
||||
try:
|
||||
with handle:
|
||||
image.save(handle, "JPEG", quality=85)
|
||||
os.replace(temp, local_path)
|
||||
except BaseException:
|
||||
try:
|
||||
os.remove(temp)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
|
||||
def sweep_stale_downloads(cache_dir: Path) -> None:
|
||||
"""Remove ``.pm-*.part`` files a killed download left in the image cache.
|
||||
|
||||
Only that exact pattern, only in ``cache_dir`` and its LoRA folders, and
|
||||
only once older than ``STALE_PART_SECONDS``; a file that cannot be read
|
||||
or removed is left for a later import.
|
||||
"""
|
||||
import time
|
||||
|
||||
cutoff = time.time() - STALE_PART_SECONDS
|
||||
try:
|
||||
folders = [cache_dir] + [p for p in cache_dir.iterdir() if p.is_dir()]
|
||||
except OSError:
|
||||
return
|
||||
for folder in folders:
|
||||
try:
|
||||
candidates = list(folder.glob(f"{_PART_PREFIX}*{_PART_SUFFIX}"))
|
||||
except OSError:
|
||||
continue
|
||||
for path in candidates:
|
||||
try:
|
||||
if path.is_file() and path.stat().st_mtime < cutoff:
|
||||
os.remove(path)
|
||||
except OSError:
|
||||
logger.debug(f"Could not remove stale download {path}")
|
||||
|
||||
|
||||
def _civitai_cache_plan(
|
||||
metadata: Dict, metadata_path: Path, cache_dir: Path
|
||||
) -> Tuple[Path, List[str], List[Tuple[str, Path]]]:
|
||||
"""``(lora_cache, cached, tasks)`` for a LoRA's CivitAI example images.
|
||||
|
||||
``cached`` are the images already in the cache, ``tasks`` the
|
||||
``(url, local_path)`` downloads still needed. Reads only; creates nothing.
|
||||
"""
|
||||
civitai = _get_civitai(metadata)
|
||||
images = civitai.get("images", []) or []
|
||||
if not images:
|
||||
return []
|
||||
|
||||
file_name = metadata.get("file_name", "")
|
||||
if not file_name:
|
||||
file_name = metadata_path.name.replace(".metadata.json", "")
|
||||
lora_cache = cache_dir / Path(file_name).stem
|
||||
lora_stem = Path(file_name).stem
|
||||
|
||||
lora_cache = cache_dir / lora_stem
|
||||
lora_cache.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Build download tasks
|
||||
cached = []
|
||||
tasks = [] # (url, local_path)
|
||||
for img in images:
|
||||
@@ -744,42 +548,9 @@ def _civitai_cache_plan(
|
||||
cached.append(str(local_path.resolve()))
|
||||
else:
|
||||
tasks.append((url, local_path))
|
||||
return lora_cache, cached, tasks
|
||||
|
||||
|
||||
def planned_civitai_downloads(
|
||||
metadata: Dict, metadata_path: Path, cache_dir: Path
|
||||
) -> int:
|
||||
"""How many example images an import would download (no request, no write)."""
|
||||
return len(_civitai_cache_plan(metadata, metadata_path, cache_dir)[2])
|
||||
|
||||
|
||||
def download_civitai_images(
|
||||
metadata: Dict,
|
||||
metadata_path: Path,
|
||||
cache_dir: Path,
|
||||
api_key: str = "",
|
||||
stats: Optional[Dict[str, int]] = None,
|
||||
) -> List[str]:
|
||||
"""Download civitai example images to a local cache directory.
|
||||
|
||||
Uses thumbnail URLs (512px) instead of full-size originals, and
|
||||
downloads in parallel (up to 8 concurrent) for speed.
|
||||
|
||||
Args:
|
||||
api_key: CivitAI API key for authenticated downloads (NSFW content).
|
||||
stats: When given, ``downloaded`` and ``download_errors`` are
|
||||
increased by the downloads that succeeded and failed.
|
||||
|
||||
Returns:
|
||||
List of absolute paths to downloaded image files.
|
||||
"""
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
lora_cache, cached, tasks = _civitai_cache_plan(metadata, metadata_path, cache_dir)
|
||||
if not tasks:
|
||||
return cached
|
||||
lora_cache.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Download in parallel
|
||||
downloaded = []
|
||||
@@ -792,46 +563,42 @@ def download_civitai_images(
|
||||
if result:
|
||||
downloaded.append(result)
|
||||
|
||||
if stats is not None:
|
||||
stats["downloaded"] = stats.get("downloaded", 0) + len(downloaded)
|
||||
stats["download_errors"] = (
|
||||
stats.get("download_errors", 0) + len(tasks) - len(downloaded)
|
||||
)
|
||||
return cached + downloaded
|
||||
|
||||
|
||||
def get_lora_image_cache_dir(create: bool = True) -> Path:
|
||||
"""Get the directory used to cache downloaded LoRA example images.
|
||||
|
||||
``<data folder>/data/lora_images`` (``utils.library_location``): the node
|
||||
folder for a library that was not moved, else ComfyUI's
|
||||
``user/__promptmanager``. ``create=False`` only names it (a preview
|
||||
writes nothing).
|
||||
"""
|
||||
try:
|
||||
from ..utils import library_location
|
||||
except ImportError:
|
||||
from utils import library_location
|
||||
|
||||
def get_lora_image_cache_dir() -> Path:
|
||||
"""Get the directory used to cache downloaded LoRA example images."""
|
||||
# Store in the extension's own directory
|
||||
ext_root = Path(__file__).resolve().parent.parent
|
||||
cache = Path(
|
||||
library_location.cache_dir("lora_images", extension_root=str(ext_root))
|
||||
)
|
||||
if create:
|
||||
cache.mkdir(parents=True, exist_ok=True)
|
||||
cache = ext_root / "data" / "lora_images"
|
||||
cache.mkdir(parents=True, exist_ok=True)
|
||||
return cache
|
||||
|
||||
|
||||
def get_example_images_dir(lora_manager_path: str) -> Optional[str]:
|
||||
"""Find the LoraManager example_images directory."""
|
||||
lm_path = Path(lora_manager_path)
|
||||
|
||||
# Direct subdirectory
|
||||
candidate = lm_path / "example_images"
|
||||
if candidate.is_dir():
|
||||
return str(candidate.resolve())
|
||||
|
||||
# Search one level in user data dirs
|
||||
for child in lm_path.iterdir():
|
||||
if child.is_dir():
|
||||
sub = child / "example_images"
|
||||
if sub.is_dir():
|
||||
return str(sub.resolve())
|
||||
|
||||
return None
|
||||
|
||||
|
||||
# ── Trigger word cache & injection ───────────────────────────────────
|
||||
|
||||
_LORA_PATTERN = re.compile(r"<lora:([^:>]+):[^>]+>", re.IGNORECASE)
|
||||
|
||||
|
||||
# A failed load is retried by a running node only after this long (an explicit
|
||||
# reload, from the settings routes, always runs).
|
||||
TRIGGER_RETRY_SECONDS = 60
|
||||
|
||||
|
||||
class TriggerWordCache:
|
||||
"""Thread-safe cache mapping LoRA names to their trigger words.
|
||||
|
||||
@@ -842,81 +609,52 @@ class TriggerWordCache:
|
||||
self._cache: Dict[str, List[str]] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._loaded = False
|
||||
self._failed_at: Optional[float] = None
|
||||
|
||||
def load(self, lora_manager_path: str) -> int:
|
||||
"""Build the trigger word map from the sidecars LoraManager uses.
|
||||
|
||||
Never raises (a node calls this while it runs): a folder that cannot
|
||||
be listed, or any other failure, is logged and the previous map is
|
||||
kept. When LoraManager's settings cannot be used, the reason is
|
||||
logged and the default LoRA folders are read in alongside mode.
|
||||
"""Scan LoRA metadata files and build the trigger word mapping.
|
||||
|
||||
Returns:
|
||||
Number of LoRAs with trigger words in the map.
|
||||
Number of LoRAs with trigger words found.
|
||||
"""
|
||||
try:
|
||||
new_cache = self._build(lora_manager_path)
|
||||
except Exception as error:
|
||||
logger.warning(f"Trigger word cache not refreshed: {error}")
|
||||
with self._lock:
|
||||
self._failed_at = time.monotonic()
|
||||
return len(self._cache)
|
||||
new_cache: Dict[str, List[str]] = {}
|
||||
|
||||
lora_dirs = find_lora_directories(lora_manager_path)
|
||||
for lora_dir in lora_dirs:
|
||||
dir_path = Path(lora_dir)
|
||||
for meta_file in dir_path.rglob("*.metadata.json"):
|
||||
metadata = read_lora_metadata(meta_file)
|
||||
if not metadata:
|
||||
continue
|
||||
|
||||
words = get_trigger_words_from_metadata(metadata)
|
||||
if not words:
|
||||
continue
|
||||
|
||||
# Key by filename stem (what appears in <lora:NAME:weight>)
|
||||
file_name = metadata.get("file_name", "")
|
||||
if file_name:
|
||||
stem = Path(file_name).stem
|
||||
new_cache[stem.lower()] = words
|
||||
|
||||
# Also key by the metadata file stem
|
||||
meta_stem = meta_file.name.replace(".metadata.json", "")
|
||||
if meta_stem.lower() not in new_cache:
|
||||
new_cache[meta_stem.lower()] = words
|
||||
|
||||
with self._lock:
|
||||
self._cache = new_cache
|
||||
self._loaded = True
|
||||
self._failed_at = None
|
||||
|
||||
logger.info(
|
||||
f"Trigger word cache loaded: {len(new_cache)} LoRAs with trigger words"
|
||||
)
|
||||
return len(new_cache)
|
||||
|
||||
@staticmethod
|
||||
def _build(lora_manager_path: str) -> Dict[str, List[str]]:
|
||||
layout = lora_manager_layout(lora_manager_path)
|
||||
if layout.error is not None:
|
||||
logger.warning(
|
||||
f"LoraManager settings not usable ({layout.error}); trigger words "
|
||||
"come from the default LoRA folders"
|
||||
)
|
||||
layout = replace(
|
||||
layout,
|
||||
error=None,
|
||||
mode=MODE_ALONGSIDE,
|
||||
sidecar_root=None,
|
||||
example_images_root=None,
|
||||
)
|
||||
new_cache: Dict[str, List[str]] = {}
|
||||
for sidecar in discover_sidecars(layout).sidecars:
|
||||
metadata = read_lora_metadata(sidecar.metadata_path)
|
||||
if not metadata:
|
||||
continue
|
||||
words = get_trigger_words_from_metadata(metadata)
|
||||
if not words:
|
||||
continue
|
||||
# Key by the model file name (what appears in <lora:NAME:weight>)
|
||||
# and by the sidecar's name.
|
||||
new_cache[sidecar.model_path.stem.lower()] = words
|
||||
sidecar_stem = sidecar.metadata_path.name[: -len(METADATA_SUFFIX)]
|
||||
new_cache.setdefault(sidecar_stem.lower(), words)
|
||||
return new_cache
|
||||
|
||||
def get_trigger_words(self, lora_name: str) -> List[str]:
|
||||
"""Look up trigger words for a LoRA by name (case-insensitive)."""
|
||||
with self._lock:
|
||||
return self._cache.get(lora_name.lower(), [])
|
||||
|
||||
def retry_due(self) -> bool:
|
||||
"""Whether a node may try loading again: never failed, or the last
|
||||
failure is TRIGGER_RETRY_SECONDS old (no rescan on every prompt run)."""
|
||||
with self._lock:
|
||||
failed_at = self._failed_at
|
||||
return failed_at is None or time.monotonic() - failed_at >= (
|
||||
TRIGGER_RETRY_SECONDS
|
||||
)
|
||||
|
||||
@property
|
||||
def is_loaded(self) -> bool:
|
||||
with self._lock:
|
||||
@@ -926,7 +664,6 @@ class TriggerWordCache:
|
||||
with self._lock:
|
||||
self._cache.clear()
|
||||
self._loaded = False
|
||||
self._failed_at = None
|
||||
|
||||
|
||||
# Module-level singleton
|
||||
|
||||
@@ -1,55 +0,0 @@
|
||||
"""The tagging configuration's constants and validators (#328).
|
||||
|
||||
Server-free on purpose: ``py.api`` must import without ComfyUI's ``server``
|
||||
module, which ``py/config.py`` needs; config re-exports all of this.
|
||||
"""
|
||||
|
||||
import math
|
||||
from typing import Any
|
||||
|
||||
# The tagging engines (``py/autotag.py`` MODELS) and the options they take.
|
||||
WD14_MODEL_TYPES = frozenset({"wd14-swinv2", "wd14-vit"})
|
||||
LLM_MODEL_TYPES = frozenset({"gguf", "hf"})
|
||||
MODEL_TYPES = WD14_MODEL_TYPES | LLM_MODEL_TYPES
|
||||
MAX_INSTRUCTION_CHARS = 4000
|
||||
MAX_INSTRUCTION_BYTES = 16 * 1024
|
||||
# What a batch does with a prompt that already has tags (#328).
|
||||
EXISTING_POLICIES = ("skip", "merge", "replace_after_review")
|
||||
|
||||
# The language-model tag template every page starts from until one is saved.
|
||||
DEFAULT_TAG_TEMPLATE = (
|
||||
"Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces).\n"
|
||||
"Include: subject_count (1girl/1boy/etc), gender, ethnicity (asian, caucasian, "
|
||||
"etc), age_approx (teen, young_woman), body_attributes, face_features, "
|
||||
"hairstyle, appearance, clothing, accessories, pose, expression, action, "
|
||||
"camera_angle, camera_shot, scene_composition (wide_shot, long_shot, "
|
||||
"dutch_angle), environment, background_elements, lighting, mood.\n"
|
||||
"Prioritize visual accuracy.\n"
|
||||
"No sentences, no commentary, no captions. Only tags. Keep length short but "
|
||||
"descriptive."
|
||||
)
|
||||
|
||||
|
||||
def validate_threshold(name: str, value: Any) -> float:
|
||||
"""A WD14 confidence threshold in [0, 1], or ValueError naming ``name``."""
|
||||
if type(value) not in (int, float) or not math.isfinite(value):
|
||||
raise ValueError(f"{name} must be a finite number")
|
||||
if not 0 <= value <= 1:
|
||||
raise ValueError(f"{name} must be between 0 and 1")
|
||||
return float(value)
|
||||
|
||||
|
||||
def validate_template(value: Any, name: str = "instructions") -> str:
|
||||
"""Non-empty language-model instructions within the character and byte
|
||||
caps, or ValueError naming ``name``."""
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise ValueError(f"{name} must be non-empty text")
|
||||
if len(value) > MAX_INSTRUCTION_CHARS:
|
||||
raise ValueError(f"{name} must be at most {MAX_INSTRUCTION_CHARS} characters")
|
||||
try:
|
||||
encoded = value.encode("utf-8")
|
||||
except UnicodeEncodeError:
|
||||
raise ValueError(f"{name} must be valid Unicode text") from None
|
||||
if len(encoded) > MAX_INSTRUCTION_BYTES:
|
||||
raise ValueError(f"{name} {'are' if name == 'instructions' else 'is'} too long")
|
||||
return value
|
||||
+2
-2
@@ -3,11 +3,11 @@ name = "promptmanager"
|
||||
description = "A powerful ComfyUI custom node that extends the standard text encoder with persistent prompt storage, advanced search capabilities, and an automatic image gallery system using SQLite."
|
||||
version = "3.3.1"
|
||||
license = {file = "LICENSE"}
|
||||
requires-python = ">=3.12"
|
||||
requires-python = ">=3.10"
|
||||
dependencies = ["watchdog>=2.1.0", "Pillow>=8.0.0"]
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=84.2.1"]
|
||||
requires = ["setuptools>=84.0.0"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project.urls]
|
||||
|
||||
@@ -1,27 +0,0 @@
|
||||
"""Final stops for the background threads a test starts (#301, #302).
|
||||
|
||||
Every real ``ImageGenerationHandler``, ``ImageMonitor`` or ``PromptTracker``
|
||||
a test constructs ends with its final stop and an assertion, before the
|
||||
test's folder and databases go: ``close()`` for a handler,
|
||||
``stop_monitoring()`` for a monitor, ``shutdown()`` for a tracker.
|
||||
"""
|
||||
|
||||
from utils.prompt_tracker import PromptTracker
|
||||
|
||||
|
||||
def stop_tracker(testcase, tracker):
|
||||
"""Shut a real tracker's cleanup thread down; mocks are left alone."""
|
||||
if isinstance(tracker, PromptTracker):
|
||||
testcase.assertTrue(tracker.shutdown(), "prompt cleanup thread still running")
|
||||
|
||||
|
||||
def close_handler(testcase, handler):
|
||||
"""Close ``handler`` for good, then stop the tracker it was built with."""
|
||||
testcase.assertTrue(handler.close(), "image worker still running")
|
||||
stop_tracker(testcase, getattr(handler, "prompt_tracker", None))
|
||||
|
||||
|
||||
def stop_monitor(testcase, monitor):
|
||||
"""Stop ``monitor`` (its observer and worker), then its tracker."""
|
||||
testcase.assertTrue(monitor.stop_monitoring(), "image monitor still running")
|
||||
stop_tracker(testcase, getattr(monitor, "prompt_tracker", None))
|
||||
@@ -149,49 +149,3 @@ test('Classic test button reports success and failure and releases expired notif
|
||||
await page.clock.fastForward(3001);
|
||||
await expect(page.locator('.pm-classic-notification')).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('Classic shows linked text, keeps the typed text through save and refuses to overwrite it (#343)',async({page})=>{
|
||||
await page.route('**/prompt_manager/recent?**',route=>route.fulfill({json:{success:true,results:[{text:'Saved prompt',tags:[]}]}}));
|
||||
const note=page.locator('.pm-classic').getByText('Text comes from a link. Your typed text is kept for when you disconnect it.',{exact:true});
|
||||
await expect(note).toBeHidden();
|
||||
await page.evaluate(()=>{
|
||||
const widget=node.widgets.find(item=>item.name==='text');widget.value=widget.inputEl.value='Typed classic';
|
||||
node.inputs=[{name:'text',type:'STRING',link:3,widget:{name:'text'}}];node.onConnectionsChange();
|
||||
});
|
||||
await expect(note).toBeVisible();
|
||||
// The link drives the run; the saved workflow keeps what the user typed, not a stale or empty value.
|
||||
expect(await page.evaluate(()=>node.serialize().widgets_values[0])).toBe('Typed classic');
|
||||
await page.getByRole('button',{name:'📋 Recent',exact:true}).click();
|
||||
await page.getByRole('button',{name:/Saved prompt/}).click();
|
||||
expect(await page.evaluate(()=>node.widgets.find(item=>item.name==='text').value)).toBe('Typed classic');
|
||||
await expect(page.getByRole('alert').filter({hasText:'Text comes from a link. Disconnect it to use a saved prompt.'})).toBeVisible();
|
||||
await page.evaluate(()=>{node.inputs[0].link=null;node.onConnectionsChange();});
|
||||
await expect(note).toBeHidden();
|
||||
await page.getByRole('button',{name:/Saved prompt/}).click();
|
||||
expect(await page.evaluate(()=>node.widgets.find(item=>item.name==='text').value)).toBe('Saved prompt');
|
||||
// A workflow load re-evaluates the link state.
|
||||
await page.evaluate(()=>{node.inputs[0].link=8;node.configure({});});
|
||||
await expect(note).toBeVisible();
|
||||
});
|
||||
|
||||
test('PromptManagerText shows linked text in its legacy controls and keeps the typed text (#343)',async({page})=>{
|
||||
await page.evaluate(async()=>{
|
||||
class TextNode {
|
||||
constructor(){this.widgets=['text','category','tags','search_text'].map(name=>({name,value:'',inputEl:document.createElement('textarea')}));this.widgets_values=[];this.inputs=[{name:'text',type:'STRING',link:null,widget:{name:'text'}}];}
|
||||
addDOMWidget(name,type,element){element.id='legacy-text-node';document.body.append(element);return {element};}
|
||||
setSize(){} setDirtyCanvas(){} serialize(){return {properties:{...this.properties}};}
|
||||
}
|
||||
await fixtureApp.extensions[0].beforeRegisterNodeDef(TextNode,{name:'PromptManagerText'},fixtureApp);
|
||||
window.textNode=new TextNode();textNode.onNodeCreated();
|
||||
const widget=textNode.widgets[0];widget.value=widget.inputEl.value='Typed text node';
|
||||
});
|
||||
const note=page.locator('#legacy-text-node').getByText('Text comes from a link. Your typed text is kept for when you disconnect it.',{exact:true});
|
||||
await expect(note).toBeHidden();
|
||||
await page.evaluate(()=>{textNode.inputs[0].link=5;textNode.onConnectionsChange();});
|
||||
await expect(note).toBeVisible();
|
||||
expect(await page.evaluate(()=>textNode.serialize().widgets_values[0])).toBe('Typed text node');
|
||||
await page.evaluate(()=>{textNode.usePrompt({text:'Replacement'},0);});
|
||||
expect(await page.evaluate(()=>textNode.widgets[0].value)).toBe('Typed text node');
|
||||
await page.evaluate(()=>{textNode.inputs[0].link=null;textNode.onConnectionsChange();textNode.onRemoved();});
|
||||
await expect(note).toBeHidden();
|
||||
});
|
||||
|
||||
@@ -26,8 +26,7 @@ test.beforeEach(async({page,browserName})=>{
|
||||
window.fixtureApp={registerExtension(extension){this.extension=extension;},ui:{settings:new EventTarget()},graph:{setDirtyCanvas(){}},canvas:{}};
|
||||
await import('/extensions/promptmanager/prompt_manager.js');
|
||||
class Node {
|
||||
constructor(){this.widgets=['text','category','tags','search_text'].map(name=>({name,type:'text',value:'',options:{},inputEl:document.createElement('textarea')}));this.inputs=[...['clip','prepend_text','append_text'].map(name=>({name,type:name==='clip'?'CLIP':'STRING',link:null})),...['text','category','tags','search_text'].map(name=>({name,type:'STRING',link:null,widget:{name}}))];this.properties={};this.size=[550,700];}
|
||||
findInputByType(type){const index=this.inputs.findIndex(input=>input.type===type&&input.link==null);return index<0?undefined:{index,slot:this.inputs[index]};}
|
||||
constructor(){this.widgets=['text','category','tags','search_text'].map(name=>({name,type:'text',value:'',options:{},inputEl:document.createElement('textarea')}));this.inputs=['clip','prepend_text','append_text'].map(name=>({name,link:null}));this.properties={};this.size=[550,700];}
|
||||
addWidget(type,name,value,callback,options){const widget={type,name,value,callback,options};this.widgets.push(widget);return widget;}
|
||||
addDOMWidget(name,type,element){const host=document.createElement('div');host.style.height='650px';host.className='fixture-dom-widget';host.append(element);document.body.append(host);return {element};}
|
||||
setSize(size){this.size=size;}setDirtyCanvas(){}
|
||||
@@ -247,39 +246,13 @@ test(`${variant} restores workflow values and disclosures and protects connected
|
||||
await page.evaluate(()=>{const saved=node.serialize();node.configure(saved);});
|
||||
await expect(page.getByLabel('Category',{exact:true})).toBeVisible();
|
||||
await expect(page.getByLabel('Category',{exact:true})).toHaveValue('Saved category');
|
||||
await page.evaluate(()=>{node.inputs.find(input=>input.name==='text').link=9;node.onConnectionsChange();});
|
||||
await page.evaluate(()=>{node.inputs.push({name:'text',link:9});node.onConnectionsChange();});
|
||||
await expect(page.getByRole('textbox',{name:'Main prompt',exact:true})).toBeHidden();
|
||||
await expect(page.getByText('Main prompt comes from a link. Your typed text is kept for when you disconnect it.',{exact:true})).toBeVisible();
|
||||
await expect(page.getByText('Text is supplied by an upstream connection.',{exact:true})).toBeVisible();
|
||||
await page.evaluate(()=>{node.inputs.find(input=>input.name==='text').link=null;node.onConnectionsChange();});
|
||||
await expect(page.getByRole('textbox',{name:'Main prompt',exact:true})).toHaveValue('Saved draft');
|
||||
});
|
||||
|
||||
test(`${variant} offers text as a link socket, keeps typed text while linked and through a reload (#343)`,async({page})=>{
|
||||
const sockets=()=>page.evaluate(()=>Object.fromEntries(node.inputs.map(input=>[input.name,input.widget?.name??null])));
|
||||
// The hidden text widget no longer owns the input, so the host draws it as a socket; the other widget inputs are unchanged.
|
||||
expect(await sockets()).toEqual({clip:null,prepend_text:null,append_text:null,text:null,category:'category',tags:'tags',search_text:'search_text'});
|
||||
expect(await page.evaluate(()=>node.findInputByType('STRING').slot.name)).toBe('text');
|
||||
expect(await page.evaluate(()=>node.findInputByType('CLIP').slot.name)).toBe('clip');
|
||||
const main=page.getByRole('textbox',{name:'Main prompt',exact:true});
|
||||
await main.fill('Typed main');
|
||||
await page.evaluate(()=>{node.inputs.find(input=>input.name==='text').link=21;node.onConnectionsChange();});
|
||||
await expect(main).toBeHidden();
|
||||
const note=page.getByText('Main prompt comes from a link. Your typed text is kept for when you disconnect it.',{exact:true});
|
||||
await expect(note).toBeVisible();
|
||||
await expect(page.locator('[data-composed-preview]')).toHaveText('[text supplied by connection]');
|
||||
// A second STRING link dropped on the editor falls back to the frontend's own choice.
|
||||
expect(await page.evaluate(()=>node.findInputByType('STRING').slot.name)).toBe('prepend_text');
|
||||
// An older saved workflow brings the widget binding back with its inputs: configure exposes the socket again.
|
||||
const saved=await page.evaluate(()=>{const saved=node.serialize();node.inputs.find(input=>input.name==='text').widget={name:'text'};node.configure(saved);return saved;});
|
||||
expect(saved.widgets_values[0]).toBe('Typed main');
|
||||
expect((await sockets()).text).toBeNull();
|
||||
await expect(note).toBeVisible();
|
||||
await page.evaluate(()=>{node.inputs.find(input=>input.name==='text').link=null;node.onConnectionsChange();});
|
||||
await expect(main).toBeVisible();
|
||||
await expect(main).toHaveValue('Typed main');
|
||||
await expect(note).toBeHidden();
|
||||
});
|
||||
|
||||
test(`${variant} searches quoted tags with Enter and explicitly loads metadata`, async ({page}) => {
|
||||
let requested;
|
||||
await page.route('**/prompt_manager/search?**', async route => {
|
||||
|
||||
@@ -1,152 +0,0 @@
|
||||
// #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`).
|
||||
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();
|
||||
}
|
||||
|
||||
async function downloads(request) {
|
||||
const response = await request.get('/prompt_manager/jobs');
|
||||
await expect(response).toBeOK();
|
||||
return (await response.json()).jobs.filter(job => job.kind === 'model_download');
|
||||
}
|
||||
|
||||
test.beforeEach(async ({ page, request, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
// The 3.2.0 update notice would cover the admin page's toolbar.
|
||||
await page.addInitScript(() => { localStorage.setItem('pm_update_notice_dismissed', '3.2.0'); });
|
||||
await post(request, '/__test__/jobs/restart');
|
||||
await post(request, '/__test__/tagging/reset');
|
||||
await post(request, '/__test__/tagging/models', { unavailable: ['wd14-vit'] });
|
||||
await post(request, '/__test__/tagging/download', { mode: 'hold' });
|
||||
});
|
||||
|
||||
test.afterEach(async ({ page, request, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
// Releases a download a failed test left held, and puts the raising default back.
|
||||
await post(request, '/__test__/tagging/reset');
|
||||
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.
|
||||
async function watchProgress(page) {
|
||||
await page.evaluate(() => {
|
||||
window.progressSeen = [];
|
||||
const node = document.getElementById('autoTagProgressPercent');
|
||||
new MutationObserver(() => window.progressSeen.push(node.textContent))
|
||||
.observe(node, { childList: true, characterData: true, subtree: true });
|
||||
});
|
||||
}
|
||||
|
||||
async function runGalleryBatch(page) {
|
||||
await page.locator('#autoTagBtn').click();
|
||||
await page.locator('input[name="autoTagModel"][value="wd14-swinv2"]').check();
|
||||
await watchProgress(page);
|
||||
const done = page.waitForResponse(response => response.url().includes('/prompt_manager/autotag/start'));
|
||||
await page.locator('#startAutoTagBtn').click();
|
||||
await done;
|
||||
await expect(page.locator('#autoTagProgressModal')).toBeHidden();
|
||||
await expect(page.getByText(/^Auto tagging complete! Applied tags to \d+ prompts\.$/)).toBeVisible();
|
||||
}
|
||||
|
||||
test('gallery batch modal counts progress and closes on complete', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await runGalleryBatch(page);
|
||||
const seen = await page.evaluate(() => window.progressSeen);
|
||||
// The fixture's prompts are already tagged, so the batch skips them; its progress frames still move the bar.
|
||||
expect(seen).toContain('95%');
|
||||
expect((await downloads(request)).length).toBe(0);
|
||||
});
|
||||
|
||||
test('admin review opens the first image', async ({ page }) => {
|
||||
await page.goto('/prompt_manager/admin');
|
||||
await page.locator('#autoTagBtn').click();
|
||||
await page.locator('input[name="autoTagModel"][value="wd14-swinv2"]').check();
|
||||
await page.locator('#startReviewBtn').click();
|
||||
// The first image's prompt is already tagged: confirm a new pass, then the review opens on that image.
|
||||
await page.locator('#retagConfirmBtn').click();
|
||||
await expect(page.locator('#autoTagReviewModal')).toBeVisible();
|
||||
await expect(page.locator('#reviewImageName')).not.toHaveText(/^(image_name\.png|Unknown|)$/);
|
||||
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();
|
||||
});
|
||||
@@ -17,8 +17,7 @@ test.beforeEach(async({page,browserName})=>{
|
||||
window.fixtureApp={registerExtension(extension){this.extension=extension;},ui:{settings},graph:{_nodes:[],setDirtyCanvas(){}},canvas:{},graphToPrompt:async()=>({})};
|
||||
await import('/extensions/promptmanager/prompt_manager.js');
|
||||
class Node {
|
||||
constructor(){this.id=1;this.type='PromptManagerPair';this.widgets=['text','negative_text','category','tags','search_text'].map(name=>({name,type:'text',value:'',options:{},inputEl:document.createElement('textarea')}));this.inputs=[...['clip','prepend_text','append_text'].map(name=>({name,type:name==='clip'?'CLIP':'STRING',link:null})),...['text','negative_text','category','tags','search_text'].map(name=>({name,type:'STRING',link:null,widget:{name}}))];this.properties={};this.size=[1100,700];}
|
||||
findInputByType(type){const index=this.inputs.findIndex(input=>input.type===type&&input.link==null);return index<0?undefined:{index,slot:this.inputs[index]};}
|
||||
constructor(){this.id=1;this.type='PromptManagerPair';this.widgets=['text','negative_text','category','tags','search_text'].map(name=>({name,type:'text',value:'',options:{},inputEl:document.createElement('textarea')}));this.inputs=['clip','prepend_text','append_text'].map(name=>({name,link:null}));this.properties={};this.size=[1100,700];}
|
||||
addWidget(type,name,value,callback,options){const widget={type,name,value,callback,options};this.widgets.push(widget);return widget;}
|
||||
addDOMWidget(name,type,element){const host=document.createElement('div');host.style.cssText='height:700px;width:1100px;max-width:100%';host.className='fixture-dom-widget';host.append(element);document.body.append(host);return {element};}
|
||||
setSize(size){this.size=size;}setDirtyCanvas(){}
|
||||
@@ -124,46 +123,6 @@ test('a superseded library response never replaces the latest pair selection',as
|
||||
await expect(page.locator('.pm-editor-result')).not.toContainText('Stale positive');
|
||||
});
|
||||
|
||||
test('paired editor offers positive and negative link sockets and keeps both typed texts (#343)',async({page})=>{
|
||||
await draft(page,'Typed positive','Typed negative');
|
||||
expect(await page.evaluate(()=>node.inputs.filter(input=>!input.widget).map(input=>input.name))).toEqual(['clip','prepend_text','append_text','text','negative_text']);
|
||||
const target=()=>page.evaluate(()=>node.findInputByType('STRING').slot.name);
|
||||
expect(await target()).toBe('text');
|
||||
const positive=page.getByRole('textbox',{name:'Positive prompt',exact:true}),negative=page.getByRole('textbox',{name:'Negative prompt',exact:true});
|
||||
const positiveNote=page.getByText('Positive prompt comes from a link. Your typed text is kept for when you disconnect it.',{exact:true});
|
||||
const negativeNote=page.getByText('Negative prompt comes from a link. Your typed text is kept for when you disconnect it.',{exact:true});
|
||||
const link=(name,value)=>page.evaluate(({name,value})=>{node.inputs.find(input=>input.name===name).link=value;node.onConnectionsChange();},{name,value});
|
||||
for(const [linked,next] of [[['text'],'negative_text'],[['negative_text'],'text'],[['text','negative_text'],'prepend_text'],[[],'text']]){
|
||||
for(const name of ['text','negative_text'])await link(name,linked.includes(name)?30:null);
|
||||
expect(await target()).toBe(next);
|
||||
await expect(positive).toBeVisible({visible:!linked.includes('text')});
|
||||
await expect(negative).toBeVisible({visible:!linked.includes('negative_text')});
|
||||
await expect(positiveNote).toBeVisible({visible:linked.includes('text')});
|
||||
await expect(negativeNote).toBeVisible({visible:linked.includes('negative_text')});
|
||||
const saved=await page.evaluate(()=>{const saved=node.serialize();for(const input of node.inputs)if(['text','negative_text'].includes(input.name))input.widget={name:input.name};node.configure(saved);return saved.widgets_values.slice(0,2);});
|
||||
expect(saved).toEqual(['Typed positive','Typed negative']);
|
||||
expect(await page.evaluate(()=>node.inputs.filter(input=>input.widget&&['text','negative_text'].includes(input.name)).length)).toBe(0);
|
||||
}
|
||||
await expect(positive).toHaveValue('Typed positive');await expect(negative).toHaveValue('Typed negative');
|
||||
});
|
||||
|
||||
test('a link dropped on either prompt box, or its note, fills that box (#343 review M1)',async({page})=>{
|
||||
const keys=await page.evaluate(()=>[...document.querySelectorAll('.pm-pair .lg-node-widget [data-slot-key]')].map(key=>key.dataset.slotKey));
|
||||
// Nodes 2.0 resolves a drop through the nearest .lg-node-widget and its slot key: positive box and note, negative box and note.
|
||||
expect(keys).toEqual(['1-in-3','1-in-3','1-in-4','1-in-4']);
|
||||
const at=(name,connecting=true)=>page.evaluate(({name,connecting})=>{
|
||||
fixtureApp.canvas={linkConnector:{state:{connectingTo:connecting?'input':undefined}},ds:{convertOffsetToCanvas:point=>point},canvas:{getBoundingClientRect:()=>({left:0,top:0})}};
|
||||
const box=[...document.querySelectorAll('.pm-pair .lg-node-widget')].find(element=>element.getClientRects().length&&element.textContent.includes(name)).getBoundingClientRect();
|
||||
return node.getInputOnPos([box.left+box.width/2,box.top+box.height/2])?.name??null;
|
||||
},{name,connecting});
|
||||
expect(await at('Negative prompt')).toBe('negative_text');
|
||||
expect(await at('Positive prompt')).toBe('text');
|
||||
expect(await at('Negative prompt',false)).toBeNull();
|
||||
await page.evaluate(()=>{node.inputs.find(input=>input.name==='negative_text').link=4;node.onConnectionsChange();});
|
||||
expect(await at('Negative prompt comes from a link')).toBe('negative_text');
|
||||
expect(await page.evaluate(()=>{fixtureApp.canvas={};return node.getInputOnPos([1,1])??null;})).toBeNull();
|
||||
});
|
||||
|
||||
test('queue hook synchronizes both editors, preserves links and forwards branch execution arguments',async({page})=>{
|
||||
await draft(page);
|
||||
const args=await page.evaluate(async()=>{
|
||||
|
||||
+2
-43
@@ -3,7 +3,6 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import socket
|
||||
@@ -23,9 +22,6 @@ from tests.browser.request_hold import register_request_holds
|
||||
from tests.browser import slow_reset_watch
|
||||
from tests.browser.workspace_activity_fixture import register_activity_commands
|
||||
from tests.browser.workspace_system_fixture import register_system_commands
|
||||
from tests.browser.workspace_integrations_fixture import (
|
||||
register_integrations_commands,
|
||||
)
|
||||
from tests.browser.workspace_concept_fixture import (
|
||||
MAX_EXTRA_IMAGES as CONCEPT_MAX_EXTRA_IMAGES,
|
||||
MAX_SEEDLESS_IMAGES as CONCEPT_MAX_SEEDLESS_IMAGES,
|
||||
@@ -72,8 +68,6 @@ class TestModel:
|
||||
# Scored names the concept comparisons answer instead of the fixed list
|
||||
# (``/__test__/tagging/suggestions``); None keeps the fixed list.
|
||||
suggestions: list | None = None
|
||||
# The engine's model table (``py/autotag.py`` MODELS), set at startup.
|
||||
models_config: dict = {}
|
||||
|
||||
def get_models_status(self) -> dict:
|
||||
return {
|
||||
@@ -81,7 +75,7 @@ class TestModel:
|
||||
"downloaded": name not in self.unavailable,
|
||||
"loaded": self.loaded == name,
|
||||
}
|
||||
for name in self.models_config
|
||||
for name in ("gguf", "hf", "wd14-vit", "wd14-convnext", "wd14-swinv2")
|
||||
}
|
||||
|
||||
def is_model_loaded(self) -> bool:
|
||||
@@ -168,20 +162,7 @@ def register_reveal_stub(routes, api, root: Path):
|
||||
return types.SimpleNamespace(reset=reset)
|
||||
|
||||
|
||||
def isolate_lora_manager_settings(root: Path) -> Path:
|
||||
"""LoraManager's settings lookup reads an empty folder under ``root`` (#355).
|
||||
|
||||
The server's temporary root is removed when it exits, and with it this
|
||||
folder: no page test reads the developer's real LoraManager settings.
|
||||
"""
|
||||
folder = root / "lora-manager-settings"
|
||||
folder.mkdir(exist_ok=True)
|
||||
os.environ["LORA_MANAGER_SETTINGS_DIR"] = str(folder)
|
||||
return folder
|
||||
|
||||
|
||||
def create_app(root: Path) -> web.Application:
|
||||
isolate_lora_manager_settings(root)
|
||||
os.environ["PROMPT_MANAGER_CONFIG_PATH"] = str(root / "config.json")
|
||||
os.chdir(root)
|
||||
app = web.Application()
|
||||
@@ -209,7 +190,6 @@ def create_app(root: Path) -> web.Application:
|
||||
from py import autotag
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
TestModel.models_config = autotag.MODELS
|
||||
model = TestModel()
|
||||
autotag.get_autotag_service = lambda: model
|
||||
api = PromptManagerAPI()
|
||||
@@ -222,7 +202,6 @@ def create_app(root: Path) -> web.Application:
|
||||
reveal_stub = register_reveal_stub(routes, api, root)
|
||||
register_activity_commands(app, routes, api)
|
||||
register_system_commands(app, routes, api, root)
|
||||
register_integrations_commands(app, routes, api, root)
|
||||
request_holds = register_request_holds(app, routes)
|
||||
slow_reset_watch.install(app, api)
|
||||
|
||||
@@ -376,12 +355,7 @@ def create_app(root: Path) -> web.Application:
|
||||
with Image.new("RGB", (64, 64), (100, 80, 120)) as image:
|
||||
image.save(root / "output" / name, pnginfo=info)
|
||||
os.utime(root / "output" / name, (time.time() - 100, time.time() - 100))
|
||||
from py.api.thumbnail_settings import ThumbnailSettings
|
||||
|
||||
# 150 px previews in the default format (#317 thumbnail settings).
|
||||
api._generate_gallery_thumbnails_sync(
|
||||
api._get_all_output_dirs(), ThumbnailSettings("webp", 80, 150)
|
||||
)
|
||||
api._generate_gallery_thumbnails_sync(api._get_all_output_dirs(), (150, 150))
|
||||
if include_video:
|
||||
# Seed after thumbnail generation to exercise a real unthumbnailed video.
|
||||
video = root / "output" / "clip.webm"
|
||||
@@ -412,8 +386,6 @@ def create_app(root: Path) -> web.Application:
|
||||
async def reset(request: web.Request) -> web.Response:
|
||||
nonlocal history_fixture
|
||||
request_holds.release()
|
||||
# No preview queued by the last test may land in the new fixture.
|
||||
await api._run_in_executor(api.close_auto_thumbnails)
|
||||
library_editor_commands.reset()
|
||||
data = await request.json() if request.can_read_body else {}
|
||||
if not isinstance(data, dict):
|
||||
@@ -682,18 +654,6 @@ def exit_when_stdin_closes(loop: asyncio.AbstractEventLoop) -> threading.Thread:
|
||||
return thread
|
||||
|
||||
|
||||
def close_owned_log_files() -> None:
|
||||
"""Close Prompt Manager's own log handlers before the temp root goes.
|
||||
|
||||
The logs live in the library's data folder (#249 S6b), which for this
|
||||
server is <root>/user/__promptmanager; Windows cannot remove a file
|
||||
this process still has open.
|
||||
"""
|
||||
for handler in logging.getLogger("prompt_manager").handlers[:]:
|
||||
if getattr(handler, "_prompt_manager_owned", False) is True:
|
||||
handler.close()
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
@@ -718,7 +678,6 @@ def main(argv: list[str] | None = None) -> None:
|
||||
finally:
|
||||
# create_app works inside the root; Windows cannot remove a process's cwd.
|
||||
os.chdir(original_cwd)
|
||||
close_owned_log_files()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,329 +0,0 @@
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from './fixtures.mjs';
|
||||
import { SIZES, THEMES, seedConcept, openConcept, openProduction, diffInventory, mutate } from './concept-inventory.mjs';
|
||||
|
||||
// #249 System S9 Task 25: System › Backup & export (app.js renderSystem 691–693), its "Restore preview" (`restore`,
|
||||
// 1472–1476) and its "Import & inspect" (`import`, 1483–1487) against the immutable concept, on the concept-shaped
|
||||
// fixture. Each concept measurement is mapped only by the recorded deviations (D103, D104, D105; the Import dialog's
|
||||
// D15 and D55), each commented where it applies; after that the difference must be empty, and a mutated label, row,
|
||||
// danger mark or token must be noticed.
|
||||
test.use({ locale: 'en-US', timezoneId: 'UTC' });
|
||||
|
||||
test.beforeEach(async ({ page, request, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
await seedConcept(request);
|
||||
await expect(await request.post('/__test__/system/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/jobs/restart')).toBeOK();
|
||||
page.conceptErrors = [];
|
||||
page.on('pageerror', error => page.conceptErrors.push(error.message));
|
||||
});
|
||||
|
||||
test.afterEach(async ({ page, request, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
await request.post('/__test__/jobs/restart');
|
||||
await request.post('/__test__/system/reset');
|
||||
expect(page.conceptErrors).toEqual([]);
|
||||
});
|
||||
|
||||
const PREVIEW_URL = '/prompt_manager/restore/preview';
|
||||
|
||||
/** Runs in the page: the section's surfaces under `mainSelector`, their text, controls, geometry and tokens. */
|
||||
function measureBackups(mainSelector) {
|
||||
const layout = document.querySelector(`${mainSelector} .settings-layout`);
|
||||
const css = (node, ...keys) => node ? Object.fromEntries(keys.map(key => [key, getComputedStyle(node)[key]])) : null;
|
||||
const name = node => [node.localName, ...[...node.classList].sort()].join('.');
|
||||
const text = node => node?.textContent.trim() ?? null;
|
||||
const button = node => ({ text: text(node), className: node.className,
|
||||
...css(node, 'backgroundColor', 'color', 'fontSize', 'paddingTop', 'paddingLeft', 'borderTopColor') });
|
||||
const base = layout.getBoundingClientRect();
|
||||
const box = node => {
|
||||
const rect = node.getBoundingClientRect();
|
||||
return { left: Math.round(rect.left - base.left), top: Math.round(rect.top - base.top), width: Math.round(rect.width) };
|
||||
};
|
||||
const control = row => {
|
||||
const select = row.querySelector(':scope > select');
|
||||
if (select) return { kind: 'select', id: select.id, options: [...select.options].map(option => [option.textContent, option.disabled]),
|
||||
selected: select.selectedOptions[0]?.textContent ?? null, disabled: select.disabled };
|
||||
const toggle = row.querySelector(':scope > * input[type="checkbox"]');
|
||||
return toggle ? { kind: 'switch', id: toggle.id, wrapper: name(toggle.parentElement), checked: toggle.checked,
|
||||
disabled: toggle.disabled, ...css(toggle.parentElement, 'width', 'height') } : null;
|
||||
};
|
||||
const columns = layout.querySelector(':scope > .two-col');
|
||||
const surfaces = [...columns.querySelectorAll(':scope > section.surface')];
|
||||
const policy = layout.querySelector(':scope > section.surface');
|
||||
const chip = policy.querySelector(':scope > h2 > .chip');
|
||||
const row = policy.querySelector('.setting-row');
|
||||
const status = surfaces[0].querySelector(':scope > p.small');
|
||||
const warning = policy.querySelector(':scope > .warning-note');
|
||||
return {
|
||||
layout: [...layout.children].map(name),
|
||||
columns: { ...css(columns, 'display', 'gridTemplateColumns', 'columnGap'), box: box(columns) },
|
||||
surfaces: surfaces.map(surface => ({ tree: [...surface.children].map(name), heading: text(surface.querySelector(':scope > h2')),
|
||||
intro: text(surface.querySelector(':scope > p:not(.small)')), box: box(surface),
|
||||
actions: [...surface.querySelectorAll(':scope > .row-actions')].map(group => ({ tree: [...group.children].map(name),
|
||||
...css(group, 'display', 'columnGap', 'flexWrap', 'marginTop'),
|
||||
buttons: [...group.children].map(button) })) })),
|
||||
status: { text: text(status), ...css(status, 'fontSize', 'color', 'marginTop') },
|
||||
surface: css(surfaces[0], 'paddingTop', 'paddingLeft', 'backgroundColor', 'borderTopWidth', 'borderTopColor', 'borderTopLeftRadius'),
|
||||
h2: css(surfaces[0].querySelector('h2'), 'fontSize', 'fontWeight', 'color', 'marginBottom'),
|
||||
intro: css(surfaces[0].querySelector(':scope > p'), 'fontSize', 'color', 'lineHeight'),
|
||||
// Its place below the columns (the concept's `.surface + .surface` 18px), whatever the columns' height.
|
||||
policy: { tree: [...policy.children].map(name), heading: text(policy.querySelector(':scope > h2')),
|
||||
box: { left: box(policy).left, width: box(policy).width,
|
||||
gap: Math.round(policy.getBoundingClientRect().top - columns.getBoundingClientRect().bottom) },
|
||||
chip: { text: text(chip), className: chip.className, ...css(chip, 'color', 'fontSize', 'borderTopColor', 'backgroundColor') },
|
||||
rows: [...policy.querySelectorAll(':scope > .setting-row')].map(item => ({ tree: [...item.children].map(name),
|
||||
text: [...item.firstElementChild.children].map(node => [name(node), text(node), node.localName === 'label' ? node.htmlFor : null]),
|
||||
control: control(item) })),
|
||||
warning: { text: text(warning), ...css(warning, 'fontSize', 'color', 'marginTop', 'marginBottom') } },
|
||||
row: css(row, 'display', 'alignItems', 'justifyContent', 'columnGap', 'paddingTop', 'paddingBottom', 'borderTopWidth', 'borderTopColor',
|
||||
'marginTop'),
|
||||
label: css(row.querySelector('label'), 'fontSize', 'fontWeight', 'color'),
|
||||
description: css(row.querySelector('p'), 'fontSize', 'color', 'marginTop', 'maxWidth'),
|
||||
select: css(policy.querySelector('.setting-row > select'), 'fontSize', 'minWidth', 'maxWidth', 'backgroundColor', 'color',
|
||||
'borderTopColor'),
|
||||
pageScroll: document.documentElement.scrollWidth > document.documentElement.clientWidth,
|
||||
};
|
||||
}
|
||||
|
||||
const SCHEDULE_REASON = 'Daily and Weekly are not available yet: back up with Back up whenever you choose.';
|
||||
const SAFETY_DESCRIPTION = 'Keep a verified copy of the database on disk before maintenance, cleanups and LoRA imports change ' +
|
||||
'records. A restore always keeps one.';
|
||||
|
||||
/**
|
||||
* D103: "Back up" (the concept's "Simulate backup"), the status line from /backup/status, "Export JSON" (its "Export
|
||||
* sample JSON") and the two secondary exports in a second `.row-actions.spaced` with the concept's plain buttons.
|
||||
* D104: the toggle's line says what the policy does, the schedule's reason replaces the prototype's proposal line,
|
||||
* Daily and Weekly are disabled, and the warning note is truthful. The fixture has no backup yet.
|
||||
*/
|
||||
function mapBackupsConcept(reference) {
|
||||
const [database, portable] = reference.surfaces;
|
||||
const [backup, restore] = database.actions[0].buttons;
|
||||
const [exportButton, importButton] = portable.actions[0].buttons;
|
||||
const secondary = ['Export all prompts JSON', 'Export saved pairs JSON'].map(label => ({ ...exportButton, text: label }));
|
||||
return {
|
||||
...reference,
|
||||
surfaces: [
|
||||
{ ...database, actions: [{ ...database.actions[0], buttons: [{ ...backup, text: 'Back up' }, restore] }] },
|
||||
{ ...portable, tree: [...portable.tree, 'div.row-actions.spaced'],
|
||||
actions: [{ ...portable.actions[0], buttons: [{ ...exportButton, text: 'Export JSON' }, importButton] },
|
||||
{ ...portable.actions[0], tree: ['button', 'button'], buttons: secondary }] },
|
||||
],
|
||||
status: { ...reference.status, text: 'No backup recorded.' },
|
||||
policy: { ...reference.policy,
|
||||
rows: reference.policy.rows.map(item => (item.control.kind !== 'select'
|
||||
// D104 (M2): what the safety-backup policy really does.
|
||||
? { ...item, text: item.text.map(([node, value, target]) => [node, node === 'p' ? SAFETY_DESCRIPTION : value, target]) }
|
||||
: { ...item,
|
||||
text: item.text.map(([node, value, target]) => [node, node === 'p' ? SCHEDULE_REASON : value, target]),
|
||||
control: { ...item.control, options: item.control.options.map(([label]) => [label, label !== 'Manual only']) } })),
|
||||
warning: { ...reference.policy.warning, text: 'Restores are staged, migrated and validated before replacing live data.' } },
|
||||
};
|
||||
}
|
||||
|
||||
async function conceptBackups(page, theme) {
|
||||
await openConcept(page, 'system', theme);
|
||||
await page.locator('#main .section-nav').getByRole('button', { name: 'Backup & export', exact: true }).click();
|
||||
await expect(page.locator('#main .settings-layout > .two-col > section.surface')).toHaveCount(2);
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureBackups, '#main');
|
||||
}
|
||||
|
||||
async function productionBackups(page, theme, status = 'No backup recorded.') {
|
||||
await openProduction(page, 'system/backups', theme);
|
||||
await expect(page.locator('#workspace-main .settings-layout > .two-col > section.surface')).toHaveCount(2);
|
||||
// The status line and the policy row are their own reads: compare once both answered.
|
||||
await expect(page.locator('#workspace-main #backup-status')).toHaveText(status);
|
||||
await expect(page.locator('#workspace-main #backup-before')).toBeEnabled();
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureBackups, '#workspace-main');
|
||||
}
|
||||
|
||||
for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
test(`backup & export matches the concept ${theme} ${width}x${height}`, async ({ page }, testInfo) => {
|
||||
await page.setViewportSize({ width, height });
|
||||
const concept = await conceptBackups(page, theme);
|
||||
await page.screenshot({ path: testInfo.outputPath(`concept-backups-${theme}-${width}.png`), fullPage: true });
|
||||
const production = await productionBackups(page, theme);
|
||||
await page.screenshot({ path: testInfo.outputPath(`production-backups-${theme}-${width}.png`), fullPage: true });
|
||||
expect(diffInventory(mapBackupsConcept(concept), production)).toEqual([]);
|
||||
expect(production.pageScroll).toBe(false);
|
||||
});
|
||||
}
|
||||
|
||||
/** Runs in the page: the open dialog's title, body tree, grid, note, backup box and actions with the concept's tokens. */
|
||||
function measureDialog() {
|
||||
const dialog = document.querySelector('dialog#modal[open]');
|
||||
const css = (node, ...keys) => node ? Object.fromEntries(keys.map(key => [key, getComputedStyle(node)[key]])) : null;
|
||||
const name = node => [node.localName, ...[...node.classList].sort()].join('.');
|
||||
const text = node => node?.textContent.trim() ?? null;
|
||||
const body = dialog.querySelector('#modal-body');
|
||||
const grid = body.querySelector(':scope > dl.meta-grid');
|
||||
const note = body.querySelector(':scope > p.small');
|
||||
const actions = body.querySelector(':scope > .dialog-actions');
|
||||
const list = body.querySelector(':scope > .operation-list');
|
||||
return {
|
||||
title: text(dialog.querySelector('#modal-title')),
|
||||
tree: [...body.children].map(name),
|
||||
grid: grid && { rows: [...grid.children].map(row => [text(row.querySelector('dt')), text(row.querySelector('dd'))]),
|
||||
...css(grid, 'display', 'gridTemplateColumns', 'columnGap', 'rowGap'),
|
||||
dt: css(grid.querySelector('dt'), 'fontSize', 'color', 'textTransform'), dd: css(grid.querySelector('dd'), 'fontSize', 'color') },
|
||||
note: note && { text: text(note), ...css(note, 'fontSize', 'color') },
|
||||
backup: body.querySelector(':scope > .notice') && { tree: [...body.querySelector(':scope > .notice').children].map(name),
|
||||
text: [...body.querySelectorAll(':scope > .notice > div:first-child > *')].map(node => [node.localName, text(node)]),
|
||||
buttons: [...body.querySelectorAll(':scope > .notice button')].map(text) },
|
||||
list: list && { ...css(list, 'display', 'rowGap'),
|
||||
rows: [...list.children].map(row => ({ tree: [...row.children].map(name), title: text(row.querySelector('h3')),
|
||||
description: text(row.querySelector('p')), buttons: [...row.querySelectorAll('button')].map(text),
|
||||
...css(row, 'display', 'alignItems', 'paddingTop', 'borderTopColor') })) },
|
||||
actions: actions && { tree: [...actions.children].map(name), ...css(actions, 'display', 'justifyContent', 'columnGap', 'marginTop'),
|
||||
buttons: [...actions.children].map(item => ({ text: text(item), className: item.className,
|
||||
...css(item, 'backgroundColor', 'color', 'fontSize', 'paddingTop', 'paddingLeft', 'borderTopColor') })) },
|
||||
};
|
||||
}
|
||||
|
||||
const RESTORE_BACKUP_BOX = Object.freeze({ tree: ['div', 'div.row-actions'], text: [['h3', 'Back up the database first (recommended)'],
|
||||
['p', 'Downloads a copy of the prompts database before the restore replaces it.']], buttons: ['Back up now'] });
|
||||
|
||||
/**
|
||||
* D105: the server's review in the concept's "Restore preview": no sample note; the four rows with the review's values,
|
||||
* then "In the backup" and "In the library now"; the hold line in the concept's dialog `p.small` (its Import dialog's
|
||||
* intro, `small`); the D55 backup box before the actions; "Restore" (the concept's "Simulate restore", still danger).
|
||||
*/
|
||||
function mapRestoreConcept(reference, small, rows) {
|
||||
const [cancel, restore] = reference.actions.buttons;
|
||||
return {
|
||||
...reference,
|
||||
tree: ['dl.meta-grid', 'p.small', 'div.notice', 'div.dialog-actions'],
|
||||
grid: { ...reference.grid, rows },
|
||||
note: { ...small, text: 'Restore replaces the whole library with this backup and runs alone in Activity. ' +
|
||||
'Nothing changes until you restore. This review holds for 15 minutes.' },
|
||||
backup: RESTORE_BACKUP_BOX,
|
||||
actions: { ...reference.actions, buttons: [cancel, { ...restore, text: 'Restore' }] },
|
||||
};
|
||||
}
|
||||
|
||||
/** D15 and D55, as on Images: Rescan for "Scan an output folder", then "Scan for metadata"; no portable recipe import. */
|
||||
function mapImportConcept(reference) {
|
||||
return { ...reference,
|
||||
note: { ...reference.note, text: 'Image inspection and output scanning have different effects. Choose the workflow you want.' },
|
||||
list: { ...reference.list, rows: reference.list.rows.filter(row => row.title !== 'Import portable recipes').flatMap(row =>
|
||||
row.title === 'Scan an output folder'
|
||||
? [{ ...row, description: 'Find files added or removed in the output folders.', buttons: ['Rescan'] },
|
||||
{ ...row, title: 'Scan for metadata' }] : [row]) } };
|
||||
}
|
||||
|
||||
async function conceptDialog(page, theme, action, title) {
|
||||
await conceptBackups(page, theme);
|
||||
await page.locator(`#main .settings-layout [data-action="${action}"]`).click();
|
||||
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText(title);
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureDialog);
|
||||
}
|
||||
|
||||
/** Production's restore review of a fresh backup of this server's library; the review's own answer too. */
|
||||
async function productionRestore(page, request, theme) {
|
||||
const copy = await request.get('/prompt_manager/backup');
|
||||
await expect(copy).toBeOK();
|
||||
// That copy is this session's download now.
|
||||
await productionBackups(page, theme, /^Downloaded backup · .+ · verified\.$/);
|
||||
const chooser = page.waitForEvent('filechooser');
|
||||
const answer = page.waitForResponse(response => response.url().endsWith(PREVIEW_URL));
|
||||
await page.locator('#workspace-main').getByRole('button', { name: 'Review restore…', exact: true }).click();
|
||||
await (await chooser).setFiles({ name: 'prompts_backup_20261025_093000.db', mimeType: 'application/octet-stream',
|
||||
buffer: await copy.body() });
|
||||
const preview = await (await answer).json();
|
||||
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Restore preview');
|
||||
await page.mouse.move(0, 0);
|
||||
return { preview, production: await page.evaluate(measureDialog) };
|
||||
}
|
||||
|
||||
const counts = summary => `${summary.prompts} ${summary.prompts === 1 ? 'prompt' : 'prompts'} · ${summary.image_links} ` +
|
||||
`${summary.image_links === 1 ? 'image link' : 'image links'} · ${summary.pairs} ${summary.pairs === 1 ? 'saved pair' : 'saved pairs'}` +
|
||||
` · ${summary.runs} ${summary.runs === 1 ? 'run' : 'runs'}`;
|
||||
|
||||
function restoreRows(preview, snapshot) {
|
||||
return [['Snapshot', snapshot], ['Integrity', 'Verified'], ['Migration', 'Current schema'],
|
||||
['Safety backup', `Created before replacing (${preview.safety_backup.destination_label})`],
|
||||
['In the backup', counts(preview.summary)], ['In the library now', counts(preview.live)]];
|
||||
}
|
||||
|
||||
for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
test(`restore preview and import dialogs match the concept ${theme} ${width}x${height}`, async ({ page, request }, testInfo) => {
|
||||
await page.setViewportSize({ width, height });
|
||||
// The concept's own dialog `p.small` (its Import dialog's intro) is the reference for the review's note.
|
||||
const conceptImport = await conceptDialog(page, theme, 'import', 'Import & inspect');
|
||||
await page.screenshot({ path: testInfo.outputPath(`concept-import-${theme}-${width}.png`) });
|
||||
const conceptRestore = await conceptDialog(page, theme, 'restore', 'Restore preview');
|
||||
expect(conceptRestore.tree).toEqual(['p.demo-note', 'dl.meta-grid', 'div.dialog-actions']);
|
||||
await page.screenshot({ path: testInfo.outputPath(`concept-restore-${theme}-${width}.png`) });
|
||||
const { preview, production } = await productionRestore(page, request, theme);
|
||||
await page.screenshot({ path: testInfo.outputPath(`production-restore-${theme}-${width}.png`) });
|
||||
const snapshot = production.grid.rows[0][1];
|
||||
expect(snapshot).toMatch(/^prompts_backup_20261025_093000.db · (\d+ KB|\d+\.\d MB) · .+$/);
|
||||
const { text: _intro, ...small } = conceptImport.note;
|
||||
expect(diffInventory(mapRestoreConcept(conceptRestore, small, restoreRows(preview, snapshot)), production)).toEqual([]);
|
||||
await page.keyboard.press('Escape');
|
||||
await expect(page.locator('dialog#modal[open]')).toHaveCount(0);
|
||||
await page.locator('#workspace-main').getByRole('button', { name: 'Import / inspect file', exact: true }).click();
|
||||
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Import & inspect');
|
||||
await page.mouse.move(0, 0);
|
||||
await page.screenshot({ path: testInfo.outputPath(`production-import-${theme}-${width}.png`) });
|
||||
expect(diffInventory(mapImportConcept(conceptImport), await page.evaluate(measureDialog))).toEqual([]);
|
||||
});
|
||||
}
|
||||
|
||||
test('backup & export comparison detects mutations', async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
const concept = await conceptBackups(page, 'dark');
|
||||
const expected = mapBackupsConcept(concept);
|
||||
await productionBackups(page, 'dark');
|
||||
const measure = () => page.evaluate(measureBackups, '#workspace-main');
|
||||
const paths = async () => diffInventory(expected, await measure()).map(item => item.path);
|
||||
expect(await paths()).toEqual([]);
|
||||
// Unmapped, the section differs exactly on the D103 and D104 paths (the added row's own paths, then the rest).
|
||||
const unmapped = diffInventory(concept, await measure()).map(item => item.path).sort();
|
||||
const added = 'surfaces.1.actions.1.';
|
||||
expect(unmapped.some(path => path.startsWith(added))).toBe(true);
|
||||
expect(unmapped.filter(path => !path.startsWith(added))).toEqual([
|
||||
'policy.rows.0.text.1.1', 'policy.rows.1.control.options.1.1', 'policy.rows.1.control.options.2.1', 'policy.rows.1.text.1.1',
|
||||
'policy.warning.text',
|
||||
'status.text', 'surfaces.0.actions.0.buttons.0.text', 'surfaces.1.actions.0.buttons.0.text', 'surfaces.1.actions.length',
|
||||
'surfaces.1.tree.3', 'surfaces.1.tree.length']);
|
||||
// A renamed label, a removed row, a lost danger mark and a swapped token are noticed.
|
||||
await page.evaluate(() => { document.querySelector('#workspace-main .setting-row label').textContent = 'Safety backup'; });
|
||||
expect(await paths()).toEqual(['policy.rows.0.text.0.1']);
|
||||
await page.evaluate(() => { document.querySelector('#workspace-main .setting-row label').textContent = 'Create a safety backup before maintenance'; });
|
||||
const undoRow = await mutate(page, { remove: ['#workspace-main .settings-layout > section.surface > .setting-row:last-of-type'] });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['policy.rows.length', 'policy.tree.length']));
|
||||
await undoRow();
|
||||
await page.evaluate(() => { document.querySelector('#workspace-main button.danger').classList.remove('danger'); });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['surfaces.0.actions.0.buttons.1.className']));
|
||||
await page.evaluate(() => { document.querySelector('#workspace-main .two-col .row-actions > button:nth-child(2)').classList.add('danger'); });
|
||||
const undoToken = await mutate(page, { css: '#workspace-main .settings-layout { --line: rgb(255, 0, 0); }' });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['row.borderTopColor']));
|
||||
await undoToken();
|
||||
const undoChip = await mutate(page, { css: '#workspace-main .chip.blue { color: rgb(255, 0, 0); }' });
|
||||
expect(await paths()).toEqual(['policy.chip.color']);
|
||||
await undoChip();
|
||||
expect(await paths()).toEqual([]);
|
||||
|
||||
// The restore review: a renamed term, a lost danger mark and the grid's token are noticed.
|
||||
const conceptImport = await conceptDialog(page, 'dark', 'import', 'Import & inspect');
|
||||
const conceptRestore = await conceptDialog(page, 'dark', 'restore', 'Restore preview');
|
||||
const { preview, production } = await productionRestore(page, request, 'dark');
|
||||
const { text: _intro, ...small } = conceptImport.note;
|
||||
const expectedRestore = mapRestoreConcept(conceptRestore, small, restoreRows(preview, production.grid.rows[0][1]));
|
||||
const restorePaths = async () => diffInventory(expectedRestore, await page.evaluate(measureDialog)).map(item => item.path);
|
||||
expect(await restorePaths()).toEqual([]);
|
||||
await page.evaluate(() => { document.querySelector('dialog#modal[open] dl.meta-grid dt').textContent = 'Backup'; });
|
||||
expect(await restorePaths()).toEqual(['grid.rows.0.0']);
|
||||
await page.evaluate(() => { document.querySelector('dialog#modal[open] dl.meta-grid dt').textContent = 'Snapshot'; });
|
||||
await page.evaluate(() => { document.querySelector('dialog#modal[open] .dialog-actions .danger').classList.remove('danger'); });
|
||||
expect(await restorePaths()).toEqual(expect.arrayContaining(['actions.buttons.1.className']));
|
||||
await page.evaluate(() => { document.querySelector('dialog#modal[open] .dialog-actions button:last-child').classList.add('danger'); });
|
||||
const undoGrid = await mutate(page, { css: 'dialog#modal dl.meta-grid { column-gap: 3px; }' });
|
||||
expect(await restorePaths()).toEqual(expect.arrayContaining(['grid.columnGap']));
|
||||
await undoGrid();
|
||||
expect(await restorePaths()).toEqual([]);
|
||||
});
|
||||
@@ -172,8 +172,8 @@ function mapOverviewConcept(reference, overview) {
|
||||
// D93: the real pending count; Review metadata opens Needs review (the concept's go-review).
|
||||
'3 ambiguous image links': row => ({ ...row,
|
||||
title: `${number(overview.needs_review.pending)} ambiguous image link${overview.needs_review.pending === 1 ? '' : 's'}` }),
|
||||
// D93 (S9): what GET /backup/status reports; the fixture has no backup yet.
|
||||
'Backup status': row => ({ ...row, description: 'No backup recorded.' }),
|
||||
// D93: no backup has been recorded (verified backups arrive with S8).
|
||||
'Backup status': row => ({ ...row, description: 'No retained backup record.' }),
|
||||
// D93: there is no paging, so the shared settings are grid and density.
|
||||
'One display configuration': row => ({ ...row, description: 'Grid and density settings are shared across views.' }),
|
||||
// D93: the real sources, not samples.
|
||||
@@ -181,17 +181,6 @@ function mapOverviewConcept(reference, overview) {
|
||||
// D96: the dialog shows runs and ratings by category and library totals, not prompt, tag or date usage.
|
||||
Statistics: row => ({ ...row, description: 'Runs and ratings by category, with library totals.' }),
|
||||
};
|
||||
const surfaces = reference.columns.surfaces.map(surface => ({ ...surface,
|
||||
rows: surface.rows.map(row => (rows[row.title] ?? (value => value))(row)) }));
|
||||
// D129: "Library location" after Backup status, in the backup row's own markup; the fixture's database is configured
|
||||
// explicitly (config.json), so it does not move and its action is Details.
|
||||
const [health] = surfaces;
|
||||
const backup = health.rows.find(row => row.title === 'Backup status');
|
||||
const location = { ...backup, title: 'Library location',
|
||||
description: `${overview.location.path} · The database path set in config.json.`,
|
||||
action: { ...backup.action, text: 'Details' } };
|
||||
surfaces[0] = { ...health, tree: [...health.tree, 'div.operation'],
|
||||
rows: [...health.rows.slice(0, 2), location, ...health.rows.slice(2)] };
|
||||
return {
|
||||
...reference,
|
||||
// D92: no "Sample configuration" chip.
|
||||
@@ -199,7 +188,8 @@ function mapOverviewConcept(reference, overview) {
|
||||
// D122: the concept's literal label is its prototype action name ("system-tab sections").
|
||||
nav: { ...reference.nav, label: 'System sections' },
|
||||
metrics: { ...reference.metrics, items: reference.metrics.items.map(item => ({ ...item, value: relabel[item.label] })) },
|
||||
columns: { ...reference.columns, surfaces },
|
||||
columns: { ...reference.columns, surfaces: reference.columns.surfaces.map(surface => ({ ...surface,
|
||||
rows: surface.rows.map(row => (rows[row.title] ?? (value => value))(row)) })) },
|
||||
save: { ...reference.save,
|
||||
// D92: truthful status copy, and "Save preferences" is aria-disabled (dimmed) while nothing is unsaved.
|
||||
status: { ...reference.save.status, text: 'No unsaved changes.' },
|
||||
@@ -219,9 +209,6 @@ async function productionOverview(page, theme) {
|
||||
await expect(page.locator('#workspace-main .summary-metrics')).toBeVisible();
|
||||
// The thumbnail row is its own read (I2): compare once it answered.
|
||||
await expect(page.locator('#workspace-main .operation h3').filter({ hasText: 'thumbnails' })).toHaveText('No missing thumbnails');
|
||||
// So is the backup row (S9).
|
||||
await expect(page.locator('#workspace-main .operation').filter({ has: page.locator('h3', { hasText: 'Backup status' }) })
|
||||
.locator('p')).toHaveText('No backup recorded.');
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureSystem, '#workspace-main');
|
||||
}
|
||||
@@ -245,21 +232,13 @@ for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
|
||||
// The recorded differences, by path: D93 the review row (its real count), the backup row, the
|
||||
// preferences row and the sources row, and metric 4; D96 the statistics row; D92 the chip, the status and the dimmed
|
||||
// Save; D122 the nav label; D129 the inserted "Library location" row, which shifts the thumbnail row to the fourth place
|
||||
// (row 2 takes its title, text and action; row 3 and the surface's fifth child are new). (The fixture's first three
|
||||
// metrics are the concept's own counts.)
|
||||
const D129_ROW_ACTION = ['backgroundColor', 'borderTopColor', 'borderTopLeftRadius', 'borderTopWidth', 'className', 'color',
|
||||
'cursor', 'fontSize', 'opacity', 'paddingLeft', 'paddingTop', 'text'].map(key => `columns.surfaces.0.rows.3.action.${key}`);
|
||||
// Save; D122 the nav label. (The fixture's first three metrics are the concept's own counts.)
|
||||
const UNMAPPED_PATHS = [
|
||||
'columns.surfaces.0.rows.0.title', 'columns.surfaces.0.rows.1.description',
|
||||
'columns.surfaces.0.rows.2.action.text', 'columns.surfaces.0.rows.2.description', 'columns.surfaces.0.rows.2.title',
|
||||
...D129_ROW_ACTION, 'columns.surfaces.0.rows.3.description', 'columns.surfaces.0.rows.3.title',
|
||||
'columns.surfaces.0.rows.3.tree.0', 'columns.surfaces.0.rows.3.tree.1', 'columns.surfaces.0.rows.3.tree.length',
|
||||
'columns.surfaces.0.rows.length', 'columns.surfaces.0.tree.4', 'columns.surfaces.0.tree.length',
|
||||
'columns.surfaces.1.rows.0.description', 'columns.surfaces.1.rows.1.title', 'columns.surfaces.1.rows.2.description',
|
||||
'heading.actions.0', 'heading.actions.length', 'metrics.items.3.value', 'nav.label',
|
||||
'save.button.cursor', 'save.button.opacity', 'save.status.text',
|
||||
].sort();
|
||||
];
|
||||
|
||||
test('system comparison detects mutations', async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
@@ -485,8 +464,6 @@ function mapAppearanceConcept(reference, settings) {
|
||||
control: { ...toggleRow.control, id, checked } }));
|
||||
return {
|
||||
...reference,
|
||||
// #317: the server's thumbnail settings follow in their own surface ("Thumbnails", README: retained in Appearance).
|
||||
layout: [reference.layout[0], 'section.surface', ...reference.layout.slice(1)],
|
||||
// D97: two added toggles and the "Reset display defaults" operation close the surface.
|
||||
surface: { ...reference.surface, tree: [...reference.surface.tree, 'div.setting-row', 'div.setting-row', 'div.operation'] },
|
||||
rows: [...reference.rows.map(row => row.control?.id === 'timeout'
|
||||
@@ -535,8 +512,7 @@ test('appearance comparison detects mutations', async ({ page, request }) => {
|
||||
// Unmapped, the concept differs only where D97 and the real timeout are recorded.
|
||||
expect(diffInventory(concept, production).map(item => item.path.split('.').slice(0, 2).join('.'))
|
||||
.filter((path, index, all) => all.indexOf(path) === index).sort())
|
||||
.toEqual(['layout.1', 'layout.2', 'layout.length', 'operations.0', 'operations.length', 'rows.11', 'rows.12', 'rows.13',
|
||||
'rows.length', 'surface.tree'].sort());
|
||||
.toEqual(['operations.0', 'operations.length', 'rows.11', 'rows.12', 'rows.13', 'rows.length', 'surface.tree'].sort());
|
||||
const measure = () => page.evaluate(measureAppearance, '#workspace-main');
|
||||
const paths = async () => diffInventory(expected, await measure()).map(item => item.path);
|
||||
// A removed row.
|
||||
@@ -848,16 +824,13 @@ function measureThumbnailReview() {
|
||||
|
||||
/**
|
||||
* D99: the planned counts line (`p.small[role=status]`) before the actions, from the server's plan for the chosen
|
||||
* options, and the primary "Start" (the concept's "Simulate thumbnail job"). D131 (#317): the Quality and Preview size
|
||||
* limit fields give way to a second `p.small` naming the server's thumbnail settings, and Scope gains "Convert
|
||||
* outdated (removes older formats)".
|
||||
* options, and the primary "Start" (the concept's "Simulate thumbnail job").
|
||||
*/
|
||||
function mapThumbnailConcept(reference, plan) {
|
||||
const [cancel, start] = reference.actions.buttons;
|
||||
const [scope] = reference.fields;
|
||||
const tree = ['div.field', 'p.small', 'p.small', 'div.dialog-actions'];
|
||||
const fields = [{ ...scope, control: { ...scope.control, options: [...scope.control.options, 'Convert outdated (removes older formats)'] } }];
|
||||
return { ...reference, tree, fields,
|
||||
const tree = [...reference.tree];
|
||||
tree.splice(tree.indexOf('div.dialog-actions'), 0, 'p.small');
|
||||
return { ...reference, tree,
|
||||
plan: { ...reference.plan, role: 'status',
|
||||
text: `${number(plan.targets)} ${plan.targets === 1 ? 'preview' : 'previews'} to generate · ${number(plan.fresh)} already up to date.` },
|
||||
actions: { ...reference.actions, buttons: [cancel, { ...start, text: 'Start' }] } };
|
||||
@@ -883,12 +856,12 @@ async function productionThumbnailReview(page, theme, entry) {
|
||||
await page.locator('#workspace-main').getByRole('button', { name: 'Review thumbnails', exact: true }).click();
|
||||
}
|
||||
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Thumbnail maintenance');
|
||||
await expect(page.locator('dialog#modal[open] #modal-body > p.small').first()).toHaveText(/ already up to date\.$/);
|
||||
await expect(page.locator('dialog#modal[open] #modal-body > p.small')).toHaveText(/ already up to date\.$/);
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureThumbnailReview);
|
||||
}
|
||||
|
||||
const thumbnailPlan = request => serverJson(request, '/prompt_manager/images/thumbnails/plan?mode=missing');
|
||||
const thumbnailPlan = request => serverJson(request, '/prompt_manager/images/thumbnails/plan?mode=missing&size=512&quality=medium');
|
||||
|
||||
for (const entry of ['sources', 'overview']) for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
test(`thumbnail dialog matches the concept from ${entry} ${theme} ${width}x${height}`, async ({ page, request }, testInfo) => {
|
||||
@@ -911,106 +884,22 @@ test('thumbnail dialog comparison detects mutations', async ({ page, request })
|
||||
const expected = mapThumbnailConcept(concept, plan);
|
||||
const production = await productionThumbnailReview(page, 'dark', 'sources');
|
||||
expect(diffInventory(expected, production)).toEqual([]);
|
||||
// Unmapped, the concept differs only where D99 and D131 are recorded.
|
||||
expect(diffInventory(concept, production).map(item => item.path.split('.').slice(0, 2).join('.'))
|
||||
.filter((path, index, all) => all.indexOf(path) === index).sort())
|
||||
.toEqual(['actions.buttons', 'fields.0', 'fields.1', 'fields.2', 'fields.length', 'plan.role', 'plan.text', 'tree.1', 'tree.2']
|
||||
.sort());
|
||||
// Unmapped, the concept differs only where D99 is recorded.
|
||||
expect(diffInventory(concept, production).map(item => item.path).sort())
|
||||
.toEqual(['actions.buttons.1.text', 'plan.role', 'plan.text', 'tree.3', 'tree.4', 'tree.length']);
|
||||
const paths = async () => diffInventory(expected, await page.evaluate(measureThumbnailReview)).map(item => item.path);
|
||||
// A changed option label, a swapped option order, a swapped action order and a changed style.
|
||||
await page.evaluate(() => { document.querySelector('#thumb-mode option').textContent = 'Missing only'; });
|
||||
expect(await paths()).toEqual(['fields.0.control.options.0']);
|
||||
await page.evaluate(() => { document.querySelector('#thumb-mode option').textContent = 'Missing thumbnails only'; });
|
||||
await page.evaluate(() => { document.querySelector('#thumb-quality option').textContent = 'Medium'; });
|
||||
expect(await paths()).toEqual(['fields.1.control.options.0']);
|
||||
await page.evaluate(() => { document.querySelector('#thumb-quality option').textContent = 'Medium · optimized'; });
|
||||
await page.evaluate(() => { const select = document.querySelector('#thumb-mode'); select.prepend(select.lastElementChild); });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['fields.0.control.options.0', 'fields.0.control.options.1']));
|
||||
await page.evaluate(() => { const select = document.querySelector('#thumb-mode'); select.append(select.firstElementChild); select.selectedIndex = 0; });
|
||||
await page.evaluate(() => { const select = document.querySelector('#thumb-mode'); select.prepend(select.lastElementChild); select.selectedIndex = 0; });
|
||||
await page.evaluate(() => { const actions = document.querySelector('dialog#modal[open] .dialog-actions'); actions.prepend(actions.lastElementChild); });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['actions.buttons.0.text', 'actions.buttons.1.text']));
|
||||
await page.evaluate(() => { const actions = document.querySelector('dialog#modal[open] .dialog-actions'); actions.prepend(actions.lastElementChild); });
|
||||
const undoStyle = await mutate(page, { css: 'dialog#modal #thumb-mode { border-color: rgb(255, 0, 0); } dialog#modal #modal-body > p.small { color: rgb(255, 0, 0); }' });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['fields.0.control.borderTopColor', 'plan.color']));
|
||||
const undoStyle = await mutate(page, { css: 'dialog#modal #thumb-limit { border-color: rgb(255, 0, 0); } dialog#modal #modal-body > p.small { color: rgb(255, 0, 0); }' });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['fields.2.control.borderTopColor', 'plan.color']));
|
||||
await undoStyle();
|
||||
expect(diffInventory(expected, await page.evaluate(measureThumbnailReview))).toEqual([]);
|
||||
});
|
||||
|
||||
/**
|
||||
* Runs in the page: one `.notice` (the concept's markup: `div > h2 + p`, then `.row-actions` buttons) with its tokens and
|
||||
* geometry, and its left edge against `edgeSelector` (the content it lines up with).
|
||||
*/
|
||||
function measureNoticeOf([selector, edgeSelector]) {
|
||||
const notice = document.querySelector(selector);
|
||||
const css = (node, ...keys) => node ? Object.fromEntries(keys.map(key => [key, getComputedStyle(node)[key]])) : null;
|
||||
const name = node => [node.localName, ...[...node.classList].sort()].join('.');
|
||||
const text = node => node?.textContent.trim() ?? null;
|
||||
const button = node => ({ text: text(node), className: node.className,
|
||||
...css(node, 'backgroundColor', 'color', 'fontSize', 'paddingTop', 'paddingLeft', 'borderTopColor', 'borderTopLeftRadius') });
|
||||
const edge = document.querySelector(edgeSelector);
|
||||
return {
|
||||
tree: [...notice.children].map(name),
|
||||
text: [...notice.querySelectorAll(':scope > div:not(.row-actions) > *')].map(node => [node.localName, text(node)]),
|
||||
actions: [...notice.querySelectorAll(':scope > .row-actions > button')].map(button),
|
||||
frame: css(notice, 'display', 'alignItems', 'justifyContent', 'columnGap', 'borderTopWidth', 'borderTopColor',
|
||||
'borderTopLeftRadius', 'backgroundColor', 'paddingTop', 'paddingLeft', 'marginBottom', 'flexDirection'),
|
||||
h2: css(notice.querySelector('h2'), 'fontSize', 'fontWeight', 'color'),
|
||||
p: css(notice.querySelector('p'), 'fontSize', 'color', 'marginTop'),
|
||||
rowActions: css(notice.querySelector(':scope > .row-actions'), 'display', 'columnGap', 'flexShrink'),
|
||||
alignedWith: Math.round(notice.getBoundingClientRect().left - edge.getBoundingClientRect().left),
|
||||
pageScroll: document.documentElement.scrollWidth > document.documentElement.clientWidth,
|
||||
};
|
||||
}
|
||||
|
||||
/** D129: the concept's Sources notice as the library notice: its own copy and actions, no bottom margin (main's padding). */
|
||||
function mapLibraryNotice(reference) {
|
||||
return {
|
||||
...reference,
|
||||
text: [['h2', 'Your library is inside the node folder'], ['p', 'Your library is stored inside the node folder, where an '
|
||||
+ 'update or reinstall can delete it. Move it to ComfyUI\'s user folder: it is copied and verified first, and nothing '
|
||||
+ 'is deleted.']],
|
||||
actions: reference.actions.map((action, index) => ({ ...action, text: ['Keep it here', 'Move library'][index] })),
|
||||
frame: { ...reference.frame, marginBottom: '0px' },
|
||||
};
|
||||
}
|
||||
|
||||
async function conceptNotice(page, theme) {
|
||||
await openConcept(page, 'system', theme);
|
||||
await page.locator('#main .section-nav').getByRole('button', { name: 'Sources & scans', exact: true }).click();
|
||||
await expect(page.locator('#main .settings-layout > .notice')).toBeVisible();
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureNoticeOf, ['#main .settings-layout > .notice', '#main .heading']);
|
||||
}
|
||||
|
||||
async function productionNotice(page, theme) {
|
||||
await openProduction(page, 'system/overview', theme);
|
||||
await expect(page.locator('#library-banner > .notice')).toBeVisible();
|
||||
await expect(page.locator('#workspace-main .summary-metrics')).toBeVisible();
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureNoticeOf, ['#library-banner > .notice', '#workspace-main .heading']);
|
||||
}
|
||||
|
||||
for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
test(`library notice matches the concept notice ${theme} ${width}x${height}`, async ({ page, request }, testInfo) => {
|
||||
await page.setViewportSize({ width, height });
|
||||
await expect(await request.post('/__test__/system/library', { data: { state: 'node_folder' } })).toBeOK();
|
||||
const concept = await conceptNotice(page, theme);
|
||||
const production = await productionNotice(page, theme);
|
||||
await page.screenshot({ path: testInfo.outputPath(`production-library-notice-${theme}-${width}.png`), fullPage: true });
|
||||
expect(diffInventory(mapLibraryNotice(concept), production)).toEqual([]);
|
||||
expect(production.alignedWith).toBe(0);
|
||||
expect(production.pageScroll).toBe(false);
|
||||
});
|
||||
}
|
||||
|
||||
test('library notice comparison detects mutations', async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
await expect(await request.post('/__test__/system/library', { data: { state: 'node_folder' } })).toBeOK();
|
||||
const expected = mapLibraryNotice(await conceptNotice(page, 'dark'));
|
||||
await productionNotice(page, 'dark');
|
||||
const paths = async () => diffInventory(expected,
|
||||
await page.evaluate(measureNoticeOf, ['#library-banner > .notice', '#workspace-main .heading'])).map(item => item.path);
|
||||
expect(await paths()).toEqual([]);
|
||||
const undoToken = await mutate(page, { css: '#library-banner { --soft-blue: rgb(255, 0, 0); }' });
|
||||
expect(await paths()).toEqual(['frame.backgroundColor']);
|
||||
await undoToken();
|
||||
await page.evaluate(() => { document.querySelector('#library-banner .notice h2').textContent = 'Library'; });
|
||||
expect(await paths()).toEqual(['text.0.1']);
|
||||
});
|
||||
|
||||
@@ -1,344 +0,0 @@
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from './fixtures.mjs';
|
||||
import { SIZES, THEMES, seedConcept, openConcept, openProduction, diffInventory, mutate } from './concept-inventory.mjs';
|
||||
|
||||
// #249 System S10 Task 28: System › Integrations (app.js renderSystem 682–684) and the LoRA import review against the
|
||||
// immutable concept, on the concept-shaped fixture with a fake LoraManager install. Each concept measurement is mapped
|
||||
// only by the recorded deviations (D106, D107, D114), each commented where it applies; after that the difference must
|
||||
// be empty, and a mutated label, option, row or token must be noticed. The concept starts its sample import without a
|
||||
// review, so the import review is compared with the concept's own record-changing review, `maintenancePlan`
|
||||
// (953–958), whose shell, grid, note and actions it shares.
|
||||
test.use({ locale: 'en-US', timezoneId: 'UTC' });
|
||||
|
||||
test.beforeEach(async ({ page, request, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
await seedConcept(request);
|
||||
await expect(await request.post('/__test__/system/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/integrations/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/jobs/restart')).toBeOK();
|
||||
page.conceptErrors = [];
|
||||
page.on('pageerror', error => page.conceptErrors.push(error.message));
|
||||
});
|
||||
|
||||
test.afterEach(async ({ page, request, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
await request.post('/__test__/jobs/restart');
|
||||
await request.post('/__test__/integrations/reset');
|
||||
await request.post('/__test__/system/reset');
|
||||
expect(page.conceptErrors).toEqual([]);
|
||||
});
|
||||
|
||||
const PREVIEW_URL = '/prompt_manager/lora/import/preview';
|
||||
|
||||
/** Runs in the page: Integrations' surfaces under `mainSelector`, every row's text and control, geometry and tokens. */
|
||||
function measureIntegrations(mainSelector) {
|
||||
const layout = document.querySelector(`${mainSelector} .settings-layout`);
|
||||
const css = (node, ...keys) => node ? Object.fromEntries(keys.map(key => [key, getComputedStyle(node)[key]])) : null;
|
||||
const name = node => [node.localName, ...[...node.classList].sort()].join('.');
|
||||
const text = node => node?.textContent.trim() ?? null;
|
||||
const control = row => {
|
||||
const select = row.querySelector(':scope > select');
|
||||
if (select) return { kind: 'select', id: select.id, options: [...select.options].map(option => [option.textContent, option.disabled]),
|
||||
selected: select.selectedOptions[0]?.textContent ?? null, disabled: select.disabled };
|
||||
const input = row.querySelector(':scope > input');
|
||||
if (input) return { kind: `input.${input.type}`, id: input.id, value: input.value, readOnly: input.readOnly };
|
||||
const toggle = row.querySelector(':scope > * input[type="checkbox"]');
|
||||
return toggle ? { kind: 'switch', id: toggle.id, wrapper: name(toggle.parentElement), checked: toggle.checked,
|
||||
disabled: toggle.disabled, ...css(toggle.parentElement, 'width', 'height') } : null;
|
||||
};
|
||||
const surfaces = [...layout.querySelectorAll(':scope > section.surface')];
|
||||
const firstRow = surfaces[0].querySelector('.setting-row');
|
||||
const note = layout.querySelector('.setting-row p.small');
|
||||
return {
|
||||
layout: [...layout.children].map(name),
|
||||
surfaces: surfaces.map(surface => ({ tree: [...surface.children].map(name), heading: text(surface.querySelector(':scope > h2')),
|
||||
intro: text(surface.querySelector(':scope > p')),
|
||||
rows: [...surface.querySelectorAll(':scope > .setting-row')].map(row => ({ tree: [...row.children].map(name),
|
||||
text: [...row.firstElementChild.children].map(node => [name(node), text(node), node.localName === 'label' ? node.htmlFor : null]),
|
||||
control: control(row) })),
|
||||
operations: [...surface.querySelectorAll(':scope > .operation')].map(row => ({ tree: [...row.children].map(name),
|
||||
title: text(row.querySelector('h3')), description: text(row.querySelector('p')),
|
||||
actions: [...row.querySelectorAll('.row-actions > button')].map(item => [text(item), item.className]) })) })),
|
||||
surface: css(surfaces[1], 'paddingTop', 'paddingLeft', 'backgroundColor', 'borderTopWidth', 'borderTopColor', 'borderTopLeftRadius',
|
||||
'marginTop'),
|
||||
row: css(firstRow, 'display', 'alignItems', 'justifyContent', 'columnGap', 'paddingTop', 'paddingBottom', 'borderTopWidth',
|
||||
'borderTopColor', 'marginTop'),
|
||||
label: css(firstRow.querySelector('label'), 'fontSize', 'fontWeight', 'color'),
|
||||
description: css(firstRow.querySelector('p'), 'fontSize', 'color', 'marginTop', 'maxWidth'),
|
||||
select: css(firstRow.querySelector('select'), 'fontSize', 'minWidth', 'maxWidth', 'backgroundColor', 'color', 'borderTopColor'),
|
||||
password: css(layout.querySelector('.setting-row input[type="password"]'), 'fontSize', 'minWidth', 'maxWidth', 'backgroundColor',
|
||||
'color', 'borderTopColor'),
|
||||
// The concept's `.small` text (its Save row status); production's row notes use the same class.
|
||||
small: css(layout.querySelector('#settings-status'), 'fontSize', 'color'),
|
||||
note: note && css(note, 'fontSize', 'color'),
|
||||
// Whether each select is as wide as the browser lays out a copy holding only its chosen option, i.e. it shows
|
||||
// that option whole (the browser's own select sizing, arrow included).
|
||||
selectsFit: [...layout.querySelectorAll('.setting-row > select')].map(select => {
|
||||
const copy = select.cloneNode(false);
|
||||
copy.removeAttribute('id');
|
||||
copy.append(new Option(select.selectedOptions[0]?.textContent ?? ''));
|
||||
Object.assign(copy.style, { position: 'absolute', visibility: 'hidden', width: 'max-content', minWidth: '0',
|
||||
maxWidth: 'none', flex: 'none' });
|
||||
select.parentElement.append(copy);
|
||||
const needed = copy.getBoundingClientRect().width;
|
||||
copy.remove();
|
||||
return select.getBoundingClientRect().width + 0.5 >= needed;
|
||||
}),
|
||||
pageScroll: document.documentElement.scrollWidth > document.documentElement.clientWidth,
|
||||
};
|
||||
}
|
||||
|
||||
const SEND_REASON = 'Sending always asks for the ComfyUI window and node, so “Ask each time” is the only option.';
|
||||
const RENDERER_REASON = 'Prompt Manager nodes always follow ComfyUI’s own renderer setting.';
|
||||
const RENDERER_NONE = 'No ComfyUI window is open in this browser, so its renderer is not known.';
|
||||
const RETURN_REASON = 'Not available: sending applies the text in ComfyUI and keeps this page open.';
|
||||
const PATH_DESCRIPTION = 'Found automatically in ComfyUI’s custom_nodes folders; read-only, as on the existing settings page.';
|
||||
|
||||
/**
|
||||
* The concept's Integrations with the recorded deviations applied, for the fixture's LoraManager (detected, disabled,
|
||||
* no key, trigger words off and not loaded).
|
||||
*/
|
||||
function mapIntegrationsConcept(reference) {
|
||||
const [comfy, lora] = reference.surfaces;
|
||||
const small = text => ['p.small', text, null];
|
||||
const firstOnly = row => ({ ...row, control: { ...row.control, options: row.control.options.map(([label], index) => [label, index > 0]) } });
|
||||
const [send, renderer, back] = comfy.rows;
|
||||
const [enable, triggers, key] = lora.rows;
|
||||
return {
|
||||
...reference,
|
||||
surfaces: [
|
||||
{ ...comfy, rows: [
|
||||
// D106: only "Ask each time" can be chosen, and the row says why.
|
||||
{ ...firstOnly(send), text: [...send.text, small(SEND_REASON)] },
|
||||
// D106: only "Follow host renderer" can be chosen; the reason, then the detected renderer.
|
||||
{ ...firstOnly(renderer), text: [...renderer.text, small(RENDERER_REASON), small(RENDERER_NONE)] },
|
||||
// D106: Return to canvas is not available: off, disabled and saying why.
|
||||
{ ...back, text: [...back.text, small(RETURN_REASON)], control: { ...back.control, checked: false, disabled: true } },
|
||||
] },
|
||||
{ ...lora,
|
||||
// D106: the LoraManager path row (legacy, read-only) follows the integration toggle.
|
||||
tree: [lora.tree[0], lora.tree[1], 'div.setting-row', 'div.setting-row', 'div.setting-row', 'div.setting-row', 'div.operation'],
|
||||
rows: [
|
||||
// The real detection replaces the concept's sample line; the integration is off in the fixture.
|
||||
{ ...enable, text: [enable.text[0], ['p', 'Detected: custom_nodes/ComfyUI-Lora-Manager.', null]],
|
||||
control: { ...enable.control, checked: false } },
|
||||
{ tree: ['div', 'input'],
|
||||
text: [['label', 'LoraManager path', 'lora-path'], ['p', PATH_DESCRIPTION, null]],
|
||||
control: { kind: 'input.text', id: 'lora-path', value: 'custom_nodes/ComfyUI-Lora-Manager', readOnly: true } },
|
||||
// The trigger cache is a separate fact from the setting: its own line.
|
||||
{ ...triggers, text: [...triggers.text, small('Trigger words are not loaded.')], control: { ...triggers.control, checked: false } },
|
||||
// D114: no prototype warning; the saved-key line under it.
|
||||
// (No key is saved in the fixture, so "Clear saved key" is not offered.)
|
||||
{ ...key, text: [key.text[0], ['p', 'Credentials stay in the server configuration.', null], small('No key saved.')],
|
||||
control: { ...key.control, value: '' } },
|
||||
],
|
||||
// D107: the import is reviewed first ("Review import"; the concept's button only simulated a job).
|
||||
operations: lora.operations.map(operation => ({ ...operation, actions: [['Review import', '']] })) },
|
||||
],
|
||||
// D106: both selects show their whole chosen option at every width (at 390 the concept cuts "Follow host
|
||||
// renderer"; production lets the two selects grow to 160px and never shrink for the reason lines).
|
||||
selectsFit: reference.selectsFit.map(() => true),
|
||||
select: reference.select.maxWidth === '120px' ? { ...reference.select, maxWidth: '160px' } : reference.select,
|
||||
// The D106/D114 row notes are `p.small` inside the row text, so they read as its description (`.setting-row p`).
|
||||
note: { fontSize: reference.description.fontSize, color: reference.description.color },
|
||||
};
|
||||
}
|
||||
|
||||
async function conceptIntegrations(page, theme) {
|
||||
await openConcept(page, 'system', theme);
|
||||
await page.locator('#main .section-nav').getByRole('button', { name: 'Integrations', exact: true }).click();
|
||||
await expect(page.locator('#main .settings-layout > section.surface')).toHaveCount(2);
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureIntegrations, '#main');
|
||||
}
|
||||
|
||||
async function productionIntegrations(page, theme) {
|
||||
await openProduction(page, 'system/integrations', theme);
|
||||
// Compare once the LoRA state and the renderer check answered.
|
||||
await expect(page.locator('#workspace-main #lora-path')).toHaveValue('custom_nodes/ComfyUI-Lora-Manager');
|
||||
await expect(page.locator('#workspace-main #renderer-status')).toHaveText(RENDERER_NONE);
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureIntegrations, '#workspace-main');
|
||||
}
|
||||
|
||||
const fakeLoraManager = async (request, data = {}) =>
|
||||
expect(await request.post('/__test__/integrations/lora-manager', { data })).toBeOK();
|
||||
|
||||
for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
test(`integrations matches the concept ${theme} ${width}x${height}`, async ({ page, request }, testInfo) => {
|
||||
await page.setViewportSize({ width, height });
|
||||
await fakeLoraManager(request);
|
||||
const concept = await conceptIntegrations(page, theme);
|
||||
await page.screenshot({ path: testInfo.outputPath(`concept-integrations-${theme}-${width}.png`), fullPage: true });
|
||||
const production = await productionIntegrations(page, theme);
|
||||
await page.screenshot({ path: testInfo.outputPath(`production-integrations-${theme}-${width}.png`), fullPage: true });
|
||||
expect(diffInventory(mapIntegrationsConcept(concept), production)).toEqual([]);
|
||||
expect(production.pageScroll).toBe(false);
|
||||
});
|
||||
}
|
||||
|
||||
// D106 below the concept's smallest size: on a 320px phone both selects still show their whole option, without page
|
||||
// scroll (S10 re-review N9).
|
||||
for (const theme of THEMES) test(`integrations selects fit at 320 wide ${theme}`, async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 320, height: 700 });
|
||||
await fakeLoraManager(request);
|
||||
const production = await productionIntegrations(page, theme);
|
||||
expect(production.selectsFit).toEqual([true, true]);
|
||||
expect(production.pageScroll).toBe(false);
|
||||
});
|
||||
|
||||
test('integrations comparison detects mutations', async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
await fakeLoraManager(request);
|
||||
const concept = await conceptIntegrations(page, 'dark');
|
||||
const expected = mapIntegrationsConcept(concept);
|
||||
const production = await productionIntegrations(page, 'dark');
|
||||
expect(diffInventory(expected, production)).toEqual([]);
|
||||
// Unmapped, the section differs only where D106, D107 and D114 are recorded.
|
||||
expect(diffInventory(concept, production).map(item => item.path.split('.').slice(0, 4).join('.'))
|
||||
.filter((path, index, all) => all.indexOf(path) === index).sort()).toEqual([
|
||||
'note', 'note.color', 'note.fontSize', 'surfaces.0.rows.0', 'surfaces.0.rows.1', 'surfaces.0.rows.2', 'surfaces.1.operations.0',
|
||||
'surfaces.1.rows.0', 'surfaces.1.rows.1', 'surfaces.1.rows.2', 'surfaces.1.rows.3', 'surfaces.1.rows.length',
|
||||
'surfaces.1.tree.5', 'surfaces.1.tree.6', 'surfaces.1.tree.length'].sort());
|
||||
const measure = () => page.evaluate(measureIntegrations, '#workspace-main');
|
||||
const paths = async () => diffInventory(expected, await measure()).map(item => item.path);
|
||||
// A removed row.
|
||||
const undoRow = await mutate(page, { remove: ['#workspace-main section.surface:nth-of-type(2) > .setting-row:nth-of-type(3)'] });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['surfaces.1.rows.length', 'surfaces.1.tree.length']));
|
||||
await undoRow();
|
||||
// A changed label and a changed option.
|
||||
await page.evaluate(() => { document.querySelector('#workspace-main label[for="renderer"]').textContent = 'Renderer'; });
|
||||
expect(await paths()).toEqual(['surfaces.0.rows.1.text.0.1']);
|
||||
await page.evaluate(() => { document.querySelector('#workspace-main label[for="renderer"]').textContent = 'Host compatibility'; });
|
||||
await page.evaluate(() => { document.querySelector('#workspace-main #send-destination option').textContent = 'Always ask'; });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['surfaces.0.rows.0.control.options.0.0']));
|
||||
await page.evaluate(() => { document.querySelector('#workspace-main #send-destination option').textContent = 'Ask each time'; });
|
||||
// An unsupported option that became selectable.
|
||||
await page.evaluate(() => { document.querySelectorAll('#workspace-main #renderer option')[1].disabled = false; });
|
||||
expect(await paths()).toEqual(['surfaces.0.rows.1.control.options.1.1']);
|
||||
await page.evaluate(() => { document.querySelectorAll('#workspace-main #renderer option')[1].disabled = true; });
|
||||
// A select squeezed until its chosen option is cut (S10 review M3).
|
||||
const undoSqueeze = await mutate(page, { css: '#workspace-main #renderer { max-width: 70px !important; }' });
|
||||
expect(await paths()).toEqual(['selectsFit.1']);
|
||||
await undoSqueeze();
|
||||
// A swapped token.
|
||||
const undoToken = await mutate(page, { css: '#workspace-main .settings-layout { --line: rgb(255, 0, 0); }' });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['row.borderTopColor']));
|
||||
await undoToken();
|
||||
expect(diffInventory(expected, await measure())).toEqual([]);
|
||||
});
|
||||
|
||||
/** Runs in the page: the open review dialog: its title, body and the concept's tokens. */
|
||||
function measureReview() {
|
||||
const dialog = document.querySelector('dialog#modal[open]');
|
||||
const css = (node, ...keys) => node ? Object.fromEntries(keys.map(key => [key, getComputedStyle(node)[key]])) : null;
|
||||
const name = node => [node.localName, ...[...node.classList].sort()].join('.');
|
||||
const text = node => node?.textContent.trim() ?? null;
|
||||
const body = dialog.querySelector('#modal-body');
|
||||
const grid = body.querySelector(':scope > dl.meta-grid');
|
||||
const note = [...body.querySelectorAll(':scope > p.small')].at(-1);
|
||||
const actions = body.querySelector(':scope > .dialog-actions');
|
||||
return {
|
||||
title: text(dialog.querySelector('#modal-title')),
|
||||
tree: [...body.children].map(name),
|
||||
heading: { text: text(body.querySelector(':scope > h3')),
|
||||
...css(body.querySelector(':scope > h3'), 'fontSize', 'fontWeight', 'color', 'marginBottom') },
|
||||
grid: { rows: [...grid.children].map(row => [text(row.querySelector('dt')), text(row.querySelector('dd'))]),
|
||||
...css(grid, 'display', 'gridTemplateColumns', 'columnGap', 'rowGap'),
|
||||
dt: css(grid.querySelector('dt'), 'fontSize', 'color', 'textTransform'), dd: css(grid.querySelector('dd'), 'fontSize', 'color') },
|
||||
details: body.querySelector(':scope > details') && { summary: text(body.querySelector(':scope > details > summary')),
|
||||
tree: [...body.querySelector(':scope > details').children].map(name) },
|
||||
backup: body.querySelector(':scope > .notice') && { tree: [...body.querySelector(':scope > .notice').children].map(name),
|
||||
text: [...body.querySelectorAll(':scope > .notice > div:first-child > *')].map(node => [node.localName, text(node)]),
|
||||
buttons: [...body.querySelectorAll(':scope > .notice button')].map(text) },
|
||||
note: { text: text(note), ...css(note, 'fontSize', 'color', 'marginTop') },
|
||||
actions: { tree: [...actions.children].map(name), ...css(actions, 'display', 'justifyContent', 'columnGap', 'marginTop'),
|
||||
buttons: [...actions.children].map(item => ({ text: text(item), className: item.className,
|
||||
...css(item, 'backgroundColor', 'color', 'fontSize', 'paddingTop', 'paddingLeft', 'borderTopColor') })) },
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* D107: the server's import plan in the concept's review: "Review LoRA import" with "LoRA reference data" as its `h3`,
|
||||
* the concept's four rows with the plan's values plus Found and Network, the LoRAs it found (`details`), the hold
|
||||
* line with the cache sentence in the concept's `p.small`, the D55 backup box, and Cancel / "Start import".
|
||||
*/
|
||||
function mapImportConcept(reference, plan) {
|
||||
const [cancel, start] = reference.actions.buttons;
|
||||
return {
|
||||
...reference,
|
||||
title: 'Review LoRA import',
|
||||
tree: ['h3', 'dl.meta-grid', 'details', 'p.small', 'div.notice', 'div.dialog-actions'],
|
||||
heading: { ...reference.heading, text: 'LoRA reference data' },
|
||||
grid: { ...reference.grid, rows: [['LoraManager', 'environment settings · alongside storage · library default'],
|
||||
['Scope', plan.scope.label], ['Found', '2 metadata files · 2 LoRAs · 0 with local images'],
|
||||
['Protected', 'Nothing is protected'], ['Network', 'Downloads up to 2 example images from CivitAI'],
|
||||
['Backup', `Create and verify before changes (${plan.backup.destination_label})`],
|
||||
['Media on disk', 'No files are removed; image links are replaced']] },
|
||||
details: { summary: 'LoRAs found (2)', tree: ['summary', 'pre'] },
|
||||
backup: { tree: ['div', 'div.row-actions'], text: [['h3', 'Back up the database first (recommended)'],
|
||||
['p', 'Downloads a copy of the prompts database before this operation changes any records.']], buttons: ['Back up now'] },
|
||||
note: { ...reference.note, text: 'Nothing changes until you apply. This review holds for 15 minutes. '
|
||||
+ 'Example images downloaded before a failure stay in the image cache.' },
|
||||
actions: { ...reference.actions, buttons: [cancel, { ...start, text: 'Start import' }] },
|
||||
};
|
||||
}
|
||||
|
||||
async function conceptReview(page, theme) {
|
||||
await openConcept(page, 'system', theme);
|
||||
await page.locator('#main .section-nav').getByRole('button', { name: 'Maintenance', exact: true }).click();
|
||||
await page.locator('#main [data-plan="duplicate prompts"]').click();
|
||||
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Review maintenance plan');
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureReview);
|
||||
}
|
||||
|
||||
async function productionReview(page, request, theme) {
|
||||
await openProduction(page, 'system/integrations', theme);
|
||||
await expect(page.locator('#workspace-main #lora-path')).toHaveValue('custom_nodes/ComfyUI-Lora-Manager');
|
||||
// The renderer line above settles first: its new text moves the rows below (a click must not land mid-shift).
|
||||
await expect(page.locator('#workspace-main #renderer-status')).toHaveText(RENDERER_NONE);
|
||||
// The import reviews what is enabled: save the integration on first.
|
||||
await page.locator('#workspace-main #lora-enable').check();
|
||||
await page.getByRole('button', { name: 'Save preferences', exact: true }).click();
|
||||
await expect(page.locator('#settings-status')).toHaveText('Saved to the server configuration.');
|
||||
const answer = page.waitForResponse(response => response.url().endsWith(PREVIEW_URL));
|
||||
await page.locator('#workspace-main').getByRole('button', { name: 'Review import', exact: true }).click();
|
||||
const plan = await (await answer).json();
|
||||
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Review LoRA import');
|
||||
await page.mouse.move(0, 0);
|
||||
return { plan, production: await page.evaluate(measureReview) };
|
||||
}
|
||||
|
||||
for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
test(`LoRA import review matches the concept ${theme} ${width}x${height}`, async ({ page, request }, testInfo) => {
|
||||
await page.setViewportSize({ width, height });
|
||||
await fakeLoraManager(request, { loras: 2, civitai: true, preview: false });
|
||||
const concept = await conceptReview(page, theme);
|
||||
expect(concept.tree).toEqual(['p.demo-note', 'h3', 'dl.meta-grid', 'p.small', 'div.dialog-actions']);
|
||||
await page.screenshot({ path: testInfo.outputPath(`concept-plan-review-${theme}-${width}.png`) });
|
||||
const { plan, production } = await productionReview(page, request, theme);
|
||||
await page.screenshot({ path: testInfo.outputPath(`production-lora-review-${theme}-${width}.png`) });
|
||||
expect(diffInventory(mapImportConcept(concept, plan), production)).toEqual([]);
|
||||
});
|
||||
}
|
||||
|
||||
test('LoRA import review comparison detects mutations', async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
await fakeLoraManager(request, { loras: 2, civitai: true, preview: false });
|
||||
const concept = await conceptReview(page, 'dark');
|
||||
const { plan } = await productionReview(page, request, 'dark');
|
||||
const expected = mapImportConcept(concept, plan);
|
||||
const paths = async () => diffInventory(expected, await page.evaluate(measureReview)).map(item => item.path);
|
||||
expect(await paths()).toEqual([]);
|
||||
await page.evaluate(() => { document.querySelectorAll('dialog#modal[open] dl.meta-grid dt')[4].textContent = 'Downloads'; });
|
||||
expect(await paths()).toEqual(['grid.rows.4.0']);
|
||||
await page.evaluate(() => { document.querySelectorAll('dialog#modal[open] dl.meta-grid dt')[4].textContent = 'Network'; });
|
||||
const undoGrid = await mutate(page, { css: 'dialog#modal dl.meta-grid { column-gap: 3px; }' });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['grid.columnGap']));
|
||||
await undoGrid();
|
||||
const undoBox = await mutate(page, { remove: ['dialog#modal[open] #modal-body > .notice'] });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['backup', 'tree.length']));
|
||||
await undoBox();
|
||||
expect(await paths()).toEqual([]);
|
||||
});
|
||||
@@ -1,332 +0,0 @@
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from './fixtures.mjs';
|
||||
import { SIZES, THEMES, seedConcept, openConcept, openProduction, diffInventory, mutate } from './concept-inventory.mjs';
|
||||
|
||||
// #249 System S7 Task 20: System › Maintenance (app.js renderSystem 688–690), its plan dialog (`maintenancePlan`,
|
||||
// 953–958) for each operation and its duplicates dialog (`duplicates`, 1457–1461) against the immutable concept, on
|
||||
// the concept-shaped real fixture. Each concept measurement is mapped only by the recorded deviations (D100, D101,
|
||||
// D102; the health report is D95, compared in system-concept.spec), each commented where it applies; after that the
|
||||
// difference must be empty, and a mutated label, row or token must be noticed.
|
||||
test.use({ locale: 'en-US', timezoneId: 'UTC' });
|
||||
|
||||
test.beforeEach(async ({ page, request, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
await seedConcept(request);
|
||||
await expect(await request.post('/__test__/system/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/jobs/restart')).toBeOK();
|
||||
page.conceptErrors = [];
|
||||
page.on('pageerror', error => page.conceptErrors.push(error.message));
|
||||
});
|
||||
|
||||
test.afterEach(async ({ page, request, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
await request.post('/__test__/jobs/restart');
|
||||
await request.post('/__test__/system/reset');
|
||||
expect(page.conceptErrors).toEqual([]);
|
||||
});
|
||||
|
||||
const PREVIEW_URL = '/prompt_manager/maintenance/preview';
|
||||
|
||||
/** Runs in the page: Maintenance's notice and surfaces under `mainSelector`, with the concept's geometry and tokens. */
|
||||
function measureMaintenance(mainSelector) {
|
||||
const layout = document.querySelector(`${mainSelector} .settings-layout`);
|
||||
const css = (node, ...keys) => node ? Object.fromEntries(keys.map(key => [key, getComputedStyle(node)[key]])) : null;
|
||||
const name = node => [node.localName, ...[...node.classList].sort()].join('.');
|
||||
const text = node => node?.textContent.trim() ?? null;
|
||||
const button = node => ({ text: text(node), className: node.className,
|
||||
...css(node, 'backgroundColor', 'color', 'fontSize', 'paddingTop', 'paddingLeft', 'borderTopColor', 'opacity', 'cursor') });
|
||||
const box = node => {
|
||||
const rect = node.getBoundingClientRect(), base = layout.getBoundingClientRect();
|
||||
return { left: Math.round(rect.left - base.left), width: Math.round(rect.width) };
|
||||
};
|
||||
const notice = layout.querySelector(':scope > .notice');
|
||||
const surfaces = [...layout.querySelectorAll(':scope > section.surface')];
|
||||
const operation = surfaces[0].querySelector('.operation');
|
||||
return {
|
||||
layout: [...layout.children].map(name),
|
||||
notice: { tree: [...notice.children].map(name), text: [...notice.querySelectorAll(':scope > div > *')].map(node => [node.localName, text(node)]),
|
||||
run: button(notice.querySelector(':scope > button')),
|
||||
...css(notice, 'display', 'alignItems', 'justifyContent', 'columnGap', 'borderTopColor', 'backgroundColor', 'paddingTop', 'marginBottom') },
|
||||
surfaces: surfaces.map(surface => ({ tree: [...surface.children].map(name), heading: text(surface.querySelector(':scope > h2')),
|
||||
box: box(surface),
|
||||
rows: [...surface.querySelectorAll(':scope > .operation')].map(row => ({ tree: [...row.children].map(name),
|
||||
title: text(row.querySelector('h3')), description: text(row.querySelector('p')),
|
||||
actions: [...row.querySelectorAll('.row-actions > button')].map(button) })) })),
|
||||
surface: css(surfaces[0], 'paddingTop', 'paddingLeft', 'backgroundColor', 'borderTopWidth', 'borderTopColor', 'borderTopLeftRadius', 'marginTop'),
|
||||
h2: css(surfaces[0].querySelector('h2'), 'fontSize', 'fontWeight', 'color', 'marginBottom'),
|
||||
operation: css(operation, 'display', 'alignItems', 'justifyContent', 'columnGap', 'paddingTop', 'paddingBottom',
|
||||
'borderTopWidth', 'borderTopColor'),
|
||||
h3: css(operation.querySelector('h3'), 'fontSize', 'fontWeight', 'color'),
|
||||
p: css(operation.querySelector('p'), 'fontSize', 'color', 'marginTop'),
|
||||
actions: css(operation.querySelector('.row-actions'), 'display', 'columnGap', 'flexShrink'),
|
||||
pageScroll: document.documentElement.scrollWidth > document.documentElement.clientWidth,
|
||||
};
|
||||
}
|
||||
|
||||
/** D102: the Metadata cache row says why its review is unavailable and its button is aria-disabled (D93's dimmed style). */
|
||||
function mapMaintenanceConcept(reference) {
|
||||
return {
|
||||
...reference,
|
||||
surfaces: reference.surfaces.map(surface => ({ ...surface, rows: surface.rows.map(row => row.title !== 'Metadata cache' ? row : {
|
||||
...row, description: `${row.description} Nothing is cached in this browser.`,
|
||||
actions: row.actions.map(action => ({ ...action, opacity: '0.55', cursor: 'not-allowed' })) }) })),
|
||||
};
|
||||
}
|
||||
|
||||
async function conceptMaintenance(page, theme) {
|
||||
await openConcept(page, 'system', theme);
|
||||
await page.locator('#main .section-nav').getByRole('button', { name: 'Maintenance', exact: true }).click();
|
||||
await expect(page.locator('#main .settings-layout > section.surface')).toHaveCount(2);
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureMaintenance, '#main');
|
||||
}
|
||||
|
||||
async function productionMaintenance(page, theme) {
|
||||
await openProduction(page, 'system/maintenance', theme);
|
||||
await expect(page.locator('#workspace-main .settings-layout > section.surface')).toHaveCount(2);
|
||||
// The Metadata cache row reads this browser's storage once: compare after it answered.
|
||||
await expect(page.locator('#workspace-main .operation').filter({ hasText: 'Metadata cache' }).locator('p'))
|
||||
.toHaveText(/Nothing is cached in this browser\.$/);
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureMaintenance, '#workspace-main');
|
||||
}
|
||||
|
||||
for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
test(`maintenance matches the concept ${theme} ${width}x${height}`, async ({ page }, testInfo) => {
|
||||
await page.setViewportSize({ width, height });
|
||||
const concept = await conceptMaintenance(page, theme);
|
||||
await page.screenshot({ path: testInfo.outputPath(`concept-maintenance-${theme}-${width}.png`), fullPage: true });
|
||||
const production = await productionMaintenance(page, theme);
|
||||
await page.screenshot({ path: testInfo.outputPath(`production-maintenance-${theme}-${width}.png`), fullPage: true });
|
||||
expect(diffInventory(mapMaintenanceConcept(concept), production)).toEqual([]);
|
||||
expect(production.pageScroll).toBe(false);
|
||||
});
|
||||
}
|
||||
|
||||
/** Runs in the page: the open review dialog (plan or duplicates): its title, body and the concept's tokens. */
|
||||
function measureReview() {
|
||||
const dialog = document.querySelector('dialog#modal[open]');
|
||||
const css = (node, ...keys) => node ? Object.fromEntries(keys.map(key => [key, getComputedStyle(node)[key]])) : null;
|
||||
const name = node => [node.localName, ...[...node.classList].sort()].join('.');
|
||||
const text = node => node?.textContent.trim() ?? null;
|
||||
const body = dialog.querySelector('#modal-body');
|
||||
const grid = body.querySelector(':scope > dl.meta-grid');
|
||||
const compare = body.querySelector('.compare-grid');
|
||||
const item = compare?.querySelector('.compare-item');
|
||||
const notes = [...body.querySelectorAll(':scope > p.small')];
|
||||
const note = notes.at(-1);
|
||||
const actions = body.querySelector(':scope > .dialog-actions');
|
||||
return {
|
||||
title: text(dialog.querySelector('#modal-title')),
|
||||
tree: [...body.children].map(name),
|
||||
heading: body.querySelector(':scope > h3') && { text: text(body.querySelector(':scope > h3')),
|
||||
...css(body.querySelector(':scope > h3'), 'fontSize', 'fontWeight', 'color', 'marginBottom') },
|
||||
grid: grid && { rows: [...grid.children].map(row => [text(row.querySelector('dt')), text(row.querySelector('dd'))]),
|
||||
...css(grid, 'display', 'gridTemplateColumns', 'columnGap', 'rowGap'),
|
||||
dt: css(grid.querySelector('dt'), 'fontSize', 'color', 'textTransform'), dd: css(grid.querySelector('dd'), 'fontSize', 'color') },
|
||||
details: body.querySelector(':scope > details') && { summary: text(body.querySelector(':scope > details > summary')),
|
||||
tree: [...body.querySelector(':scope > details').children].map(name) },
|
||||
backup: body.querySelector(':scope > .notice') && { tree: [...body.querySelector(':scope > .notice').children].map(name),
|
||||
text: [...body.querySelectorAll(':scope > .notice > div:first-child > *')].map(node => [node.localName, text(node)]),
|
||||
buttons: [...body.querySelectorAll(':scope > .notice button')].map(text) },
|
||||
compare: compare && { tree: [...compare.children].map(name), ...css(compare, 'display', 'gridTemplateColumns', 'columnGap'),
|
||||
items: [...compare.children].map(entry => ({ tree: [...entry.children].map(name), label: text(entry.querySelector('label')),
|
||||
checked: entry.querySelector('input[type=radio]').checked, image: entry.querySelector('img')?.className ?? null })),
|
||||
item: css(item, 'paddingTop', 'paddingLeft', 'backgroundColor', 'borderTopWidth', 'borderTopColor', 'borderTopLeftRadius'),
|
||||
image: css(item.querySelector('img'), 'width', 'objectFit'),
|
||||
label: css(item.querySelector('label'), 'fontSize', 'color'),
|
||||
radio: { name: item.querySelector('input[type=radio]').name.startsWith('keep-') ? 'keep-*' : 'other',
|
||||
type: item.querySelector('input').type } },
|
||||
note: note && { text: text(note), ...css(note, 'fontSize', 'color', 'marginTop') },
|
||||
actions: { tree: [...actions.children].map(name), ...css(actions, 'display', 'justifyContent', 'columnGap', 'marginTop'),
|
||||
buttons: [...actions.children].map(item => ({ text: text(item), className: item.className,
|
||||
...css(item, 'backgroundColor', 'color', 'fontSize', 'paddingTop', 'paddingLeft', 'borderTopColor') })) },
|
||||
};
|
||||
}
|
||||
|
||||
// The concept's plan names (`data-plan`) and production's operation, row and button for each.
|
||||
const PLANS = Object.freeze([
|
||||
['duplicate prompts', 'duplicate_prompts', 'Duplicate prompt text'],
|
||||
['database optimization', 'optimize', 'Optimize database'],
|
||||
['unused prompt components', 'unused_components', 'Unused prompt components'],
|
||||
['missing-file links', 'missing_links', 'Missing-file links'],
|
||||
]);
|
||||
const BACKUP_BOX = Object.freeze({ tree: ['div', 'div.row-actions'], text: [['h3', 'Back up the database first (recommended)'],
|
||||
['p', 'Downloads a copy of the prompts database before this operation changes any records.']], buttons: ['Back up now'] });
|
||||
|
||||
/**
|
||||
* D100: the server's plan in the concept's dialog: the operation's row title as its `h3` (the concept prints its
|
||||
* `data-plan` name), the four rows with the plan's values, the listed records (`details`), the hold line in the
|
||||
* concept's `p.small`, the D55 backup box before the actions, and "Apply reviewed operation". No sample note.
|
||||
*/
|
||||
function mapPlanConcept(reference, title, plan) {
|
||||
const backup = plan.backup.policy === 'off' ? 'No safety backup (the policy is off)'
|
||||
: `Create and verify before changes (${plan.backup.destination_label})`;
|
||||
const tree = ['h3', 'dl.meta-grid', ...(plan.details.length ? ['details'] : []), 'p.small', 'div.notice', 'div.dialog-actions'];
|
||||
const [cancel, apply] = reference.actions.buttons;
|
||||
return {
|
||||
...reference,
|
||||
tree,
|
||||
heading: { ...reference.heading, text: title },
|
||||
grid: { ...reference.grid, rows: [['Scope', plan.scope.label],
|
||||
['Protected', plan.protected.reasons.length ? plan.protected.reasons.join(' · ') : 'Nothing is protected'],
|
||||
['Backup', backup], ['Media on disk', plan.media.label]] },
|
||||
details: plan.details.length ? { summary: `Listed records (${plan.details.length})`, tree: ['summary', 'pre'] } : null,
|
||||
backup: BACKUP_BOX,
|
||||
note: { ...reference.note, text: 'Nothing changes until you apply. This review holds for 15 minutes.' },
|
||||
actions: { ...reference.actions, buttons: [cancel, { ...apply, text: 'Apply reviewed operation' }] },
|
||||
};
|
||||
}
|
||||
|
||||
async function conceptPlan(page, theme, name) {
|
||||
await conceptMaintenance(page, theme);
|
||||
await page.locator(`#main [data-plan="${name}"]`).click();
|
||||
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Review maintenance plan');
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureReview);
|
||||
}
|
||||
|
||||
/** Production's plan dialog for the row `title`, with the server's answer its review read. */
|
||||
async function productionPlan(page, title) {
|
||||
const answer = page.waitForResponse(response => response.url().endsWith(PREVIEW_URL));
|
||||
await page.locator('#workspace-main .operation').filter({ has: page.locator('h3', { hasText: title }) })
|
||||
.getByRole('button', { name: 'Review plan', exact: true }).click();
|
||||
const plan = await (await answer).json();
|
||||
await expect(page.locator('dialog#modal[open] #modal-body > h3')).toHaveText(title);
|
||||
await page.mouse.move(0, 0);
|
||||
return { plan, production: await page.evaluate(measureReview) };
|
||||
}
|
||||
|
||||
for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
test(`plan dialogs match the concept ${theme} ${width}x${height}`, async ({ page, request }, testInfo) => {
|
||||
await page.setViewportSize({ width, height });
|
||||
// Something to merge, and a missing-file link, so no plan is empty.
|
||||
await expect(await request.post('/__test__/system/duplicate-prompts')).toBeOK();
|
||||
await expect(await request.post('/__test__/system/missing-file')).toBeOK();
|
||||
const concepts = [];
|
||||
for (const [name] of PLANS) {
|
||||
concepts.push(await conceptPlan(page, theme, name));
|
||||
await page.screenshot({ path: testInfo.outputPath(`concept-plan-${name.replace(/\W+/g, '-')}-${theme}-${width}.png`) });
|
||||
}
|
||||
expect(concepts[0].tree).toEqual(['p.demo-note', 'h3', 'dl.meta-grid', 'p.small', 'div.dialog-actions']);
|
||||
await productionMaintenance(page, theme);
|
||||
for (const [index, [, operation, title]] of PLANS.entries()) {
|
||||
const { plan, production } = await productionPlan(page, title);
|
||||
expect(plan.operation).toBe(operation);
|
||||
await page.screenshot({ path: testInfo.outputPath(`production-plan-${operation}-${theme}-${width}.png`) });
|
||||
expect(diffInventory(mapPlanConcept(concepts[index], title, plan), production)).toEqual([]);
|
||||
await page.keyboard.press('Escape');
|
||||
await expect(page.locator('dialog#modal[open]')).toHaveCount(0);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* D101: the real group in the concept's markup: the status line (the concept's sample note's place), a group's
|
||||
* `p.small` and `.compare-grid` in one wrapper, a "Keep <source>/<path>" radio per copy, the oldest checked, and the
|
||||
* production note in the concept's `p.small`.
|
||||
*/
|
||||
function mapDuplicatesConcept(reference) {
|
||||
const items = reference.compare.items.map((item, index) => ({ ...item,
|
||||
label: `Keep output/s7-dupes/group_00_${index ? 'b' : 'a'}.png` }));
|
||||
return {
|
||||
...reference,
|
||||
tree: ['p.small', 'div.duplicate-groups', 'p.small', 'div.dialog-actions'],
|
||||
compare: { ...reference.compare, items },
|
||||
note: { ...reference.note, text: 'Removing a copy deletes its file from disk; its database links move to the kept copy. '
|
||||
+ 'History keeps the removed paths in its run records. Nothing is deleted until you review and apply the cleanup.' },
|
||||
};
|
||||
}
|
||||
|
||||
async function conceptDuplicates(page, theme) {
|
||||
await conceptMaintenance(page, theme);
|
||||
await page.locator('#main').getByRole('button', { name: 'Find duplicates', exact: true }).click();
|
||||
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Review duplicate image files');
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureReview);
|
||||
}
|
||||
|
||||
async function productionDuplicates(page, theme) {
|
||||
await productionMaintenance(page, theme);
|
||||
await page.locator('#workspace-main').getByRole('button', { name: 'Find duplicates', exact: true }).click();
|
||||
await page.locator('#workspace-main').getByRole('button', { name: 'Review groups', exact: true }).click();
|
||||
await expect(page.locator('dialog#modal[open] .compare-item img')).toHaveCount(2);
|
||||
await page.locator('dialog#modal[open] .compare-item img').evaluateAll(images => Promise.all(images.map(image => image.decode())));
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureReview);
|
||||
}
|
||||
|
||||
for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
test(`duplicates dialog matches the concept ${theme} ${width}x${height}`, async ({ page, request }, testInfo) => {
|
||||
await page.setViewportSize({ width, height });
|
||||
await expect(await request.post('/__test__/system/duplicate-files', { data: { groups: 1, trio: false } })).toBeOK();
|
||||
const concept = await conceptDuplicates(page, theme);
|
||||
expect(concept.tree).toEqual(['p.demo-note', 'div.compare-grid', 'p.small', 'div.dialog-actions']);
|
||||
await page.screenshot({ path: testInfo.outputPath(`concept-duplicates-${theme}-${width}.png`) });
|
||||
const production = await productionDuplicates(page, theme);
|
||||
await page.screenshot({ path: testInfo.outputPath(`production-duplicates-${theme}-${width}.png`) });
|
||||
expect(diffInventory(mapDuplicatesConcept(concept), production)).toEqual([]);
|
||||
});
|
||||
}
|
||||
|
||||
test('maintenance comparison detects mutations', async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
const concept = await conceptMaintenance(page, 'dark');
|
||||
const expected = mapMaintenanceConcept(concept);
|
||||
const production = await productionMaintenance(page, 'dark');
|
||||
expect(diffInventory(expected, production)).toEqual([]);
|
||||
// Unmapped, the section differs only on the D102 row.
|
||||
expect(diffInventory(concept, production).map(item => item.path).sort()).toEqual([
|
||||
'surfaces.1.rows.3.actions.0.cursor', 'surfaces.1.rows.3.actions.0.opacity', 'surfaces.1.rows.3.description']);
|
||||
const measure = () => page.evaluate(measureMaintenance, '#workspace-main');
|
||||
const paths = async () => diffInventory(expected, await measure()).map(item => item.path);
|
||||
// A removed row, a changed label, a lost danger mark and a swapped token.
|
||||
const undoRow = await mutate(page, { remove: ['#workspace-main section.surface:nth-of-type(1) > .operation:nth-of-type(2)'] });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['surfaces.0.rows.length', 'surfaces.0.tree.length']));
|
||||
await undoRow();
|
||||
await page.evaluate(() => { document.querySelector('#workspace-main section.surface h3').textContent = 'Integrity checks'; });
|
||||
expect(await paths()).toEqual(['surfaces.0.rows.0.title']);
|
||||
await page.evaluate(() => { document.querySelector('#workspace-main section.surface h3').textContent = 'Integrity & relationship checks'; });
|
||||
await page.evaluate(() => { document.querySelector('#workspace-main button.danger').classList.remove('danger'); });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['surfaces.0.rows.3.actions.0.className']));
|
||||
await page.evaluate(() => { document.querySelectorAll('#workspace-main section.surface')[0].querySelectorAll('.row-actions > button')[3].classList.add('danger'); });
|
||||
const undoToken = await mutate(page, { css: '#workspace-main .settings-layout { --line: rgb(255, 0, 0); }' });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['operation.borderTopColor']));
|
||||
await undoToken();
|
||||
expect(diffInventory(expected, await measure())).toEqual([]);
|
||||
|
||||
// The plan dialog: a changed grid term and the duplicates dialog's grid token are noticed.
|
||||
await expect(await request.post('/__test__/system/duplicate-prompts')).toBeOK();
|
||||
const conceptPlanDialog = await conceptPlan(page, 'dark', 'duplicate prompts');
|
||||
await productionMaintenance(page, 'dark');
|
||||
const { plan } = await productionPlan(page, 'Duplicate prompt text');
|
||||
const expectedPlan = mapPlanConcept(conceptPlanDialog, 'Duplicate prompt text', plan);
|
||||
const planPaths = async () => diffInventory(expectedPlan, await page.evaluate(measureReview)).map(item => item.path);
|
||||
expect(await planPaths()).toEqual([]);
|
||||
await page.evaluate(() => { document.querySelector('dialog#modal[open] dl.meta-grid dt').textContent = 'Impact'; });
|
||||
expect(await planPaths()).toEqual(['grid.rows.0.0']);
|
||||
await page.evaluate(() => { document.querySelector('dialog#modal[open] dl.meta-grid dt').textContent = 'Scope'; });
|
||||
const undoGrid = await mutate(page, { css: 'dialog#modal dl.meta-grid { column-gap: 3px; }' });
|
||||
expect(await planPaths()).toEqual(expect.arrayContaining(['grid.columnGap']));
|
||||
await undoGrid();
|
||||
expect(await planPaths()).toEqual([]);
|
||||
});
|
||||
|
||||
test('duplicates comparison detects mutations', async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
await expect(await request.post('/__test__/system/duplicate-files', { data: { groups: 1, trio: false } })).toBeOK();
|
||||
const expected = mapDuplicatesConcept(await conceptDuplicates(page, 'dark'));
|
||||
await productionDuplicates(page, 'dark');
|
||||
const paths = async () => diffInventory(expected, await page.evaluate(measureReview)).map(item => item.path);
|
||||
expect(await paths()).toEqual([]);
|
||||
const undoToken = await mutate(page, { css: 'dialog#modal .compare-item { background: rgb(255, 0, 0); }' });
|
||||
expect(await paths()).toEqual(['compare.item.backgroundColor']);
|
||||
await undoToken();
|
||||
await page.evaluate(() => { document.querySelector('dialog#modal[open] .dialog-actions button:last-child').classList.add('ca-primary'); });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['actions.buttons.1.className']));
|
||||
await page.evaluate(() => { document.querySelector('dialog#modal[open] .dialog-actions button:last-child').classList.remove('ca-primary'); });
|
||||
const undoItem = await mutate(page, { remove: ['dialog#modal[open] .compare-item:last-child'] });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['compare.items.length']));
|
||||
await undoItem();
|
||||
expect(await paths()).toEqual([]);
|
||||
});
|
||||
@@ -1,385 +0,0 @@
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from './fixtures.mjs';
|
||||
import { SIZES, THEMES, seedConcept, openConcept, openProduction, diffInventory, mutate } from './concept-inventory.mjs';
|
||||
|
||||
// #328 System S11 Task 8: System › Auto-tagging (app.js renderSystem 685–686) and the Tag output images review against
|
||||
// the immutable concept, on the concept-shaped fixture with the fake tagging model (every engine downloaded, none
|
||||
// loaded). Each concept measurement is mapped only by the recorded deviations (D132, D133, D134), each commented where
|
||||
// it applies; after that the difference must be empty, and a mutated label, option, order, class, select width or
|
||||
// token must be noticed. The concept has no batch review, so the Tag output images review is compared with the
|
||||
// concept's own record-changing review, `maintenancePlan` (953–958), whose shell, grid, note and actions it shares.
|
||||
test.use({ locale: 'en-US', timezoneId: 'UTC' });
|
||||
|
||||
test.beforeEach(async ({ page, request, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
// The tagging reset restores the fake model (nothing loaded, every engine on disk) and the saved tagging defaults;
|
||||
// the concept fixture then replaces its rows and files.
|
||||
await expect(await request.post('/__test__/tagging/reset', { data: {} })).toBeOK();
|
||||
await seedConcept(request);
|
||||
await expect(await request.post('/__test__/system/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/jobs/restart')).toBeOK();
|
||||
page.conceptErrors = [];
|
||||
page.on('pageerror', error => page.conceptErrors.push(error.message));
|
||||
});
|
||||
|
||||
test.afterEach(async ({ page, request, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
await request.post('/__test__/jobs/restart');
|
||||
await request.post('/__test__/system/reset');
|
||||
await request.post('/__test__/tagging/reset', { data: {} });
|
||||
expect(page.conceptErrors).toEqual([]);
|
||||
});
|
||||
|
||||
const PREVIEW_URL = '/prompt_manager/autotag/batch/preview';
|
||||
const MODEL_OPTIONS = ['WD14 SwinV2 · ~400 MB', 'WD14 ViT · ~400 MB', 'GGUF JoyCaption · ~4.5 GB', 'HuggingFace JoyCaption · ~16 GB'];
|
||||
const SUGGESTIONS = 'From 1 to 100: the most suggestions Generate shows for one image.';
|
||||
const BATCH = 'Add tags to each prompt from its linked output images with the model above. Review the scope first.';
|
||||
|
||||
/** Runs in the page: Auto-tagging's surface under `mainSelector`: rows, field, actions, operation, geometry and tokens. */
|
||||
function measureTagging(mainSelector) {
|
||||
const layout = document.querySelector(`${mainSelector} .settings-layout`);
|
||||
const css = (node, ...keys) => node ? Object.fromEntries(keys.map(key => [key, getComputedStyle(node)[key]])) : null;
|
||||
const name = node => [node.localName, ...[...node.classList].sort()].join('.');
|
||||
const text = node => node?.textContent.trim() ?? null;
|
||||
const control = row => {
|
||||
const select = row.querySelector(':scope > select');
|
||||
if (select) return { kind: 'select', id: select.id, options: [...select.options].map(option => [option.textContent, option.disabled]),
|
||||
selected: select.selectedOptions[0]?.textContent ?? null, disabled: select.disabled };
|
||||
const input = row.querySelector(':scope > input');
|
||||
return input ? { kind: `input.${input.type}`, id: input.id, value: input.value, disabled: input.disabled } : null;
|
||||
};
|
||||
const surface = layout.querySelector(':scope > section.surface');
|
||||
const firstRow = surface.querySelector('.setting-row');
|
||||
const field = surface.querySelector(':scope > .field');
|
||||
const actions = surface.querySelector(':scope > .row-actions');
|
||||
return {
|
||||
layout: [...layout.children].map(name),
|
||||
tree: [...surface.children].map(name),
|
||||
heading: text(surface.querySelector(':scope > h2')),
|
||||
intro: text(surface.querySelector(':scope > p')),
|
||||
// Per row: whether it may wrap, and whether its control sits under the text (a wrapped row) or beside it.
|
||||
rows: [...surface.querySelectorAll(':scope > .setting-row')].map(row => ({ tree: [...row.children].map(name),
|
||||
text: [...row.firstElementChild.children].map(node => [name(node), text(node), node.localName === 'label' ? node.htmlFor : null]),
|
||||
control: control(row), flexWrap: getComputedStyle(row).flexWrap,
|
||||
controlBelow: row.lastElementChild.getBoundingClientRect().top >= row.firstElementChild.getBoundingClientRect().bottom - 0.5 })),
|
||||
field: { tree: [...field.children].map(name), label: [text(field.querySelector('label')), field.querySelector('label').htmlFor],
|
||||
textarea: { id: field.querySelector('textarea').id, value: field.querySelector('textarea').value,
|
||||
...css(field.querySelector('textarea'), 'fontSize', 'minHeight', 'backgroundColor', 'color', 'borderTopColor') },
|
||||
...css(field, 'display', 'rowGap', 'fontSize', 'marginTop', 'marginBottom') },
|
||||
actions: { buttons: [...actions.children].map(item => [text(item), item.className]),
|
||||
...css(actions, 'display', 'flexWrap', 'columnGap', 'justifyContent'),
|
||||
primary: css(actions.querySelector('.ca-primary'), 'backgroundColor', 'color', 'fontSize', 'paddingTop', 'paddingLeft'),
|
||||
plain: css(actions.querySelector('button:not(.ca-primary)'), 'backgroundColor', 'color', 'fontSize', 'borderTopColor') },
|
||||
operations: [...surface.querySelectorAll(':scope > .operation')].map(row => ({ tree: [...row.children].map(name),
|
||||
title: text(row.querySelector('h3')), description: text(row.querySelector('p')),
|
||||
actions: [...row.querySelectorAll('.row-actions > button')].map(item => [text(item), item.className]) })),
|
||||
surface: css(surface, 'paddingTop', 'paddingLeft', 'backgroundColor', 'borderTopWidth', 'borderTopColor', 'borderTopLeftRadius'),
|
||||
row: css(firstRow, 'display', 'alignItems', 'justifyContent', 'columnGap', 'paddingTop', 'paddingBottom', 'borderTopWidth',
|
||||
'borderTopColor', 'marginTop'),
|
||||
label: css(firstRow.querySelector('label'), 'fontSize', 'fontWeight', 'color'),
|
||||
description: css(firstRow.querySelector('p'), 'fontSize', 'color', 'marginTop', 'maxWidth'),
|
||||
select: css(firstRow.querySelector('select'), 'fontSize', 'minWidth', 'backgroundColor', 'color', 'borderTopColor'),
|
||||
number: css(surface.querySelector('.setting-row input[type="number"]'), 'fontSize', 'minWidth', 'maxWidth', 'backgroundColor',
|
||||
'color', 'borderTopColor'),
|
||||
// Whether each select is as wide as the browser lays out a copy holding only its chosen option, i.e. it shows
|
||||
// that option whole (the browser's own select sizing, arrow included).
|
||||
selectsFit: [...surface.querySelectorAll('.setting-row > select')].map(select => {
|
||||
const copy = select.cloneNode(false);
|
||||
copy.removeAttribute('id');
|
||||
copy.append(new Option(select.selectedOptions[0]?.textContent ?? ''));
|
||||
// Important inline values: the phone cap is an !important rule (concept.css), which a plain inline style loses to.
|
||||
for (const [key, value] of [['position', 'absolute'], ['visibility', 'hidden'], ['width', 'max-content'], ['min-width', '0'],
|
||||
['max-width', 'none'], ['flex', 'none']]) copy.style.setProperty(key, value, 'important');
|
||||
select.parentElement.append(copy);
|
||||
const needed = copy.getBoundingClientRect().width;
|
||||
copy.remove();
|
||||
return select.getBoundingClientRect().width + 0.5 >= needed;
|
||||
}),
|
||||
pageScroll: document.documentElement.scrollWidth > document.documentElement.clientWidth,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* The concept's Auto-tagging with the recorded deviations applied, for the fixture (all engines on disk, none loaded).
|
||||
* `phone` is a viewport at the concept's 700px breakpoint or below.
|
||||
*/
|
||||
function mapTaggingConcept(reference, { template, phone }) {
|
||||
const [model, existing, general, character] = reference.rows;
|
||||
const [preview] = reference.actions.buttons;
|
||||
return {
|
||||
...reference,
|
||||
// D133 adds a row after the thresholds; D134 adds the operation after the actions.
|
||||
tree: [...reference.tree.slice(0, 6), 'div.setting-row', ...reference.tree.slice(6), 'div.operation'],
|
||||
rows: [
|
||||
// D132 (OQ5, OQ6): the four real engines with their download size, and the chosen one's live status.
|
||||
{ ...model, text: [model.text[0], ['p', 'Downloaded. Not loaded.', null]],
|
||||
control: { ...model.control, options: MODEL_OPTIONS.map(label => [label, false]), selected: MODEL_OPTIONS[0] },
|
||||
// D132: at phone width the row wraps, so the engine select sits whole under its text.
|
||||
...(phone ? { flexWrap: 'wrap', controlBelow: true } : {}) },
|
||||
existing, general, character,
|
||||
// D133: the persisted suggestion count, on the concept's number row.
|
||||
{ ...general, text: [['label', 'Suggestions per image', 'tag-suggestions'], ['p', SUGGESTIONS, null]],
|
||||
control: { ...general.control, id: 'tag-suggestions', value: '10' } },
|
||||
],
|
||||
// The saved template (the server's default until one is saved) replaces the concept's sample text.
|
||||
field: { ...reference.field, textarea: { ...reference.field.textarea, value: template } },
|
||||
// D132: Download (for an engine already on disk, "Download again") and Unload replace the two simulated jobs.
|
||||
actions: { ...reference.actions, buttons: [preview, ['Download again', ''], ['Unload', '']] },
|
||||
// D134: the prompt-tag batch, reviewed first.
|
||||
operations: [{ tree: ['div', 'div.row-actions'], title: 'Tag output images', description: BATCH, actions: [['Review', '']] }],
|
||||
// D132: the engine select shows its whole chosen option at every width (at 390 the concept's 120px cap cuts it;
|
||||
// `#tag-model` alone may grow there). Existing tags keeps the concept's own fit, cut or not.
|
||||
selectsFit: reference.selectsFit.map((fit, index) => (index === 0 ? true : fit)),
|
||||
};
|
||||
}
|
||||
|
||||
async function conceptTagging(page, theme) {
|
||||
await openConcept(page, 'system', theme);
|
||||
await page.locator('#main .section-nav').getByRole('button', { name: 'Auto-tagging', exact: true }).click();
|
||||
await expect(page.locator('#main #tag-template')).toBeVisible();
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureTagging, '#main');
|
||||
}
|
||||
|
||||
async function productionTagging(page, theme) {
|
||||
await openProduction(page, 'system/tagging', theme);
|
||||
// Compare once the settings and the model status answered.
|
||||
await expect(page.locator('#workspace-main #tag-model')).toBeEnabled();
|
||||
await expect(page.locator('#workspace-main .setting-row').first().locator('p')).toHaveText('Downloaded. Not loaded.');
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureTagging, '#workspace-main');
|
||||
}
|
||||
|
||||
async function savedTemplate(request, page) {
|
||||
const response = await request.get('/prompt_manager/settings');
|
||||
await expect(response).toBeOK();
|
||||
return { template: (await response.json()).settings.autotag_template, phone: page.viewportSize().width <= 700 };
|
||||
}
|
||||
|
||||
for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
test(`auto-tagging matches the concept ${theme} ${width}x${height}`, async ({ page, request }, testInfo) => {
|
||||
await page.setViewportSize({ width, height });
|
||||
const concept = await conceptTagging(page, theme);
|
||||
await page.screenshot({ path: testInfo.outputPath(`concept-tagging-${theme}-${width}.png`), fullPage: true });
|
||||
const production = await productionTagging(page, theme);
|
||||
await page.screenshot({ path: testInfo.outputPath(`production-tagging-${theme}-${width}.png`), fullPage: true });
|
||||
expect(diffInventory(mapTaggingConcept(concept, await savedTemplate(request, page)), production)).toEqual([]);
|
||||
expect(production.pageScroll).toBe(false);
|
||||
});
|
||||
}
|
||||
|
||||
// D132 below the concept's smallest size: on a 320px phone the engine select still shows its whole option.
|
||||
for (const theme of THEMES) test(`auto-tagging engine select fits at 320 wide ${theme}`, async ({ page }) => {
|
||||
await page.setViewportSize({ width: 320, height: 700 });
|
||||
const production = await productionTagging(page, theme);
|
||||
for (const model of ['hf', 'gguf', 'wd14-swinv2']) {
|
||||
await page.locator('#workspace-main #tag-model').selectOption(model);
|
||||
expect((await page.evaluate(measureTagging, '#workspace-main')).selectsFit[0], model).toBe(true);
|
||||
}
|
||||
expect(production.pageScroll).toBe(false);
|
||||
});
|
||||
|
||||
test('auto-tagging comparison detects mutations', async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
const concept = await conceptTagging(page, 'dark');
|
||||
const saved = await savedTemplate(request, page);
|
||||
const expected = mapTaggingConcept(concept, saved);
|
||||
const production = await productionTagging(page, 'dark');
|
||||
expect(diffInventory(expected, production)).toEqual([]);
|
||||
// Unmapped, the section differs only where D132, D133 and D134 are recorded (and the saved template's text).
|
||||
expect(diffInventory(concept, production).map(item => item.path.split('.').slice(0, 3).join('.'))
|
||||
.filter((path, index, all) => all.indexOf(path) === index).sort()).toEqual([
|
||||
'actions.buttons.1', 'actions.buttons.2', 'field.textarea.value', 'operations.0.actions', 'operations.0.description',
|
||||
'operations.0.title', 'operations.0.tree', 'operations.length', 'rows.0.control', 'rows.0.text', 'rows.4.control', 'rows.4.controlBelow',
|
||||
'rows.4.flexWrap', 'rows.4.text', 'rows.4.tree', 'rows.length', 'tree.6', 'tree.7', 'tree.8', 'tree.9', 'tree.length'].sort());
|
||||
const measure = () => page.evaluate(measureTagging, '#workspace-main');
|
||||
const paths = async () => diffInventory(expected, await measure()).map(item => item.path);
|
||||
// A removed row.
|
||||
const undoRow = await mutate(page, { remove: ['#workspace-main section.surface > .setting-row:nth-of-type(2)'] });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['rows.length', 'tree.length']));
|
||||
await undoRow();
|
||||
// A renamed label.
|
||||
await page.evaluate(() => { document.querySelector('#workspace-main label[for="tag-existing"]').textContent = 'Existing'; });
|
||||
expect(await paths()).toEqual(['rows.1.text.0.1']);
|
||||
await page.evaluate(() => { document.querySelector('#workspace-main label[for="tag-existing"]').textContent = 'Existing tags'; });
|
||||
// A reordered option.
|
||||
await page.evaluate(() => { const select = document.querySelector('#workspace-main #tag-existing'); select.append(select.options[0]); });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['rows.1.control.options.0.0']));
|
||||
await page.evaluate(() => { const select = document.querySelector('#workspace-main #tag-existing'); select.prepend(select.options[2]); });
|
||||
// ca-primary moved to another action.
|
||||
await page.evaluate(() => {
|
||||
const [first, second] = document.querySelectorAll('#workspace-main section.surface > .row-actions > button');
|
||||
first.classList.remove('ca-primary');
|
||||
second.classList.add('ca-primary');
|
||||
});
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['actions.buttons.0.1', 'actions.buttons.1.1']));
|
||||
await page.evaluate(() => {
|
||||
const [first, second] = document.querySelectorAll('#workspace-main section.surface > .row-actions > button');
|
||||
second.classList.remove('ca-primary');
|
||||
first.classList.add('ca-primary');
|
||||
});
|
||||
// A select squeezed until its chosen option is cut.
|
||||
const undoSqueeze = await mutate(page, { css: '#workspace-main #tag-model { max-width: 60px !important; min-width: 0 !important; }' });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['selectsFit.0']));
|
||||
await undoSqueeze();
|
||||
// A row that wraps where the concept's does not.
|
||||
const undoWrap = await mutate(page, { css: '#workspace-main .setting-row:has(> #tag-existing) { flex-wrap: wrap; }' });
|
||||
expect(await paths()).toEqual(['rows.1.flexWrap']);
|
||||
await undoWrap();
|
||||
// A swapped token.
|
||||
const undoToken = await mutate(page, { css: '#workspace-main .settings-layout { --line: rgb(255, 0, 0); }' });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['row.borderTopColor']));
|
||||
await undoToken();
|
||||
expect(diffInventory(expected, await measure())).toEqual([]);
|
||||
});
|
||||
|
||||
// D132 at phone width, unmapped: only the engine row wraps; Existing tags keeps the concept's cut select.
|
||||
test('auto-tagging at 390 wraps only the engine row', async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 390, height: 844 });
|
||||
const concept = await conceptTagging(page, 'dark');
|
||||
const production = await productionTagging(page, 'dark');
|
||||
expect(diffInventory(mapTaggingConcept(concept, await savedTemplate(request, page)), production)).toEqual([]);
|
||||
expect(production.rows.map(row => row.flexWrap)).toEqual(['wrap', 'nowrap', 'nowrap', 'nowrap', 'nowrap']);
|
||||
expect(production.selectsFit[1]).toBe(concept.selectsFit[1]);
|
||||
});
|
||||
|
||||
/** Runs in the page: the open review dialog: its title, body and the concept's tokens. */
|
||||
function measureReview() {
|
||||
const dialog = document.querySelector('dialog#modal[open]');
|
||||
const css = (node, ...keys) => node ? Object.fromEntries(keys.map(key => [key, getComputedStyle(node)[key]])) : null;
|
||||
const name = node => [node.localName, ...[...node.classList].sort()].join('.');
|
||||
const text = node => node?.textContent.trim() ?? null;
|
||||
const body = dialog.querySelector('#modal-body');
|
||||
const grid = body.querySelector(':scope > dl.meta-grid');
|
||||
const note = [...body.querySelectorAll(':scope > p.small')].at(-1);
|
||||
const actions = body.querySelector(':scope > .dialog-actions');
|
||||
return {
|
||||
title: text(dialog.querySelector('#modal-title')),
|
||||
tree: [...body.children].map(name),
|
||||
heading: { text: text(body.querySelector(':scope > h3')),
|
||||
...css(body.querySelector(':scope > h3'), 'fontSize', 'fontWeight', 'color', 'marginBottom') },
|
||||
grid: { rows: [...grid.children].map(row => [text(row.querySelector('dt')), text(row.querySelector('dd'))]),
|
||||
...css(grid, 'display', 'gridTemplateColumns', 'columnGap', 'rowGap'),
|
||||
dt: css(grid.querySelector('dt'), 'fontSize', 'color', 'textTransform'), dd: css(grid.querySelector('dd'), 'fontSize', 'color') },
|
||||
backup: body.querySelector(':scope > .notice') && { tree: [...body.querySelector(':scope > .notice').children].map(name),
|
||||
text: [...body.querySelectorAll(':scope > .notice > div:first-child > *')].map(node => [node.localName, text(node)]),
|
||||
buttons: [...body.querySelectorAll(':scope > .notice button')].map(text) },
|
||||
note: { text: text(note), ...css(note, 'fontSize', 'color', 'marginTop') },
|
||||
actions: { tree: [...actions.children].map(name), ...css(actions, 'display', 'justifyContent', 'columnGap', 'marginTop'),
|
||||
buttons: [...actions.children].map(item => ({ text: text(item), className: item.className,
|
||||
...css(item, 'backgroundColor', 'color', 'fontSize', 'paddingTop', 'paddingLeft', 'borderTopColor') })) },
|
||||
};
|
||||
}
|
||||
|
||||
const number = value => value.toLocaleString('en-US');
|
||||
|
||||
/**
|
||||
* D134: the server's batch preview in the concept's review: "Review tag output images" with "Tag output images" as its
|
||||
* `h3`, the model, policy, counts and settings in the concept's `dl.meta-grid` (Backup in the concept's wording), the
|
||||
* upper-bound sentence in the concept's `p.small`, the D55 backup box and Cancel / "Start tagging".
|
||||
*/
|
||||
function mapBatchConcept(reference, plan) {
|
||||
const [cancel, start] = reference.actions.buttons;
|
||||
return {
|
||||
...reference,
|
||||
title: 'Review tag output images',
|
||||
tree: ['h3', 'dl.meta-grid', 'p.small', 'div.notice', 'div.dialog-actions'],
|
||||
heading: { ...reference.heading, text: 'Tag output images' },
|
||||
grid: { ...reference.grid, rows: [['Model', 'WD14 SwinV2 · ~400 MB'], ['Existing tags', 'Merge with existing'],
|
||||
['Images linked', number(plan.images)], ['Prompts', number(plan.prompts)], ['Already tagged', number(plan.already_tagged)],
|
||||
['Prompts to tag', `up to ${number(plan.to_tag)}`], ['Settings', 'General 0.35 · Character 0.85'],
|
||||
['Backup', 'Create and verify before changes']] },
|
||||
backup: { tree: ['div', 'div.row-actions'], text: [['h3', 'Back up the database first (recommended)'],
|
||||
['p', 'Downloads a copy of the prompts database before this operation changes any records.']], buttons: ['Back up now'] },
|
||||
note: { ...reference.note, text: 'Nothing is written until you start. Missing files and files outside the output folders '
|
||||
+ 'are counted, so the prompts to tag are an upper bound; Activity reports what was tagged.' },
|
||||
actions: { ...reference.actions, buttons: [cancel, { ...start, text: 'Start tagging' }] },
|
||||
};
|
||||
}
|
||||
|
||||
async function conceptReview(page, theme) {
|
||||
await openConcept(page, 'system', theme);
|
||||
await page.locator('#main .section-nav').getByRole('button', { name: 'Maintenance', exact: true }).click();
|
||||
await page.locator('#main [data-plan="duplicate prompts"]').click();
|
||||
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Review maintenance plan');
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureReview);
|
||||
}
|
||||
|
||||
async function productionReview(page, theme) {
|
||||
await productionTagging(page, theme);
|
||||
// Every concept-fixture prompt with an image is already tagged, so under the saved Skip there is nothing to start
|
||||
// (the Cancel-only review). Merge reviews a batch with work: the full review the concept's shell is compared with.
|
||||
await page.locator('#workspace-main #tag-existing').selectOption('merge');
|
||||
const answer = page.waitForResponse(response => response.url().endsWith(PREVIEW_URL));
|
||||
await page.locator('#workspace-main .operation').getByRole('button', { name: 'Review', exact: true }).click();
|
||||
const plan = await (await answer).json();
|
||||
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Review tag output images');
|
||||
await page.mouse.move(0, 0);
|
||||
return { plan, production: await page.evaluate(measureReview) };
|
||||
}
|
||||
|
||||
for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
test(`tag output images review matches the concept ${theme} ${width}x${height}`, async ({ page }, testInfo) => {
|
||||
await page.setViewportSize({ width, height });
|
||||
const concept = await conceptReview(page, theme);
|
||||
expect(concept.tree).toEqual(['p.demo-note', 'h3', 'dl.meta-grid', 'p.small', 'div.dialog-actions']);
|
||||
await page.screenshot({ path: testInfo.outputPath(`concept-plan-review-${theme}-${width}.png`) });
|
||||
const { plan, production } = await productionReview(page, theme);
|
||||
// Under Merge every prompt with a linked image is tagged again: the batch has work to review.
|
||||
expect(plan.to_tag).toBeGreaterThan(0);
|
||||
await page.screenshot({ path: testInfo.outputPath(`production-tagging-review-${theme}-${width}.png`) });
|
||||
expect(diffInventory(mapBatchConcept(concept, plan), production)).toEqual([]);
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* D134 (CUJ 5): Replace after review never runs as a batch: the same review shell with the counts, the reason as its
|
||||
* `p.small`, no backup box and only Cancel (the shared close-only review Integrations also uses).
|
||||
*/
|
||||
function mapCancelOnlyConcept(reference, plan) {
|
||||
const [cancel] = reference.actions.buttons;
|
||||
const mapped = mapBatchConcept(reference, plan);
|
||||
return {
|
||||
...mapped,
|
||||
tree: ['h3', 'dl.meta-grid', 'p.small', 'div.dialog-actions'],
|
||||
grid: { ...mapped.grid, rows: mapped.grid.rows.map(([term, value]) => [term, term === 'Existing tags' ? 'Replace after review' : value]) },
|
||||
backup: null,
|
||||
note: { ...reference.note, text: 'Replacing tags needs a review of each image. Use Preview tagging review, or choose Skip or '
|
||||
+ 'Merge for a batch run.' },
|
||||
actions: { ...reference.actions, tree: [reference.actions.tree[0]], buttons: [cancel] },
|
||||
};
|
||||
}
|
||||
|
||||
test('the Cancel-only review matches the concept review shell', async ({ page }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
const concept = await conceptReview(page, 'dark');
|
||||
await productionTagging(page, 'dark');
|
||||
await page.locator('#workspace-main #tag-existing').selectOption('replace_after_review');
|
||||
const answer = page.waitForResponse(response => response.url().endsWith(PREVIEW_URL));
|
||||
await page.locator('#workspace-main .operation').getByRole('button', { name: 'Review', exact: true }).click();
|
||||
const plan = await (await answer).json();
|
||||
expect(plan.review_required).toBe(true);
|
||||
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Review tag output images');
|
||||
await page.mouse.move(0, 0);
|
||||
expect(diffInventory(mapCancelOnlyConcept(concept, plan), await page.evaluate(measureReview))).toEqual([]);
|
||||
});
|
||||
|
||||
test('tag output images review comparison detects mutations', async ({ page }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
const concept = await conceptReview(page, 'dark');
|
||||
const { plan } = await productionReview(page, 'dark');
|
||||
const expected = mapBatchConcept(concept, plan);
|
||||
const paths = async () => diffInventory(expected, await page.evaluate(measureReview)).map(item => item.path);
|
||||
expect(await paths()).toEqual([]);
|
||||
await page.evaluate(() => { document.querySelectorAll('dialog#modal[open] dl.meta-grid dt')[5].textContent = 'To tag'; });
|
||||
expect(await paths()).toEqual(['grid.rows.5.0']);
|
||||
await page.evaluate(() => { document.querySelectorAll('dialog#modal[open] dl.meta-grid dt')[5].textContent = 'Prompts to tag'; });
|
||||
const undoGrid = await mutate(page, { css: 'dialog#modal dl.meta-grid { column-gap: 3px; }' });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['grid.columnGap']));
|
||||
await undoGrid();
|
||||
const undoBox = await mutate(page, { remove: ['dialog#modal[open] #modal-body > .notice'] });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['backup', 'tree.length']));
|
||||
await undoBox();
|
||||
expect(await paths()).toEqual([]);
|
||||
});
|
||||
@@ -236,7 +236,7 @@ test('paired Tag dialog captures', async ({ page, request }, testInfo) => {
|
||||
await expect(await request.post('/__test__/tagging/models', { data: { unavailable: ['wd14-swinv2'] } })).toBeOK();
|
||||
await openTagDialog(page, 'production', 'dark');
|
||||
await generateWith(page, 'production').catch(() => {});
|
||||
await expect(status()).toHaveText('WD14 SwinV2 is not downloaded. Download it in System › Auto-tagging.');
|
||||
await expect(status()).toHaveText('WD14 SwinV2 is not downloaded. Download it from the legacy Auto-tag page.');
|
||||
await shot('production-unavailable');
|
||||
await expect(await request.post('/__test__/tagging/models', { data: { unavailable: [] } })).toBeOK();
|
||||
|
||||
|
||||
@@ -19,13 +19,6 @@ test.afterEach(async ({ page, browserName }) => {
|
||||
expect(page.workflowErrors).toEqual([]);
|
||||
});
|
||||
|
||||
// #249 S6b: the legacy admin page says where the library lives (the fixture's database is configured explicitly).
|
||||
test('the admin page names the library location', async ({ page }) => {
|
||||
await page.goto('/prompt_manager/admin');
|
||||
await expect(page.locator('#libraryLocation')).toHaveText('Library: browser.db (set in config.json)');
|
||||
await expect(page.locator('#totalPrompts')).not.toHaveText('-');
|
||||
});
|
||||
|
||||
test('actual gallery form saves a prompt that the admin page can edit and reload', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await page.locator('#addPromptBtn').click();
|
||||
@@ -99,34 +92,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'] });
|
||||
});
|
||||
|
||||
test('actual gallery review generates and persists tags for the selected root', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(2);
|
||||
@@ -257,18 +222,6 @@ async function chooseRestore(page, file) {
|
||||
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');
|
||||
@@ -299,7 +252,12 @@ test('downloaded backup restores deleted prompts, tags, and image links through
|
||||
expect(await promptSnapshot(request)).not.toEqual(before);
|
||||
await chooseRestore(page, backup);
|
||||
confirmNextDialog(page, true, backup.name);
|
||||
await confirmRestoreAndWait(page);
|
||||
const restored = page.waitForResponse(response => response.url().endsWith('/prompt_manager/restore'));
|
||||
await page.locator('#confirmRestore').click();
|
||||
const response = await restored;
|
||||
expect(response.ok()).toBe(true);
|
||||
expect(await response.json()).toMatchObject({ success: true, prompt_count: 2, backup_created: expect.any(String) });
|
||||
await expect(page.locator('#restoreModal')).toBeHidden();
|
||||
await expect(page.locator('.prompt-text')).toHaveCount(2);
|
||||
await page.reload();
|
||||
await expect(page.locator('.prompt-text').filter({ hasText: 'Cat portrait' })).toBeVisible();
|
||||
@@ -342,7 +300,8 @@ test('restore validation and cancellation preserve live data and allow a valid r
|
||||
expect(await promptSnapshot(request)).toEqual(before);
|
||||
await page.locator('#restoreFileInput').setInputFiles(backup);
|
||||
confirmNextDialog(page, true, backup.name);
|
||||
await confirmRestoreAndWait(page);
|
||||
await page.locator('#confirmRestore').click();
|
||||
await expect(page.locator('#restoreModal')).toBeHidden();
|
||||
expect(await promptSnapshot(request)).toEqual(before);
|
||||
});
|
||||
|
||||
@@ -362,7 +321,8 @@ test('backup and restore recover from interrupted network requests without losin
|
||||
await expect(page.locator('#restoreModal')).toBeVisible();
|
||||
expect(await promptSnapshot(request)).toEqual(before);
|
||||
confirmNextDialog(page, true, backup.name);
|
||||
await confirmRestoreAndWait(page);
|
||||
await page.locator('#confirmRestore').click();
|
||||
await expect(page.locator('#restoreModal')).toBeHidden();
|
||||
expect(await promptSnapshot(request)).toEqual(before);
|
||||
});
|
||||
|
||||
|
||||
@@ -1088,36 +1088,6 @@ test('a held legacy job and the health check show no tiles, no Results and no ma
|
||||
await expect(toast(page)).toHaveText('Health check finished.');
|
||||
});
|
||||
|
||||
// #328: a tagging model download is a job of its own kind with no byte progress and no Cancel; Activity shows it
|
||||
// through the generic held view (its title, status line and an indeterminate bar), then its end toast.
|
||||
test('a held model download shows its title, an indeterminate bar and no Cancel; its end toast follows', async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
const title = 'Tagging model download · WD14 ViT';
|
||||
await expect(await request.post('/__test__/tagging/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/tagging/download', { data: { mode: 'hold' } })).toBeOK();
|
||||
try {
|
||||
const response = await request.post('/prompt_manager/jobs/model-download', { data: { model_type: 'wd14-vit' } });
|
||||
expect(response.status()).toBe(202);
|
||||
const job = (await response.json()).job;
|
||||
expect([job.kind, job.title, job.cancellable]).toEqual(['model_download', title, false]);
|
||||
await page.goto(`${workspace}#/activity`);
|
||||
const item = items(page).filter({ has: page.getByRole('link', { name: new RegExp(`${title}$`) }) });
|
||||
await expect(item.locator('progress')).not.toHaveAttribute('value');
|
||||
await expect(item.getByRole('button', { name: 'Cancel job', exact: true })).toHaveCount(0);
|
||||
await item.getByRole('link').click();
|
||||
await expect(heading(page).locator('h1')).toHaveText(title);
|
||||
await expect(progressSection(page).locator('.panel-heading > span.small')).toHaveText('Downloading WD14 ViT (~400 MB)');
|
||||
await expect(progressSection(page).locator('progress')).not.toHaveAttribute('value');
|
||||
await expect(progressSection(page).locator('dl.stat-row')).toHaveCount(0);
|
||||
await expect(heading(page).getByRole('button', { name: /^Cancel/ })).toHaveCount(0);
|
||||
await expect(await request.post('/__test__/tagging/download', { data: { mode: 'succeed' } })).toBeOK();
|
||||
await expect(toast(page)).toHaveText(`${title} finished.`);
|
||||
await expect(heading(page).locator('p.muted')).toHaveText(/^Finished \d\d:\d\d$/);
|
||||
} finally {
|
||||
await request.post('/__test__/tagging/reset');
|
||||
}
|
||||
});
|
||||
|
||||
// S5 review I3: a thumbnail job with nothing to generate ends on the server's final line ("No previews to generate · N
|
||||
// already up to date"), never on a stale "Processing file 0 of 0" under the Completed chip.
|
||||
test('a finished thumbnail job with nothing to do says so', async ({ page, request }) => {
|
||||
|
||||
@@ -398,20 +398,18 @@ test('preview uses actual Viewer.js and cleans up when closing or leaving Images
|
||||
test('real video media has a usable card and native playback cleans up on close and navigation', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { include_video: true } })).toBeOK();
|
||||
const media = await (await request.get('/prompt_manager/images/output?limit=40')).json();
|
||||
// #249: a video with no preview yet still lists its preview URL; the server makes the frame on first view.
|
||||
expect(media.images.find(image => image.filename === 'clip.webm').thumbnail_url).toMatch(/\/thumbnails\/clip_thumb\.webp\?root=0/);
|
||||
expect(media.images.find(image => image.filename === 'clip.webm').thumbnail_url).toBeNull();
|
||||
await page.goto(workspace);
|
||||
await refine(page, { Media: 'video' });
|
||||
await expect(cards(page)).toHaveCount(1);
|
||||
const card = cards(page).first();
|
||||
await expect(card).toHaveAttribute('data-image-key', '0:clip.webm');
|
||||
await expect(card.locator('img')).toHaveCount(0);
|
||||
await expect(card.getByRole('button', { name: 'Open clip.webm', exact: true })).toBeVisible();
|
||||
// The card shows the preview frame, never the video itself; D11 names the video's own extension.
|
||||
await expect(card.locator('video')).toHaveCount(0);
|
||||
await expect(card.locator('.media-art > .image-open > img')).toHaveCount(1);
|
||||
// Legacy video art inside the concept card; D11 names the video's own extension.
|
||||
await expect(card.locator('.media-art > .image-open > video')).toHaveCount(1);
|
||||
await expect(card.locator('.media-info p')).toHaveText(/(?:^|· )WEBM$/);
|
||||
await expect.poll(() => card.locator('img').evaluate(node => node.complete && node.naturalHeight > 0)).toBe(true);
|
||||
expect(await card.locator('img').evaluate(node => node.getBoundingClientRect().height)).toBeGreaterThan(20);
|
||||
expect(await card.locator('video').evaluate(node => node.getBoundingClientRect().height)).toBeGreaterThan(20);
|
||||
await card.getByRole('button', { name: 'Open clip.webm', exact: true }).click();
|
||||
const video = page.locator('dialog[open] video');
|
||||
await expect(video).toBeVisible();
|
||||
@@ -715,45 +713,16 @@ test('selected layout stays blue with readable text on hover in either theme', a
|
||||
}
|
||||
});
|
||||
|
||||
test('refresh releases the previous inspector cover source', async ({ page, request }) => {
|
||||
test('refresh releases the previous inspector video source', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { include_video: true } })).toBeOK();
|
||||
await page.goto(workspace);
|
||||
const card = cards(page).filter({ hasText: 'clip.webm' });
|
||||
await card.getByRole('button', { name: 'Inspect clip.webm', exact: true }).click();
|
||||
// #249: a video's cover is its preview frame, not the video.
|
||||
const cover = details(page).locator('img.detail-cover');
|
||||
await expect(cover).toHaveAttribute('src', /\/thumbnails\/clip_thumb\.webp\?/);
|
||||
await expect(details(page).locator('video')).toHaveCount(0);
|
||||
await cover.evaluate(element => { window.__oldInspectorCover = element; });
|
||||
await rescan(page);
|
||||
await expect.poll(() => page.evaluate(() => !window.__oldInspectorCover.isConnected)).toBe(true);
|
||||
expect(await page.evaluate(() => window.__oldInspectorCover.getAttribute('src'))).toBeNull();
|
||||
});
|
||||
|
||||
/** A listing that names no preview for its videos (the fallback a card and the cover keep for one). */
|
||||
async function videosWithoutPreview(page) {
|
||||
await page.route(url => url.pathname === '/prompt_manager/images/output', async route => {
|
||||
const response = await route.fetch();
|
||||
const body = await response.json();
|
||||
if (Array.isArray(body.images)) body.images = body.images.map(image => image.is_video ? { ...image, thumbnail_url: null } : image);
|
||||
await route.fulfill({ response, json: body });
|
||||
});
|
||||
}
|
||||
|
||||
test('a video listed without a preview keeps its own element, released on refresh', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { include_video: true } })).toBeOK();
|
||||
await videosWithoutPreview(page);
|
||||
await page.goto(workspace);
|
||||
const card = cards(page).filter({ hasText: 'clip.webm' });
|
||||
await expect(card.locator('.media-art > .image-open > video')).toHaveCount(1);
|
||||
await expect(card.locator('img')).toHaveCount(0);
|
||||
await card.getByRole('button', { name: 'Inspect clip.webm', exact: true }).click();
|
||||
await expect(details(page).locator('video')).toBeVisible();
|
||||
await details(page).locator('video').evaluate(element => { window.__oldInspectorVideo = element; });
|
||||
await card.locator('video').evaluate(element => { window.__oldCardVideo = element; });
|
||||
await rescan(page);
|
||||
await expect.poll(() => page.evaluate(() => !window.__oldInspectorVideo.isConnected && !window.__oldCardVideo.isConnected)).toBe(true);
|
||||
expect(await page.evaluate(() => [window.__oldInspectorVideo.getAttribute('src'), window.__oldCardVideo.getAttribute('src')])).toEqual([null, null]);
|
||||
await expect.poll(() => page.evaluate(() => !window.__oldInspectorVideo.isConnected)).toBe(true);
|
||||
expect(await page.evaluate(() => window.__oldInspectorVideo.getAttribute('src'))).toBeNull();
|
||||
});
|
||||
|
||||
test('Import Rescan rescans cached files and preserves the collection when rescan fails', async ({ page, request }) => {
|
||||
@@ -2184,8 +2153,6 @@ test('scan for metadata opens the review dialog from Import', async ({ page, req
|
||||
await dialog.getByRole('button', { name: 'Cancel', exact: true }).click();
|
||||
await expect(scanReview(page)).toHaveCount(0);
|
||||
await expect(page.getByRole('button', { name: 'Import', exact: true })).toBeFocused();
|
||||
// No scan started: the one job is the backup's finished registry hold
|
||||
// (#249 S8: a backup download holds a guarded slot while it streams).
|
||||
await expect.poll(async () => (await (await request.get('/prompt_manager/jobs')).json()).jobs
|
||||
.map(job => [job.kind, job.state])).toEqual([['backup', 'finished']]);
|
||||
const jobs = await (await request.get('/prompt_manager/jobs')).json();
|
||||
expect(jobs.jobs).toEqual([]);
|
||||
});
|
||||
|
||||
@@ -1,490 +0,0 @@
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from './fixtures.mjs';
|
||||
import { seedConcept } from './concept-inventory.mjs';
|
||||
|
||||
// System › Backup & export (#249 System S9 Task 24) on the S8 backend and the concept-shaped fixture: Back up and its
|
||||
// status line (D103), the restore review (D105) and its job, the exports, Import / inspect and the recovery policy
|
||||
// (D104). A restore here replaces this worker's own temp database; every test starts from a fresh concept fixture.
|
||||
const workspace = '/prompt_manager/workspace';
|
||||
const STATUS_URL = '/prompt_manager/backup/status';
|
||||
const PREVIEW_URL = '/prompt_manager/restore/preview';
|
||||
const APPLY_URL = '/prompt_manager/restore/apply';
|
||||
const WHEN = { dateStyle: 'medium', timeStyle: 'short' };
|
||||
const SAFETY_DESCRIPTION = 'Keep a verified copy of the database on disk before maintenance, cleanups and LoRA imports change ' +
|
||||
'records. A restore always keeps one.';
|
||||
|
||||
test.use({ locale: 'en-US', timezoneId: 'UTC' });
|
||||
|
||||
test.beforeEach(async ({ page, request, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
await seedConcept(request);
|
||||
await expect(await request.post('/__test__/system/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/jobs/restart')).toBeOK();
|
||||
page.systemErrors = [];
|
||||
page.on('pageerror', error => page.systemErrors.push(error.message));
|
||||
});
|
||||
|
||||
test.afterEach(async ({ page, request, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
await request.post('/__test__/release');
|
||||
await request.post('/__test__/jobs/restart');
|
||||
await request.post('/__test__/system/reset');
|
||||
expect(page.systemErrors).toEqual([]);
|
||||
});
|
||||
|
||||
const toast = page => page.locator('#workspace-toast');
|
||||
const layoutOf = page => page.locator('#workspace-main .settings-layout');
|
||||
const dialog = page => page.locator('dialog#modal[open]');
|
||||
const statusLine = page => layoutOf(page).locator('#backup-status');
|
||||
const settingsStatus = page => page.locator('#settings-status');
|
||||
const save = page => page.getByRole('button', { name: 'Save preferences', exact: true });
|
||||
const buttonOf = (page, name) => layoutOf(page).getByRole('button', { name, exact: true });
|
||||
const settingOf = (page, label) => layoutOf(page).locator('.setting-row').filter({ has: page.locator(`label:text-is("${label}")`) });
|
||||
|
||||
async function serverJson(request, url) {
|
||||
const response = await request.get(url);
|
||||
await expect(response).toBeOK();
|
||||
return response.json();
|
||||
}
|
||||
|
||||
async function openBackups(page) {
|
||||
// A fresh page each time: the drafts live for the page session (Q10), and a hash-only goto would not reload.
|
||||
if (page.url().endsWith('#/system/backups')) await page.reload();
|
||||
else await page.goto(`${workspace}#/system/backups`);
|
||||
await expect(layoutOf(page).locator(':scope > .two-col > section.surface')).toHaveCount(2);
|
||||
await expect(statusLine(page)).not.toHaveText('Checking backups…');
|
||||
// The policy row is painted once the server settings were read.
|
||||
await expect(settingOf(page, 'Create a safety backup before maintenance').locator('input')).toBeEnabled();
|
||||
return layoutOf(page);
|
||||
}
|
||||
|
||||
/** The status line the page should show for the server's current status (its own wording, in this locale). */
|
||||
async function expectedStatus(page, request) {
|
||||
const status = await serverJson(request, STATUS_URL);
|
||||
return page.evaluate(([answer, options]) => {
|
||||
const when = value => new Date(value).toLocaleString(undefined, options);
|
||||
const parts = [];
|
||||
if (answer.last_download) parts.push(`Downloaded backup · ${when(answer.last_download.at)} · ${answer.last_download.verified ? 'verified' : 'not verified'}.`);
|
||||
if (answer.safety_backups.length) parts.push(`Safety backup · ${when(answer.safety_backups[0].modified_at)} · ${answer.safety_backups.length} on disk.`);
|
||||
return parts.join(' ') || 'No backup recorded.';
|
||||
}, [status, WHEN]);
|
||||
}
|
||||
|
||||
/** A verified copy of the library as it is now (this records the session's last download, as Back up does). */
|
||||
async function backupBytes(request) {
|
||||
const response = await request.get('/prompt_manager/backup');
|
||||
await expect(response).toBeOK();
|
||||
expect(response.headers()['x-promptmanager-verify']).toBe('ok');
|
||||
return response.body();
|
||||
}
|
||||
|
||||
/** Pick `buffer` as the backup file through Review restore…'s file chooser; resolves with the preview's answer. */
|
||||
async function chooseBackup(page, buffer, name = 'prompts_backup_20261025_093000.db') {
|
||||
const chooser = page.waitForEvent('filechooser');
|
||||
const answer = page.waitForResponse(response => response.url().endsWith(PREVIEW_URL) && response.request().method() === 'POST');
|
||||
await buttonOf(page, 'Review restore…').click();
|
||||
await (await chooser).setFiles({ name, mimeType: 'application/octet-stream', buffer });
|
||||
return answer;
|
||||
}
|
||||
|
||||
const gridOf = page => dialog(page).locator('dl.meta-grid > div').evaluateAll(items => items.map(item =>
|
||||
[item.querySelector('dt').textContent, item.querySelector('dd').textContent]));
|
||||
|
||||
const libraryTotal = async request => (await serverJson(request, '/prompt_manager/library?limit=1')).total;
|
||||
|
||||
test('Back up downloads through the api client and revokes the object URL (D103)', async ({ page, request }) => {
|
||||
await page.addInitScript(() => {
|
||||
window.objectUrls = { created: [], revoked: [] };
|
||||
const create = URL.createObjectURL.bind(URL), revoke = URL.revokeObjectURL.bind(URL);
|
||||
URL.createObjectURL = blob => { const url = create(blob); window.objectUrls.created.push(url); return url; };
|
||||
URL.revokeObjectURL = url => { window.objectUrls.revoked.push(url); revoke(url); };
|
||||
});
|
||||
await openBackups(page);
|
||||
await expect(statusLine(page)).toHaveText('No backup recorded.');
|
||||
const download = page.waitForEvent('download');
|
||||
const answer = page.waitForResponse(response => new URL(response.url()).pathname === '/prompt_manager/backup');
|
||||
await buttonOf(page, 'Back up').click();
|
||||
expect((await answer).headers()['x-promptmanager-verify']).toBe('ok');
|
||||
expect((await download).suggestedFilename()).toMatch(/^prompts_backup_\d{8}_\d{6}\.db$/);
|
||||
await expect(toast(page)).toHaveText('Database backup downloaded.');
|
||||
// The status line is read again: this session's verified download, as the server records it.
|
||||
await expect(statusLine(page)).toHaveText(/^Downloaded backup · .+ · verified\.$/);
|
||||
await expect(statusLine(page)).toHaveText(await expectedStatus(page, request));
|
||||
await expect(buttonOf(page, 'Back up')).not.toHaveAttribute('aria-disabled', 'true');
|
||||
const urls = await page.evaluate(() => window.objectUrls.created);
|
||||
expect(urls).toHaveLength(1);
|
||||
await expect.poll(() => page.evaluate(() => window.objectUrls.revoked)).toEqual(urls);
|
||||
});
|
||||
|
||||
test('status line shows only server-reported backups (D103)', async ({ page, request }) => {
|
||||
await openBackups(page);
|
||||
await expect(statusLine(page)).toHaveText('No backup recorded.');
|
||||
await expect(await request.post('/__test__/system/safety-backup')).toBeOK();
|
||||
await openBackups(page);
|
||||
await expect(statusLine(page)).toHaveText(/^Safety backup · .+ · 1 on disk\.$/);
|
||||
await expect(statusLine(page)).toHaveText(await expectedStatus(page, request));
|
||||
// A status the server could not report is never shown as "no backup".
|
||||
await page.route(`**${STATUS_URL}`, route => route.fulfill({ status: 500, json: { success: false, error: 'Failed' } }));
|
||||
await page.reload();
|
||||
await expect(statusLine(page)).toHaveText('The backup status could not be read.');
|
||||
});
|
||||
|
||||
test('Overview backup row uses the same status (D93)', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/safety-backup')).toBeOK();
|
||||
await backupBytes(request);
|
||||
await page.goto(`${workspace}#/system/overview`);
|
||||
const row = layoutOf(page).locator('.operation').filter({ has: page.locator('h3', { hasText: 'Backup status' }) });
|
||||
const expected = await expectedStatus(page, request);
|
||||
expect(expected).toMatch(/^Downloaded backup · .+ · verified\. Safety backup · .+ · 1 on disk\.$/);
|
||||
await expect(row.locator('p')).toHaveText(expected);
|
||||
await row.getByRole('button', { name: 'View backups', exact: true }).click();
|
||||
await expect(page).toHaveURL(`${workspace}#/system/backups`);
|
||||
await expect(statusLine(page)).toHaveText(expected);
|
||||
});
|
||||
|
||||
test('Review restore stages the chosen file and shows the real summary with the backup box (D105)', async ({ page, request }) => {
|
||||
const bytes = await backupBytes(request);
|
||||
await expect(await request.post('/prompt_manager/save', { data: { text: 'S9 prompt saved after the backup' } })).toBeOK();
|
||||
const before = await serverJson(request, '/__test__/system/digest');
|
||||
await openBackups(page);
|
||||
const answer = await chooseBackup(page, bytes);
|
||||
const preview = await answer.json();
|
||||
expect(preview.integrity).toBe('ok');
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Restore preview');
|
||||
const counts = summary => `${summary.prompts.toLocaleString('en-US')} ${summary.prompts === 1 ? 'prompt' : 'prompts'} · ` +
|
||||
`${summary.image_links} ${summary.image_links === 1 ? 'image link' : 'image links'} · ` +
|
||||
`${summary.pairs} ${summary.pairs === 1 ? 'saved pair' : 'saved pairs'} · ${summary.runs} ${summary.runs === 1 ? 'run' : 'runs'}`;
|
||||
const rows = await gridOf(page);
|
||||
expect(rows.map(([term]) => term)).toEqual(['Snapshot', 'Integrity', 'Migration', 'Safety backup', 'In the backup', 'In the library now']);
|
||||
expect(rows[0][1]).toMatch(/^prompts_backup_20261025_093000.db · (\d+ KB|\d+\.\d MB) · .+$/);
|
||||
expect(rows.slice(1)).toEqual([['Integrity', 'Verified'], ['Migration', 'Current schema'],
|
||||
['Safety backup', `Created before replacing (${preview.safety_backup.destination_label})`],
|
||||
['In the backup', counts(preview.summary)], ['In the library now', counts(preview.live)]]);
|
||||
// The prompt saved after the backup is the difference the review shows.
|
||||
expect(preview.live.prompts).toBe(preview.summary.prompts + 1);
|
||||
await expect(dialog(page).locator('#modal-body > p.small')).toHaveText('Restore replaces the whole library with this backup ' +
|
||||
'and runs alone in Activity. Nothing changes until you restore. This review holds for 15 minutes.');
|
||||
await expect(dialog(page).locator('.notice h3')).toHaveText('Back up the database first (recommended)');
|
||||
await expect(dialog(page).locator('.notice p')).toHaveText('Downloads a copy of the prompts database before the restore replaces it.');
|
||||
const actions = dialog(page).locator('.dialog-actions > button');
|
||||
await expect(actions).toHaveText(['Cancel', 'Restore']);
|
||||
await expect(actions.nth(1)).toHaveClass('danger');
|
||||
// The backup box downloads a copy and keeps the review open.
|
||||
const download = page.waitForEvent('download');
|
||||
await dialog(page).getByRole('button', { name: 'Back up now', exact: true }).click();
|
||||
await download;
|
||||
await expect(toast(page)).toHaveText('Database backup downloaded.');
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Restore preview');
|
||||
// A review writes nothing: the rows and the media are as they were.
|
||||
expect(await serverJson(request, '/__test__/system/digest')).toEqual(before);
|
||||
expect((await serverJson(request, '/__test__/system/backups')).stages).toBe(1);
|
||||
});
|
||||
|
||||
test('Cancel discards the stage; Escape too', async ({ page, request }) => {
|
||||
const bytes = await backupBytes(request);
|
||||
const before = await serverJson(request, '/__test__/system/digest');
|
||||
await openBackups(page);
|
||||
await chooseBackup(page, bytes);
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Restore preview');
|
||||
expect((await serverJson(request, '/__test__/system/backups')).stages).toBe(1);
|
||||
let discarded = page.waitForResponse(response => response.url().includes(`${PREVIEW_URL}/`) && response.request().method() === 'DELETE');
|
||||
await dialog(page).getByRole('button', { name: 'Cancel', exact: true }).click();
|
||||
expect((await discarded).status()).toBe(200);
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
await expect(buttonOf(page, 'Review restore…')).toBeFocused();
|
||||
expect((await serverJson(request, '/__test__/system/backups')).stages).toBe(0);
|
||||
await chooseBackup(page, bytes);
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Restore preview');
|
||||
discarded = page.waitForResponse(response => response.url().includes(`${PREVIEW_URL}/`) && response.request().method() === 'DELETE');
|
||||
await page.keyboard.press('Escape');
|
||||
expect((await discarded).status()).toBe(200);
|
||||
expect((await serverJson(request, '/__test__/system/backups')).stages).toBe(0);
|
||||
expect(await serverJson(request, '/__test__/system/digest')).toEqual(before);
|
||||
});
|
||||
|
||||
test('Restore starts the job, opens Activity; views reload after it finishes', async ({ page, request }) => {
|
||||
const bytes = await backupBytes(request);
|
||||
const restoredTotal = await libraryTotal(request);
|
||||
await expect(await request.post('/prompt_manager/save', { data: { text: 'S9 prompt the restore removes' } })).toBeOK();
|
||||
expect(await libraryTotal(request)).toBe(restoredTotal + 1);
|
||||
await openBackups(page);
|
||||
const libraryCount = page.locator('aside.sidebar a[data-route="library"] .nav-count');
|
||||
await expect(libraryCount).toHaveText((restoredTotal + 1).toLocaleString('en-US'));
|
||||
await chooseBackup(page, bytes);
|
||||
const applied = page.waitForResponse(response => response.url().endsWith(APPLY_URL));
|
||||
await dialog(page).getByRole('button', { name: 'Restore', exact: true }).click();
|
||||
const { job } = await (await applied).json();
|
||||
expect([job.kind, job.title, job.cancellable]).toEqual(['restore', 'Database restore', false]);
|
||||
await expect(page).toHaveURL(`${workspace}#/activity/${job.id}`);
|
||||
await expect(toast(page)).toHaveText(/^Database restore finished · Restored \d+ prompts?, /);
|
||||
// The library is the backup's again, and the page says so without a reload: the sidebar counts were read again.
|
||||
expect(await libraryTotal(request)).toBe(restoredTotal);
|
||||
await expect(libraryCount).toHaveText(restoredTotal.toLocaleString('en-US'));
|
||||
expect((await serverJson(request, '/__test__/system/backups')).safety).toHaveLength(1);
|
||||
// The Activity job view stayed, and Backup & export now names the safety backup the restore kept.
|
||||
await expect(page).toHaveURL(`${workspace}#/activity/${job.id}`);
|
||||
await page.goto(`${workspace}#/system/backups`);
|
||||
await expect(statusLine(page)).toHaveText(/Safety backup · .+ · 1 on disk\.$/);
|
||||
});
|
||||
|
||||
test('a finished restore shows the open page again from the restored data', async ({ page, request }) => {
|
||||
const bytes = await backupBytes(request);
|
||||
const restoredTotal = await libraryTotal(request);
|
||||
await expect(await request.post('/prompt_manager/save', { data: { text: 'S9 prompt the restore removes' } })).toBeOK();
|
||||
await openBackups(page);
|
||||
await chooseBackup(page, bytes);
|
||||
const preview = page.locator('dialog#modal[open] #modal-title');
|
||||
await expect(preview).toHaveText('Restore preview');
|
||||
// The restore is held at its start, so the reader is on the Library (showing the prompt) when it finishes.
|
||||
await expect(await request.post('/__test__/hold', { data: { path: APPLY_URL } })).toBeOK();
|
||||
const applied = page.waitForResponse(response => response.url().endsWith(APPLY_URL));
|
||||
await dialog(page).getByRole('button', { name: 'Restore', exact: true }).click();
|
||||
await expect.poll(async () => (await serverJson(request, '/__test__/hold')).held[APPLY_URL]).toBe(1);
|
||||
await page.evaluate(() => { window.location.hash = '#/library'; });
|
||||
const card = page.locator('#workspace-main article.prompt-card').filter({ hasText: 'S9 prompt the restore removes' });
|
||||
await expect(card).toHaveCount(1);
|
||||
await expect(await request.post('/__test__/release')).toBeOK();
|
||||
await applied;
|
||||
await expect(toast(page)).toHaveText(/^Database restore finished · /);
|
||||
// The shelf was read again from the restored library, on the same route.
|
||||
await expect(card).toHaveCount(0);
|
||||
await expect(page).toHaveURL(`${workspace}#/library`);
|
||||
expect(await libraryTotal(request)).toBe(restoredTotal);
|
||||
});
|
||||
|
||||
test('invalid file shows the validator message; nothing is staged', async ({ page, request }) => {
|
||||
const before = await serverJson(request, '/__test__/system/digest');
|
||||
await openBackups(page);
|
||||
const answer = await chooseBackup(page, Buffer.from('this is not a database'), 'notes.db');
|
||||
expect(answer.status()).toBe(400);
|
||||
const body = await answer.json();
|
||||
expect(body.code).toBe('invalid-backup');
|
||||
await expect(toast(page)).toHaveText(`This file can't be restored: ${body.error.replace(/\.?$/, '.')}`);
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
await expect(buttonOf(page, 'Review restore…')).not.toHaveAttribute('aria-disabled', 'true');
|
||||
expect((await serverJson(request, '/__test__/system/backups')).stages).toBe(0);
|
||||
expect(await serverJson(request, '/__test__/system/digest')).toEqual(before);
|
||||
});
|
||||
|
||||
test('Export JSON downloads /library/export; secondary buttons download all prompts (/export) and /pairs/export (D103)', async ({ page }) => {
|
||||
const layout = await openBackups(page);
|
||||
const portable = layout.locator('section.surface').filter({ has: page.locator('h2', { hasText: 'Portable export' }) });
|
||||
expect(await portable.locator('.row-actions > button').allTextContents())
|
||||
.toEqual(['Export JSON', 'Import / inspect file', 'Export all prompts JSON', 'Export saved pairs JSON']);
|
||||
for (const [label, url, name, kind] of [
|
||||
['Export JSON', '/prompt_manager/library/export', /^prompt-manager-library-v1\.json$/, 'prompt-manager-library'],
|
||||
['Export all prompts JSON', '/prompt_manager/export', /^prompt_manager_\d{8}_\d{6}\.json$/, null],
|
||||
['Export saved pairs JSON', '/prompt_manager/pairs/export', /^prompt-pairs\.json$/, 'prompt-manager-pairs']]) {
|
||||
const download = page.waitForEvent('download');
|
||||
const answer = page.waitForResponse(response => new URL(response.url()).pathname === url);
|
||||
await portable.getByRole('button', { name: label, exact: true }).click();
|
||||
expect((await answer).status()).toBe(200);
|
||||
const file = await download;
|
||||
expect(file.suggestedFilename()).toMatch(name);
|
||||
const parsed = JSON.parse(await (await file.createReadStream()).toArray().then(chunks => Buffer.concat(chunks).toString('utf8')));
|
||||
if (kind) expect(parsed.kind).toBe(kind);
|
||||
else expect(Array.isArray(parsed.prompts)).toBe(true);
|
||||
}
|
||||
// A failed export says so and saves nothing.
|
||||
await page.route('**/prompt_manager/pairs/export', route => route.fulfill({ status: 500, json: { success: false } }));
|
||||
await portable.getByRole('button', { name: 'Export saved pairs JSON', exact: true }).click();
|
||||
await expect(toast(page)).toHaveText('The export could not be downloaded. Nothing was saved to this device.');
|
||||
});
|
||||
|
||||
test('Import / inspect lists Inspect, Rescan and the shared scan review (D15, D55)', async ({ page, request }) => {
|
||||
await openBackups(page);
|
||||
const open = async () => {
|
||||
await buttonOf(page, 'Import / inspect file').click();
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Import & inspect');
|
||||
};
|
||||
await open();
|
||||
const rows = dialog(page).locator('.operation-list > .operation');
|
||||
expect(await rows.locator('h3').allTextContents()).toEqual(['Inspect a local image', 'Scan an output folder', 'Scan for metadata']);
|
||||
expect(await rows.locator('button').allTextContents()).toEqual(['Open inspector', 'Rescan', 'Review scan']);
|
||||
// Rescan reads the output folders again and says what it found; while it runs, a reopened dialog says so.
|
||||
await expect(await request.post('/__test__/hold', { data: { path: '/prompt_manager/gallery/rescan' } })).toBeOK();
|
||||
const rescan = page.waitForResponse(response => response.url().endsWith('/prompt_manager/gallery/rescan'));
|
||||
await rows.getByRole('button', { name: 'Rescan', exact: true }).click();
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
await expect(toast(page)).toHaveText('Rescanning output folders…');
|
||||
await open();
|
||||
await expect(rows.nth(1).locator('button')).toHaveText('Rescanning…');
|
||||
await expect(rows.nth(1).locator('button')).toHaveAttribute('aria-disabled', 'true');
|
||||
await expect(await request.post('/__test__/release')).toBeOK();
|
||||
expect((await rescan).status()).toBe(200);
|
||||
await expect(toast(page)).toHaveText(/^Rescan finished\. \d+ files? in \d+ output folders?\.$/);
|
||||
await expect(rows.nth(1).locator('button')).toHaveText('Rescan');
|
||||
// Scan for metadata opens the shared review over the server's output folders; Cancel starts nothing.
|
||||
await dialog(page).getByRole('button', { name: 'Review scan', exact: true }).click();
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Review metadata scan');
|
||||
await expect(dialog(page).locator('dl.meta-grid dd').first()).toHaveText(/^\d+ output folders?$/);
|
||||
await dialog(page).getByRole('button', { name: 'Cancel', exact: true }).click();
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
// Open inspector goes to Inspect a file.
|
||||
await open();
|
||||
await dialog(page).getByRole('button', { name: 'Open inspector', exact: true }).click();
|
||||
await expect(page).toHaveURL(`${workspace}#/inspect`);
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('safety backup toggle saves through Save preferences with revision (D104)', async ({ page, request }) => {
|
||||
const before = await serverJson(request, '/prompt_manager/settings');
|
||||
expect(before.settings.safety_backup).toBe(true);
|
||||
await openBackups(page);
|
||||
const row = settingOf(page, 'Create a safety backup before maintenance');
|
||||
// D104: what the policy really does (the concept's "Preview a verified backup…" line does not say it).
|
||||
await expect(row.locator('p').first()).toHaveText(SAFETY_DESCRIPTION);
|
||||
const toggle = row.locator('input');
|
||||
await expect(toggle).toBeChecked();
|
||||
await toggle.uncheck();
|
||||
await expect(settingsStatus(page)).toHaveText('Unsaved changes.');
|
||||
const saved = page.waitForRequest(item => item.url().endsWith('/prompt_manager/settings') && item.method() === 'POST');
|
||||
await save(page).click();
|
||||
const body = (await saved).postDataJSON();
|
||||
expect(body).toEqual({ safety_backup: false, expected_revision: before.revision });
|
||||
await expect(settingsStatus(page)).toHaveText('Saved to the server configuration.');
|
||||
expect((await serverJson(request, '/prompt_manager/settings')).settings.safety_backup).toBe(false);
|
||||
// A maintenance review now says so.
|
||||
const plan = await request.post('/prompt_manager/maintenance/preview', { data: { operation: 'optimize', options: {} } });
|
||||
expect((await plan.json()).backup.policy).toBe('off');
|
||||
await openBackups(page);
|
||||
await expect(toggle).not.toBeChecked();
|
||||
// A configuration changed elsewhere is a conflict, never a blind write.
|
||||
await toggle.check();
|
||||
await expect(await request.post('/prompt_manager/settings', { data: { result_timeout: 7,
|
||||
expected_revision: (await serverJson(request, '/prompt_manager/settings')).revision } })).toBeOK();
|
||||
await save(page).click();
|
||||
await expect(settingsStatus(page)).toHaveText('Settings changed elsewhere. Reload this section to see them.');
|
||||
expect((await serverJson(request, '/prompt_manager/settings')).settings.safety_backup).toBe(false);
|
||||
});
|
||||
|
||||
test('schedule offers Manual only; Daily/Weekly disabled with reason (D104)', async ({ page }) => {
|
||||
await openBackups(page);
|
||||
const schedule = settingOf(page, 'Scheduled database backups');
|
||||
expect(await schedule.locator('option').evaluateAll(options => options.map(option => [option.textContent, option.disabled])))
|
||||
.toEqual([['Manual only', false], ['Daily', true], ['Weekly', true]]);
|
||||
await expect(schedule.locator('select')).toHaveValue('Manual only');
|
||||
await expect(schedule.locator('p')).toHaveText('Daily and Weekly are not available yet: back up with Back up whenever you choose.');
|
||||
const policy = layoutOf(page).locator(':scope > section.surface');
|
||||
await expect(policy.locator('h2')).toHaveText('Recovery policy Proposed');
|
||||
await expect(policy.locator('h2 > span.chip.blue')).toHaveText('Proposed');
|
||||
await expect(policy.locator('p.warning-note')).toHaveText('Restores are staged, migrated and validated before replacing live data.');
|
||||
// Nothing on the schedule is ever drafted.
|
||||
await expect(settingsStatus(page)).toHaveText('No unsaved changes.');
|
||||
});
|
||||
|
||||
test('a failed settings read leaves the policy disabled with the reason', async ({ page }) => {
|
||||
await page.route('**/prompt_manager/settings', route => (route.request().method() === 'GET'
|
||||
? route.fulfill({ status: 500, json: { success: false, error: 'Failed' } }) : route.continue()));
|
||||
await page.goto(`${workspace}#/system/backups`);
|
||||
const row = settingOf(page, 'Create a safety backup before maintenance');
|
||||
await expect(row.locator('p.small')).toHaveText('Server settings could not be read.');
|
||||
await expect(row.locator('input')).toBeDisabled();
|
||||
});
|
||||
|
||||
test('the legacy admin Backup and Restore keep working', 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);
|
||||
await page.goto('/prompt_manager/admin');
|
||||
await expect(page.locator('#restoreFileInput')).toHaveCount(1);
|
||||
});
|
||||
|
||||
/** Start a fake running job (a metadata scan another tab started), which an exclusive restore must wait for. */
|
||||
async function runningScan(request) {
|
||||
const response = await request.post('/__test__/jobs/start', { data: { kind: 'scan', title: 'Metadata scan', total: 4 } });
|
||||
await expect(response).toBeOK();
|
||||
return (await response.json()).job;
|
||||
}
|
||||
|
||||
const stagesOf = async request => (await serverJson(request, '/__test__/system/backups')).stages;
|
||||
const deleted = page => page.waitForResponse(response => response.url().includes(`${PREVIEW_URL}/`) &&
|
||||
response.request().method() === 'DELETE');
|
||||
|
||||
test('a refused restore reopens its review: Restore retries it, Discard releases it (I1)', async ({ page, request }) => {
|
||||
const bytes = await backupBytes(request);
|
||||
await openBackups(page);
|
||||
await chooseBackup(page, bytes);
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Restore preview');
|
||||
const scan = await runningScan(request);
|
||||
const refused = page.waitForResponse(response => response.url().endsWith(APPLY_URL));
|
||||
await dialog(page).getByRole('button', { name: 'Restore', exact: true }).click();
|
||||
expect((await refused).status()).toBe(409);
|
||||
await expect(toast(page)).toHaveText('Another job is running: Metadata scan.');
|
||||
// The same review is open again, its stage still held: Restore retries, Discard releases it.
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Restore preview');
|
||||
await expect(dialog(page).locator('.dialog-actions > button')).toHaveText(['Discard', 'Restore']);
|
||||
expect(await stagesOf(request)).toBe(1);
|
||||
await expect(await request.post(`/__test__/jobs/${scan.id}/finish`, { data: { state: 'finished' } })).toBeOK();
|
||||
const applied = page.waitForResponse(response => response.url().endsWith(APPLY_URL));
|
||||
await dialog(page).getByRole('button', { name: 'Restore', exact: true }).click();
|
||||
const { job } = await (await applied).json();
|
||||
await expect(page).toHaveURL(`${workspace}#/activity/${job.id}`);
|
||||
await expect(toast(page)).toHaveText(/^Database restore finished · /);
|
||||
expect(await stagesOf(request)).toBe(0);
|
||||
});
|
||||
|
||||
test('two refused restores in a row, then a third Review restore… still opens the preview (I1)', async ({ page, request }) => {
|
||||
const bytes = await backupBytes(request);
|
||||
await openBackups(page);
|
||||
await runningScan(request);
|
||||
for (const leave of ['discard', 'escape']) {
|
||||
await chooseBackup(page, bytes);
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Restore preview');
|
||||
const refused = page.waitForResponse(response => response.url().endsWith(APPLY_URL));
|
||||
await dialog(page).getByRole('button', { name: 'Restore', exact: true }).click();
|
||||
expect((await refused).status()).toBe(409);
|
||||
await expect(dialog(page).locator('.dialog-actions > button')).toHaveText(['Discard', 'Restore']);
|
||||
const released = deleted(page);
|
||||
if (leave === 'discard') await dialog(page).getByRole('button', { name: 'Discard', exact: true }).click();
|
||||
else await page.keyboard.press('Escape');
|
||||
expect((await released).status()).toBe(200);
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
expect(await stagesOf(request)).toBe(0);
|
||||
}
|
||||
const answer = await chooseBackup(page, bytes);
|
||||
expect((await answer).status()).toBe(200);
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Restore preview');
|
||||
expect(await stagesOf(request)).toBe(1);
|
||||
});
|
||||
|
||||
test('× and leaving the section or the page discard the stage (M1)', async ({ page, request }) => {
|
||||
const bytes = await backupBytes(request);
|
||||
await openBackups(page);
|
||||
const leaves = [
|
||||
['×', () => page.locator('#modal-close').click()],
|
||||
// The modal covers the page, so the section changes as a link or the Back button changes it: by its hash.
|
||||
['section', () => page.evaluate(() => { window.location.hash = '#/system/maintenance'; })],
|
||||
['route', () => page.evaluate(() => { window.location.hash = '#/images'; })],
|
||||
];
|
||||
for (const [name, leave] of leaves) {
|
||||
if (!page.url().endsWith('#/system/backups')) {
|
||||
await page.goto(`${workspace}#/system/backups`);
|
||||
await expect(statusLine(page)).not.toHaveText('Checking backups…');
|
||||
}
|
||||
await chooseBackup(page, bytes);
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Restore preview');
|
||||
expect(await stagesOf(request), name).toBe(1);
|
||||
// A keepalive DELETE (it may outlive the page) is asserted by its request and the server's stage count.
|
||||
const released = page.waitForRequest(item => item.url().includes(`${PREVIEW_URL}/`) && item.method() === 'DELETE');
|
||||
await leave();
|
||||
await released;
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
await expect.poll(() => stagesOf(request), { message: name }).toBe(0);
|
||||
}
|
||||
});
|
||||
|
||||
test('a safety backup file (no .db extension) can be picked and reviewed (M3)', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/safety-backup')).toBeOK();
|
||||
const bytes = await backupBytes(request);
|
||||
await openBackups(page);
|
||||
// The picker lists every file; the server decides what is a database.
|
||||
await expect(page.locator('#restore-file')).not.toHaveAttribute('accept', /./);
|
||||
const answer = await chooseBackup(page, bytes, 'prompts.db.backup_20261025_093000_000000');
|
||||
expect((await answer).status()).toBe(200);
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Restore preview');
|
||||
// The copy's size depends on what earlier tests on this server wrote: KB or MB.
|
||||
expect((await gridOf(page))[0][1]).toMatch(/^prompts\.db\.backup_20261025_093000_000000 · [\d,.]+ (KB|MB) · .+$/);
|
||||
});
|
||||
@@ -1,482 +0,0 @@
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from './fixtures.mjs';
|
||||
import { seedConcept } from './concept-inventory.mjs';
|
||||
|
||||
// System › Integrations (#249 System S10 Task 27) on the concept-shaped fixture with a fake LoraManager install
|
||||
// (workspace_integrations_fixture.py): the D106 ComfyUI rows and renderer status, the LoRA rows and the D114 key, saved
|
||||
// through the one Save row, and the D107 import review run as an Activity job. No test reaches the network: CivitAI
|
||||
// downloads are answered by the fixture's stub opener.
|
||||
const workspace = '/prompt_manager/workspace';
|
||||
const PREVIEW_URL = '/prompt_manager/lora/import/preview';
|
||||
const APPLY_URL = '/prompt_manager/lora/import/apply';
|
||||
const ENABLE_URL = '/prompt_manager/lora/enable';
|
||||
const KEPT = 'The import failed; the previous LoRA imports were kept. Images already downloaded stay in the image cache.';
|
||||
|
||||
test.beforeEach(async ({ page, request, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
await seedConcept(request);
|
||||
await expect(await request.post('/__test__/system/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/integrations/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/jobs/restart')).toBeOK();
|
||||
page.systemErrors = [];
|
||||
page.on('pageerror', error => page.systemErrors.push(error.message));
|
||||
});
|
||||
|
||||
test.afterEach(async ({ page, request, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
await request.post('/__test__/integrations/gate', { data: { hold: false } });
|
||||
await request.post('/__test__/release');
|
||||
await request.post('/__test__/jobs/restart');
|
||||
await request.post('/__test__/integrations/reset');
|
||||
await request.post('/__test__/system/reset');
|
||||
expect(page.systemErrors).toEqual([]);
|
||||
});
|
||||
|
||||
const toast = page => page.locator('#workspace-toast');
|
||||
const layoutOf = page => page.locator('#workspace-main .settings-layout');
|
||||
const dialog = page => page.locator('dialog#modal[open]');
|
||||
const status = page => page.locator('#settings-status');
|
||||
const save = page => page.getByRole('button', { name: 'Save preferences', exact: true });
|
||||
const settingOf = (page, label) => layoutOf(page).locator('.setting-row').filter({ has: page.locator(`label:text-is("${label}")`) });
|
||||
|
||||
async function serverJson(request, url) {
|
||||
const response = await request.get(url);
|
||||
await expect(response).toBeOK();
|
||||
return response.json();
|
||||
}
|
||||
|
||||
async function lora(request, data = {}) {
|
||||
await expect(await request.post('/__test__/integrations/lora-manager', { data })).toBeOK();
|
||||
}
|
||||
|
||||
async function openIntegrations(page) {
|
||||
// A fresh page each time: the drafts live for the page session (Q10), and a hash-only goto would not reload.
|
||||
if (page.url().endsWith('#/system/integrations')) await page.reload();
|
||||
else await page.goto(`${workspace}#/system/integrations`);
|
||||
await expect(layoutOf(page).locator(':scope > section.surface')).toHaveCount(2);
|
||||
// The LoRA rows are painted once the server's LoRA state was read.
|
||||
await expect(settingOf(page, 'LoraManager integration').locator('input')).toBeEnabled();
|
||||
// The renderer line settles first: its new text moves the rows below, so no click lands mid-shift.
|
||||
await expect(settingOf(page, 'Host compatibility').getByRole('status')).not.toHaveText('Looking for open ComfyUI windows…');
|
||||
return layoutOf(page);
|
||||
}
|
||||
|
||||
/** Enable the integration (and optionally its trigger words) through the Save row. */
|
||||
async function enable(page, { triggers = false } = {}) {
|
||||
await settingOf(page, 'LoraManager integration').locator('input').check();
|
||||
if (triggers) await settingOf(page, 'Automatically include trigger words').locator('input').check();
|
||||
await save(page).click();
|
||||
await expect(status(page)).toHaveText('Saved to the server configuration.');
|
||||
}
|
||||
|
||||
/** The open review's `[dt, dd]` rows. */
|
||||
const gridOf = page => dialog(page).locator('dl.meta-grid > div').evaluateAll(items => items.map(item =>
|
||||
[item.querySelector('dt').textContent, item.querySelector('dd').textContent]));
|
||||
|
||||
async function ended(request, id) {
|
||||
await expect.poll(async () => (await serverJson(request, `/prompt_manager/jobs/${id}`)).job.state).not.toBe('running');
|
||||
return (await serverJson(request, `/prompt_manager/jobs/${id}`)).job;
|
||||
}
|
||||
|
||||
/** A second page of this origin that serves a fake ComfyUI window on the send bridge. */
|
||||
async function fakeComfyWindow(context, renderer) {
|
||||
const host = await context.newPage();
|
||||
await host.goto(`${workspace}#/overview`);
|
||||
await host.evaluate(async drawn => {
|
||||
const { createBridgeHost } = await import('/prompt_manager/js/shared/toolkit/prompt-bridge.js');
|
||||
window.fakeHost = createBridgeHost({ namespace: '/', origin: location.origin, label: 'PromptManager — ComfyUI',
|
||||
adapter: { listTargets: () => (drawn ? [{ id: '7', label: 'Prompt Manager', type: 'PromptManager', roles: ['positive'],
|
||||
renderer: drawn }] : []), preview() { throw new Error('unused'); }, apply() { throw new Error('unused'); } } });
|
||||
}, renderer);
|
||||
return host;
|
||||
}
|
||||
|
||||
test('send destination offers Ask each time; unsupported options disabled with the reason (D106)', async ({ page }) => {
|
||||
const layout = await openIntegrations(page);
|
||||
expect(await layout.locator(':scope > section.surface > h2').allTextContents()).toEqual(['ComfyUI', 'LoRA & resources']);
|
||||
const send = settingOf(page, 'Default send destination');
|
||||
expect(await send.locator('option').evaluateAll(options => options.map(option => [option.textContent, option.disabled])))
|
||||
.toEqual([['Ask each time', false], ['Selected Prompt Manager node', true], ['Create a Prompt Manager (+/-) node', true]]);
|
||||
await expect(send.locator('select')).toHaveValue('Ask each time');
|
||||
await expect(send.locator('p.small')).toHaveText(
|
||||
'Sending always asks for the ComfyUI window and node, so “Ask each time” is the only option.');
|
||||
const renderer = settingOf(page, 'Host compatibility');
|
||||
expect(await renderer.locator('option').evaluateAll(options => options.map(option => [option.textContent, option.disabled])))
|
||||
.toEqual([['Follow host renderer', false], ['Classic', true], ['Nodes 2.0', true]]);
|
||||
await expect(renderer.locator('p.small').first()).toHaveText('Prompt Manager nodes always follow ComfyUI’s own renderer setting.');
|
||||
const back = settingOf(page, 'Return to canvas after send');
|
||||
await expect(back.locator('input')).toBeDisabled();
|
||||
await expect(back.locator('input')).not.toBeChecked();
|
||||
await expect(back.locator('p.small')).toHaveText('Not available: sending applies the text in ComfyUI and keeps this page open.');
|
||||
// Nothing on this surface is ever drafted: the Save row stays clean.
|
||||
await expect(status(page)).toHaveText('No unsaved changes.');
|
||||
await expect(save(page)).toHaveAttribute('aria-disabled', 'true');
|
||||
});
|
||||
|
||||
test('host compatibility shows the detected renderer as status (D106)', async ({ page, context }) => {
|
||||
await openIntegrations(page);
|
||||
const line = settingOf(page, 'Host compatibility').getByRole('status');
|
||||
await expect(line).toHaveText('No ComfyUI window is open in this browser, so its renderer is not known.');
|
||||
const host = await fakeComfyWindow(context, 'nodes2');
|
||||
try {
|
||||
await page.reload();
|
||||
await expect(line).toHaveText('Detected: Nodes 2.0 (1 ComfyUI window).');
|
||||
await host.evaluate(() => window.fakeHost.dispose());
|
||||
await host.evaluate(async () => {
|
||||
const { createBridgeHost } = await import('/prompt_manager/js/shared/toolkit/prompt-bridge.js');
|
||||
window.fakeHost = createBridgeHost({ namespace: '/', origin: location.origin, label: 'PromptManager — ComfyUI',
|
||||
adapter: { listTargets: () => [], preview() {}, apply() {} } });
|
||||
});
|
||||
await page.reload();
|
||||
await expect(line).toHaveText('ComfyUI is open; add a Prompt Manager node to see which renderer it uses.');
|
||||
} finally {
|
||||
await host.close();
|
||||
}
|
||||
});
|
||||
|
||||
test('LoRA rows show detected / enabled / cache-loaded separately', async ({ page, request }) => {
|
||||
await openIntegrations(page);
|
||||
const integration = settingOf(page, 'LoraManager integration');
|
||||
const triggers = settingOf(page, 'Automatically include trigger words');
|
||||
await expect(integration.locator('p').first()).toHaveText('Not detected in ComfyUI’s custom_nodes folders.');
|
||||
await expect(integration.locator('input')).not.toBeChecked();
|
||||
await expect(triggers.locator('p.small')).toHaveText('Trigger words are not loaded.');
|
||||
// Enabling what is not installed is refused by the server: the toggle stays highlighted, nothing was saved.
|
||||
await integration.locator('input').check();
|
||||
await save(page).click();
|
||||
await expect(status(page)).toHaveText('Some changes were not saved. Review the highlighted settings.');
|
||||
await expect(integration.locator('input')).toHaveAttribute('aria-invalid', 'true');
|
||||
expect((await serverJson(request, '/__test__/integrations/state')).enabled).toBe(false);
|
||||
|
||||
await lora(request);
|
||||
await openIntegrations(page);
|
||||
await expect(integration.locator('p').first()).toHaveText('Detected: custom_nodes/ComfyUI-Lora-Manager.');
|
||||
await expect(integration.locator('input')).not.toBeChecked();
|
||||
await enable(page, { triggers: true });
|
||||
await expect(integration.locator('input')).toBeChecked();
|
||||
await expect(triggers.locator('p.small')).toHaveText('Trigger words are loaded.');
|
||||
expect(await serverJson(request, '/__test__/integrations/state')).toMatchObject({ enabled: true, trigger_words: true });
|
||||
// A reload reads the saved state back.
|
||||
await openIntegrations(page);
|
||||
await expect(integration.locator('input')).toBeChecked();
|
||||
await expect(triggers.locator('input')).toBeChecked();
|
||||
await expect(status(page)).toHaveText('No unsaved changes.');
|
||||
});
|
||||
|
||||
test('LoraManager path row shows the detected path read-only, as legacy (D106)', async ({ page, request }) => {
|
||||
await openIntegrations(page);
|
||||
const row = settingOf(page, 'LoraManager path');
|
||||
const input = row.locator('input');
|
||||
await expect(input).toHaveAttribute('readonly', '');
|
||||
await expect(input).toHaveValue('');
|
||||
await expect(input).toHaveAttribute('placeholder', 'Not detected');
|
||||
await lora(request);
|
||||
await openIntegrations(page);
|
||||
await expect(input).toHaveValue('custom_nodes/ComfyUI-Lora-Manager');
|
||||
await input.focus();
|
||||
await page.keyboard.type('elsewhere');
|
||||
await expect(input).toHaveValue('custom_nodes/ComfyUI-Lora-Manager');
|
||||
await expect(status(page)).toHaveText('No unsaved changes.');
|
||||
});
|
||||
|
||||
test('key never prefilled; blank keeps; Clear saved key clears; saved through Save preferences (D114)', async ({ page, request }) => {
|
||||
const enables = [];
|
||||
page.on('request', item => { if (item.url().endsWith(ENABLE_URL)) enables.push(item.postDataJSON()); });
|
||||
await lora(request);
|
||||
await openIntegrations(page);
|
||||
const key = settingOf(page, 'CivitAI API key');
|
||||
const field = key.locator('input[type="password"]');
|
||||
const clear = key.getByRole('button', { name: 'Clear saved key', exact: true });
|
||||
await expect(field).toHaveValue('');
|
||||
await expect(key.locator('p.small')).toHaveText('No key saved.');
|
||||
await expect(clear).toHaveCount(0);
|
||||
// A typed key is dropped when the section is left: the field comes back empty with nothing hidden to save (N4).
|
||||
await field.fill('sk-dropped');
|
||||
await expect(status(page)).toHaveText('Unsaved changes.');
|
||||
await page.locator('#workspace-main .section-nav').getByRole('button', { name: 'Appearance', exact: true }).click();
|
||||
await page.locator('#workspace-main .section-nav').getByRole('button', { name: 'Integrations', exact: true }).click();
|
||||
await expect(status(page)).toHaveText('No unsaved changes.');
|
||||
await expect(key.locator('p.small')).toHaveText('No key saved.');
|
||||
// Typing a key drafts it; Save sends it once through /lora/enable with the settings revision.
|
||||
await field.fill('sk-test-key');
|
||||
await expect(key.locator('p.small')).toHaveText('The new key replaces the saved one when you save.');
|
||||
await expect(status(page)).toHaveText('Unsaved changes.');
|
||||
await save(page).click();
|
||||
await expect(status(page)).toHaveText('Saved to the server configuration.');
|
||||
expect(enables).toHaveLength(1);
|
||||
expect(enables[0]).toMatchObject({ civitai_api_key: 'sk-test-key', enabled: false });
|
||||
expect(enables[0].expected_revision).toMatch(/^(sha256:|absent)/);
|
||||
await expect(field).toHaveValue('');
|
||||
await expect(key.locator('p.small')).toHaveText('A key is saved.');
|
||||
expect((await serverJson(request, '/__test__/integrations/state')).key_saved).toBe(true);
|
||||
// Never prefilled, even after a reload; the page never receives the key.
|
||||
const statusBodies = [];
|
||||
page.on('response', async response => {
|
||||
if (response.url().endsWith('/prompt_manager/lora/status')) statusBodies.push(await response.text());
|
||||
});
|
||||
await openIntegrations(page);
|
||||
await expect(field).toHaveValue('');
|
||||
await expect(key.locator('p.small')).toHaveText('A key is saved.');
|
||||
expect(statusBodies.join('')).not.toContain('sk-test-key');
|
||||
// The announced line is rewritten only when its text changes, not on every edit elsewhere (S10 re-review N8).
|
||||
await page.evaluate(() => {
|
||||
window.__keyNoteWrites = 0;
|
||||
new MutationObserver(records => { window.__keyNoteWrites += records.length; })
|
||||
.observe(document.querySelector('#civitai-key-status'), { childList: true, characterData: true, subtree: true });
|
||||
});
|
||||
const triggers = settingOf(page, 'Automatically include trigger words').locator('input');
|
||||
await triggers.check();
|
||||
await expect(status(page)).toHaveText('Unsaved changes.');
|
||||
await triggers.uncheck();
|
||||
await expect(status(page)).toHaveText('No unsaved changes.');
|
||||
expect(await page.evaluate(() => window.__keyNoteWrites)).toBe(0);
|
||||
// A blank field keeps the saved key: nothing to save.
|
||||
await field.fill(' ');
|
||||
await expect(status(page)).toHaveText('No unsaved changes.');
|
||||
await field.fill('');
|
||||
// Clear saved key drafts the removal; Save removes it. The keyboard stays on the key field, and the line that
|
||||
// says so is announced (S10 review M2).
|
||||
await expect(key.locator('p.small')).toHaveAttribute('role', 'status');
|
||||
await expect(field).toHaveAttribute('aria-describedby', 'civitai-key-status');
|
||||
await clear.focus();
|
||||
await page.keyboard.press('Enter');
|
||||
await expect(key.locator('p.small')).toHaveText('The saved key is removed when you save.');
|
||||
await expect(clear).toHaveCount(0);
|
||||
await expect(field).toBeFocused();
|
||||
await save(page).click();
|
||||
await expect(status(page)).toHaveText('Saved to the server configuration.');
|
||||
expect(enables.at(-1)).toMatchObject({ clear_civitai_api_key: true });
|
||||
expect(enables.at(-1)).not.toHaveProperty('civitai_api_key');
|
||||
await expect(key.locator('p.small')).toHaveText('No key saved.');
|
||||
expect((await serverJson(request, '/__test__/integrations/state')).key_saved).toBe(false);
|
||||
});
|
||||
|
||||
test('a LoRA save over settings changed elsewhere is a conflict, never a blind write', async ({ page, request }) => {
|
||||
await lora(request);
|
||||
await openIntegrations(page);
|
||||
await settingOf(page, 'Automatically include trigger words').locator('input').check();
|
||||
// Another page saves the settings meanwhile: the revision this page read is stale.
|
||||
await expect(await request.post('/prompt_manager/settings', { data: { result_timeout: 42 } })).toBeOK();
|
||||
await save(page).click();
|
||||
await expect(status(page)).toHaveText('Settings changed elsewhere. Reload this section to see them.');
|
||||
expect((await serverJson(request, '/__test__/integrations/state')).trigger_words).toBe(false);
|
||||
// The draft stays; the next Save carries the reread revision and succeeds.
|
||||
await save(page).click();
|
||||
await expect(status(page)).toHaveText('Saved to the server configuration.');
|
||||
expect((await serverJson(request, '/__test__/integrations/state')).trigger_words).toBe(true);
|
||||
});
|
||||
|
||||
test('a save where only the LoRA part is written says so; the next save needs no conflict (S10 review I3)',
|
||||
async ({ page, request }) => {
|
||||
await lora(request);
|
||||
await expect(await request.post('/__test__/system/missing-source', { data: { inside: true } })).toBeOK();
|
||||
await openIntegrations(page);
|
||||
const triggers = settingOf(page, 'Automatically include trigger words').locator('input');
|
||||
await triggers.check();
|
||||
// In Sources, the drive is unplugged after the page read the list, and another source is removed.
|
||||
const nav = page.locator('#workspace-main .section-nav');
|
||||
await nav.getByRole('button', { name: 'Sources & scans', exact: true }).click();
|
||||
const sourceRow = title => layoutOf(page).locator('#source-list .operation').filter({ has: page.locator('h3', { hasText: title }) });
|
||||
await expect(sourceRow('unplugged')).toBeVisible();
|
||||
await expect(await request.post('/__test__/system/unplug')).toBeOK();
|
||||
await sourceRow('ComfyUI input').getByRole('button', { name: 'Remove', exact: true }).click();
|
||||
await save(page).click();
|
||||
// The LoRA part was saved, the sources were refused: the status gives the server's reason for the source, and
|
||||
// only the source row is highlighted.
|
||||
await expect(status(page)).toHaveText("output/unplugged can't be used: Gallery root does not exist. Remove it or reconnect it, then save.");
|
||||
await expect(sourceRow('unplugged')).toHaveAttribute('aria-invalid', 'true');
|
||||
expect((await serverJson(request, '/__test__/integrations/state')).trigger_words).toBe(true);
|
||||
await nav.getByRole('button', { name: 'Integrations', exact: true }).click();
|
||||
await expect(triggers).toBeChecked();
|
||||
await expect(triggers).not.toHaveAttribute('aria-invalid', 'true');
|
||||
// The next save carries the revision the saved part left: it saves, with no "changed elsewhere".
|
||||
await nav.getByRole('button', { name: 'Sources & scans', exact: true }).click();
|
||||
await sourceRow('unplugged').getByRole('button', { name: 'Remove', exact: true }).click();
|
||||
await save(page).click();
|
||||
await expect(status(page)).toHaveText('Saved to the server configuration.');
|
||||
expect((await serverJson(request, '/prompt_manager/settings')).sources.configured.map(source => source.path)).toEqual(['output']);
|
||||
});
|
||||
|
||||
test('a partial save made from Integrations keeps its outcome when Sources opens (S10 re-review R1)',
|
||||
async ({ page, request }) => {
|
||||
await lora(request);
|
||||
await expect(await request.post('/__test__/system/missing-source', { data: { inside: true } })).toBeOK();
|
||||
await openIntegrations(page);
|
||||
const nav = page.locator('#workspace-main .section-nav');
|
||||
const sourceRow = title => layoutOf(page).locator('#source-list .operation').filter({ has: page.locator('h3', { hasText: title }) });
|
||||
await nav.getByRole('button', { name: 'Sources & scans', exact: true }).click();
|
||||
await expect(sourceRow('unplugged')).toBeVisible();
|
||||
await expect(await request.post('/__test__/system/unplug')).toBeOK();
|
||||
await sourceRow('ComfyUI input').getByRole('button', { name: 'Remove', exact: true }).click();
|
||||
// Back in Integrations, a LoRA change, then Save from there.
|
||||
await nav.getByRole('button', { name: 'Integrations', exact: true }).click();
|
||||
await settingOf(page, 'Automatically include trigger words').locator('input').check();
|
||||
await save(page).click();
|
||||
const reason = "output/unplugged can't be used: Gallery root does not exist. Remove it or reconnect it, then save.";
|
||||
await expect(status(page)).toHaveText(reason);
|
||||
// Leaving Integrations (no key was typed) keeps the outcome: Sources shows the reason on the refused row.
|
||||
await nav.getByRole('button', { name: 'Sources & scans', exact: true }).click();
|
||||
await expect(sourceRow('unplugged')).toHaveAttribute('aria-invalid', 'true');
|
||||
await expect(status(page)).toHaveText(reason);
|
||||
});
|
||||
|
||||
test('Review import shows replacement, blocked items and network effects with the backup box; Start opens the job (D107)',
|
||||
async ({ page, request }) => {
|
||||
await lora(request, { loras: 3, civitai: true, preview: false, previous: 2 });
|
||||
const previews = [];
|
||||
const applies = [];
|
||||
page.on('request', item => {
|
||||
if (item.url().endsWith(PREVIEW_URL)) previews.push(item.postDataJSON());
|
||||
if (item.url().endsWith(APPLY_URL)) applies.push(item.postDataJSON());
|
||||
});
|
||||
await openIntegrations(page);
|
||||
await enable(page);
|
||||
const before = await serverJson(request, '/__test__/integrations/state');
|
||||
const review = layoutOf(page).locator('.operation').filter({ hasText: 'LoRA reference data' })
|
||||
.getByRole('button', { name: 'Review import', exact: true });
|
||||
await review.click();
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Review LoRA import');
|
||||
await expect(dialog(page).locator('#modal-body > h3')).toHaveText('LoRA reference data');
|
||||
expect(previews).toEqual([{}]);
|
||||
expect(await gridOf(page)).toEqual([
|
||||
['LoraManager', 'environment settings · alongside storage · library default'],
|
||||
['Scope', 'Replace 2 imported prompts with 3 LoRAs'],
|
||||
['Found', '3 metadata files · 3 LoRAs · 0 with local images'],
|
||||
['Protected', 'Nothing is protected'],
|
||||
['Network', 'Downloads up to 3 example images from CivitAI'],
|
||||
['Backup', 'Create and verify before changes (browser.db.backup_<date and time>)'],
|
||||
['Media on disk', 'No files are removed; image links are replaced'],
|
||||
]);
|
||||
// #355: the settings file the server read is the LoraManager row's hover title.
|
||||
await expect(dialog(page).locator('dl.meta-grid > div').first().locator('dd'))
|
||||
.toHaveAttribute('title', /lora-manager-settings[\/\\]settings\.json$/);
|
||||
await expect(dialog(page).locator('details > summary')).toHaveText('LoRAs found (3)');
|
||||
await expect(dialog(page).locator('.notice h3')).toHaveText('Back up the database first (recommended)');
|
||||
await expect(dialog(page).locator('p.small').last()).toContainText('stay in the image cache');
|
||||
// The review wrote and downloaded nothing; Cancel neither.
|
||||
await dialog(page).getByRole('button', { name: 'Cancel', exact: true }).click();
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
await expect(review).toBeFocused();
|
||||
expect(await serverJson(request, '/__test__/integrations/state')).toEqual(before);
|
||||
|
||||
await review.click();
|
||||
await dialog(page).getByRole('button', { name: 'Start import', exact: true }).click();
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
await expect(page).toHaveURL(/#\/activity\/[0-9a-f]{32}$/);
|
||||
await expect(toast(page)).toHaveText(/LoRA reference import started\.|LoRA import finished · 3 imported · 0 errors/);
|
||||
expect(applies).toHaveLength(1);
|
||||
expect(Object.keys(applies[0])).toEqual(['token']);
|
||||
const job = await ended(request, page.url().split('/').pop());
|
||||
expect([job.kind, job.title, job.state]).toEqual(['lora_import', 'LoRA reference import', 'finished']);
|
||||
expect(job.counters).toMatchObject({ imported: 3, downloaded: 3, download_errors: 0, errors: 0 });
|
||||
const main = page.locator('#workspace-main');
|
||||
await expect(main.locator('h1')).toHaveText('LoRA reference import');
|
||||
await expect(main.locator('.stat-row .stat dt')).toHaveText(['Imported', 'Linked to existing', 'Skipped', 'Images downloaded',
|
||||
'Download errors', 'Errors']);
|
||||
await expect(main.getByRole('button', { name: 'Cancel import', exact: true })).toHaveCount(0);
|
||||
const after = await serverJson(request, '/__test__/integrations/state');
|
||||
expect(after.imports).toEqual(['style_01 example prompt', 'style_02 example prompt', 'style_03 example prompt']);
|
||||
expect(after.downloads).toBe(3);
|
||||
});
|
||||
|
||||
test('saved pairs block the import: the review explains it and offers only Close', async ({ page, request }) => {
|
||||
await lora(request, { previous: 1 });
|
||||
await expect(await request.post('/__test__/integrations/pair')).toBeOK();
|
||||
await openIntegrations(page);
|
||||
await enable(page);
|
||||
await layoutOf(page).getByRole('button', { name: 'Review import', exact: true }).click();
|
||||
// gridOf reads once; poll until the review has rendered its rows.
|
||||
await expect.poll(async () => (await gridOf(page))[3]).toEqual(['Protected', '1 imported prompt used by saved pairs']);
|
||||
await expect(dialog(page).locator('p.small')).toHaveText(
|
||||
'Delete or change the saved pairs that use these prompts, then review the import again. Nothing was changed.');
|
||||
expect(await dialog(page).locator('.dialog-actions button').allTextContents()).toEqual(['Close']);
|
||||
await dialog(page).getByRole('button', { name: 'Close', exact: true }).click();
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('unusable LoraManager settings: the review says why and offers only Close (#355)', async ({ page, request }) => {
|
||||
await lora(request, { previous: 1, settings: 'invalid' });
|
||||
await openIntegrations(page);
|
||||
await enable(page);
|
||||
const before = await serverJson(request, '/__test__/integrations/state');
|
||||
await layoutOf(page).getByRole('button', { name: 'Review import', exact: true }).click();
|
||||
await expect.poll(async () => (await gridOf(page))[0])
|
||||
.toEqual(['LoraManager', 'Not usable: the settings file is not valid JSON']);
|
||||
await expect(dialog(page).locator('p.small')).toHaveText(
|
||||
'LoraManager settings: not usable: the settings file is not valid JSON. Nothing was changed.');
|
||||
expect(await dialog(page).locator('.dialog-actions button').allTextContents()).toEqual(['Close']);
|
||||
await dialog(page).getByRole('button', { name: 'Close', exact: true }).click();
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
expect(await serverJson(request, '/__test__/integrations/state')).toEqual(before);
|
||||
});
|
||||
|
||||
test('imported LoRA previews show on the Library card and open in the inspector and viewer (#355 G6)', async ({ page, request }) => {
|
||||
await lora(request, { loras: 1 });
|
||||
await openIntegrations(page);
|
||||
await enable(page);
|
||||
await layoutOf(page).getByRole('button', { name: 'Review import', exact: true }).click();
|
||||
await dialog(page).getByRole('button', { name: 'Start import', exact: true }).click();
|
||||
await expect(page).toHaveURL(/#\/activity\/[0-9a-f]{32}$/);
|
||||
expect((await ended(request, page.url().split('/').pop())).state).toBe('finished');
|
||||
await page.goto(`${workspace}#/library`);
|
||||
const card = page.locator('#workspace-main article.prompt-card').filter({ hasText: 'Style 01' });
|
||||
const cover = card.locator('.card-art img');
|
||||
await expect(cover).toHaveAttribute('src', /\/prompt_manager\/images\/serve\/style_01\.png\?root=\d+&root_scope=[0-9a-f]{64}$/);
|
||||
await expect.poll(() => cover.evaluate(node => node.complete && node.naturalWidth)).toBeGreaterThan(0);
|
||||
await expect(card.locator('.card-art')).not.toContainText('No available image previews');
|
||||
// The prompt's gallery: Info inspects the file and the viewer opens it; no Tag action or star for LoRA media.
|
||||
await card.getByRole('button', { name: 'Explore history', exact: true }).click();
|
||||
const item = page.getByRole('region', { name: 'Shared gallery', exact: true }).locator('article[data-image-key]').first();
|
||||
await expect(item.getByRole('button', { name: 'Tag or auto-tag style_01.png', exact: true })).toHaveCount(0);
|
||||
await expect(item.locator('.star-button')).toHaveCount(0);
|
||||
await item.getByRole('button', { name: 'Inspect style_01.png', exact: true }).click();
|
||||
const inspector = page.getByRole('complementary', { name: 'Selected image details', exact: true });
|
||||
await expect(inspector).toContainText('style_01.png');
|
||||
await item.getByRole('button', { name: 'Open style_01.png', exact: true }).click();
|
||||
await expect(page.locator('dialog[open] .viewer-canvas img')).toBeVisible();
|
||||
await expect.poll(() => page.locator('dialog[open] .viewer-canvas img').evaluate(node => node.complete && node.naturalWidth))
|
||||
.toBeGreaterThan(0);
|
||||
await page.keyboard.press('Escape');
|
||||
});
|
||||
|
||||
test('a failed import keeps old imports and its job says so', async ({ page, request }) => {
|
||||
await lora(request, { loras: 2, previous: 2 });
|
||||
await openIntegrations(page);
|
||||
await enable(page);
|
||||
const before = await serverJson(request, '/__test__/integrations/state');
|
||||
await expect(await request.post('/__test__/integrations/fail-import')).toBeOK();
|
||||
await layoutOf(page).getByRole('button', { name: 'Review import', exact: true }).click();
|
||||
await dialog(page).getByRole('button', { name: 'Start import', exact: true }).click();
|
||||
await expect(page).toHaveURL(/#\/activity\/[0-9a-f]{32}$/);
|
||||
await expect(toast(page)).toHaveText(`LoRA import failed · ${KEPT}`);
|
||||
const job = await ended(request, page.url().split('/').pop());
|
||||
expect([job.state, job.scope.failure, job.counters.errors]).toEqual(['failed', 'import-failed', 1]);
|
||||
await expect(page.locator('#workspace-main')).toContainText(KEPT);
|
||||
expect((await serverJson(request, '/__test__/integrations/state')).imports).toEqual(before.imports);
|
||||
expect(before.imports).toEqual(['previous lora import 1', 'previous lora import 2']);
|
||||
});
|
||||
|
||||
test('the import can be cancelled while it downloads; then nothing is written', async ({ page, request }) => {
|
||||
await lora(request, { loras: 2, civitai: true, preview: false, previous: 1 });
|
||||
await openIntegrations(page);
|
||||
await enable(page);
|
||||
const before = await serverJson(request, '/__test__/integrations/state');
|
||||
await expect(await request.post('/__test__/integrations/gate', { data: { hold: true } })).toBeOK();
|
||||
await layoutOf(page).getByRole('button', { name: 'Review import', exact: true }).click();
|
||||
await dialog(page).getByRole('button', { name: 'Start import', exact: true }).click();
|
||||
await expect(page).toHaveURL(/#\/activity\/[0-9a-f]{32}$/);
|
||||
const id = page.url().split('/').pop();
|
||||
await expect.poll(async () => (await serverJson(request, '/__test__/integrations/state')).downloads).toBeGreaterThan(0);
|
||||
const cancel = page.locator('#workspace-main').getByRole('button', { name: 'Cancel import', exact: true });
|
||||
await expect(cancel).toBeVisible();
|
||||
// The cancel must have reached the server before the held download goes on (S10 review M1).
|
||||
await Promise.all([
|
||||
page.waitForResponse(response => /\/prompt_manager\/jobs\/[0-9a-f]+\/cancel$/.test(response.url()) && response.ok()),
|
||||
cancel.click(),
|
||||
]);
|
||||
await expect(await request.post('/__test__/integrations/gate', { data: { hold: false } })).toBeOK();
|
||||
const job = await ended(request, id);
|
||||
expect(job.state).toBe('cancelled');
|
||||
expect((await serverJson(request, '/__test__/integrations/state')).imports).toEqual(before.imports);
|
||||
await expect(cancel).toHaveCount(0);
|
||||
});
|
||||
@@ -1,483 +0,0 @@
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from './fixtures.mjs';
|
||||
import { seedConcept } from './concept-inventory.mjs';
|
||||
|
||||
// System › Maintenance (#249 System S7 Task 19) on the concept-shaped real fixture and the S6 backend: the concept's
|
||||
// rows, the read-only health check, the D100 plan review applied as an Activity job, the D101 duplicate search and
|
||||
// review, the S5 thumbnail dialog and the D102 metadata cache row. The hooks are workspace_system_fixture.py's.
|
||||
const workspace = '/prompt_manager/workspace';
|
||||
const PREVIEW_URL = '/prompt_manager/maintenance/preview';
|
||||
const APPLY_URL = '/prompt_manager/maintenance/apply';
|
||||
const CHANGED = 'The library changed since the review; nothing was changed. Review the plan again.';
|
||||
|
||||
test.beforeEach(async ({ page, request, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
await seedConcept(request);
|
||||
await expect(await request.post('/__test__/system/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/jobs/restart')).toBeOK();
|
||||
page.systemErrors = [];
|
||||
page.on('pageerror', error => page.systemErrors.push(error.message));
|
||||
});
|
||||
|
||||
test.afterEach(async ({ page, request, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
await request.post('/__test__/release');
|
||||
await request.post('/__test__/jobs/restart');
|
||||
await request.post('/__test__/system/reset');
|
||||
expect(page.systemErrors).toEqual([]);
|
||||
});
|
||||
|
||||
const toast = page => page.locator('#workspace-toast');
|
||||
const layoutOf = page => page.locator('#workspace-main .settings-layout');
|
||||
const dialog = page => page.locator('dialog#modal[open]');
|
||||
const row = (page, title) => layoutOf(page).locator('.operation').filter({ has: page.locator('h3', { hasText: title }) });
|
||||
|
||||
async function serverJson(request, url) {
|
||||
const response = await request.get(url);
|
||||
await expect(response).toBeOK();
|
||||
return response.json();
|
||||
}
|
||||
|
||||
async function openMaintenance(page) {
|
||||
await page.goto(`${workspace}#/system/maintenance`);
|
||||
await expect(layoutOf(page).locator(':scope > section.surface')).toHaveCount(2);
|
||||
return layoutOf(page);
|
||||
}
|
||||
|
||||
/** The open plan dialog's `[dt, dd]` rows. */
|
||||
const gridOf = page => dialog(page).locator('dl.meta-grid > div').evaluateAll(items => items.map(item =>
|
||||
[item.querySelector('dt').textContent, item.querySelector('dd').textContent]));
|
||||
|
||||
/** Wait until the job `id` ended on the server; answers its summary. */
|
||||
async function ended(request, id) {
|
||||
await expect.poll(async () => (await serverJson(request, `/prompt_manager/jobs/${id}`)).job.state).not.toBe('running');
|
||||
return (await serverJson(request, `/prompt_manager/jobs/${id}`)).job;
|
||||
}
|
||||
|
||||
test('the section shows the concept rows; Metadata cache follows Q5 (D102)', async ({ page }) => {
|
||||
const layout = await openMaintenance(page);
|
||||
await expect(layout.locator(':scope > .notice h2')).toHaveText('Check first. Review changes. Then apply.');
|
||||
await expect(layout.locator(':scope > section.handoff')).toHaveCount(0);
|
||||
expect(await layout.locator(':scope > section.surface').evaluateAll(surfaces => surfaces.map(surface => [
|
||||
surface.querySelector('h2').textContent,
|
||||
[...surface.querySelectorAll('.operation')].map(item => [item.querySelector('h3').textContent,
|
||||
item.querySelector('.row-actions > button').textContent, item.querySelector('.row-actions > button').className]),
|
||||
]))).toEqual([
|
||||
['Database', [['Integrity & relationship checks', 'Check', ''], ['Duplicate prompt text', 'Review plan', ''],
|
||||
['Optimize database', 'Review plan', ''], ['Unused prompt components', 'Review plan', 'danger']]],
|
||||
['Images & cache', [['Missing-file links', 'Review plan', 'danger'], ['Duplicate image files', 'Find duplicates', ''],
|
||||
['Thumbnail cache', 'Review options', ''], ['Metadata cache', 'Review plan', '']]],
|
||||
]);
|
||||
// Nothing of Prompt Manager's is cached in a browser (#222): the row says so and its review is unavailable.
|
||||
const cache = row(page, 'Metadata cache');
|
||||
await expect(cache.locator('p')).toHaveText('Clear only Prompt Manager metadata cached by this browser. Nothing is cached in this browser.');
|
||||
const review = cache.getByRole('button', { name: 'Review plan', exact: true });
|
||||
await expect(review).toHaveAttribute('aria-disabled', 'true');
|
||||
// An aria-disabled button stays focusable and clickable; its press opens nothing.
|
||||
await review.dispatchEvent('click');
|
||||
await expect(review).toHaveAttribute('aria-disabled', 'true');
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('Metadata cache reviews and clears only owned browser entries when there are some (D102)', async ({ page }) => {
|
||||
// The shipped inventory is empty (#222); this test names one owned entry through the module's inventory seam.
|
||||
await page.addInitScript(() => {
|
||||
localStorage.setItem('pm-owned-test-entry', 'cached');
|
||||
localStorage.setItem('pm_workspace_theme', 'dark');
|
||||
let wrapped;
|
||||
Object.defineProperty(window, 'CacheOwnership', { configurable: true, get: () => wrapped, set(api) {
|
||||
const owned = { storageKeys: ['pm-owned-test-entry'], cacheNames: [] };
|
||||
wrapped = { ...api, ownedEntries: options => api.ownedEntries({ ...options, owned }),
|
||||
clearOwned: options => api.clearOwned({ ...options, owned }) };
|
||||
} });
|
||||
});
|
||||
await openMaintenance(page);
|
||||
const cache = row(page, 'Metadata cache');
|
||||
const review = cache.getByRole('button', { name: 'Review plan', exact: true });
|
||||
await expect(review).not.toHaveAttribute('aria-disabled', 'true');
|
||||
await expect(cache.locator('p')).toHaveText('Clear only Prompt Manager metadata cached by this browser.');
|
||||
await review.click();
|
||||
await expect(dialog(page).locator('#modal-body > h3')).toHaveText('Metadata cache');
|
||||
expect(await gridOf(page)).toEqual([['Scope', '1 cached entry in this browser'],
|
||||
['Protected', 'Preferences and other tools’ data in this browser'], ['Backup', 'Not needed: the database is not changed'],
|
||||
['Media on disk', 'No files are removed']]);
|
||||
await dialog(page).getByRole('button', { name: 'Clear browser cache', exact: true }).click();
|
||||
await expect(toast(page)).toHaveText('Cleared 1 cached entry from this browser.');
|
||||
await expect(review).toHaveAttribute('aria-disabled', 'true');
|
||||
expect(await page.evaluate(() => [localStorage.getItem('pm-owned-test-entry'), localStorage.getItem('pm_workspace_theme')]))
|
||||
.toEqual([null, 'dark']);
|
||||
});
|
||||
|
||||
test('Run health checks and Check open the Health report; Review missing links starts a fresh preview, never a cleanup',
|
||||
async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/missing-file')).toBeOK();
|
||||
const before = await serverJson(request, '/__test__/system/digest');
|
||||
const previews = [];
|
||||
page.on('request', item => { if (item.url().endsWith(PREVIEW_URL)) previews.push(item.postDataJSON()); });
|
||||
await openMaintenance(page);
|
||||
const check = row(page, 'Integrity & relationship checks').getByRole('button', { name: 'Check', exact: true });
|
||||
await check.click();
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Health report');
|
||||
await expect(dialog(page).locator('.operation').filter({ hasText: 'Missing files · output' }).locator('.chip'))
|
||||
.toHaveText('1 link needs review');
|
||||
await expect(check).toHaveText('Check');
|
||||
await dialog(page).getByRole('button', { name: 'Review missing links', exact: true }).click();
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Review maintenance plan');
|
||||
await expect(dialog(page).locator('#modal-body > h3')).toHaveText('Missing-file links');
|
||||
expect(previews).toEqual([{ operation: 'missing_links', options: {} }]);
|
||||
expect((await gridOf(page))[0]).toEqual(['Scope', '1 link to a missing file']);
|
||||
await dialog(page).getByRole('button', { name: 'Cancel', exact: true }).click();
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
await expect(check).toBeFocused();
|
||||
// The notice's Run health checks opens the same report.
|
||||
await layoutOf(page).getByRole('button', { name: 'Run health checks', exact: true }).click();
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Health report');
|
||||
expect(await serverJson(request, '/__test__/system/digest')).toEqual(before);
|
||||
});
|
||||
|
||||
test("plan dialog shows the server's scope, protected, backup policy and media lines with the backup box (D100); Cancel writes nothing",
|
||||
async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/duplicate-prompts')).toBeOK();
|
||||
const before = await serverJson(request, '/__test__/system/digest');
|
||||
const answers = [];
|
||||
page.on('response', response => { if (response.url().endsWith(PREVIEW_URL)) answers.push(response.json()); });
|
||||
await openMaintenance(page);
|
||||
const reviewButton = row(page, 'Duplicate prompt text').getByRole('button', { name: 'Review plan', exact: true });
|
||||
await reviewButton.click();
|
||||
const plan = dialog(page);
|
||||
await expect(plan.locator('#modal-title')).toHaveText('Review maintenance plan');
|
||||
const answer = await answers[0];
|
||||
expect(answer.scope.count).toBe(1);
|
||||
await expect(plan.locator('#modal-body > h3')).toHaveText('Duplicate prompt text');
|
||||
expect(await gridOf(page)).toEqual([
|
||||
['Scope', answer.scope.label],
|
||||
['Protected', answer.protected.reasons.length ? answer.protected.reasons.join(' · ') : 'Nothing is protected'],
|
||||
['Backup', `Create and verify before changes (${answer.backup.destination_label})`],
|
||||
['Media on disk', answer.media.label],
|
||||
]);
|
||||
expect(await plan.locator('#modal-body').evaluate(body => [...body.children].map(node =>
|
||||
[node.localName, ...node.classList].join('.')))).toEqual(['h3', 'dl.meta-grid', 'details', 'p.small', 'div.notice',
|
||||
'div.dialog-actions']);
|
||||
await expect(plan.locator('details > summary')).toHaveText('Listed records (1)');
|
||||
await expect(plan.locator('#modal-body > p.small')).toHaveText(/^Nothing changes until you apply\. This review holds for 1[45] minutes\.$/);
|
||||
await expect(plan.locator('.notice h3')).toHaveText('Back up the database first (recommended)');
|
||||
await expect(plan.locator('.dialog-actions button')).toHaveText(['Cancel', 'Apply reviewed operation']);
|
||||
await plan.getByRole('button', { name: 'Cancel', exact: true }).click();
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
await expect(reviewButton).toBeFocused();
|
||||
// Optimize and Unused prompt components are reviewed the same way.
|
||||
for (const [title, scope] of [['Optimize database', 'The whole database'], ['Unused prompt components', null]]) {
|
||||
await row(page, title).getByRole('button', { name: 'Review plan', exact: true }).click();
|
||||
await expect(dialog(page).locator('#modal-body > h3')).toHaveText(title);
|
||||
const value = (await answers.at(-1)).scope.label;
|
||||
expect((await gridOf(page))[0]).toEqual(['Scope', scope ?? value]);
|
||||
await page.keyboard.press('Escape');
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
}
|
||||
expect(await serverJson(request, '/__test__/system/digest')).toEqual(before);
|
||||
});
|
||||
|
||||
test('a plan with nothing to do shows only Close, never Apply', async ({ page, request }) => {
|
||||
const applies = [];
|
||||
page.on('request', item => { if (item.url().endsWith(APPLY_URL)) applies.push(item.url()); });
|
||||
await openMaintenance(page);
|
||||
// The concept fixture has no duplicate prompt text.
|
||||
const reviewButton = row(page, 'Duplicate prompt text').getByRole('button', { name: 'Review plan', exact: true });
|
||||
await reviewButton.click();
|
||||
await expect(dialog(page).locator('#modal-body > h3')).toHaveText('Duplicate prompt text');
|
||||
expect((await gridOf(page))[0]).toEqual(['Scope', 'No duplicate prompt text']);
|
||||
await expect(dialog(page).locator('#modal-body > p.small')).toHaveText('There is nothing to apply; nothing was changed.');
|
||||
await expect(dialog(page).locator('.dialog-actions button')).toHaveText(['Close']);
|
||||
await expect(dialog(page).locator('.notice')).toHaveCount(0);
|
||||
await dialog(page).getByRole('button', { name: 'Close', exact: true }).click();
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
await expect(reviewButton).toBeFocused();
|
||||
expect(applies).toEqual([]);
|
||||
});
|
||||
|
||||
test('Apply reviewed operation closes, toasts, opens the job in Activity', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/duplicate-prompts')).toBeOK();
|
||||
const applies = [];
|
||||
page.on('request', item => { if (item.url().endsWith(APPLY_URL)) applies.push(item.postDataJSON()); });
|
||||
await openMaintenance(page);
|
||||
await row(page, 'Duplicate prompt text').getByRole('button', { name: 'Review plan', exact: true }).click();
|
||||
await dialog(page).getByRole('button', { name: 'Apply reviewed operation', exact: true }).click();
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
await expect(page).toHaveURL(/#\/activity\/[0-9a-f]{32}$/);
|
||||
await expect(toast(page)).toHaveText(/Duplicate prompt text (started|finished · 1 changed)/);
|
||||
expect(applies).toHaveLength(1);
|
||||
expect(Object.keys(applies[0])).toEqual(['token']);
|
||||
const job = await ended(request, page.url().split('/').pop());
|
||||
expect([job.kind, job.state, job.counters.changed]).toEqual(['maintenance', 'finished', 1]);
|
||||
const main = page.locator('#workspace-main');
|
||||
await expect(main.locator('h1')).toHaveText('Duplicate prompt text');
|
||||
await expect(main.locator('.stat-row .stat dt')).toHaveText(['Changed', 'Protected', 'Errors']);
|
||||
await expect(main.locator('.stat-row .stat dd').first()).toHaveText('1');
|
||||
// Not cancellable: no Cancel, a report like every job.
|
||||
await expect(main.getByRole('button', { name: 'Cancel job', exact: true })).toHaveCount(0);
|
||||
await expect(main.getByRole('button', { name: 'Download report', exact: true })).toBeVisible();
|
||||
});
|
||||
|
||||
test('plan-changed asks for a new review', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/duplicate-prompts')).toBeOK();
|
||||
await openMaintenance(page);
|
||||
await row(page, 'Duplicate prompt text').getByRole('button', { name: 'Review plan', exact: true }).click();
|
||||
await expect(dialog(page).locator('#modal-body > h3')).toHaveText('Duplicate prompt text');
|
||||
// The library changes after the review: the apply's recheck refuses it and nothing is changed.
|
||||
await expect(await request.post('/__test__/system/duplicate-prompts', { data: { again: true } })).toBeOK();
|
||||
const before = await serverJson(request, '/__test__/system/digest');
|
||||
await dialog(page).getByRole('button', { name: 'Apply reviewed operation', exact: true }).click();
|
||||
await expect(page).toHaveURL(/#\/activity\/[0-9a-f]{32}$/);
|
||||
await expect(toast(page)).toHaveText(`Duplicate prompt text failed · ${CHANGED}`);
|
||||
await expect(page.locator('#workspace-main')).toContainText(CHANGED);
|
||||
expect((await serverJson(request, '/__test__/system/digest')).rows).toEqual(before.rows);
|
||||
// S7 review M2: the failed job offers Review again, which reviews the same operation afresh.
|
||||
const previews = [];
|
||||
page.on('request', item => { if (item.url().endsWith(PREVIEW_URL)) previews.push(item.postDataJSON()); });
|
||||
await page.locator('#workspace-main').getByRole('button', { name: 'Review again', exact: true }).click();
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Review maintenance plan');
|
||||
await expect(dialog(page).locator('#modal-body > h3')).toHaveText('Duplicate prompt text');
|
||||
expect(previews).toEqual([{ operation: 'duplicate_prompts', options: {} }]);
|
||||
});
|
||||
|
||||
test('a review button pressed twice sends one preview (S7 review M4)', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/missing-file')).toBeOK();
|
||||
await expect(await request.post('/__test__/system/duplicate-files')).toBeOK();
|
||||
const previews = [];
|
||||
page.on('request', item => { if (item.url().endsWith(PREVIEW_URL)) previews.push(item.postDataJSON().operation); });
|
||||
await openMaintenance(page);
|
||||
await row(page, 'Integrity & relationship checks').getByRole('button', { name: 'Check', exact: true }).click();
|
||||
// Hold the preview at the server so the second press lands while it is read.
|
||||
await expect(await request.post('/__test__/hold', { data: { path: PREVIEW_URL } })).toBeOK();
|
||||
// The pressed button reads "Reviewing…" while held, so it is found by its place, not its name.
|
||||
const missing = dialog(page).locator('.dialog-actions > button').last();
|
||||
await expect(missing).toHaveText('Review missing links');
|
||||
await missing.click();
|
||||
await expect(missing).toHaveAttribute('aria-disabled', 'true');
|
||||
await expect(missing).toHaveText('Reviewing…');
|
||||
await missing.dispatchEvent('click');
|
||||
await expect.poll(async () => (await serverJson(request, '/__test__/hold')).held[PREVIEW_URL]).toBe(1);
|
||||
await expect(await request.post('/__test__/release')).toBeOK();
|
||||
await expect(dialog(page).locator('#modal-body > h3')).toHaveText('Missing-file links');
|
||||
expect(previews).toEqual(['missing_links']);
|
||||
await page.keyboard.press('Escape');
|
||||
|
||||
await row(page, 'Duplicate image files').getByRole('button', { name: 'Find duplicates', exact: true }).click();
|
||||
await page.locator('#workspace-main').getByRole('button', { name: 'Review groups', exact: true }).click();
|
||||
await expect(await request.post('/__test__/hold', { data: { path: PREVIEW_URL } })).toBeOK();
|
||||
const cleanup = dialog(page).locator('.dialog-actions > button').last();
|
||||
await expect(cleanup).toHaveText('Review duplicate cleanup');
|
||||
await cleanup.click();
|
||||
await expect(cleanup).toHaveAttribute('aria-disabled', 'true');
|
||||
await cleanup.dispatchEvent('click');
|
||||
await expect.poll(async () => (await serverJson(request, '/__test__/hold')).held[PREVIEW_URL]).toBe(1);
|
||||
await expect(await request.post('/__test__/release')).toBeOK();
|
||||
await expect(dialog(page).locator('#modal-body > h3')).toHaveText('Duplicate image files');
|
||||
expect(previews).toEqual(['missing_links', 'duplicate_files']);
|
||||
});
|
||||
|
||||
test('plan-used links the job that applied it', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/duplicate-prompts')).toBeOK();
|
||||
// A double-submitted apply: the first request applies the plan, the second one's answer reaches the page.
|
||||
await page.route(`**${APPLY_URL}`, async route => {
|
||||
await route.fetch();
|
||||
await route.fulfill({ response: await route.fetch() });
|
||||
});
|
||||
await openMaintenance(page);
|
||||
await row(page, 'Duplicate prompt text').getByRole('button', { name: 'Review plan', exact: true }).click();
|
||||
await dialog(page).getByRole('button', { name: 'Apply reviewed operation', exact: true }).click();
|
||||
await expect(toast(page)).toHaveText('This plan was already applied: Duplicate prompt text.');
|
||||
const link = toast(page).getByRole('link', { name: 'Duplicate prompt text', exact: true });
|
||||
const { jobs } = await serverJson(request, '/prompt_manager/jobs');
|
||||
const applied = jobs.find(job => job.kind === 'maintenance');
|
||||
await expect(link).toHaveAttribute('href', `#/activity/${applied.id}`);
|
||||
await expect(page).toHaveURL(`${workspace}#/system/maintenance`);
|
||||
});
|
||||
|
||||
/**
|
||||
* Record every toast the page shows (`[text, [link hrefs]]`): a job's finish toast may replace an earlier one before a
|
||||
* test reads it.
|
||||
*/
|
||||
async function recordToasts(page) {
|
||||
await page.addInitScript(() => {
|
||||
window.toastLog = [];
|
||||
const record = node => window.toastLog.push([node.textContent, [...node.querySelectorAll('a')].map(a => a.getAttribute('href'))]);
|
||||
new MutationObserver(() => {
|
||||
const node = document.getElementById('workspace-toast');
|
||||
const last = window.toastLog.at(-1);
|
||||
if (node && node.textContent && last?.[0] !== node.textContent) record(node);
|
||||
}).observe(document, { subtree: true, childList: true, characterData: true });
|
||||
});
|
||||
return () => page.evaluate(() => window.toastLog);
|
||||
}
|
||||
|
||||
const NOT_CONFIRMED = ['Outcome not confirmed. Check Activity before applying it again.', ['#/activity']];
|
||||
|
||||
/** The maintenance jobs the server lists (the apply ran even though its answer was lost). */
|
||||
async function maintenanceJobs(request) {
|
||||
return (await serverJson(request, '/prompt_manager/jobs')).jobs.filter(job => job.kind === 'maintenance');
|
||||
}
|
||||
|
||||
test('lost apply response says Outcome not confirmed and points to Activity', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/duplicate-prompts')).toBeOK();
|
||||
// The server applies the plan; its answer never reaches the page.
|
||||
const reads = [];
|
||||
await page.route(`**${APPLY_URL}`, async route => {
|
||||
await route.fetch();
|
||||
page.on('request', item => { if (item.url().endsWith('/prompt_manager/jobs')) reads.push(item.url()); });
|
||||
await route.abort('failed');
|
||||
});
|
||||
const toasts = await recordToasts(page);
|
||||
await openMaintenance(page);
|
||||
await row(page, 'Duplicate prompt text').getByRole('button', { name: 'Review plan', exact: true }).click();
|
||||
await dialog(page).getByRole('button', { name: 'Apply reviewed operation', exact: true }).click();
|
||||
await expect.poll(toasts).toContainEqual(NOT_CONFIRMED);
|
||||
await expect(page).toHaveURL(`${workspace}#/system/maintenance`);
|
||||
// The Activity list is read again at once, and the job the lost answer started is in it.
|
||||
await expect.poll(() => reads.length).toBeGreaterThan(0);
|
||||
expect((await maintenanceJobs(request)).map(job => job.title)).toEqual(['Duplicate prompt text']);
|
||||
await page.evaluate(() => { location.hash = '#/activity'; });
|
||||
await expect(page.locator('#workspace-main')).toContainText('Duplicate prompt text');
|
||||
});
|
||||
|
||||
test('a proxy error or an unreadable answer is not a refusal: Outcome not confirmed (S7 review I1)', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/duplicate-prompts')).toBeOK();
|
||||
const reads = [];
|
||||
// A reverse proxy times out after the server started the job.
|
||||
await page.route(`**${APPLY_URL}`, async route => {
|
||||
await route.fetch();
|
||||
page.on('request', item => { if (item.url().endsWith('/prompt_manager/jobs')) reads.push(item.url()); });
|
||||
await route.fulfill({ status: 504, contentType: 'text/html', body: '<html><body>Gateway Timeout</body></html>' });
|
||||
});
|
||||
const toasts = await recordToasts(page);
|
||||
await openMaintenance(page);
|
||||
await row(page, 'Duplicate prompt text').getByRole('button', { name: 'Review plan', exact: true }).click();
|
||||
await dialog(page).getByRole('button', { name: 'Apply reviewed operation', exact: true }).click();
|
||||
await expect.poll(toasts).toContainEqual(NOT_CONFIRMED);
|
||||
await expect.poll(() => reads.length).toBeGreaterThan(0);
|
||||
expect(await maintenanceJobs(request)).toHaveLength(1);
|
||||
});
|
||||
|
||||
test('a refused apply names the running job', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/duplicate-prompts')).toBeOK();
|
||||
await openMaintenance(page);
|
||||
await row(page, 'Duplicate prompt text').getByRole('button', { name: 'Review plan', exact: true }).click();
|
||||
await expect(dialog(page).locator('#modal-body > h3')).toHaveText('Duplicate prompt text');
|
||||
const hold = await (await request.post('/__test__/jobs/hold', { data: { title: 'Database restore (admin page)' } })).json();
|
||||
await dialog(page).getByRole('button', { name: 'Apply reviewed operation', exact: true }).click();
|
||||
await expect(toast(page)).toHaveText('Another job is running: Database restore (admin page).');
|
||||
await expect(toast(page).getByRole('link')).toHaveAttribute('href', `#/activity/${hold.job.id}`);
|
||||
await expect(page).toHaveURL(`${workspace}#/system/maintenance`);
|
||||
});
|
||||
|
||||
test('Find duplicates starts a job; its view links Review groups; Keep radios build a duplicate_files plan (D101)',
|
||||
async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/duplicate-files')).toBeOK();
|
||||
const previews = [];
|
||||
page.on('request', item => { if (item.url().endsWith(PREVIEW_URL)) previews.push(item.postDataJSON()); });
|
||||
await openMaintenance(page);
|
||||
await row(page, 'Duplicate image files').getByRole('button', { name: 'Find duplicates', exact: true }).click();
|
||||
await expect(page).toHaveURL(/#\/activity\/[0-9a-f]{32}$/);
|
||||
const jobId = page.url().split('/').pop();
|
||||
const main = page.locator('#workspace-main');
|
||||
await expect(main.locator('h1')).toHaveText('Find duplicate image files');
|
||||
const reviewGroups = main.getByRole('button', { name: 'Review groups', exact: true });
|
||||
await expect(reviewGroups).toBeVisible();
|
||||
await expect(main.locator('.stat-row .stat dt')).toHaveText(['Groups found']);
|
||||
await reviewGroups.click();
|
||||
const review = dialog(page);
|
||||
await expect(review.locator('#modal-title')).toHaveText('Review duplicate image files');
|
||||
const group = review.locator('.compare-grid').filter({ hasText: 's7-dupes/copy.png' });
|
||||
await expect(group.locator('.compare-item img')).toHaveCount(3);
|
||||
const labels = (await group.locator('.compare-item label').allTextContents()).map(text => text.trim());
|
||||
expect([...labels].sort()).toEqual(['Keep input/s7-dupes/same.png', 'Keep output/s7-dupes/copy.png',
|
||||
'Keep output/s7-dupes/same.png']);
|
||||
await expect(review.locator('#modal-body > p.small').last())
|
||||
.toHaveText(/History keeps the removed paths in its run records\./);
|
||||
// Keep the input copy: the same relative path as the output copy, told apart by its source index.
|
||||
await group.locator('.compare-item').filter({ hasText: 'input/s7-dupes/same.png' }).getByRole('radio').check();
|
||||
await review.getByRole('button', { name: 'Review duplicate cleanup', exact: true }).click();
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Review maintenance plan');
|
||||
await expect(dialog(page).locator('#modal-body > h3')).toHaveText('Duplicate image files');
|
||||
expect(previews).toHaveLength(1);
|
||||
expect(previews[0].operation).toBe('duplicate_files');
|
||||
expect(previews[0].options.job_id).toBe(jobId);
|
||||
expect(previews[0].options.keep).toContainEqual(expect.objectContaining({ root_index: 1, path: 's7-dupes/same.png' }));
|
||||
for (const entry of previews[0].options.keep) expect(Object.keys(entry).sort()).toEqual(['group_id', 'path', 'root_index']);
|
||||
expect((await gridOf(page))[3][1]).toMatch(/^Deletes 2 files \(.+\) from disk; the kept copies stay\./);
|
||||
await dialog(page).getByRole('button', { name: 'Apply reviewed operation', exact: true }).click();
|
||||
await expect(page).not.toHaveURL(new RegExp(jobId));
|
||||
await expect(page).toHaveURL(/#\/activity\/[0-9a-f]{32}$/);
|
||||
const cleanup = await ended(request, page.url().split('/').pop());
|
||||
expect([cleanup.kind, cleanup.state, cleanup.counters.removed]).toEqual(['duplicate_cleanup', 'finished', 2]);
|
||||
await expect(main.locator('.stat-row .stat dt')).toHaveText(['Removed', 'Kept', 'Errors']);
|
||||
expect((await request.get('/prompt_manager/images/serve/s7-dupes/same.png?root=1')).status()).toBe(200);
|
||||
expect((await request.get('/prompt_manager/images/serve/s7-dupes/same.png?root=0')).status()).toBe(404);
|
||||
expect((await request.get('/prompt_manager/images/serve/s7-dupes/copy.png?root=0')).status()).toBe(404);
|
||||
});
|
||||
|
||||
test('the duplicates review pages ten groups at a time and keeps each choice (D101)', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/duplicate-files', { data: { groups: 11 } })).toBeOK();
|
||||
const groupReads = [];
|
||||
page.on('request', item => { if (item.url().includes('/maintenance/duplicates/')) groupReads.push(new URL(item.url()).search); });
|
||||
const previews = [];
|
||||
page.on('request', item => { if (item.url().endsWith(PREVIEW_URL)) previews.push(item.postDataJSON()); });
|
||||
await openMaintenance(page);
|
||||
await row(page, 'Duplicate image files').getByRole('button', { name: 'Find duplicates', exact: true }).click();
|
||||
const main = page.locator('#workspace-main');
|
||||
await main.getByRole('button', { name: 'Review groups', exact: true }).click();
|
||||
const review = dialog(page);
|
||||
const status = review.locator('#modal-body > p.small[role=status]');
|
||||
await expect(status).toHaveText('Groups 1–10 of 12');
|
||||
await expect(review.locator('.compare-grid')).toHaveCount(10);
|
||||
await expect(review.getByRole('button', { name: 'Previous groups', exact: true })).toBeDisabled();
|
||||
// The three-copy group sorts first (the most bytes to free): keep its input copy.
|
||||
const first = review.locator('.compare-grid').first();
|
||||
await first.locator('.compare-item').filter({ hasText: 'input/s7-dupes/same.png' }).getByRole('radio').check();
|
||||
await review.getByRole('button', { name: 'Next groups', exact: true }).click();
|
||||
await expect(status).toHaveText('Groups 11–12 of 12');
|
||||
await expect(review.locator('.compare-grid')).toHaveCount(2);
|
||||
await expect(review.getByRole('button', { name: 'Next groups', exact: true })).toBeDisabled();
|
||||
await review.getByRole('button', { name: 'Previous groups', exact: true }).click();
|
||||
await expect(status).toHaveText('Groups 1–10 of 12');
|
||||
await expect(first.locator('.compare-item').filter({ hasText: 'input/s7-dupes/same.png' }).getByRole('radio')).toBeChecked();
|
||||
expect(groupReads).toEqual(['?offset=0&limit=10', '?offset=10&limit=10', '?offset=0&limit=10']);
|
||||
await review.getByRole('button', { name: 'Review duplicate cleanup', exact: true }).click();
|
||||
await expect(dialog(page).locator('#modal-body > h3')).toHaveText('Duplicate image files');
|
||||
// Every group shown so far is reviewed, each with its chosen copy.
|
||||
expect(previews[0].options.keep).toHaveLength(12);
|
||||
expect(previews[0].options.keep).toContainEqual(expect.objectContaining({ root_index: 1, path: 's7-dupes/same.png' }));
|
||||
await dialog(page).getByRole('button', { name: 'Cancel', exact: true }).click();
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
});
|
||||
|
||||
for (const [width, height] of [[1600, 1100], [390, 844]]) {
|
||||
test(`long copy paths wrap; the duplicates review never scrolls sideways ${width}x${height}`, async ({ page, request }) => {
|
||||
await page.setViewportSize({ width, height });
|
||||
await expect(await request.post('/__test__/system/duplicate-files', { data: { long: true, trio: false } })).toBeOK();
|
||||
await openMaintenance(page);
|
||||
await row(page, 'Duplicate image files').getByRole('button', { name: 'Find duplicates', exact: true }).click();
|
||||
await page.locator('#workspace-main').getByRole('button', { name: 'Review groups', exact: true }).click();
|
||||
const review = dialog(page);
|
||||
await expect(review.locator('.compare-item')).toHaveCount(2);
|
||||
const overflow = await review.evaluate(node => {
|
||||
const body = node.querySelector('#modal-body');
|
||||
// An item wider than its box, or a label (inline: measured by its boxes) past its item's right edge.
|
||||
const wide = [...node.querySelectorAll('.compare-item')].filter(item => item.scrollWidth > item.clientWidth + 1
|
||||
|| [...item.querySelector('label').getClientRects()].some(rect => rect.right > item.getBoundingClientRect().right + 1)).length;
|
||||
return { dialog: node.scrollWidth > node.clientWidth + 1, body: body.scrollWidth > body.clientWidth + 1, wide,
|
||||
page: document.documentElement.scrollWidth > document.documentElement.clientWidth };
|
||||
});
|
||||
expect(overflow).toEqual({ dialog: false, body: false, wide: 0, page: false });
|
||||
});
|
||||
}
|
||||
|
||||
test('Thumbnail cache opens the S5 dialog', async ({ page }) => {
|
||||
await openMaintenance(page);
|
||||
const options = row(page, 'Thumbnail cache').getByRole('button', { name: 'Review options', exact: true });
|
||||
await options.click();
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Thumbnail maintenance');
|
||||
await expect(dialog(page).locator('.field > label')).toHaveText(['Scope']);
|
||||
await page.keyboard.press('Escape');
|
||||
await expect(options).toBeFocused();
|
||||
});
|
||||
@@ -206,9 +206,10 @@ test('Remove marks the draft; removing the last entry states auto-detect', async
|
||||
expect(await rowsOf(page)).toEqual([['p', 'small', 'No source is configured. After saving, ComfyUI’s output folder is used.']]);
|
||||
await saveButton(page).click();
|
||||
await expect(statusOf(page)).toHaveText('Saved to the server configuration.');
|
||||
// Saved: ComfyUI's output folder is detected and listed without Remove, its row saying it is the default.
|
||||
// Saved: ComfyUI's output folder is detected and listed without Remove.
|
||||
await expect.poll(rowsOf.bind(null, page)).toEqual([
|
||||
['ComfyUI output', 'ComfyUI/output · manual scan · default', ['Scan folder']],
|
||||
['ComfyUI output', 'output · manual scan', ['Scan folder']],
|
||||
['p', 'small', 'No source is configured; ComfyUI’s output folder is used.'],
|
||||
]);
|
||||
const { sources } = await serverJson(request, SETTINGS_URL);
|
||||
expect([sources.configured, sources.auto_detected]).toEqual([[], true]);
|
||||
@@ -331,25 +332,27 @@ function thumbnailRequests(page) {
|
||||
const missingLine = plan => `${planCount(plan.targets)} ${plan.targets === 1 ? 'preview' : 'previews'} to generate · ` +
|
||||
`${planCount(plan.fresh)} already up to date.`;
|
||||
|
||||
const SETTINGS_LINE = 'Previews are written as WebP · quality 80 · 512 px. Change this in System › Appearance.';
|
||||
|
||||
test('Review options shows the plan counts for the chosen scope and the server settings before Start', async ({ page, request }) => {
|
||||
test('Review options shows the plan counts for the chosen scope/quality/size before Start', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/unthumbnailed', { data: { count: 3 } })).toBeOK();
|
||||
// Image loading quality (Appearance, Q3) preselects the dialog's quality.
|
||||
await page.addInitScript(() => { if (location.protocol.startsWith('http')) localStorage.setItem('gallerySettings', JSON.stringify({ imageQuality: 'high' })); });
|
||||
const seen = thumbnailRequests(page);
|
||||
const plan = await serverJson(request, `${THUMBNAIL_PLAN_URL}?mode=missing`);
|
||||
const plan = await serverJson(request, `${THUMBNAIL_PLAN_URL}?mode=missing&size=512&quality=high`);
|
||||
expect(plan.targets).toBe(3);
|
||||
const review = await openThumbnailReview(page);
|
||||
// D131: Scope, the plan line, the server's settings line, the actions (no quality or size field of its own).
|
||||
await expect(review.locator('#modal-body > *')).toHaveCount(4);
|
||||
await expect(review.locator('.field > label')).toHaveText(['Scope']);
|
||||
await expect(review.locator('#thumb-mode option')).toHaveText(['Missing thumbnails only', 'Rebuild all (explicit opt-in)',
|
||||
'Convert outdated (removes older formats)']);
|
||||
await expect(review.locator('#modal-body > *')).toHaveCount(5);
|
||||
await expect(review.locator('.field > label')).toHaveText(['Scope', 'Quality', 'Preview size limit']);
|
||||
await expect(review.locator('#thumb-mode option')).toHaveText(['Missing thumbnails only', 'Rebuild all (explicit opt-in)']);
|
||||
await expect(review.locator('#thumb-quality option')).toHaveText(['Medium · optimized', 'High · original', 'Low · fastest']);
|
||||
await expect(review.locator('#thumb-mode')).toHaveValue('missing');
|
||||
const [line, settings] = [review.locator('#modal-body > p.small').first(), review.locator('#modal-body > p.small').last()];
|
||||
await expect(review.locator('#thumb-quality')).toHaveValue('high');
|
||||
await expect(review.locator('#thumb-limit')).toHaveValue('512');
|
||||
await expect(review.locator('#thumb-limit')).toHaveAttribute('min', '128');
|
||||
await expect(review.locator('#thumb-limit')).toHaveAttribute('max', '2048');
|
||||
const line = review.locator('#modal-body > p.small');
|
||||
await expect(line).toHaveAttribute('role', 'status');
|
||||
await expect(line).toHaveText(missingLine(plan));
|
||||
await expect(settings).toHaveText(SETTINGS_LINE);
|
||||
expect(seen.plans).toEqual([{ mode: 'missing' }]);
|
||||
expect(seen.plans).toEqual([{ mode: 'missing', size: '512', quality: 'high' }]);
|
||||
await expect(review.locator('.dialog-actions > button')).toHaveText(['Cancel', 'Start']);
|
||||
// Reviewing starts nothing; Cancel closes.
|
||||
await review.getByRole('button', { name: 'Cancel', exact: true }).click();
|
||||
@@ -358,67 +361,63 @@ test('Review options shows the plan counts for the chosen scope and the server s
|
||||
expect((await serverJson(request, '/prompt_manager/jobs')).jobs).toEqual([]);
|
||||
});
|
||||
|
||||
test('Convert outdated replaces older-format previews from their originals and removes the old files (#317)', async ({ page, request }) => {
|
||||
const { files } = await (await request.post('/__test__/system/legacy-previews', { data: { count: 3 } })).json();
|
||||
const before = await serverJson(request, '/prompt_manager/images/thumbnails/status');
|
||||
expect(before.outdated).toBe(3);
|
||||
test('size outside 128–2048 reports validity and requests nothing', async ({ page, request }) => {
|
||||
const seen = thumbnailRequests(page);
|
||||
const review = await openThumbnailReview(page);
|
||||
const line = review.locator('#modal-body > p.small').first();
|
||||
const line = review.locator('#modal-body > p.small');
|
||||
await expect(line).toHaveText(/ already up to date\.$/);
|
||||
await review.locator('#thumb-mode').selectOption('convert');
|
||||
await expect(line).toHaveText('3 previews in an older format to convert.');
|
||||
// Choosing the scope converts nothing: only Start does.
|
||||
expect(seen.starts).toEqual([]);
|
||||
await review.getByRole('button', { name: 'Start', exact: true }).click();
|
||||
await expect(page).toHaveURL(/#\/activity\/[0-9a-f]{32}$/);
|
||||
expect(seen.starts).toEqual([{ mode: 'convert' }]);
|
||||
const id = page.url().split('/').pop();
|
||||
await expect.poll(async () => (await serverJson(request, `/prompt_manager/jobs/${id}`)).job.state).toBe('finished');
|
||||
const { job } = await serverJson(request, `/prompt_manager/jobs/${id}`);
|
||||
expect(job).toMatchObject({ kind: 'thumbnails', title: 'Convert outdated thumbnails', scope: { mode: 'convert', format: 'webp' } });
|
||||
expect(job.status_line).toMatch(/^3 converted · 0 failed · .+ → .+$/);
|
||||
await expect(page.locator('#workspace-main h1')).toHaveText('Convert outdated thumbnails');
|
||||
const after = await serverJson(request, '/prompt_manager/images/thumbnails/status');
|
||||
expect(after.outdated).toBe(0);
|
||||
// The listing now names the WebP previews; the old PNG previews are gone.
|
||||
for (const file of files) {
|
||||
const name = file.replace(/\.png$/, '');
|
||||
expect((await request.get(`/prompt_manager/images/serve/thumbnails/${name}_thumb.webp?root=0`)).headers()['content-type']).toBe('image/webp');
|
||||
expect((await request.get(`/prompt_manager/images/serve/${file}?root=0`)).ok()).toBe(true);
|
||||
const size = review.locator('#thumb-limit');
|
||||
for (const value of ['4096', '64', '']) {
|
||||
await size.fill(value);
|
||||
await expect(size).toHaveAttribute('aria-invalid', 'true');
|
||||
await expect(line).toHaveText('Enter a size from 128 to 2048 pixels.');
|
||||
expect(await size.evaluate(input => input.validity.valid)).toBe(false);
|
||||
// Start keeps the dialog open and sends nothing; neither does a scope change.
|
||||
await review.getByRole('button', { name: 'Start', exact: true }).click();
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Thumbnail maintenance');
|
||||
await review.locator('#thumb-mode').selectOption('rebuild');
|
||||
await expect(line).toHaveText('Enter a size from 128 to 2048 pixels.');
|
||||
await review.locator('#thumb-mode').selectOption('missing');
|
||||
}
|
||||
expect(seen.starts).toEqual([]);
|
||||
expect(seen.plans).toEqual([{ mode: 'missing', size: '512', quality: 'medium' }]);
|
||||
await size.fill('1024');
|
||||
await expect(size).not.toHaveAttribute('aria-invalid');
|
||||
await expect(line).toHaveText(/ already up to date\.$/);
|
||||
expect((await serverJson(request, '/prompt_manager/jobs')).jobs).toEqual([]);
|
||||
});
|
||||
|
||||
test('Start closes the dialog, toasts, and opens the job in Activity', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/unthumbnailed', { data: { count: 2 } })).toBeOK();
|
||||
const seen = thumbnailRequests(page);
|
||||
const review = await openThumbnailReview(page);
|
||||
await review.locator('#thumb-quality').selectOption('low');
|
||||
await review.locator('#thumb-limit').fill('256');
|
||||
await review.getByRole('button', { name: 'Start', exact: true }).click();
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
await expect(page).toHaveURL(/#\/activity\/[0-9a-f]{32}$/);
|
||||
// The toast says it started (or, for a run this quick, that it already finished); Activity shows its progress.
|
||||
await expect(toast(page)).toHaveText(/^(Thumbnail generation started\.|Thumbnails finished · 2 generated · 0 errors)$/);
|
||||
expect(seen.starts).toEqual([{ mode: 'missing' }]);
|
||||
expect(seen.starts).toEqual([{ mode: 'missing', size: 256, quality: 'low' }]);
|
||||
const id = page.url().split('/').pop();
|
||||
const { job } = await serverJson(request, `/prompt_manager/jobs/${id}`);
|
||||
expect(job).toMatchObject({ kind: 'thumbnails', title: 'Thumbnail generation',
|
||||
scope: { mode: 'missing', format: 'webp', size: 512, quality: 80 } });
|
||||
expect(job).toMatchObject({ kind: 'thumbnails', title: 'Thumbnail generation', scope: { mode: 'missing', size: 256, quality: 'low' } });
|
||||
await expect(page.locator('#workspace-main h1')).toHaveText('Thumbnail generation');
|
||||
await expect.poll(async () => (await serverJson(request, `/prompt_manager/jobs/${id}`)).job.state).toBe('finished');
|
||||
});
|
||||
|
||||
test('Rebuild all is explicit and the plan line says every preview is regenerated', async ({ page, request }) => {
|
||||
const seen = thumbnailRequests(page);
|
||||
const rebuild = await serverJson(request, `${THUMBNAIL_PLAN_URL}?mode=rebuild`);
|
||||
const rebuild = await serverJson(request, `${THUMBNAIL_PLAN_URL}?mode=rebuild&size=512&quality=medium`);
|
||||
expect(rebuild.targets).toBeGreaterThan(0);
|
||||
const review = await openThumbnailReview(page);
|
||||
// Missing thumbnails only is the default; rebuilding is the explicit second option.
|
||||
await expect(review.locator('#thumb-mode')).toHaveValue('missing');
|
||||
const line = review.locator('#modal-body > p.small').first();
|
||||
const line = review.locator('#modal-body > p.small');
|
||||
await expect(line).toHaveText(/ already up to date\.$/);
|
||||
await review.locator('#thumb-mode').selectOption('rebuild');
|
||||
await expect(line).toHaveText(`Every preview is regenerated: ${planCount(rebuild.targets)} files.`);
|
||||
expect(seen.plans).toEqual([{ mode: 'missing' }, { mode: 'rebuild' }]);
|
||||
expect(seen.plans).toEqual([{ mode: 'missing', size: '512', quality: 'medium' }, { mode: 'rebuild', size: '512', quality: 'medium' }]);
|
||||
// Back and forth reads each scope once.
|
||||
await review.locator('#thumb-mode').selectOption('missing');
|
||||
await expect(line).toHaveText(/ already up to date\.$/);
|
||||
@@ -427,7 +426,7 @@ test('Rebuild all is explicit and the plan line says every preview is regenerate
|
||||
expect(seen.plans).toHaveLength(2);
|
||||
await review.getByRole('button', { name: 'Start', exact: true }).click();
|
||||
await expect(page).toHaveURL(/#\/activity\/[0-9a-f]{32}$/);
|
||||
expect(seen.starts).toEqual([{ mode: 'rebuild' }]);
|
||||
expect(seen.starts).toEqual([{ mode: 'rebuild', size: 512, quality: 'medium' }]);
|
||||
const id = page.url().split('/').pop();
|
||||
await expect.poll(async () => (await serverJson(request, `/prompt_manager/jobs/${id}`)).job.state).toBe('finished');
|
||||
});
|
||||
@@ -435,7 +434,7 @@ test('Rebuild all is explicit and the plan line says every preview is regenerate
|
||||
test('a plan that cannot be read says so; Start still starts', async ({ page }) => {
|
||||
await page.route(`**${THUMBNAIL_PLAN_URL}?*`, route => route.fulfill({ status: 500, json: { success: false, error: 'disk' } }));
|
||||
const review = await openThumbnailReview(page);
|
||||
await expect(review.locator('#modal-body > p.small').first()).toHaveText('The previews could not be counted.');
|
||||
await expect(review.locator('#modal-body > p.small')).toHaveText('The previews could not be counted.');
|
||||
await expect(review.getByRole('button', { name: 'Start', exact: true })).not.toHaveAttribute('aria-disabled');
|
||||
await review.getByRole('button', { name: 'Cancel', exact: true }).click();
|
||||
});
|
||||
|
||||
@@ -1,356 +0,0 @@
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from './fixtures.mjs';
|
||||
|
||||
// System › Auto-tagging (#328, System S11 Tasks 6–7) on the tagging fixture (workspace_tagging_fixture.py) with the
|
||||
// injected fake model: no model file exists and nothing is downloaded or run for real. The fake download waits on the
|
||||
// fixture (`/__test__/tagging/download {mode}`), and every reset releases it. Every test asserts that no legacy
|
||||
// download stream (`/autotag/download/`) was requested: the section downloads only through the Activity job.
|
||||
const workspace = '/prompt_manager/workspace';
|
||||
const MODELS_URL = '/prompt_manager/autotag/models';
|
||||
const PREVIEW_URL = '/prompt_manager/autotag/batch/preview';
|
||||
const SETTINGS_URL = '/prompt_manager/settings';
|
||||
const MODEL_LABELS = ['WD14 SwinV2 · ~400 MB', 'WD14 ViT · ~400 MB', 'GGUF JoyCaption · ~4.5 GB', 'HuggingFace JoyCaption · ~16 GB'];
|
||||
const FIRST = 'tagging/dunes_00001.png';
|
||||
const SECOND = 'tagging/dunes_00002.png';
|
||||
|
||||
test.beforeEach(async ({ page, request, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
await expect(await request.post('/__test__/tagging/reset', { data: {} })).toBeOK();
|
||||
await expect(await request.post('/__test__/system/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/jobs/restart')).toBeOK();
|
||||
page.systemErrors = [];
|
||||
page.on('pageerror', error => page.systemErrors.push(error.message));
|
||||
page.downloads = [];
|
||||
page.modelCalls = [];
|
||||
page.on('request', item => {
|
||||
const { pathname } = new URL(item.url());
|
||||
if (pathname.includes('/autotag/download/')) page.downloads.push(pathname);
|
||||
if (pathname.startsWith('/prompt_manager/autotag/') || pathname.startsWith('/prompt_manager/image-annotations/suggest')
|
||||
|| pathname.startsWith('/prompt_manager/jobs/model-download')) page.modelCalls.push(`${item.method()} ${pathname}`);
|
||||
});
|
||||
});
|
||||
|
||||
test.afterEach(async ({ page, request, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
// The reset releases a held fake download, so no runner outlives the test.
|
||||
await request.post('/__test__/tagging/reset', { data: {} });
|
||||
await request.post('/__test__/jobs/restart');
|
||||
await request.post('/__test__/system/reset');
|
||||
expect(page.systemErrors).toEqual([]);
|
||||
expect(page.downloads).toEqual([]);
|
||||
});
|
||||
|
||||
const toast = page => page.locator('#workspace-toast');
|
||||
const layoutOf = page => page.locator('#workspace-main .settings-layout');
|
||||
const dialog = page => page.locator('dialog#modal[open]');
|
||||
const status = page => page.locator('#settings-status');
|
||||
const save = page => page.getByRole('button', { name: 'Save preferences', exact: true });
|
||||
const surface = page => layoutOf(page).locator(':scope > section.surface');
|
||||
const settingOf = (page, label) => layoutOf(page).locator('.setting-row').filter({ has: page.locator(`label:text-is("${label}")`) });
|
||||
const modelRow = page => settingOf(page, 'Tagging model');
|
||||
const actionOf = (page, name) => surface(page).locator(':scope > .row-actions').getByRole('button', { name, exact: true });
|
||||
const reviewButton = page => surface(page).locator('.operation').getByRole('button', { name: 'Review', exact: true });
|
||||
const card = (page, path) => page.locator(`article[data-image-key="0:${path}"]`);
|
||||
const tagDialog = page => page.getByRole('dialog', { name: 'Tag image', exact: true });
|
||||
const gridOf = page => dialog(page).locator('dl.meta-grid > div').evaluateAll(items => items.map(item =>
|
||||
[item.querySelector('dt').textContent, item.querySelector('dd').textContent]));
|
||||
|
||||
async function serverJson(request, url) {
|
||||
const response = await request.get(url);
|
||||
await expect(response).toBeOK();
|
||||
return response.json();
|
||||
}
|
||||
|
||||
async function fakeDownload(request, mode) {
|
||||
await expect(await request.post('/__test__/tagging/download', { data: { mode } })).toBeOK();
|
||||
}
|
||||
|
||||
async function unavailable(request, models) {
|
||||
await expect(await request.post('/__test__/tagging/models', { data: { unavailable: models } })).toBeOK();
|
||||
}
|
||||
|
||||
async function ended(request, id) {
|
||||
await expect.poll(async () => (await serverJson(request, `/prompt_manager/jobs/${id}`)).job.state).not.toBe('running');
|
||||
return (await serverJson(request, `/prompt_manager/jobs/${id}`)).job;
|
||||
}
|
||||
|
||||
/** System › Auto-tagging on a fresh page (the drafts live for the page session, Q10), settled on the model status. */
|
||||
async function openTagging(page) {
|
||||
if (page.url().endsWith('#/system/tagging')) await page.reload();
|
||||
else await page.goto(`${workspace}#/system/tagging`);
|
||||
await expect(surface(page)).toHaveCount(1);
|
||||
await expect(page.locator('#tag-model')).toBeEnabled();
|
||||
await expect(modelRow(page).locator('p').first()).not.toHaveText('Loading model status…');
|
||||
}
|
||||
|
||||
/** Select `path` in Images (the concept's `state.selectedImage`), then open System › Auto-tagging from the sidebar. */
|
||||
async function selectThenOpen(page, path) {
|
||||
await page.goto(`${workspace}#/images`);
|
||||
await card(page, path).getByRole('button', { name: `Inspect ${path.split('/').pop()}`, exact: true }).click();
|
||||
await expect(card(page, path)).toHaveClass(/\bselected\b/);
|
||||
await page.locator('aside.sidebar a[href="#/system"]').click();
|
||||
await page.locator('#workspace-main .section-nav').getByRole('button', { name: 'Auto-tagging', exact: true }).click();
|
||||
await expect(page.locator('#tag-model')).toBeEnabled();
|
||||
await expect(modelRow(page).locator('p').first()).not.toHaveText('Loading model status…');
|
||||
}
|
||||
|
||||
test('model options are the four real engines with size and the description shows live status', async ({ page, request }) => {
|
||||
await unavailable(request, ['wd14-vit', 'hf']);
|
||||
await openTagging(page);
|
||||
expect(await surface(page).locator(':scope > h2').textContent()).toBe('Auto-tagging');
|
||||
expect(await page.locator('#tag-model option').allTextContents()).toEqual(MODEL_LABELS);
|
||||
await expect(page.locator('#tag-model')).toHaveValue('wd14-swinv2');
|
||||
const line = modelRow(page).locator('p').first();
|
||||
await expect(line).toHaveText('Downloaded. Not loaded.');
|
||||
await expect(actionOf(page, 'Download again')).toBeVisible();
|
||||
// Choosing another engine drafts it and tells its own status; nothing is downloaded or loaded.
|
||||
await page.locator('#tag-model').selectOption('wd14-vit');
|
||||
await expect(line).toHaveText('Not downloaded. Download size ~400 MB.');
|
||||
await expect(actionOf(page, 'Download')).toBeVisible();
|
||||
await page.locator('#tag-model').selectOption('hf');
|
||||
await expect(line).toHaveText('Not downloaded. Download size ~16 GB.');
|
||||
await expect(status(page)).toHaveText('Unsaved changes.');
|
||||
// The other rows keep the concept's copy and the saved values (D133 adds Suggestions per image).
|
||||
expect(await surface(page).locator(':scope > .setting-row label').allTextContents()).toEqual(['Tagging model', 'Existing tags',
|
||||
'General-tag confidence', 'Character-tag confidence', 'Suggestions per image']);
|
||||
await expect(page.locator('#tag-existing')).toHaveValue('skip');
|
||||
await expect(page.locator('#threshold-general')).toHaveValue('0.35');
|
||||
await expect(page.locator('#threshold-character')).toHaveValue('0.85');
|
||||
await expect(page.locator('#tag-suggestions')).toHaveValue('10');
|
||||
const settings = await serverJson(request, SETTINGS_URL);
|
||||
await expect(page.locator('#tag-template')).toHaveValue(settings.settings.autotag_template);
|
||||
// Unload is offered only while a model is loaded (none is).
|
||||
await expect(actionOf(page, 'Unload')).toBeDisabled();
|
||||
});
|
||||
|
||||
test('entering the section makes no model call besides status', async ({ page }) => {
|
||||
await openTagging(page);
|
||||
await page.locator('#workspace-main .section-nav').getByRole('button', { name: 'Integrations', exact: true }).click();
|
||||
await page.locator('#workspace-main .section-nav').getByRole('button', { name: 'Auto-tagging', exact: true }).click();
|
||||
await expect(page.locator('#tag-model')).toBeEnabled();
|
||||
await expect(modelRow(page).locator('p').first()).toHaveText('Downloaded. Not loaded.');
|
||||
// One status read per mount (two mounts), and nothing else that could load, run or download a model.
|
||||
expect(page.modelCalls).toEqual([`GET ${MODELS_URL}`, `GET ${MODELS_URL}`]);
|
||||
const models = await (await page.request.get(MODELS_URL)).json();
|
||||
expect(models.model_loaded).toBe(false);
|
||||
});
|
||||
|
||||
test('defaults save through Save preferences and the Tag image dialog uses them', async ({ page, request }) => {
|
||||
await openTagging(page);
|
||||
await page.locator('#tag-existing').selectOption('merge');
|
||||
await page.locator('#threshold-general').fill('0.4');
|
||||
await page.locator('#threshold-general').press('Tab');
|
||||
await page.locator('#tag-suggestions').fill('15');
|
||||
await page.locator('#tag-suggestions').press('Tab');
|
||||
const posted = page.waitForRequest(item => item.method() === 'POST' && new URL(item.url()).pathname === SETTINGS_URL);
|
||||
await save(page).click();
|
||||
const body = (await posted).postDataJSON();
|
||||
expect(body).toMatchObject({ autotag_existing: 'merge', autotag_general_threshold: 0.4, autotag_suggestion_limit: 15 });
|
||||
expect(typeof body.expected_revision).toBe('string');
|
||||
await expect(status(page)).toHaveText('Saved to the server configuration.');
|
||||
const saved = (await serverJson(request, SETTINGS_URL)).settings;
|
||||
expect([saved.autotag_existing, saved.autotag_general_threshold, saved.autotag_character_threshold, saved.autotag_suggestion_limit])
|
||||
.toEqual(['merge', 0.4, 0.85, 15]);
|
||||
const models = await serverJson(request, MODELS_URL);
|
||||
expect([models.existing, models.wd14_general_threshold, models.suggestion_limit]).toEqual(['merge', 0.4, 15]);
|
||||
// The Tag image dialog in a fresh page starts from them.
|
||||
await page.goto(`${workspace}#/images`);
|
||||
await card(page, FIRST).getByRole('button', { name: `Tag or auto-tag ${FIRST.split('/').pop()}`, exact: true }).click();
|
||||
await expect(tagDialog(page)).toBeVisible();
|
||||
await expect(tagDialog(page).locator('#image-general-threshold')).toHaveValue('0.4');
|
||||
await expect(tagDialog(page).locator('#image-character-threshold')).toHaveValue('0.85');
|
||||
await expect(tagDialog(page).locator('#image-suggestion-limit')).toHaveValue('15');
|
||||
});
|
||||
|
||||
test('a refused value shows the server\'s message', async ({ page, request }) => {
|
||||
const before = (await serverJson(request, SETTINGS_URL)).settings;
|
||||
// An image is selected, so Preview tagging review is held back only by the invalid value below.
|
||||
await selectThenOpen(page, FIRST);
|
||||
await page.locator('#threshold-general').fill('1.5');
|
||||
await page.locator('#threshold-general').press('Tab');
|
||||
await save(page).click();
|
||||
await expect(status(page)).toHaveText('General-tag confidence must be between 0 and 1.');
|
||||
await expect(page.locator('#threshold-general')).toHaveAttribute('aria-invalid', 'true');
|
||||
// Nothing was written; the draft is kept for the next Save.
|
||||
expect((await serverJson(request, SETTINGS_URL)).settings.autotag_general_threshold).toBe(before.autotag_general_threshold);
|
||||
await expect(page.locator('#threshold-general')).toHaveValue('1.5');
|
||||
// While the value is invalid, Preview tagging review is aria-disabled and explains instead of opening (Playwright
|
||||
// waits on an aria-disabled control, so it is pressed from the keyboard, as a user can).
|
||||
await expect(actionOf(page, 'Preview tagging review')).toHaveAttribute('aria-disabled', 'true');
|
||||
await actionOf(page, 'Preview tagging review').focus();
|
||||
await page.keyboard.press('Enter');
|
||||
await expect(toast(page)).toHaveText('General-tag confidence must be between 0 and 1.');
|
||||
await expect(tagDialog(page)).toHaveCount(0);
|
||||
await page.locator('#threshold-general').fill('0.5');
|
||||
await page.locator('#threshold-general').press('Tab');
|
||||
await expect(page.locator('#threshold-general')).not.toHaveAttribute('aria-invalid', 'true');
|
||||
await expect(status(page)).toHaveText('Unsaved changes.');
|
||||
await expect(actionOf(page, 'Preview tagging review')).not.toHaveAttribute('aria-disabled', 'true');
|
||||
});
|
||||
|
||||
test('Download starts the job and Activity shows it without Cancel', async ({ page, request }) => {
|
||||
await unavailable(request, ['wd14-vit']);
|
||||
await fakeDownload(request, 'hold');
|
||||
await openTagging(page);
|
||||
await page.locator('#tag-model').selectOption('wd14-vit');
|
||||
const started = page.waitForResponse(item => new URL(item.url()).pathname === '/prompt_manager/jobs/model-download');
|
||||
await actionOf(page, 'Download').click();
|
||||
const answer = await started;
|
||||
expect(answer.status()).toBe(202);
|
||||
expect(answer.request().postDataJSON()).toEqual({ model_type: 'wd14-vit' });
|
||||
const { job } = await answer.json();
|
||||
expect(job.title).toBe('Tagging model download · WD14 ViT');
|
||||
await expect(toast(page)).toHaveText('Tagging model download · WD14 ViT started.');
|
||||
await expect(page).toHaveURL(`${workspace}#/activity/${job.id}`);
|
||||
const main = page.locator('#workspace-main');
|
||||
await expect(main.locator('h1')).toHaveText('Tagging model download · WD14 ViT');
|
||||
await expect(main.getByRole('button', { name: /^Cancel/ })).toHaveCount(0);
|
||||
// Indeterminate: the server knows no byte progress.
|
||||
await expect(main.locator('progress').first()).toBeVisible();
|
||||
expect(await main.locator('progress').first().evaluate(node => node.hasAttribute('value'))).toBe(false);
|
||||
await fakeDownload(request, 'succeed');
|
||||
expect((await ended(request, job.id)).state).toBe('finished');
|
||||
await expect(toast(page)).toHaveText('Tagging model download · WD14 ViT finished.');
|
||||
// Back in the section, the engine reads as on disk (the fake marks it downloaded on success).
|
||||
await page.goto(`${workspace}#/system/tagging`);
|
||||
await expect(page.locator('#tag-model')).toBeEnabled();
|
||||
await page.locator('#tag-model').selectOption('wd14-vit');
|
||||
await expect(modelRow(page).locator('p').first()).toHaveText('Downloaded. Not loaded.');
|
||||
expect(page.modelCalls.filter(call => call.startsWith('POST'))).toEqual(['POST /prompt_manager/jobs/model-download']);
|
||||
});
|
||||
|
||||
test('Unload is enabled only while a model is loaded, and a held gate\'s refusal names the operation', async ({ page, request }) => {
|
||||
await openTagging(page);
|
||||
await expect(actionOf(page, 'Unload')).toBeDisabled();
|
||||
// Load WD14 SwinV2 the way a user does: Generate suggestions in the Tag image dialog.
|
||||
await page.goto(`${workspace}#/images`);
|
||||
await card(page, FIRST).getByRole('button', { name: `Tag or auto-tag ${FIRST.split('/').pop()}`, exact: true }).click();
|
||||
await expect(tagDialog(page)).toBeVisible();
|
||||
const suggested = page.waitForResponse(item => item.url().endsWith('/image-annotations/suggest'));
|
||||
await tagDialog(page).getByRole('button', { name: 'Generate suggestions', exact: true }).click();
|
||||
expect((await suggested).status()).toBe(200);
|
||||
await tagDialog(page).getByRole('button', { name: 'Cancel', exact: true }).click();
|
||||
await openTagging(page);
|
||||
await expect(modelRow(page).locator('p').first()).toHaveText('Downloaded and loaded.');
|
||||
await expect(actionOf(page, 'Unload')).toBeEnabled();
|
||||
// A download holds the model gate: Unload is refused and the toast names it; nothing is unloaded.
|
||||
await unavailable(request, ['wd14-vit']);
|
||||
await fakeDownload(request, 'hold');
|
||||
const held = await request.post('/prompt_manager/jobs/model-download', { data: { model_type: 'wd14-vit' } });
|
||||
expect(held.status()).toBe(202);
|
||||
const { job } = await held.json();
|
||||
const refused = page.waitForResponse(item => item.url().endsWith('/autotag/unload'));
|
||||
await actionOf(page, 'Unload').click();
|
||||
const refusal = await refused;
|
||||
expect(refusal.status()).toBe(409);
|
||||
expect(await refusal.json()).toMatchObject({ code: 'model-busy', operation: 'Tagging model download · WD14 ViT' });
|
||||
await expect(toast(page)).toHaveText('The tagging model is busy: Tagging model download · WD14 ViT. Try again when it finishes.');
|
||||
expect((await serverJson(request, MODELS_URL)).loaded_model_type).toBe('wd14-swinv2');
|
||||
await expect(actionOf(page, 'Unload')).toBeEnabled();
|
||||
// Once the download ended, Unload works, says so and the row follows.
|
||||
await fakeDownload(request, 'succeed');
|
||||
await ended(request, job.id);
|
||||
await actionOf(page, 'Unload').click();
|
||||
await expect(toast(page)).toHaveText('WD14 SwinV2 unloaded.');
|
||||
await expect(modelRow(page).locator('p').first()).toHaveText('Downloaded. Not loaded.');
|
||||
await expect(actionOf(page, 'Unload')).toBeDisabled();
|
||||
});
|
||||
|
||||
test('Preview tagging review opens the Tag image dialog for the selected image', async ({ page }) => {
|
||||
await selectThenOpen(page, SECOND);
|
||||
// The section's current (unsaved) values are what the dialog starts from.
|
||||
await page.locator('#threshold-character').fill('0.6');
|
||||
await page.locator('#threshold-character').press('Tab');
|
||||
await page.locator('#tag-model').selectOption('wd14-vit');
|
||||
const preview = actionOf(page, 'Preview tagging review');
|
||||
await expect(preview).toHaveClass(/ca-primary/);
|
||||
await expect(preview).not.toHaveAttribute('aria-disabled', 'true');
|
||||
const observed = page.waitForResponse(item => item.url().endsWith('/image-annotations/observe'));
|
||||
await preview.click();
|
||||
expect((await observed).status()).toBe(200);
|
||||
await expect(tagDialog(page)).toBeVisible();
|
||||
await expect(tagDialog(page).locator('.tag-image-heading p.small')).toHaveText('dunes_00002.png');
|
||||
await expect(tagDialog(page).getByRole('combobox', { name: 'Tagging model', exact: true })).toHaveValue('wd14-vit');
|
||||
await expect(tagDialog(page).locator('#image-character-threshold')).toHaveValue('0.6');
|
||||
await expect(tagDialog(page).locator('#image-general-threshold')).toHaveValue('0.35');
|
||||
});
|
||||
|
||||
test('Preview tagging review without a selection toasts the D115 sentence', async ({ page }) => {
|
||||
await openTagging(page);
|
||||
const preview = actionOf(page, 'Preview tagging review');
|
||||
await expect(preview).toHaveAttribute('aria-disabled', 'true');
|
||||
// An aria-disabled control still takes a press; it only explains.
|
||||
await preview.focus();
|
||||
await page.keyboard.press('Enter');
|
||||
await expect(toast(page)).toHaveText('Select an image in Images first.');
|
||||
await expect(tagDialog(page)).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('Tag output images reviews the scope and starts a batch job', async ({ page, request }) => {
|
||||
const before = await serverJson(request, '/__test__/tagging/snapshot');
|
||||
await openTagging(page);
|
||||
// The fixture's one prompt is already tagged: under Merge it is tagged again with what its images give.
|
||||
await page.locator('#tag-existing').selectOption('merge');
|
||||
const previewed = page.waitForResponse(item => new URL(item.url()).pathname === PREVIEW_URL);
|
||||
await reviewButton(page).click();
|
||||
const answer = await previewed;
|
||||
expect(answer.request().postDataJSON()).toEqual({ model_type: 'wd14-swinv2', existing: 'merge' });
|
||||
const plan = await answer.json();
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Review tag output images');
|
||||
await expect(dialog(page).locator('#modal-body > h3')).toHaveText('Tag output images');
|
||||
expect(await gridOf(page)).toEqual([['Model', 'WD14 SwinV2 · ~400 MB'], ['Existing tags', 'Merge with existing'],
|
||||
['Images linked', String(plan.images)], ['Prompts', String(plan.prompts)], ['Already tagged', String(plan.already_tagged)],
|
||||
['Prompts to tag', `up to ${plan.to_tag}`], ['Settings', 'General 0.35 · Character 0.85'],
|
||||
['Backup', 'Create and verify before changes']]);
|
||||
await expect(dialog(page).locator('#modal-body > .notice h3')).toHaveText('Back up the database first (recommended)');
|
||||
// The review wrote nothing.
|
||||
expect(await serverJson(request, '/__test__/tagging/snapshot')).toEqual(before);
|
||||
await expect(dialog(page).locator('.dialog-actions button')).toHaveText(['Cancel', 'Start tagging']);
|
||||
const started = page.waitForResponse(item => new URL(item.url()).pathname === '/prompt_manager/jobs/autotag');
|
||||
await dialog(page).getByRole('button', { name: 'Start tagging', exact: true }).click();
|
||||
const start = await started;
|
||||
expect(start.request().postDataJSON()).toEqual({ model_type: 'wd14-swinv2', existing: 'merge',
|
||||
options: { general_threshold: 0.35, character_threshold: 0.85 } });
|
||||
const { job } = await start.json();
|
||||
await expect(page).toHaveURL(`${workspace}#/activity/${job.id}`);
|
||||
await expect(page.locator('#workspace-main h1')).toHaveText('Tag output images');
|
||||
const finished = await ended(request, job.id);
|
||||
expect(finished.state).toBe('finished');
|
||||
const after = await serverJson(request, '/__test__/tagging/snapshot');
|
||||
expect(after.prompt_tags.length).toBeGreaterThan(before.prompt_tags.length);
|
||||
});
|
||||
|
||||
test('Replace after review shows only Cancel with the reason', async ({ page, request }) => {
|
||||
const before = await serverJson(request, '/__test__/tagging/snapshot');
|
||||
await openTagging(page);
|
||||
await page.locator('#tag-existing').selectOption('replace_after_review');
|
||||
await reviewButton(page).click();
|
||||
await expect(dialog(page).locator('#modal-title')).toHaveText('Review tag output images');
|
||||
expect((await gridOf(page)).map(([term]) => term)).toEqual(['Model', 'Existing tags', 'Images linked', 'Prompts', 'Already tagged',
|
||||
'Prompts to tag', 'Settings', 'Backup']);
|
||||
await expect(dialog(page).locator('#modal-body > p.small')).toHaveText('Replacing tags needs a review of each image. '
|
||||
+ 'Use Preview tagging review, or choose Skip or Merge for a batch run.');
|
||||
await expect(dialog(page).locator('.dialog-actions button')).toHaveText(['Cancel']);
|
||||
await dialog(page).getByRole('button', { name: 'Cancel', exact: true }).click();
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
// Under Skip the fixture's only prompt is already tagged: nothing to start either.
|
||||
await page.locator('#tag-existing').selectOption('skip');
|
||||
await reviewButton(page).click();
|
||||
await expect(dialog(page).locator('#modal-body > p.small')).toHaveText('Every prompt with an output image is already tagged.');
|
||||
await expect(dialog(page).locator('.dialog-actions button')).toHaveText(['Cancel']);
|
||||
expect(await serverJson(request, '/__test__/tagging/snapshot')).toEqual(before);
|
||||
expect((await serverJson(request, '/prompt_manager/jobs')).jobs.filter(item => item.kind === 'autotag')).toEqual([]);
|
||||
});
|
||||
|
||||
test('a model that is not downloaded cannot be reviewed for a batch; the button says why', async ({ page, request }) => {
|
||||
await unavailable(request, ['gguf']);
|
||||
await openTagging(page);
|
||||
await page.locator('#tag-model').selectOption('gguf');
|
||||
await expect(reviewButton(page)).toHaveAttribute('aria-disabled', 'true');
|
||||
// Playwright waits on an aria-disabled control: press it from the keyboard, as a user can.
|
||||
await reviewButton(page).focus();
|
||||
await page.keyboard.press('Enter');
|
||||
await expect(toast(page)).toHaveText('GGUF JoyCaption is not downloaded. Download it first.');
|
||||
await expect(dialog(page)).toHaveCount(0);
|
||||
expect(page.modelCalls.filter(call => call.includes('/batch/preview'))).toEqual([]);
|
||||
});
|
||||
@@ -36,9 +36,6 @@ const layoutOf = page => page.locator('#workspace-main .settings-layout');
|
||||
async function openOverview(page) {
|
||||
await page.goto(`${workspace}#/system/overview`);
|
||||
await expect(layoutOf(page).locator(':scope > .summary-metrics')).toBeVisible();
|
||||
// The Backup status row is its own read of /backup/status (S9): compare once it answered.
|
||||
await expect(layoutOf(page).locator('.operation').filter({ has: page.locator('h3', { hasText: 'Backup status' }) })
|
||||
.locator('p')).not.toHaveText('Checking backups…');
|
||||
return layoutOf(page);
|
||||
}
|
||||
|
||||
@@ -70,9 +67,7 @@ test('overview shows real counts and Not checked', async ({ page, request }) =>
|
||||
]);
|
||||
expect(await rowsOf(layout, 'Library health')).toEqual([
|
||||
['0 ambiguous image links', 'Review metadata before assigning a positive / negative pair.', 'Review metadata'],
|
||||
['Backup status', 'No backup recorded.', 'View backups'],
|
||||
// D129: the fixture's database is configured explicitly, so it reads as config.json's choice and does not move.
|
||||
['Library location', 'browser.db · The database path set in config.json.', 'Details'],
|
||||
['Backup status', 'No retained backup record.', 'View backups'],
|
||||
['No missing thumbnails', 'Generate only missing previews when necessary.', 'Review thumbnails'],
|
||||
]);
|
||||
expect(await rowsOf(layout, 'Workspace preferences')).toEqual([
|
||||
@@ -83,101 +78,12 @@ test('overview shows real counts and Not checked', async ({ page, request }) =>
|
||||
expect(writes).toEqual([]);
|
||||
});
|
||||
|
||||
// D129 (#249 S6b): where the library lives. The hook makes the fixture's database read as a node-folder library; its
|
||||
// Move library is a stand-in that touches no file (workspace_system_fixture.py).
|
||||
const banner = page => page.locator('#library-banner');
|
||||
const LIBRARY_URL = '/prompt_manager/system/library';
|
||||
|
||||
test('a node-folder library is offered a move above every page; Keep it here stops the offer', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/library', { data: { state: 'node_folder' } })).toBeOK();
|
||||
await page.goto(`${workspace}#/images`);
|
||||
await expect(banner(page).locator('.notice h2')).toHaveText('Your library is inside the node folder');
|
||||
await expect(banner(page).locator('.notice p')).toHaveText(/^Your library is stored inside the node folder, where an update or reinstall can delete it\./);
|
||||
await expect(banner(page).locator('.notice .row-actions button')).toHaveText(['Keep it here', 'Move library']);
|
||||
await expect(banner(page).getByRole('button', { name: 'Move library', exact: true })).toHaveClass(/ca-primary/);
|
||||
await page.goto(`${workspace}#/library`);
|
||||
await expect(banner(page).locator('.notice')).toBeVisible();
|
||||
await banner(page).getByRole('button', { name: 'Keep it here', exact: true }).click();
|
||||
await expect(banner(page)).toBeHidden();
|
||||
const { location } = await serverJson(request, LIBRARY_URL);
|
||||
expect([location.keep, location.offer, location.movable]).toEqual([true, false, true]);
|
||||
await page.reload();
|
||||
const layout = await openOverview(page);
|
||||
await expect(banner(page)).toBeHidden();
|
||||
const row = layout.locator('.operation').filter({ hasText: 'Library location' });
|
||||
await expect(row.locator('p')).toHaveText('browser.db · The node folder, where an update or reinstall can delete it. '
|
||||
+ 'You chose to keep it here.');
|
||||
// The row stays: the move is still one press away.
|
||||
await expect(row.locator('.row-actions button')).toHaveText('Move library');
|
||||
});
|
||||
|
||||
test('Move library is reviewed, runs as an Activity job and ends the offer', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/library', { data: { state: 'node_folder' } })).toBeOK();
|
||||
const layout = await openOverview(page);
|
||||
const row = layout.locator('.operation').filter({ hasText: 'Library location' });
|
||||
await expect(row.locator('p')).toHaveText('browser.db · The node folder, where an update or reinstall can delete it.');
|
||||
const move = banner(page).getByRole('button', { name: 'Move library', exact: true });
|
||||
await move.click();
|
||||
const dialog = page.locator('dialog#modal[open]');
|
||||
await expect(dialog.locator('#modal-title')).toHaveText('Move library');
|
||||
expect(await dialog.locator('dl.meta-grid > div').evaluateAll(items => items.map(item =>
|
||||
[item.querySelector('dt').textContent, item.querySelector('dd').textContent]))).toEqual([
|
||||
['From', 'browser.db'], ['To', 'user/__promptmanager/prompts.db'], ['Size', expect.stringMatching(/ (KB|MB)$/)],
|
||||
['Moves with it', 'config.json and the database backups'], ['Old files', 'Renamed *.moved-to-…, never deleted'],
|
||||
]);
|
||||
await expect(dialog.locator('.dialog-actions button')).toHaveText(['Cancel', 'Move library']);
|
||||
await dialog.getByRole('button', { name: 'Cancel', exact: true }).click();
|
||||
await expect(page.locator('dialog#modal[open]')).toHaveCount(0);
|
||||
await expect(move).toBeFocused();
|
||||
await row.getByRole('button', { name: 'Move library', exact: true }).click();
|
||||
await dialog.getByRole('button', { name: 'Move library', exact: true }).click();
|
||||
await expect(page).toHaveURL(/#\/activity\/[^/]+$/);
|
||||
await expect(page.locator('#workspace-main')).toContainText(
|
||||
/Moved the library to user\/__promptmanager: \d+ prompts? and 0 images verified\. The old files were renamed browser\.db\.moved-to-__promptmanager; nothing was deleted\./);
|
||||
// The offer ends with the job; the Overview says where the library lives now.
|
||||
await expect(banner(page)).toBeHidden();
|
||||
const after = await openOverview(page);
|
||||
await expect(after.locator('.operation').filter({ hasText: 'Library location' }).locator('p'))
|
||||
.toHaveText('browser.db · ComfyUI\'s user folder, which updates and reinstalls of the node leave alone.');
|
||||
});
|
||||
|
||||
test('a move that fails says so and keeps the offer; a refused start says why', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/library', { data: { state: 'failing' } })).toBeOK();
|
||||
await page.goto(`${workspace}#/images`);
|
||||
await banner(page).getByRole('button', { name: 'Move library', exact: true }).click();
|
||||
await page.locator('dialog#modal[open]').getByRole('button', { name: 'Move library', exact: true }).click();
|
||||
await expect(page).toHaveURL(/#\/activity\/[^/]+$/);
|
||||
await expect(page.locator('#workspace-main')).toContainText(
|
||||
'The copy did not pass verification. Nothing was moved. The library is still in the node folder, unchanged.');
|
||||
await expect(banner(page).locator('.notice')).toBeVisible();
|
||||
await expect(await request.post('/__test__/system/fail', { data: { path: `${LIBRARY_URL}/move` } })).toBeOK();
|
||||
await banner(page).getByRole('button', { name: 'Move library', exact: true }).click();
|
||||
await page.locator('dialog#modal[open]').getByRole('button', { name: 'Move library', exact: true }).click();
|
||||
await expect(page.locator('#workspace-toast')).toHaveText('Move library could not start. Try again.');
|
||||
});
|
||||
|
||||
test('a library that does not move shows its details', async ({ page }) => {
|
||||
const layout = await openOverview(page);
|
||||
await expect(banner(page)).toBeHidden();
|
||||
const details = layout.locator('.operation').filter({ hasText: 'Library location' }).getByRole('button', { name: 'Details', exact: true });
|
||||
await details.click();
|
||||
const dialog = page.locator('dialog#modal[open]');
|
||||
await expect(dialog.locator('#modal-title')).toHaveText('Library location');
|
||||
expect(await dialog.locator('dl.meta-grid > div').evaluateAll(items => items.map(item =>
|
||||
[item.querySelector('dt').textContent, item.querySelector('dd').textContent.replace(/[\d.,]+ (KB|MB)$/, 'N')]))).toEqual([
|
||||
['Database', 'browser.db'], ['Why here', 'The database path set in config.json'],
|
||||
['Beside it', 'config.json, the database backups and the logs'], ['Size', 'N'],
|
||||
]);
|
||||
await dialog.getByRole('button', { name: 'Close', exact: true }).click();
|
||||
await expect(details).toBeFocused();
|
||||
});
|
||||
|
||||
test('Run health checks opens Health report and updates metric 4', async ({ page, request }) => {
|
||||
await page.goto(`${workspace}#/system/maintenance`);
|
||||
const layout = layoutOf(page);
|
||||
await expect(layout.locator(':scope > .notice h2')).toHaveText('Check first. Review changes. Then apply.');
|
||||
// Maintenance is native (S7): its two surfaces follow the notice.
|
||||
await expect(layout.locator(':scope > section.surface > h2')).toHaveText(['Database', 'Images & cache']);
|
||||
// The rest of Maintenance is still the existing tools' (D94).
|
||||
await expect(layout.locator(':scope > section.handoff h2')).toHaveText('Your current tools are ready');
|
||||
const run = layout.getByRole('button', { name: 'Run health checks', exact: true });
|
||||
await run.click();
|
||||
const dialog = page.locator('dialog#modal[open]');
|
||||
@@ -207,11 +113,8 @@ test('Run health checks opens Health report and updates metric 4', async ({ page
|
||||
.toHaveText(['1 link needs review', 'Passed', 'Not available']);
|
||||
await expect(dialog.locator('.operation h3').filter({ hasText: 'Missing files' }))
|
||||
.toHaveText(['Missing files · output', 'Missing files · input', 'Missing files · unmounted-drive']);
|
||||
// Review missing links reviews a fresh Missing-file links plan (S7); nothing is cleaned up.
|
||||
// Review missing links goes to Maintenance, where the cleanup's existing tool is linked.
|
||||
await dialog.getByRole('button', { name: 'Review missing links', exact: true }).click();
|
||||
await expect(dialog.locator('#modal-title')).toHaveText('Review maintenance plan');
|
||||
await expect(dialog.locator('#modal-body > h3')).toHaveText('Missing-file links');
|
||||
await dialog.getByRole('button', { name: 'Cancel', exact: true }).click();
|
||||
await expect(page.locator('dialog#modal[open]')).toHaveCount(0);
|
||||
await expect(page).toHaveURL(`${workspace}#/system/maintenance`);
|
||||
await page.locator('#workspace-main .section-nav').getByRole('button', { name: 'Overview', exact: true }).click();
|
||||
@@ -241,7 +144,7 @@ test('thumbnail row reflects the status endpoint', async ({ page, request }) =>
|
||||
|
||||
test('Overview Review thumbnails opens the same dialog', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/unthumbnailed', { data: { count: 2 } })).toBeOK();
|
||||
const plan = await serverJson(request, '/prompt_manager/images/thumbnails/plan?mode=missing');
|
||||
const plan = await serverJson(request, '/prompt_manager/images/thumbnails/plan?mode=missing&size=512&quality=medium');
|
||||
const starts = [];
|
||||
page.on('request', item => { if (item.method() === 'POST' && item.url().endsWith('/prompt_manager/jobs/thumbnails')) starts.push(item.postDataJSON()); });
|
||||
const layout = await openOverview(page);
|
||||
@@ -251,8 +154,8 @@ test('Overview Review thumbnails opens the same dialog', async ({ page, request
|
||||
const review = page.locator('dialog#modal[open]');
|
||||
await expect(review.locator('#modal-title')).toHaveText('Thumbnail maintenance');
|
||||
await expect(page).toHaveURL(`${workspace}#/system/overview`);
|
||||
await expect(review.locator('.field > label')).toHaveText(['Scope']);
|
||||
await expect(review.locator('#modal-body > p.small').first()).toHaveText(`2 previews to generate · ${number(plan.fresh)} already up to date.`);
|
||||
await expect(review.locator('.field > label')).toHaveText(['Scope', 'Quality', 'Preview size limit']);
|
||||
await expect(review.locator('#modal-body > p.small')).toHaveText(`2 previews to generate · ${number(plan.fresh)} already up to date.`);
|
||||
// Escape closes it and focus returns to the row's button.
|
||||
await page.keyboard.press('Escape');
|
||||
await expect(review).toHaveCount(0);
|
||||
@@ -260,61 +163,11 @@ test('Overview Review thumbnails opens the same dialog', async ({ page, request
|
||||
await row.getByRole('button', { name: 'Review thumbnails', exact: true }).click();
|
||||
await review.getByRole('button', { name: 'Start', exact: true }).click();
|
||||
await expect(page).toHaveURL(/#\/activity\/[0-9a-f]{32}$/);
|
||||
expect(starts).toEqual([{ mode: 'missing' }]);
|
||||
expect(starts).toEqual([{ mode: 'missing', size: 512, quality: 'medium' }]);
|
||||
const id = page.url().split('/').pop();
|
||||
await expect.poll(async () => (await serverJson(request, `/prompt_manager/jobs/${id}`)).job.state).toBe('finished');
|
||||
});
|
||||
|
||||
test('the thumbnail row counts outdated previews and offers Convert outdated (#317)', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/legacy-previews', { data: { count: 2 } })).toBeOK();
|
||||
const layout = await openOverview(page);
|
||||
const row = layout.locator('.operation').filter({ has: page.getByRole('button', { name: 'Review thumbnails', exact: true }) });
|
||||
await expect(row.locator('h3')).toHaveText('No missing thumbnails · 2 outdated');
|
||||
await expect(row.locator('.row-actions > button')).toHaveText(['Review thumbnails', 'Convert outdated']);
|
||||
await row.getByRole('button', { name: 'Convert outdated', exact: true }).click();
|
||||
const review = page.locator('dialog#modal[open]');
|
||||
await expect(review.locator('#modal-title')).toHaveText('Thumbnail maintenance');
|
||||
// It opens on the Convert outdated scope; nothing starts until Start.
|
||||
await expect(review.locator('#thumb-mode')).toHaveValue('convert');
|
||||
await expect(review.locator('#modal-body > p.small').first()).toHaveText('2 previews in an older format to convert.');
|
||||
expect((await serverJson(request, '/prompt_manager/jobs')).jobs).toEqual([]);
|
||||
await page.keyboard.press('Escape');
|
||||
await expect(row.getByRole('button', { name: 'Convert outdated', exact: true })).toBeFocused();
|
||||
});
|
||||
|
||||
test('Appearance › Thumbnails saves the server thumbnail settings and never converts by itself (#317)', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/system/jxl-encoder', { data: { available: false } })).toBeOK();
|
||||
await expect(await request.post('/__test__/system/legacy-previews', { data: { count: 1 } })).toBeOK();
|
||||
await page.goto(`${workspace}#/system/appearance`);
|
||||
const surface = page.locator('#workspace-main .settings-layout > section.surface').nth(1);
|
||||
await expect(surface.locator(':scope > h2')).toHaveText('Thumbnails');
|
||||
await expect(surface.locator('.setting-row label')).toHaveText(['Thumbnail format', 'Thumbnail quality', 'Thumbnail size limit']);
|
||||
const format = surface.locator('#thumbnail-format');
|
||||
await expect(format).toBeEnabled();
|
||||
await expect(format).toHaveValue('webp');
|
||||
await expect(surface.locator('#thumbnail-quality')).toHaveValue('80');
|
||||
await expect(surface.locator('#thumbnail-edge')).toHaveValue('512');
|
||||
// D130: JPEG XL is offered only with an encoder; without one the row says why, in the server's words.
|
||||
await expect(format.locator('option[value="jxl"]')).toBeDisabled();
|
||||
await expect(surface.locator('.setting-row').first().locator('p.small').first()).toContainText('No JPEG XL encoder found.');
|
||||
await format.selectOption('jpeg');
|
||||
await surface.locator('#thumbnail-quality').fill('65');
|
||||
await surface.locator('#thumbnail-quality').dispatchEvent('change');
|
||||
await page.getByRole('button', { name: 'Save preferences', exact: true }).click();
|
||||
await expect(page.locator('#settings-status')).toHaveText('Saved to the server configuration.');
|
||||
const { settings } = await serverJson(request, '/prompt_manager/settings');
|
||||
expect([settings.thumbnail_format, settings.thumbnail_quality, settings.thumbnail_edge]).toEqual(['jpeg', 65, 512]);
|
||||
// Changing the format converts nothing: the row says so and the old preview is still listed.
|
||||
await expect(surface.locator('.operation p')).toHaveText('Format saved. Existing previews keep their old format until you convert them.');
|
||||
expect((await serverJson(request, '/prompt_manager/jobs')).jobs).toEqual([]);
|
||||
expect((await serverJson(request, '/prompt_manager/images/thumbnails/status')).outdated).toBeGreaterThan(0);
|
||||
// With an encoder JPEG XL can be chosen.
|
||||
await expect(await request.post('/__test__/system/jxl-encoder', { data: { available: true } })).toBeOK();
|
||||
await page.reload();
|
||||
await expect(page.locator('#workspace-main #thumbnail-format option[value="jxl"]')).toBeEnabled();
|
||||
await expect(page.locator('#workspace-main #thumbnail-format')).toHaveValue('jpeg');
|
||||
});
|
||||
|
||||
test('the thumbnail row paints when its own status read answers', async ({ page, request }) => {
|
||||
const held = async () => (await serverJson(request, '/__test__/hold')).held[THUMBNAILS_URL] ?? 0;
|
||||
await expect(await request.post('/__test__/hold', { data: { path: THUMBNAILS_URL } })).toBeOK();
|
||||
@@ -412,11 +265,11 @@ test('out-of-order overview responses never repaint a newer section', async ({ p
|
||||
await nav.getByRole('button', { name: 'Appearance', exact: true }).click();
|
||||
await nav.getByRole('button', { name: 'Overview', exact: true }).click();
|
||||
await expect.poll(held).toBe(2);
|
||||
// Diagnostics still hands off (S12), so its body is one known panel.
|
||||
await nav.getByRole('button', { name: 'Diagnostics & logs', exact: true }).click();
|
||||
// Integrations still hands off (S10), so its body is one known panel.
|
||||
await nav.getByRole('button', { name: 'Integrations', exact: true }).click();
|
||||
await expect(await request.post('/__test__/release')).toBeOK();
|
||||
await expect.poll(held).toBe(0);
|
||||
// Both late answers reached a section that had left: Diagnostics keeps its own body.
|
||||
// Both late answers reached a section that had left: Integrations keeps its own body.
|
||||
await expect(layout.locator(':scope > section.handoff')).toBeVisible();
|
||||
await page.evaluate(() => new Promise(resolve => requestAnimationFrame(() => requestAnimationFrame(resolve))));
|
||||
await expect(layout.locator('.summary-metrics')).toHaveCount(0);
|
||||
@@ -439,7 +292,7 @@ const panel = page => page.getByRole('complementary', { name: 'Selected image de
|
||||
|
||||
async function openAppearance(page) {
|
||||
await page.goto(`${workspace}#/system/appearance`);
|
||||
await expect(page.locator('#workspace-main .settings-layout > section.surface > h2').first()).toHaveText('Browsing & display');
|
||||
await expect(page.locator('#workspace-main .settings-layout > section.surface > h2')).toHaveText('Browsing & display');
|
||||
// The host rows are enabled once the server settings answered.
|
||||
await expect(page.locator('#timeout')).toBeEnabled();
|
||||
}
|
||||
@@ -656,16 +509,14 @@ test('server settings that cannot be read say so in their rows', async ({ page }
|
||||
? route.fulfill({ status: 500, contentType: 'application/json', body: JSON.stringify({ success: false, error: 'Injected failure' }) })
|
||||
: route.continue()));
|
||||
await page.goto(`${workspace}#/system/appearance`);
|
||||
// #317: the Thumbnails rows are server settings too.
|
||||
for (const id of ['open-mode', 'timeout', 'thumbnail-format', 'thumbnail-quality', 'thumbnail-edge']) {
|
||||
for (const id of ['open-mode', 'timeout']) {
|
||||
const row = page.locator('#workspace-main .setting-row').filter({ has: page.locator(`#${id}`) });
|
||||
await expect(row.locator('p.small').first()).toHaveText('Server settings could not be read.');
|
||||
await expect(row.locator('p.small')).toHaveText('Server settings could not be read.');
|
||||
await expect(page.locator(`#${id}`)).toBeDisabled();
|
||||
}
|
||||
// The browser rows are this browser's, and work as before.
|
||||
await expect(page.locator('#default-view')).toBeEnabled();
|
||||
await expect(page.locator('#workspace-main .setting-row p.small', { hasText: 'Server settings could not be read.' })).toHaveCount(5);
|
||||
await expect(page.locator('#workspace-main section.surface').first().locator('.setting-row p.small')).toHaveCount(2);
|
||||
await expect(page.locator('#workspace-main .setting-row p.small')).toHaveCount(2);
|
||||
});
|
||||
|
||||
test('stale revision shows the conflict status', async ({ page, request }) => {
|
||||
@@ -727,7 +578,7 @@ test('a change another tab saved while System was closed is a conflict on return
|
||||
await other.close();
|
||||
// Back in System: the draft is kept and flagged, and nothing was written over the other tab's value.
|
||||
await page.evaluate(() => { location.hash = '#/system/appearance'; });
|
||||
await expect(page.locator('#workspace-main .settings-layout > section.surface > h2').first()).toHaveText('Browsing & display');
|
||||
await expect(page.locator('#workspace-main .settings-layout > section.surface > h2')).toHaveText('Browsing & display');
|
||||
await expect(statusOf(page)).toHaveText('Settings changed elsewhere. Reload this section to see them.');
|
||||
await expect(page.locator('#density')).toHaveAttribute('aria-invalid', 'true');
|
||||
await expect(page.locator('#density')).toHaveValue('compact');
|
||||
|
||||
@@ -17,7 +17,7 @@ const chips = page => modal(page).locator('#tag-suggestions .chip.tag-suggestion
|
||||
const actionButtons = page => modal(page).locator('.dialog-actions button');
|
||||
const panel = page => page.getByRole('complementary', { name: 'Selected image details', exact: true });
|
||||
const tagRow = page => panel(page).locator('.detail-inner > .tags').last();
|
||||
const NOT_DOWNLOADED = 'WD14 SwinV2 is not downloaded. Download it in System › Auto-tagging.';
|
||||
const NOT_DOWNLOADED = 'WD14 SwinV2 is not downloaded. Download it from the legacy Auto-tag page.';
|
||||
|
||||
test.beforeEach(async ({ page, request, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
|
||||
@@ -246,7 +246,7 @@ test('unavailable model shows inline status and never downloads', async ({ page,
|
||||
await expect(await request.post('/__test__/tagging/models', { data: { unavailable: ['wd14-swinv2'] } })).toBeOK();
|
||||
await openTag(page);
|
||||
expect((await generated(page)).status()).toBe(409);
|
||||
await expect(status(page)).toHaveText('WD14 SwinV2 is not downloaded. Download it in System › Auto-tagging.');
|
||||
await expect(status(page)).toHaveText('WD14 SwinV2 is not downloaded. Download it from the legacy Auto-tag page.');
|
||||
await expect(chips(page)).toHaveCount(0);
|
||||
await expect(dialog(page).locator('#suggestion-count')).toHaveText('Not generated');
|
||||
// D116: Retry leads the actions while the controller allows it, and generates again.
|
||||
@@ -602,10 +602,6 @@ test('a save under a Tags filter re-reads the listing: the selection stays while
|
||||
|
||||
test('a save under a search keeps the scroll position and the loaded range', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/tagging/reset', { data: { extra_images: 140 } })).toBeOK();
|
||||
// Listing pages still in flight: one that lands after the scroll below would grow the loaded range.
|
||||
const inFlight = new Set();
|
||||
page.on('request', item => { if (item.url().includes('/images/output?')) inFlight.add(item); });
|
||||
for (const done of ['requestfinished', 'requestfailed']) page.on(done, item => inFlight.delete(item));
|
||||
await page.goto(images);
|
||||
const searched = listingOf(page, query => query.get('q') === 'bulk');
|
||||
await page.getByRole('searchbox', { name: 'Search workspace', exact: true }).fill('bulk');
|
||||
@@ -619,9 +615,6 @@ test('a save under a search keeps the scroll position and the loaded range', asy
|
||||
await expect.poll(async () => { await continuation.scrollIntoViewIfNeeded(); return loadedCount(); }).toBeGreaterThanOrEqual(80);
|
||||
await page.evaluate(() => window.scrollTo(0, Math.round(document.documentElement.scrollHeight / 3)));
|
||||
await expect(continuation).not.toBeInViewport();
|
||||
// A load the continuation scheduled for the next frame has started by now; wait for every page to land.
|
||||
await page.evaluate(() => new Promise(done => requestAnimationFrame(() => requestAnimationFrame(done))));
|
||||
await expect.poll(() => inFlight.size).toBe(0);
|
||||
const loaded = await continuation.textContent();
|
||||
await expect.poll(() => continuation.textContent()).toBe(loaded);
|
||||
const y = await page.evaluate(() => window.scrollY);
|
||||
|
||||
@@ -1,170 +0,0 @@
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from './fixtures.mjs';
|
||||
|
||||
/**
|
||||
* #249: cards and strips load previews, never the full file. The approved mock has no repeated thumbnail prompt
|
||||
* (docs/concepts/workspace/README.md), so a preview that was not made yet is made by the server on first view:
|
||||
* the listing still names the `_thumb` URL. Only the media viewer's main image (and downloads) load the original.
|
||||
*/
|
||||
const images = '/prompt_manager/workspace#/images';
|
||||
const library = '/prompt_manager/workspace#/library';
|
||||
const cards = page => page.locator('article[data-image-key]');
|
||||
const details = page => page.getByRole('complementary', { name: 'Selected image details', exact: true });
|
||||
const SERVE = '/prompt_manager/images/serve/';
|
||||
|
||||
/** Every media file the page asks the server for, as `{path, preview}` (path relative to its root). */
|
||||
function recordMedia(page) {
|
||||
const seen = [];
|
||||
page.on('request', request => {
|
||||
const url = new URL(request.url());
|
||||
if (!url.pathname.startsWith(SERVE)) return;
|
||||
const path = decodeURIComponent(url.pathname.slice(SERVE.length));
|
||||
seen.push({ path, preview: /^thumbnails\/.+_thumb\.[a-z0-9]+$/.test(path) });
|
||||
});
|
||||
return seen;
|
||||
}
|
||||
/** Every card image has loaded (or failed): what it asked for is all in `seen`. */
|
||||
const settled = page => page.evaluate(() => Promise.all([...document.querySelectorAll('article img')].map(image => image.decode().catch(() => {}))));
|
||||
const originals = seen => seen.filter(item => !item.preview).map(item => item.path);
|
||||
|
||||
test.beforeEach(async ({ page, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
page.thumbErrors = [];
|
||||
page.on('pageerror', error => page.thumbErrors.push(error.message));
|
||||
});
|
||||
test.afterEach(async ({ page, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
expect(page.thumbErrors).toEqual([]);
|
||||
});
|
||||
|
||||
test('a page of the Images grid requests previews, never originals, in every layout and size', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset')).toBeOK();
|
||||
// Newest outputs with no preview yet: the live finding (844 such files listed with a null thumbnail).
|
||||
await expect(await request.post('/__test__/system/unthumbnailed', { data: { count: 6 } })).toBeOK();
|
||||
const listing = await (await request.get('/prompt_manager/images/output?limit=40')).json();
|
||||
const fresh = listing.images.filter(image => image.path.startsWith('unthumbnailed/'));
|
||||
expect(fresh).toHaveLength(6);
|
||||
for (const image of fresh) expect(image.thumbnail_url).toMatch(/\/thumbnails\/unthumbnailed\/plain_\d{3}_thumb\.webp\?root=0/);
|
||||
const seen = recordMedia(page);
|
||||
await page.goto(images);
|
||||
await expect(cards(page)).toHaveCount(8);
|
||||
await settled(page);
|
||||
expect(originals(seen)).toEqual([]);
|
||||
expect(seen.filter(item => item.path.startsWith('thumbnails/unthumbnailed/'))).toHaveLength(6);
|
||||
// The server made them on that first view.
|
||||
await expect.poll(async () => (await (await request.get('/prompt_manager/images/thumbnails/status')).json()).pending).toBe(0);
|
||||
for (const layout of ['List', 'Masonry', 'Grid']) {
|
||||
await page.getByRole('button', { name: layout, exact: true }).click();
|
||||
await expect(page.getByRole('button', { name: layout, exact: true })).toHaveAttribute('aria-pressed', 'true');
|
||||
for (const density of ['large', 'compact', 'comfortable']) {
|
||||
await page.getByRole('combobox', { name: 'Thumbnail size', exact: true }).selectOption(density);
|
||||
await expect(cards(page)).toHaveCount(8);
|
||||
await settled(page);
|
||||
}
|
||||
}
|
||||
expect(originals(seen)).toEqual([]);
|
||||
});
|
||||
|
||||
test('the inspector cover is the preview and only the media viewer loads the original', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/system/unthumbnailed', { data: { count: 1 } })).toBeOK();
|
||||
const seen = recordMedia(page);
|
||||
await page.goto(images);
|
||||
const card = cards(page).filter({ hasText: 'plain_001.png' });
|
||||
await card.getByRole('button', { name: 'Inspect plain_001.png', exact: true }).click();
|
||||
const cover = details(page).locator('img.detail-cover');
|
||||
await expect(cover).toHaveAttribute('src', /\/thumbnails\/unthumbnailed\/plain_001_thumb\.webp\?/);
|
||||
await settled(page);
|
||||
await cover.evaluate(image => image.decode().catch(() => {}));
|
||||
expect(originals(seen)).toEqual([]);
|
||||
await card.getByRole('button', { name: 'Open plain_001.png', exact: true }).click();
|
||||
await expect(page.locator('dialog#viewer-dialog')).toBeVisible();
|
||||
await expect.poll(() => originals(seen)).toEqual(['unthumbnailed/plain_001.png']);
|
||||
await page.getByRole('button', { name: 'Close image viewer', exact: true }).click();
|
||||
});
|
||||
|
||||
test('a video card shows its preview frame instead of loading the video', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { include_video: true } })).toBeOK();
|
||||
const seen = recordMedia(page);
|
||||
await page.goto(images);
|
||||
const card = cards(page).filter({ hasText: 'clip.webm' });
|
||||
await expect(card.locator('.media-art > .image-open > img')).toHaveAttribute('src', /\/thumbnails\/clip_thumb\.webp\?/);
|
||||
await expect(card.locator('video')).toHaveCount(0);
|
||||
await settled(page);
|
||||
expect(originals(seen)).toEqual([]);
|
||||
});
|
||||
|
||||
test('Library cards and the shared gallery request previews only', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { native_library: true } })).toBeOK();
|
||||
const seen = recordMedia(page);
|
||||
await page.goto(library);
|
||||
const cat = page.locator('article[data-prompt-id]').filter({ has: page.getByText('Cat portrait', { exact: true }) });
|
||||
await expect(cat.locator('.card-art img')).not.toHaveCount(0);
|
||||
await settled(page);
|
||||
expect(seen.length).toBeGreaterThan(0);
|
||||
expect(originals(seen)).toEqual([]);
|
||||
});
|
||||
|
||||
test('a video whose preview cannot be loaded swaps back to its own element', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset', { data: { include_video: true } })).toBeOK();
|
||||
await page.route(url => url.pathname.endsWith('/thumbnails/clip_thumb.webp'), route => route.fulfill({ status: 404, body: '' }));
|
||||
await page.goto(images);
|
||||
const card = cards(page).filter({ hasText: 'clip.webm' });
|
||||
await expect(card.locator('.media-art > .image-open > video')).toHaveCount(1);
|
||||
await expect(card.locator('.media-art > .image-open > img')).toHaveCount(0);
|
||||
await card.getByRole('button', { name: 'Inspect clip.webm', exact: true }).click();
|
||||
await expect(details(page).locator('video.detail-cover')).toHaveCount(1);
|
||||
await expect(details(page).locator('img.detail-cover')).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('an image listed without a preview (a nested thumbnails folder) shows its original', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset')).toBeOK();
|
||||
// M10: such files are listed with a null thumbnail_url; the card falls back to the original as before.
|
||||
await page.route(url => url.pathname === '/prompt_manager/images/output', async route => {
|
||||
const response = await route.fetch();
|
||||
const body = await response.json();
|
||||
if (Array.isArray(body.images)) body.images = body.images.map(image => ({ ...image, thumbnail_url: null }));
|
||||
await route.fulfill({ response, json: body });
|
||||
});
|
||||
const seen = recordMedia(page);
|
||||
await page.goto(images);
|
||||
await expect(cards(page)).toHaveCount(2);
|
||||
await settled(page);
|
||||
expect(seen.length).toBeGreaterThan(0);
|
||||
expect(seen.every(item => !item.preview)).toBe(true);
|
||||
});
|
||||
|
||||
test('a JPEG XL thumbnail this browser cannot show loads the original; one it can show is used (#317)', async ({ page, request }) => {
|
||||
await expect(await request.post('/__test__/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/system/reset')).toBeOK();
|
||||
try {
|
||||
await expect(await request.post('/__test__/system/jxl-encoder', { data: { available: true } })).toBeOK();
|
||||
const saved = await request.post('/prompt_manager/settings', { data: { thumbnail_format: 'jxl' } });
|
||||
await expect(saved).toBeOK();
|
||||
await expect(await request.post('/__test__/system/unthumbnailed', { data: { count: 1 } })).toBeOK();
|
||||
const listing = await (await request.get('/prompt_manager/images/output?limit=40')).json();
|
||||
const fresh = listing.images.find(image => image.path === 'unthumbnailed/plain_001.png');
|
||||
expect(fresh.thumbnail_url).toMatch(/\/thumbnails\/unthumbnailed\/plain_001_thumb\.jxl\?root=0/);
|
||||
await page.goto(images);
|
||||
// What this browser does with JPEG XL, decided the way the workspace decides it (an embedded probe).
|
||||
const decodes = await page.evaluate(() => new Promise(resolve => {
|
||||
const probe = new Image();
|
||||
probe.onload = () => resolve(true);
|
||||
probe.onerror = () => resolve(false);
|
||||
probe.src = 'data:image/jxl;base64,/woAkAEAE4gCALQAtZ8gAAAVKqOMG7yc6/nyQ4fFtI3rDG21bWEJY7O9MEhIOGwY4s3xwRCTShIA';
|
||||
}));
|
||||
const card = cards(page).filter({ hasText: 'plain_001.png' });
|
||||
const media = card.locator('.media-art > .image-open > img');
|
||||
// A browser without JPEG XL shows the original; one with it shows the preview, or the original when the
|
||||
// preview cannot be made here (no real encoder in the test server): never a broken image.
|
||||
if (!decodes) await expect(media).toHaveAttribute('src', /\/unthumbnailed\/plain_001\.png\?root=0/);
|
||||
await expect.poll(() => media.evaluate(image => image.complete && image.naturalWidth > 0)).toBe(true);
|
||||
// The inspector cover follows the same rule.
|
||||
await card.getByRole('button', { name: 'Inspect plain_001.png', exact: true }).click();
|
||||
const cover = details(page).locator('img.detail-cover');
|
||||
if (!decodes) await expect(cover).toHaveAttribute('src', /\/unthumbnailed\/plain_001\.png\?/);
|
||||
await expect.poll(() => cover.evaluate(image => image.complete && image.naturalWidth > 0)).toBe(true);
|
||||
} finally {
|
||||
await request.post('/__test__/system/reset');
|
||||
}
|
||||
});
|
||||
@@ -173,8 +173,7 @@ test('visual fixture exposes real stable media and library records', async ({ re
|
||||
await expect(await request.get(image.url)).toBeOK();
|
||||
const thumbnail = await request.get(image.thumbnail_url);
|
||||
await expect(thumbnail).toBeOK();
|
||||
// #317: previews are written as WebP by default.
|
||||
expect(thumbnail.headers()['content-type']).toBe('image/webp');
|
||||
expect(thumbnail.headers()['content-type']).toBe('image/png');
|
||||
const inspection = await (await request.get(`/prompt_manager/images/inspect/${encodeURIComponent(image.path)}?root=${image.root_index}`)).json();
|
||||
expect(inspection.success).toBe(true);
|
||||
expect(inspection.image.path).toBe(image.path);
|
||||
|
||||
@@ -413,7 +413,7 @@ for (const width of [1600, 390]) test(`overview and handoff render with concept
|
||||
|
||||
await sidebarLink(page, 'System').click();
|
||||
// System's Overview is native (S2 Task 6); a section not delivered yet still hands off (D94).
|
||||
await main.locator('.section-nav').getByRole('button', { name: 'Diagnostics & logs', exact: true }).click();
|
||||
await main.locator('.section-nav').getByRole('button', { name: 'Integrations', exact: true }).click();
|
||||
const handoff = main.locator('section.handoff');
|
||||
await expect(handoff).toHaveClass(/\bsurface\b/);
|
||||
expect(await handoff.evaluate(surfaceStyle)).toMatchObject(SURFACE);
|
||||
@@ -427,6 +427,7 @@ for (const width of [1600, 390]) test(`overview and handoff render with concept
|
||||
const SYSTEM_SECTION_LABELS = ['Overview', 'Appearance', 'Sources & scans', 'Integrations', 'Auto-tagging', 'Maintenance',
|
||||
'Backup & export', 'Diagnostics & logs'];
|
||||
const SYSTEM_HASHES = ['overview', 'appearance', 'sources', 'integrations', 'tagging', 'maintenance', 'backups', 'diagnostics'];
|
||||
const SYSTEM_HANDOFFS = { maintenance: ['/prompt_manager/admin', '/prompt_manager/gallery'] };
|
||||
|
||||
test('System is native: the concept shell, section navigation, sections, hand-offs and the Save row', async ({ page }) => {
|
||||
const mutations = [];
|
||||
@@ -463,26 +464,20 @@ test('System is native: the concept shell, section navigation, sections, hand-of
|
||||
// The section changes in place: focus stays on the pressed section, not the page title.
|
||||
await expect(button).toBeFocused();
|
||||
// S2 Task 6: Overview is native (workspace-system.spec.mjs), and so are Appearance (S3 Task 8) and Sources & scans
|
||||
// (S4 Task 11, workspace-system-sources.spec.mjs, its update notice shown in a fresh browser), Integrations (S10
|
||||
// Task 27, workspace-system-integrations.spec.mjs), Maintenance (S7 Task 19,
|
||||
// workspace-system-maintenance.spec.mjs), Backup & export (S9 Task 24, workspace-system-backups.spec.mjs) and
|
||||
// Auto-tagging (S11 Task 6, workspace-system-tagging.spec.mjs); every other section hands off.
|
||||
if (hash === 'tagging') await expect(layout.locator('#tag-model')).toBeVisible();
|
||||
// (S4 Task 11, workspace-system-sources.spec.mjs, its update notice shown in a fresh browser); Maintenance shows
|
||||
// its health-check notice above the hand-off (D123); every other section hands off.
|
||||
if (hash === 'overview') await expect(layout.locator(':scope > .summary-metrics')).toBeVisible();
|
||||
if (hash === 'sources') await expect(layout.locator('#source-list .operation')).toHaveCount(2);
|
||||
expect(await layout.evaluate(node => [...node.children].map(child => `${child.localName}.${child.className}`)))
|
||||
.toEqual({ overview: ['div.summary-metrics', 'div.two-col', 'div.save-row'], appearance: ['section.surface', 'section.surface', 'div.save-row'],
|
||||
.toEqual({ overview: ['div.summary-metrics', 'div.two-col', 'div.save-row'], appearance: ['section.surface', 'div.save-row'],
|
||||
sources: ['div.notice', 'section.surface', 'div.save-row'],
|
||||
integrations: ['section.surface', 'section.surface', 'div.save-row'],
|
||||
maintenance: ['div.notice', 'section.surface', 'section.surface', 'div.save-row'],
|
||||
backups: ['div.two-col', 'section.surface', 'div.save-row'],
|
||||
tagging: ['section.surface', 'div.save-row'] }[hash]
|
||||
maintenance: ['div.notice', 'section.surface handoff', 'div.save-row'] }[hash]
|
||||
?? ['section.surface handoff', 'div.save-row']);
|
||||
if (['overview', 'appearance', 'sources', 'integrations', 'tagging', 'maintenance', 'backups'].includes(hash)) continue;
|
||||
if (['overview', 'appearance', 'sources'].includes(hash)) continue;
|
||||
const handoff = layout.locator('section.handoff');
|
||||
await expect(handoff.locator('h2')).toHaveText('Your current tools are ready');
|
||||
expect(await handoff.locator('a.workspace-link').evaluateAll(links => links.map(link => link.getAttribute('href'))))
|
||||
.toEqual(['/prompt_manager/admin']);
|
||||
.toEqual(SYSTEM_HANDOFFS[hash] ?? ['/prompt_manager/admin']);
|
||||
await expect(layout.locator('#settings-status')).toHaveText('No unsaved changes.');
|
||||
await expect(page.locator('aside.sidebar a[href="#/system"]')).toHaveAttribute('aria-current', 'page');
|
||||
}
|
||||
|
||||
@@ -513,12 +513,7 @@ def seed_workspace_concept(
|
||||
tags = [row[0] for row in conn.execute("SELECT name FROM tags ORDER BY id")]
|
||||
if tuple(counts) != (len(CONCEPT_PROMPTS), CONCEPT_RUN_COUNT, CONCEPT_RUN_COUNT):
|
||||
raise RuntimeError(f"Concept fixture run counts are not real: {tuple(counts)}")
|
||||
from py.api.thumbnail_settings import ThumbnailSettings
|
||||
|
||||
# 150 px previews in the default format (#317 thumbnail settings).
|
||||
api._generate_gallery_thumbnails_sync(
|
||||
api._get_all_output_dirs(), ThumbnailSettings("webp", 80, 150)
|
||||
)
|
||||
api._generate_gallery_thumbnails_sync(api._get_all_output_dirs(), (150, 150))
|
||||
api.invalidate_gallery_cache()
|
||||
return {
|
||||
"epoch": epoch,
|
||||
|
||||
@@ -1,216 +0,0 @@
|
||||
"""System › Integrations test hooks for the browser tests (#249 System S10).
|
||||
|
||||
A fake LoraManager install lives under this server's temp root; nothing ever
|
||||
reaches a real LoraManager, ComfyUI install or the network:
|
||||
|
||||
- ``POST /__test__/integrations/reset`` turns the integration off (no path,
|
||||
no key, no trigger words), removes the fake install, its LoRA folder and
|
||||
the image cache, forgets every LoRA import and injected failure, and opens
|
||||
the download gate;
|
||||
- ``POST /__test__/integrations/lora-manager``
|
||||
``{loras, civitai, preview, previous, settings}`` writes
|
||||
``custom_nodes/ComfyUI-Lora-Manager`` (detected through the
|
||||
``folder_paths`` stub's base path) and ``loras`` LoRAs under
|
||||
``models/loras``: each a ``<stem>.safetensors`` model beside its
|
||||
``<stem>.metadata.json`` sidecar (LoraManager's alongside layout, #355)
|
||||
and, unless ``preview`` is false, a local preview PNG; ``civitai`` adds one
|
||||
CivitAI example image per LoRA (served by the stub opener below; an import
|
||||
downloads it only for a LoRA without a local preview); ``previous`` saves
|
||||
that many earlier ``lora-manager`` imports with tags and image links;
|
||||
``settings: "invalid"`` makes LoraManager's settings file unreadable JSON
|
||||
(the server's own isolated settings folder, ``LORA_MANAGER_SETTINGS_DIR``);
|
||||
- ``POST /__test__/integrations/gate`` ``{hold}`` holds every CivitAI
|
||||
download until it is released (``{"hold": false}``), so a test sees the job
|
||||
while it downloads;
|
||||
- ``POST /__test__/integrations/fail-import`` makes the next collection
|
||||
replacement fail (a database error after staging);
|
||||
- ``POST /__test__/integrations/pair`` saves a pair whose positive is the first
|
||||
imported prompt (a saved pair that blocks the replacement);
|
||||
- ``GET /__test__/integrations/state`` answers the imported prompt texts,
|
||||
whether a CivitAI key is saved and the downloads served.
|
||||
|
||||
CivitAI downloads are always answered by a stub opener (a small PNG), so no
|
||||
test ever opens a network connection.
|
||||
"""
|
||||
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
from PIL import Image
|
||||
|
||||
MAX_LORAS = 20
|
||||
CATEGORY = "lora-manager"
|
||||
|
||||
|
||||
def _png():
|
||||
buffer = io.BytesIO()
|
||||
Image.new("RGB", (8, 8), "#3a6ea5").save(buffer, "PNG")
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
class _Response:
|
||||
def __init__(self, body):
|
||||
self._body = io.BytesIO(body)
|
||||
self.headers = {"Content-Length": str(len(body))}
|
||||
|
||||
def read(self, amount=-1):
|
||||
return self._body.read(amount)
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc):
|
||||
return False
|
||||
|
||||
|
||||
def register_integrations_commands(app, routes, api, root: Path):
|
||||
from py import lora_utils
|
||||
from py.config import IntegrationConfig
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
lora_manager = root / "custom_nodes" / "ComfyUI-Lora-Manager"
|
||||
loras = root / "models" / "loras"
|
||||
gate = threading.Event()
|
||||
gate.set()
|
||||
served = []
|
||||
failures = {"replace": 0}
|
||||
pairs = [] # the blocking pairs this fixture saved
|
||||
real_replace = type(api.db).replace_category_prompts
|
||||
|
||||
def stub_open(request, timeout=10):
|
||||
served.append(request.full_url)
|
||||
gate.wait(30)
|
||||
return _Response(_png())
|
||||
|
||||
lora_utils._open_url = stub_open
|
||||
|
||||
def replace(db, *args, **kwargs):
|
||||
if failures["replace"] > 0:
|
||||
failures["replace"] -= 1
|
||||
raise RuntimeError("Injected replacement failure")
|
||||
return real_replace(db, *args, **kwargs)
|
||||
|
||||
type(api.db).replace_category_prompts = replace
|
||||
|
||||
def settings_file():
|
||||
return Path(os.environ["LORA_MANAGER_SETTINGS_DIR"]) / "settings.json"
|
||||
|
||||
def reset():
|
||||
settings_file().unlink(missing_ok=True)
|
||||
IntegrationConfig.LORA_MANAGER_ENABLED = False
|
||||
IntegrationConfig.LORA_MANAGER_PATH = ""
|
||||
IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED = False
|
||||
IntegrationConfig.CIVITAI_API_KEY = ""
|
||||
lora_utils.get_trigger_cache().clear()
|
||||
gate.set()
|
||||
served.clear()
|
||||
failures["replace"] = 0
|
||||
cache = Path(lora_utils.get_lora_image_cache_dir(create=False))
|
||||
for folder in (lora_manager, loras, cache):
|
||||
shutil.rmtree(folder, ignore_errors=True)
|
||||
for pair_id in pairs:
|
||||
api.db.delete_prompt_pair(pair_id)
|
||||
pairs.clear()
|
||||
real_replace(api.db, CATEGORY, [], allow_empty=True)
|
||||
|
||||
@routes.post("/__test__/integrations/reset")
|
||||
async def integrations_reset(request):
|
||||
reset()
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.post("/__test__/integrations/lora-manager")
|
||||
async def make_lora_manager(request):
|
||||
data = await request.json() if request.can_read_body else {}
|
||||
count = data.get("loras", 2)
|
||||
previous = data.get("previous", 0)
|
||||
preview = data.get("preview", True)
|
||||
settings = data.get("settings")
|
||||
for value in (count, previous):
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise web.HTTPBadRequest()
|
||||
if not 0 <= count <= MAX_LORAS or not 0 <= previous <= MAX_LORAS:
|
||||
raise web.HTTPBadRequest()
|
||||
if not isinstance(preview, bool) or settings not in (None, "invalid"):
|
||||
raise web.HTTPBadRequest()
|
||||
if settings == "invalid":
|
||||
settings_file().write_text("{not json", encoding="utf-8")
|
||||
(lora_manager / "py").mkdir(parents=True, exist_ok=True)
|
||||
(lora_manager / "__init__.py").write_text("")
|
||||
loras.mkdir(parents=True, exist_ok=True)
|
||||
for number in range(count):
|
||||
stem = f"style_{number + 1:02}"
|
||||
civitai = {"trainedWords": [f"trigger{number + 1}"]}
|
||||
if data.get("civitai") is True:
|
||||
civitai["images"] = [
|
||||
{
|
||||
"url": f"https://image.civitai.com/x/{stem}.jpeg",
|
||||
"meta": {"prompt": f"{stem} example prompt"},
|
||||
}
|
||||
]
|
||||
meta = {
|
||||
"file_name": f"{stem}.safetensors",
|
||||
"model_name": f"Style {number + 1:02}",
|
||||
"civitai": civitai,
|
||||
}
|
||||
(loras / f"{stem}.safetensors").write_bytes(b"model")
|
||||
(loras / f"{stem}.metadata.json").write_text(json.dumps(meta))
|
||||
if preview:
|
||||
(loras / f"{stem}.png").write_bytes(_png())
|
||||
for number in range(previous):
|
||||
text = f"previous lora import {number + 1}"
|
||||
prompt_id = api.db.save_prompt(
|
||||
text=text,
|
||||
category=CATEGORY,
|
||||
tags=[CATEGORY, f"lora:previous-{number + 1}"],
|
||||
prompt_hash=generate_prompt_hash(text),
|
||||
)
|
||||
api.db.link_image_to_prompt(str(prompt_id), str(loras / "old.png"))
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.post("/__test__/integrations/gate")
|
||||
async def download_gate(request):
|
||||
data = await request.json()
|
||||
if not isinstance(data.get("hold"), bool):
|
||||
raise web.HTTPBadRequest()
|
||||
if data["hold"]:
|
||||
gate.clear()
|
||||
else:
|
||||
gate.set()
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.post("/__test__/integrations/fail-import")
|
||||
async def fail_import(request):
|
||||
failures["replace"] = 1
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.post("/__test__/integrations/pair")
|
||||
async def pair(request):
|
||||
with api.db.model.get_connection() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT text FROM prompts WHERE category = ? ORDER BY id LIMIT 1",
|
||||
(CATEGORY,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
raise web.HTTPConflict()
|
||||
pairs.append(api.db.save_prompt_pair(row[0], "blurry"))
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.get("/__test__/integrations/state")
|
||||
async def state(request):
|
||||
imports = api.db.search_prompts(category=CATEGORY, limit=MAX_LORAS * 2)
|
||||
return web.json_response(
|
||||
{
|
||||
"imports": sorted(prompt["text"] for prompt in imports),
|
||||
"key_saved": bool(IntegrationConfig.CIVITAI_API_KEY),
|
||||
"enabled": IntegrationConfig.LORA_MANAGER_ENABLED,
|
||||
"trigger_words": IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED,
|
||||
"downloads": len(served),
|
||||
}
|
||||
)
|
||||
|
||||
return reset
|
||||
@@ -23,47 +23,11 @@ read. These per-server hooks only change what a System test needs on top:
|
||||
image monitor whose observer is ``alive`` (default true) and schedules the
|
||||
named folders of this server's root (``["output"]``), so Sources reads
|
||||
watching, manual scan and restart pending states (S4).
|
||||
- ``POST /__test__/system/library`` ``{state}`` (S6b): ``node_folder`` makes
|
||||
the live database read as a node-folder library that can move to
|
||||
``user/__promptmanager``; Move library then runs a stand-in move that
|
||||
touches no file and ends with the library in the default folder.
|
||||
``failing`` is the same library whose move fails verification. Reset
|
||||
restores the real resolver, the real move and "Keep it here".
|
||||
- ``POST /__test__/system/legacy-previews`` ``{count}`` (#317) writes ``count``
|
||||
PNGs into ``output/legacy`` with an older release's preview beside each
|
||||
(``thumbnails/legacy/<name>_thumb.png``, newer than its file): outdated
|
||||
previews for Convert outdated;
|
||||
- ``POST /__test__/system/jxl-encoder`` ``{available}`` (#317) makes the
|
||||
server read as having (``cjxl``) or lacking a JPEG XL encoder without
|
||||
probing; reset restores the thumbnail settings and the real detection.
|
||||
|
||||
- ``POST /__test__/system/duplicate-prompts`` (S7) adds two prompts whose
|
||||
text differs only in case and spacing (a duplicate prompt text group);
|
||||
``{"again": true}`` adds a third, so a reviewed plan no longer matches;
|
||||
- ``POST /__test__/system/duplicate-files`` (S7) writes one PNG's bytes three
|
||||
times: ``output/s7-dupes/same.png``, ``input/s7-dupes/same.png`` (the same
|
||||
relative path in the other source) and ``output/s7-dupes/copy.png``;
|
||||
``{"groups": n}`` adds ``n`` more groups of two copies each, and
|
||||
``{"trio": false}`` leaves the three ``same.png``/``copy.png`` copies out,
|
||||
and ``{"long": true}`` adds a group under long names with no spaces;
|
||||
- ``GET /__test__/system/digest`` (S7) answers the sha256 of the database's
|
||||
rows (a read-only dump) and of every media file under ``output`` and
|
||||
``input`` (not the automatic previews in ``thumbnails``), so a test proves
|
||||
a Cancel wrote nothing.
|
||||
- ``POST /__test__/system/safety-backup`` (S9) writes one verified safety
|
||||
backup beside the database, as a maintenance apply does;
|
||||
``GET /__test__/system/backups`` answers ``{safety: [names], stages: n}``:
|
||||
the ``<db>.backup_*`` files and the held ``.pm-restore-*`` stage folders.
|
||||
Reset drops every held restore review, removes the safety backups and
|
||||
forgets the session's last download.
|
||||
|
||||
Delays use the shared request holds (``request_hold.py``).
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import shutil
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
@@ -90,34 +54,14 @@ def register_system_commands(app, routes, api, root: Path):
|
||||
from py.config import GalleryConfig, PromptManagerConfig
|
||||
from utils import image_monitor
|
||||
|
||||
from database import library_move
|
||||
from utils import library_location
|
||||
|
||||
configured = list(GalleryConfig.MONITORING_DIRECTORIES)
|
||||
real_move = library_move.move_library
|
||||
keep = PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER
|
||||
config_bytes = config_path.read_bytes() if config_path.exists() else None
|
||||
from py.api.settings_routes import (
|
||||
_AUTOTAG_SETTINGS,
|
||||
_GALLERY_SETTINGS,
|
||||
_SETTINGS_FILE_KEYS,
|
||||
host_values = (
|
||||
PromptManagerConfig.RESULT_TIMEOUT,
|
||||
PromptManagerConfig.WEBUI_DISPLAY_MODE,
|
||||
PromptManagerConfig.WORKER_THREADS,
|
||||
GalleryConfig.MONITORING_ENABLED,
|
||||
)
|
||||
|
||||
# Every value POST /settings can change, so a reset restores each one
|
||||
# (a saved safety_backup=false once leaked into a later spec's review).
|
||||
# The tagging settings are the tagging fixture's to reset.
|
||||
host_values = {
|
||||
(GalleryConfig if name in _GALLERY_SETTINGS else PromptManagerConfig, name): (
|
||||
getattr(
|
||||
GalleryConfig if name in _GALLERY_SETTINGS else PromptManagerConfig,
|
||||
name,
|
||||
)
|
||||
)
|
||||
for name in _SETTINGS_FILE_KEYS
|
||||
if name not in _AUTOTAG_SETTINGS
|
||||
}
|
||||
from py.api import thumbnail_encode
|
||||
|
||||
monitor = getattr(image_monitor, "_monitor_instance", None)
|
||||
failures: dict[str, int] = {}
|
||||
|
||||
@@ -135,23 +79,19 @@ def register_system_commands(app, routes, api, root: Path):
|
||||
|
||||
unplugged = root / "output" / "unplugged"
|
||||
|
||||
dupes = [root / "output" / "s7-dupes", root / "input" / "s7-dupes"]
|
||||
|
||||
def reset():
|
||||
failures.clear()
|
||||
for folder in dupes:
|
||||
shutil.rmtree(folder, ignore_errors=True)
|
||||
if unplugged.is_dir():
|
||||
unplugged.rmdir()
|
||||
GalleryConfig.MONITORING_DIRECTORIES = list(configured)
|
||||
api._health_report = None
|
||||
for (owner, name), value in host_values.items():
|
||||
setattr(owner, name, value)
|
||||
thumbnail_encode.reset_encoders()
|
||||
(
|
||||
PromptManagerConfig.RESULT_TIMEOUT,
|
||||
PromptManagerConfig.WEBUI_DISPLAY_MODE,
|
||||
PromptManagerConfig.WORKER_THREADS,
|
||||
GalleryConfig.MONITORING_ENABLED,
|
||||
) = host_values
|
||||
image_monitor._monitor_instance = monitor
|
||||
api.__dict__.pop("_library_resolver", None)
|
||||
library_move.move_library = real_move
|
||||
PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER = keep
|
||||
if config_bytes is None:
|
||||
config_path.unlink(missing_ok=True)
|
||||
else:
|
||||
@@ -160,45 +100,11 @@ def register_system_commands(app, routes, api, root: Path):
|
||||
staged.write_bytes(config_bytes)
|
||||
os.replace(staged, config_path)
|
||||
|
||||
def backup_files():
|
||||
db_path = Path(api.db.model.db_path)
|
||||
safety = sorted(
|
||||
path.name
|
||||
for path in db_path.parent.glob(db_path.name + ".backup_*")
|
||||
if path.is_file()
|
||||
)
|
||||
stages = [
|
||||
path
|
||||
for path in db_path.parent.glob(".pm-restore-*")
|
||||
if path.is_dir() and not path.name.startswith(".pm-restore-upload-")
|
||||
]
|
||||
return db_path, safety, stages
|
||||
|
||||
async def reset_backups():
|
||||
await api._discard_held(api._restore_plans.pop_all())
|
||||
db_path, safety, _stages = backup_files()
|
||||
for name in safety:
|
||||
(db_path.parent / name).unlink(missing_ok=True)
|
||||
api._last_backup_download = None
|
||||
|
||||
@routes.post("/__test__/system/reset")
|
||||
async def system_reset(request):
|
||||
reset()
|
||||
await reset_backups()
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.post("/__test__/system/safety-backup")
|
||||
async def safety_backup(request):
|
||||
from py.api.maintenance import take_safety_backup
|
||||
|
||||
name = await api._run_in_executor(take_safety_backup, api.db.model)
|
||||
return web.json_response({"success": True, "name": name})
|
||||
|
||||
@routes.get("/__test__/system/backups")
|
||||
async def backups(request):
|
||||
_db_path, safety, stages = backup_files()
|
||||
return web.json_response({"safety": safety, "stages": len(stages)})
|
||||
|
||||
@routes.post("/__test__/system/fail")
|
||||
async def system_fail(request):
|
||||
data = await request.json()
|
||||
@@ -266,43 +172,6 @@ def register_system_commands(app, routes, api, root: Path):
|
||||
api.invalidate_gallery_cache()
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.post("/__test__/system/legacy-previews")
|
||||
async def legacy_previews(request):
|
||||
data = await request.json()
|
||||
count = data.get("count")
|
||||
if (
|
||||
isinstance(count, bool)
|
||||
or not isinstance(count, int)
|
||||
or not 1 <= count <= MAX_UNTHUMBNAILED
|
||||
):
|
||||
raise web.HTTPBadRequest()
|
||||
folder = root / "output" / "legacy"
|
||||
previews = root / "output" / "thumbnails" / "legacy"
|
||||
folder.mkdir(exist_ok=True)
|
||||
previews.mkdir(parents=True, exist_ok=True)
|
||||
names = []
|
||||
for number in range(count):
|
||||
name = f"old_{number + 1:03}"
|
||||
with Image.effect_noise((256, 256), 50) as noise:
|
||||
noise.convert("RGB").save(folder / f"{name}.png")
|
||||
noise.convert("RGB").save(previews / f"{name}_thumb.png")
|
||||
later = (folder / f"{name}.png").stat().st_mtime + 5
|
||||
os.utime(previews / f"{name}_thumb.png", (later, later))
|
||||
names.append(f"legacy/{name}.png")
|
||||
api.invalidate_gallery_cache()
|
||||
return web.json_response({"success": True, "files": names})
|
||||
|
||||
@routes.post("/__test__/system/jxl-encoder")
|
||||
async def jxl_encoder(request):
|
||||
data = await request.json()
|
||||
available = data.get("available")
|
||||
if not isinstance(available, bool):
|
||||
raise web.HTTPBadRequest()
|
||||
# Test-only: stand in the cached detection (no probe).
|
||||
with thumbnail_encode._detect_lock:
|
||||
thumbnail_encode._detected = "cjxl" if available else None
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.post("/__test__/system/monitor")
|
||||
async def fake_monitor(request):
|
||||
data = await request.json()
|
||||
@@ -318,112 +187,6 @@ def register_system_commands(app, routes, api, root: Path):
|
||||
)
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.post("/__test__/system/duplicate-prompts")
|
||||
async def duplicate_prompts(request):
|
||||
data = await request.json() if request.can_read_body else {}
|
||||
texts = [("S7 duplicate prompt", "2001-01-01 00:00:00")]
|
||||
texts.append((" s7 DUPLICATE prompt", "2002-01-01 00:00:00"))
|
||||
if data.get("again") is True:
|
||||
texts = [("S7 Duplicate Prompt ", "2003-01-01 00:00:00")]
|
||||
with api.db.model.get_connection() as conn:
|
||||
conn.executemany(
|
||||
"INSERT INTO prompts(text, created_at) VALUES(?, ?)", texts
|
||||
)
|
||||
conn.commit()
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.post("/__test__/system/duplicate-files")
|
||||
async def duplicate_files(request):
|
||||
body = await request.json() if request.can_read_body else {}
|
||||
for folder in dupes:
|
||||
folder.mkdir(exist_ok=True)
|
||||
extra = body.get("groups", 0)
|
||||
if (
|
||||
isinstance(extra, bool)
|
||||
or not isinstance(extra, int)
|
||||
or not 0 <= extra <= 20
|
||||
):
|
||||
raise web.HTTPBadRequest()
|
||||
for number in range(extra):
|
||||
# One more group each: two copies of a picture no other file has.
|
||||
with Image.new("RGB", (16, 16), (number * 9, 40, 90)) as image:
|
||||
image.save(dupes[0] / f"group_{number:02}_a.png")
|
||||
copy = (dupes[0] / f"group_{number:02}_a.png").read_bytes()
|
||||
(dupes[0] / f"group_{number:02}_b.png").write_bytes(copy)
|
||||
if body.get("long") is True:
|
||||
# One group whose copies sit under long folder and file names with
|
||||
# no space to break at, as real dataset folders do.
|
||||
deep = dupes[0] / ("mabes9sans_DATASET_ideogram_v1_" * 3).rstrip("_")
|
||||
deep.mkdir(exist_ok=True)
|
||||
with Image.new("RGB", (16, 16), "#204060") as image:
|
||||
image.save(deep / ("mabes9sans_DATASET072_" * 4 + "a.png"))
|
||||
copy = (deep / ("mabes9sans_DATASET072_" * 4 + "a.png")).read_bytes()
|
||||
(deep / ("mabes9sans_DATASET072_" * 4 + "b.png")).write_bytes(copy)
|
||||
if body.get("trio", True) is not False:
|
||||
with Image.new("RGB", (48, 48), "#7a3b2e") as image:
|
||||
image.save(dupes[0] / "same.png")
|
||||
data = (dupes[0] / "same.png").read_bytes()
|
||||
(dupes[1] / "same.png").write_bytes(data)
|
||||
(dupes[0] / "copy.png").write_bytes(data)
|
||||
api.invalidate_gallery_cache()
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.get("/__test__/system/digest")
|
||||
async def digest(request):
|
||||
uri = Path(api.db.model.db_path).as_uri() + "?mode=ro"
|
||||
rows = hashlib.sha256()
|
||||
conn = sqlite3.connect(uri, uri=True)
|
||||
try:
|
||||
for line in conn.iterdump():
|
||||
rows.update(line.encode("utf-8") + b"\n")
|
||||
finally:
|
||||
conn.close()
|
||||
media = hashlib.sha256()
|
||||
for name in ("output", "input"):
|
||||
for path in sorted((root / name).rglob("*")):
|
||||
# Automatic previews are written in the background (#316): media only.
|
||||
if path.is_file() and "thumbnails" not in path.relative_to(root).parts:
|
||||
media.update(str(path.relative_to(root)).encode("utf-8"))
|
||||
media.update(hashlib.sha256(path.read_bytes()).digest())
|
||||
return web.json_response({"rows": rows.hexdigest(), "media": media.hexdigest()})
|
||||
|
||||
@routes.post("/__test__/system/library")
|
||||
async def library_state(request):
|
||||
data = await request.json()
|
||||
state = data.get("state")
|
||||
if state not in ("node_folder", "failing"):
|
||||
raise web.HTTPBadRequest()
|
||||
target = str(root / "user" / "__promptmanager")
|
||||
|
||||
def located(source):
|
||||
return library_location.LibraryLocation(
|
||||
source,
|
||||
api.db.model.db_path,
|
||||
str(config_path),
|
||||
str(root),
|
||||
target if source == library_location.SOURCE_NODE_FOLDER else None,
|
||||
target_source=library_location.SOURCE_DEFAULT,
|
||||
)
|
||||
|
||||
current = {"source": library_location.SOURCE_NODE_FOLDER}
|
||||
api._library_resolver = lambda: located(current["source"])
|
||||
|
||||
def stand_in_move(plan, *, progress, **_kwargs):
|
||||
if state == "failing":
|
||||
raise library_move.MoveFailed("verify-failed")
|
||||
progress("Copying the database")
|
||||
counts = api.db.library_counts()
|
||||
current["source"] = library_location.SOURCE_DEFAULT
|
||||
return library_move.MoveResult(
|
||||
located(library_location.SOURCE_DEFAULT),
|
||||
{"prompts": counts["positive_prompts"], "generated_images": 0},
|
||||
["prompts.db"],
|
||||
["browser.db.moved-to-__promptmanager"],
|
||||
)
|
||||
|
||||
library_move.move_library = stand_in_move
|
||||
return web.json_response({"success": True})
|
||||
|
||||
return reset
|
||||
|
||||
|
||||
|
||||
@@ -12,7 +12,6 @@ import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
import threading
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
@@ -149,12 +148,7 @@ def seed_tagging(api: PromptManagerAPI, root: Path, extra_images: int = 0) -> di
|
||||
image_id = db.link_image_to_prompt(positive_id, str(output / relative))
|
||||
if kind == "linked":
|
||||
db.link_image_to_pair(pair_id, image_id, {"source": "fixture"})
|
||||
from py.api.thumbnail_settings import ThumbnailSettings
|
||||
|
||||
# 150 px previews in the default format (#317 thumbnail settings).
|
||||
api._generate_gallery_thumbnails_sync(
|
||||
api._get_all_output_dirs(), ThumbnailSettings("webp", 80, 150)
|
||||
)
|
||||
api._generate_gallery_thumbnails_sync(api._get_all_output_dirs(), (150, 150))
|
||||
api.invalidate_gallery_cache()
|
||||
return {
|
||||
"positive_id": positive_id,
|
||||
@@ -182,28 +176,6 @@ def register_tagging_commands(
|
||||
routes: web.RouteTableDef, api: PromptManagerAPI, root: Path, model
|
||||
) -> None:
|
||||
"""Register fixed seed, snapshot and fake-model commands; no paths or SQL input."""
|
||||
# The fake download (``/__test__/tagging/download``): waits on ``event``,
|
||||
# then answers ``result``. Every reset sets the event and puts back the
|
||||
# model's raising default, so no test can leave a worker hanging.
|
||||
# ``generation`` grows on every reset: a download a reset released (or one
|
||||
# that timed out) fails without touching the next test's model state.
|
||||
download = {"event": threading.Event(), "result": True, "generation": 0}
|
||||
|
||||
def restore_download() -> None:
|
||||
download["generation"] += 1
|
||||
download["event"].set()
|
||||
model.__dict__.pop("download_model", None)
|
||||
|
||||
def fake_download(model_type: str, progress_callback=None) -> bool:
|
||||
event, generation = download["event"], download["generation"]
|
||||
released = event.wait(DOWNLOAD_HOLD_LIMIT_S)
|
||||
if not released or download["generation"] != generation:
|
||||
return False
|
||||
if download["result"]:
|
||||
model.unavailable = model.unavailable - {model_type}
|
||||
return download["result"]
|
||||
|
||||
_set_defaults(model, DEFAULTS) # the server starts on the standard ones
|
||||
|
||||
@routes.post("/__test__/tagging/reset")
|
||||
async def reset(request: web.Request) -> web.Response:
|
||||
@@ -218,7 +190,6 @@ def register_tagging_commands(
|
||||
model.unavailable = set()
|
||||
model.suggestions = None
|
||||
model.loaded = None
|
||||
restore_download()
|
||||
_set_defaults(model, DEFAULTS)
|
||||
# A test that ends with the dialog open leaves its capture behind; drop
|
||||
# every capture so the next test starts below the store's limits.
|
||||
@@ -247,33 +218,6 @@ def register_tagging_commands(
|
||||
model.unavailable = set(data["unavailable"])
|
||||
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).
|
||||
|
||||
``{"mode": "hold"}``: the next download waits until a later
|
||||
``succeed`` or ``fail`` (or a reset) releases it. ``succeed`` and
|
||||
``fail`` release a held download with that result and answer every
|
||||
later one at once. A successful fake marks the model downloaded.
|
||||
"""
|
||||
data = await request.json() if request.can_read_body else None
|
||||
if (
|
||||
request.query
|
||||
or not isinstance(data, dict)
|
||||
or set(data) != {"mode"}
|
||||
or data["mode"] not in ("hold", "succeed", "fail")
|
||||
):
|
||||
raise web.HTTPBadRequest()
|
||||
if data["mode"] == "hold":
|
||||
download["event"].set()
|
||||
download["event"] = threading.Event()
|
||||
download["result"] = True
|
||||
else:
|
||||
download["result"] = data["mode"] == "succeed"
|
||||
download["event"].set()
|
||||
model.download_model = fake_download
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.post("/__test__/tagging/suggestions")
|
||||
async def suggestions(request: web.Request) -> web.Response:
|
||||
"""Answer Generate with fixed scored names (the concept's own sample).
|
||||
@@ -324,28 +268,18 @@ def register_tagging_commands(
|
||||
return web.json_response({"captures": count})
|
||||
|
||||
|
||||
# A held fake download gives up after this long, whatever the test does.
|
||||
DOWNLOAD_HOLD_LIMIT_S = 60
|
||||
|
||||
# The standard tagging defaults: WD14 general and character thresholds, and
|
||||
# the suggestion count (``AutotagConfig``).
|
||||
DEFAULTS = (0.35, 0.85, 10)
|
||||
|
||||
|
||||
def _set_defaults(model, values) -> None:
|
||||
"""Set the saved tagging defaults (``AutotagConfig``, #328); a reset also
|
||||
puts back the saved model, batch policy and template (the fake model's
|
||||
short instructions, which the concept comparisons show)."""
|
||||
from py.config import AutotagConfig
|
||||
|
||||
general, character, limit = values
|
||||
AutotagConfig.GENERAL_THRESHOLD = general
|
||||
AutotagConfig.CHARACTER_THRESHOLD = character
|
||||
model.wd14_general_threshold = general
|
||||
model.wd14_character_threshold = character
|
||||
AutotagConfig.SUGGESTION_LIMIT = limit
|
||||
if values == DEFAULTS:
|
||||
AutotagConfig.MODEL_TYPE = "wd14-swinv2"
|
||||
AutotagConfig.EXISTING = "skip"
|
||||
AutotagConfig.TEMPLATE = model.default_prompt
|
||||
|
||||
|
||||
def _valid_defaults(data: object) -> bool:
|
||||
|
||||
@@ -164,10 +164,5 @@ def seed_workspace_visual(api: PromptManagerAPI, root: Path) -> None:
|
||||
path = root / directory / category.lower() / filename
|
||||
_write_art(path, index, positive, negatives[offset % len(negatives)])
|
||||
api.db.link_image_to_prompt(prompt_id, str(path))
|
||||
from py.api.thumbnail_settings import ThumbnailSettings
|
||||
|
||||
# 150 px previews in the default format (#317 thumbnail settings).
|
||||
api._generate_gallery_thumbnails_sync(
|
||||
api._get_all_output_dirs(), ThumbnailSettings("webp", 80, 150)
|
||||
)
|
||||
api._generate_gallery_thumbnails_sync(api._get_all_output_dirs(), (150, 150))
|
||||
api.invalidate_gallery_cache()
|
||||
|
||||
@@ -1,195 +0,0 @@
|
||||
// #343: a link drives a PromptManager node's text (and the Pair's negative) on the real pinned host.
|
||||
// The workflow is built by the real frontend through the nodes' sockets and queued as the frontend serializes it;
|
||||
// only inference is replaced (fixtures/ledger_nodes.py). The linked text must be saved, counted, recorded in the
|
||||
// execution history, and the run's image must link to it.
|
||||
import { execFile } from 'node:child_process';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import { promisify } from 'node:util';
|
||||
import { expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs';
|
||||
import { test, readExecutionRuns } from './execution-fixture.mjs';
|
||||
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']);
|
||||
|
||||
// `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
|
||||
// TextGenerate -> PreviewAny chain, whose text is only known when the node runs.
|
||||
function build(page, { kind, positive, negative, prefix }) {
|
||||
return page.evaluate(async ({ kind, positive, negative, prefix }) => {
|
||||
const { app } = await import('/scripts/app.js');
|
||||
app.graph.clear();
|
||||
const add = (type, x, y) => { const node = LiteGraph.createNode(type); app.graph.add(node); node.pos = [x, y]; return node; };
|
||||
const slot = (node, name) => { const index = node.inputs.findIndex(input => input.name === name); if (index < 0) throw new Error(`${node.type} has no ${name} socket`); return index; };
|
||||
const connect = (from, output, to, name) => { if (!from.connect(output, to, slot(to, name))) throw new Error(`${from.type} could not link to ${to.type}.${name}`); };
|
||||
const inputs = add('PMHistoryInputs', 0, 0);
|
||||
const prompt = add(kind, 700, 0);
|
||||
const set = (name, value) => { const widget = prompt.widgets.find(item => item.name === name); widget.value = value; if (widget.inputEl) widget.inputEl.value = value; };
|
||||
set('text', positive.typed);
|
||||
if (negative) set('negative_text', negative.typed);
|
||||
const drive = (name, side, y) => {
|
||||
if (!side.via) return;
|
||||
if (side.via === 'legacy') {
|
||||
// The frontend's own Primitive: a virtual node whose value is applied when the workflow is queued.
|
||||
const primitive = add('PrimitiveNode', 0, y);
|
||||
connect(primitive, 0, prompt, name);
|
||||
primitive.widgets[0].value = side.text;
|
||||
return;
|
||||
}
|
||||
const text = add('PrimitiveString', 0, y);
|
||||
text.widgets.find(widget => widget.name === 'value').value = side.text;
|
||||
let source = text;
|
||||
if (side.via === 'preview') { source = add('PreviewAny', 350, y); connect(text, 0, source, 'source'); }
|
||||
connect(source, 0, prompt, name);
|
||||
};
|
||||
drive('text', positive, 300);
|
||||
if (negative) drive('negative_text', negative, 500);
|
||||
if (kind !== 'PromptManagerText') connect(inputs, 0, prompt, 'clip');
|
||||
const sampler = add('KSampler', 1400, 0);
|
||||
connect(inputs, 1, sampler, 'model');
|
||||
connect(inputs, 2, sampler, 'latent_image');
|
||||
if (kind === 'PromptManagerText') {
|
||||
const encode = add('CLIPTextEncode', 1100, 0);
|
||||
encode.convertWidgetToInput?.(encode.widgets.find(widget => widget.name === 'text'));
|
||||
connect(inputs, 0, encode, 'clip');
|
||||
connect(prompt, 0, encode, 'text');
|
||||
connect(encode, 0, sampler, 'positive');
|
||||
} else connect(prompt, 0, sampler, 'positive');
|
||||
if (kind === 'PromptManagerPair') connect(prompt, 1, sampler, 'negative');
|
||||
else {
|
||||
const empty = add('CLIPTextEncode', 1100, 300);
|
||||
connect(inputs, 0, empty, 'clip');
|
||||
connect(empty, 0, sampler, 'negative');
|
||||
}
|
||||
const images = add('PMHistoryImages', 1800, 0);
|
||||
images.widgets.find(widget => widget.name === 'value').value = 99;
|
||||
connect(sampler, 0, images, 'latent');
|
||||
const save = add('SaveImage', 2100, 0);
|
||||
save.widgets.find(widget => widget.name === 'filename_prefix').value = prefix;
|
||||
connect(images, 0, save, 'images');
|
||||
// The Classic nodes install their own text serializer shortly after creation; queue after it, as a user does.
|
||||
await prompt._pmTextSerializerReady;
|
||||
const { output, workflow } = await app.graphToPrompt();
|
||||
return { output, workflow, promptNode: String(prompt.id) };
|
||||
}, { kind, positive, negative, prefix });
|
||||
}
|
||||
|
||||
async function run(execution, page, testCase) {
|
||||
const prefix = `textlink/${randomUUID()}`;
|
||||
const { output, workflow, promptNode } = 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);
|
||||
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('success');
|
||||
let result;
|
||||
await expect.poll(async () => {
|
||||
result = await readExecutionRuns(execution.databasePath, id);
|
||||
const run = result.rows[0];
|
||||
return run ? result.captures[run.run_id].state : null;
|
||||
}, { timeout: 10000, message: 'the run and its prompt context reach SQLite' }).toBe('complete');
|
||||
const run = result.rows[0];
|
||||
const files = Object.values(item.outputs).flatMap(value => value.images || []).map(image => image.filename);
|
||||
return { output, promptNode, context: result.captures[run.run_id].contexts[0], usage: result.usage[run.run_id], files };
|
||||
}
|
||||
|
||||
async function scan(execution) {
|
||||
const response = await execution.request.post('/prompt_manager/scan', { timeout: 60000 });
|
||||
expect(response.ok()).toBe(true);
|
||||
const events = (await response.text()).split('\n').filter(line => line.startsWith('data: ')).map(line => JSON.parse(line.slice(6)));
|
||||
expect(events.at(-1)?.type, JSON.stringify(events.at(-1))).toBe('complete');
|
||||
}
|
||||
|
||||
// filename -> linked prompt id (null when the image is not linked), read-only from the owned host's library.
|
||||
const LINKS = `import json, sqlite3, sys, pathlib
|
||||
db = sqlite3.connect(pathlib.Path(sys.argv[1]).resolve().as_uri() + "?mode=ro", uri=True)
|
||||
links = {name: None for name in sys.argv[2:]}
|
||||
for name, prompt in db.execute("SELECT filename, prompt_id FROM generated_images WHERE filename IN (%s)" % ",".join("?" * (len(sys.argv) - 2)), sys.argv[2:]):
|
||||
links[name] = prompt
|
||||
print(json.dumps(links))`;
|
||||
async function linkedPrompts(execution, files) {
|
||||
const { stdout } = await promisify(execFile)(process.env.PM_TEST_PYTHON || 'python', ['-c', LINKS, execution.databasePath, ...files], { timeout: 5000 });
|
||||
return JSON.parse(stdout);
|
||||
}
|
||||
|
||||
// Text a link resolves from the queued graph (a primitive) links its images through their embedded metadata.
|
||||
// Text only known at run time (PreviewAny) cannot be read from the image: its attribution from the recorded run is
|
||||
// #342's. Here it must never be linked to another prompt.
|
||||
async function expectImagesLinked(execution, evidence, viaRuntime) {
|
||||
await scan(execution);
|
||||
const links = await linkedPrompts(execution, evidence.files);
|
||||
expect(evidence.files.length).toBeGreaterThan(0);
|
||||
for (const [file, prompt] of Object.entries(links)) {
|
||||
if (viaRuntime) expect([null, evidence.context.positive_prompt_id], file).toContain(prompt);
|
||||
else expect(prompt, file).toBe(evidence.context.positive_prompt_id);
|
||||
}
|
||||
}
|
||||
|
||||
const unique = label => `${label} ${randomUUID()}`;
|
||||
|
||||
const sources = { primitive: 'PrimitiveString', legacy: 'a legacy Primitive', preview: 'PreviewAny' };
|
||||
for (const kind of ['PromptManager', 'PromptManagerText', 'PromptManagerFocused', 'PromptManagerSplit']) {
|
||||
for (const via of Object.keys(sources)) {
|
||||
test(`${kind} saves, counts, records and links images to text linked from ${sources[via]}`, async ({ execution, page }) => {
|
||||
test.setTimeout(90000);
|
||||
const text = unique(`Linked ${kind} ${via}`);
|
||||
const evidence = await run(execution, page, { kind, positive: { typed: 'Typed under the link', text, via } });
|
||||
const source = evidence.output[evidence.promptNode].inputs.text;
|
||||
// A real node is queued as a link; the virtual Primitive as its own value, never the hidden typed text.
|
||||
if (via === 'legacy') expect(source).toBe(text);
|
||||
else expect(Array.isArray(source), 'the frontend queued the link, not the typed text').toBe(true);
|
||||
expect(evidence.context.positive.raw).toBe(text);
|
||||
expect(evidence.context.positive.effective).toBe(text);
|
||||
expect(evidence.usage).toEqual([{ role: 'positive', run_count: 1 }]);
|
||||
const saved = await (await execution.request.get(`/prompt_manager/library/${evidence.context.positive_prompt_id}`)).json();
|
||||
expect(saved.positive.text).toBe(text);
|
||||
await expectImagesLinked(execution, evidence, via === 'preview');
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// A legacy Primitive's wildcard is drawn once when the workflow is queued: the run saves and counts what was queued.
|
||||
for (const kind of ['PromptManager', 'PromptManagerFocused']) {
|
||||
test(`${kind} saves and counts the text a wildcard legacy Primitive queued (#343 re-review R1)`, async ({ execution, page }) => {
|
||||
test.setTimeout(90000);
|
||||
const stem = unique(`Wildcard ${kind}`);
|
||||
const evidence = await run(execution, page, { kind, positive: { typed: 'Typed under the link', text: `${stem} {alpha|beta}`, via: 'legacy' } });
|
||||
const queuedText = evidence.output[evidence.promptNode].inputs.text;
|
||||
// The Classic node keeps its widget binding: the host applies the Primitive, and its own text serializer
|
||||
// queues the text as written, as on main. The Focused editor queues one choice.
|
||||
if (kind === 'PromptManager') expect(queuedText).toBe(`${stem} {alpha|beta}`);
|
||||
else expect([`${stem} alpha`, `${stem} beta`]).toContain(queuedText);
|
||||
expect(evidence.context.positive.raw).toBe(queuedText);
|
||||
expect(evidence.usage).toEqual([{ role: 'positive', run_count: 1 }]);
|
||||
const saved = await (await execution.request.get(`/prompt_manager/library/${evidence.context.positive_prompt_id}`)).json();
|
||||
expect(saved.positive.text).toBe(queuedText);
|
||||
});
|
||||
}
|
||||
|
||||
for (const [label, linked, negativeVia] of [['positive', ['text'], 'primitive'], ['negative', ['negative_text'], 'primitive'],
|
||||
['both', ['text', 'negative_text'], 'primitive'], ['neither', [], 'primitive'], ['negative from a legacy Primitive', ['negative_text'], 'legacy']]) {
|
||||
test(`PromptManagerPair with ${label} linked saves, counts and records the pair and links its images`, async ({ execution, page }) => {
|
||||
test.setTimeout(90000);
|
||||
const sides = {
|
||||
text: { typed: unique('Typed positive'), text: unique('Linked positive'), via: linked.includes('text') ? 'preview' : null },
|
||||
negative_text: { typed: unique('Typed negative'), text: unique('Linked negative'), via: linked.includes('negative_text') ? negativeVia : null },
|
||||
};
|
||||
const evidence = await run(execution, page, { kind: 'PromptManagerPair', positive: sides.text, negative: sides.negative_text });
|
||||
const expected = name => sides[name].via ? sides[name].text : sides[name].typed;
|
||||
for (const name of ['text', 'negative_text']) {
|
||||
const queuedValue = evidence.output[evidence.promptNode].inputs[name];
|
||||
if (sides[name].via === 'legacy') expect(queuedValue, name).toBe(sides[name].text);
|
||||
else expect(Array.isArray(queuedValue), name).toBe(linked.includes(name));
|
||||
}
|
||||
expect(evidence.context.positive.raw).toBe(expected('text'));
|
||||
expect(evidence.context.negative.raw).toBe(expected('negative_text'));
|
||||
expect(evidence.usage).toEqual([{ role: 'positive', run_count: 1 }, { role: 'negative', run_count: 0 }, { role: 'pair', run_count: 1 }]);
|
||||
const pair = await (await execution.request.get(`/prompt_manager/pairs/${evidence.context.pair_id}`)).json();
|
||||
expect(pair.pair).toMatchObject({ positive_text: expected('text'), negative_text: expected('negative_text') });
|
||||
await expectImagesLinked(execution, evidence, sides.text.via === 'preview');
|
||||
});
|
||||
}
|
||||
@@ -1,7 +1,5 @@
|
||||
"""Deterministic nodes and queue controls installed only in an owned test host."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
|
||||
@@ -257,51 +255,3 @@ async def cache_mode(request: web.Request) -> web.Response:
|
||||
raise web.HTTPBadRequest()
|
||||
_CACHE_MODE = mode
|
||||
return web.json_response({"configured": True})
|
||||
|
||||
|
||||
def _plugin_module(suffix: str):
|
||||
"""A loaded PromptManager module (ComfyUI names the package by its path)."""
|
||||
import sys
|
||||
|
||||
for name, module in list(sys.modules.items()):
|
||||
if name.endswith("promptmanager." + suffix):
|
||||
return module
|
||||
raise web.HTTPServiceUnavailable(text=f"PromptManager {suffix} not loaded")
|
||||
|
||||
|
||||
@PromptServer.instance.routes.post("/pm_execution_fixture/monitor")
|
||||
async def image_monitor(request: web.Request) -> web.Response:
|
||||
"""Start or stop PromptManager's image monitor on the host output folder.
|
||||
|
||||
The managed host's config disables monitoring; only the attribution cases
|
||||
(#342) turn it on, and they turn it off again.
|
||||
"""
|
||||
enabled = (await request.json())["enabled"]
|
||||
if type(enabled) is not bool:
|
||||
raise web.HTTPBadRequest()
|
||||
import folder_paths
|
||||
|
||||
monitor_module = _plugin_module("utils.image_monitor")
|
||||
config = _plugin_module("py.config").GalleryConfig
|
||||
monitor = monitor_module._monitor_instance
|
||||
if monitor is None:
|
||||
database = _plugin_module("database.operations").PromptDatabase()
|
||||
tracker = _plugin_module("utils.prompt_tracker").get_prompt_tracker(database)
|
||||
monitor = monitor_module.get_image_monitor(database, tracker)
|
||||
if enabled:
|
||||
config.MONITORING_ENABLED = True
|
||||
monitor.start_monitoring([folder_paths.get_output_directory()])
|
||||
else:
|
||||
# Joins the image worker and the observer (up to 5 s each): off the loop.
|
||||
await asyncio.to_thread(monitor.stop_monitoring)
|
||||
config.MONITORING_ENABLED = False
|
||||
return web.json_response({"monitoring": monitor.is_monitoring})
|
||||
|
||||
|
||||
@PromptServer.instance.routes.get("/pm_execution_fixture/monitor")
|
||||
async def image_monitor_state(request: web.Request) -> web.Response:
|
||||
"""File names of keyless images waiting for their run's recorded outputs."""
|
||||
monitor = _plugin_module("utils.image_monitor")._monitor_instance
|
||||
handler = getattr(monitor, "handler", None)
|
||||
waiting = handler.awaiting_history() if handler is not None else []
|
||||
return web.json_response({"awaiting": [os.path.basename(p) for p in waiting]})
|
||||
|
||||
@@ -1,379 +0,0 @@
|
||||
{
|
||||
"id": "37b6baef-94f5-49d2-9b42-1ffc70e7da92",
|
||||
"revision": 0,
|
||||
"last_node_id": 4,
|
||||
"last_link_id": 2,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 1,
|
||||
"type": "PromptManagerFocused",
|
||||
"pos": [
|
||||
600,
|
||||
0
|
||||
],
|
||||
"size": [
|
||||
550,
|
||||
650
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "clip",
|
||||
"name": "clip",
|
||||
"type": "CLIP",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "prepend_text",
|
||||
"name": "prepend_text",
|
||||
"shape": 7,
|
||||
"type": "STRING",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "append_text",
|
||||
"name": "append_text",
|
||||
"shape": 7,
|
||||
"type": "STRING",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "text",
|
||||
"name": "text",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "text"
|
||||
},
|
||||
"link": 1
|
||||
},
|
||||
{
|
||||
"localized_name": "category",
|
||||
"name": "category",
|
||||
"shape": 7,
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "category"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "tags",
|
||||
"name": "tags",
|
||||
"shape": 7,
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "tags"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "search_text",
|
||||
"name": "search_text",
|
||||
"shape": 7,
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "search_text"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "CONDITIONING",
|
||||
"name": "CONDITIONING",
|
||||
"type": "CONDITIONING",
|
||||
"links": null
|
||||
},
|
||||
{
|
||||
"localized_name": "STRING",
|
||||
"name": "STRING",
|
||||
"type": "STRING",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"title": "Focused fed by a Primitive",
|
||||
"properties": {
|
||||
"Node name for S&R": "PromptManagerFocused",
|
||||
"focusedUI": {
|
||||
"libraryOpen": false,
|
||||
"libraryFilters": {
|
||||
"category": "",
|
||||
"tags": ""
|
||||
},
|
||||
"composition": false,
|
||||
"metadata": false,
|
||||
"preview": false
|
||||
}
|
||||
},
|
||||
"widgets_values": [
|
||||
"Primitive focused text",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
""
|
||||
],
|
||||
"widgets_values_named": {
|
||||
"text": "Primitive focused text",
|
||||
"category": "",
|
||||
"tags": "",
|
||||
"search_text": "",
|
||||
"prepend_text": "",
|
||||
"append_text": "",
|
||||
"prompt_manager_focused_ui": ""
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"type": "PromptManagerPair",
|
||||
"pos": [
|
||||
600,
|
||||
1200
|
||||
],
|
||||
"size": [
|
||||
1100,
|
||||
700
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "clip",
|
||||
"name": "clip",
|
||||
"type": "CLIP",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "prepend_text",
|
||||
"name": "prepend_text",
|
||||
"shape": 7,
|
||||
"type": "STRING",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "append_text",
|
||||
"name": "append_text",
|
||||
"shape": 7,
|
||||
"type": "STRING",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "text",
|
||||
"name": "text",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "text"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "negative_text",
|
||||
"name": "negative_text",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "negative_text"
|
||||
},
|
||||
"link": 2
|
||||
},
|
||||
{
|
||||
"localized_name": "category",
|
||||
"name": "category",
|
||||
"shape": 7,
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "category"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "tags",
|
||||
"name": "tags",
|
||||
"shape": 7,
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "tags"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "search_text",
|
||||
"name": "search_text",
|
||||
"shape": 7,
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "search_text"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "positive",
|
||||
"name": "positive",
|
||||
"type": "CONDITIONING",
|
||||
"links": null
|
||||
},
|
||||
{
|
||||
"localized_name": "negative",
|
||||
"name": "negative",
|
||||
"type": "CONDITIONING",
|
||||
"links": null
|
||||
},
|
||||
{
|
||||
"localized_name": "positive_text",
|
||||
"name": "positive_text",
|
||||
"type": "STRING",
|
||||
"links": null
|
||||
},
|
||||
{
|
||||
"localized_name": "negative_text",
|
||||
"name": "negative_text",
|
||||
"type": "STRING",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"title": "Pair negative fed by a Primitive",
|
||||
"properties": {
|
||||
"Node name for S&R": "PromptManagerPair",
|
||||
"pairUI": {
|
||||
"libraryOpen": true,
|
||||
"libraryFilters": {
|
||||
"category": "",
|
||||
"tags": ""
|
||||
},
|
||||
"composition": false,
|
||||
"metadata": false,
|
||||
"preview": false
|
||||
}
|
||||
},
|
||||
"widgets_values": [
|
||||
"Typed pair positive",
|
||||
"Primitive pair negative",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
""
|
||||
],
|
||||
"widgets_values_named": {
|
||||
"text": "Typed pair positive",
|
||||
"negative_text": "Primitive pair negative",
|
||||
"category": "",
|
||||
"tags": "",
|
||||
"search_text": "",
|
||||
"prepend_text": "",
|
||||
"append_text": "",
|
||||
"prompt_manager_focused_ui": ""
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"type": "PrimitiveNode",
|
||||
"pos": [
|
||||
0,
|
||||
0
|
||||
],
|
||||
"size": [
|
||||
210,
|
||||
88
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "STRING",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "text"
|
||||
},
|
||||
"links": [
|
||||
1
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Run widget replace on values": false
|
||||
},
|
||||
"widgets_values": [
|
||||
"Primitive focused text"
|
||||
],
|
||||
"widgets_values_named": {
|
||||
"text": "Primitive focused text"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "PrimitiveNode",
|
||||
"pos": [
|
||||
0,
|
||||
1200
|
||||
],
|
||||
"size": [
|
||||
210,
|
||||
88
|
||||
],
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "STRING",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "negative_text"
|
||||
},
|
||||
"links": [
|
||||
2
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Run widget replace on values": false
|
||||
},
|
||||
"widgets_values": [
|
||||
"Primitive pair negative"
|
||||
],
|
||||
"widgets_values_named": {
|
||||
"negative_text": "Primitive pair negative"
|
||||
}
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
1,
|
||||
3,
|
||||
0,
|
||||
1,
|
||||
3,
|
||||
"STRING"
|
||||
],
|
||||
[
|
||||
2,
|
||||
4,
|
||||
0,
|
||||
2,
|
||||
4,
|
||||
"STRING"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.9,
|
||||
"offset": [
|
||||
416,
|
||||
110
|
||||
]
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,142 +0,0 @@
|
||||
// #342 on the real pinned host: the image monitor links an output only to a prompt
|
||||
// that provably produced it. An unknown node (PMHistoryCached) stands in for the
|
||||
// maintainer's LLM; KSampler is the fixture's deterministic stand-in.
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import { execFile } from 'node:child_process';
|
||||
import { promisify } from 'node:util';
|
||||
import { fileURLToPath } from 'node:url';
|
||||
import { expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs';
|
||||
import { test } from './execution-fixture.mjs';
|
||||
|
||||
const run = promisify(execFile);
|
||||
const reader = fileURLToPath(new URL('./read_execution_history.py', import.meta.url));
|
||||
|
||||
async function readLinks(databasePath, filenames, texts = []) {
|
||||
const { stdout } = await run(process.env.PM_TEST_PYTHON || 'python',
|
||||
[reader, databasePath, '--links', ...filenames, '--texts', ...texts], { timeout: 5000, maxBuffer: 1024 * 1024 });
|
||||
return JSON.parse(stdout);
|
||||
}
|
||||
|
||||
// kind: the positive node; text: a typed editor or 'llm' (linked from the unknown
|
||||
// node); prepend: 'llm' links the unknown node into prepend_text.
|
||||
function graph({ kind = 'PromptManagerPair', text = '', prepend = '', llm = '', seed = 1, prefix }) {
|
||||
const linked = value => value === 'llm' ? ['11', 0] : value;
|
||||
const nodes = {
|
||||
'10': { class_type: 'PMHistoryInputs', inputs: {} },
|
||||
'20': { class_type: kind, inputs: { text: linked(text), clip: ['10', 0], append_text: '',
|
||||
...(prepend ? { prepend_text: linked(prepend) } : {}),
|
||||
...(kind === 'PromptManagerPair' ? { negative_text: '' } : {}) } },
|
||||
'30': { class_type: 'CLIPTextEncode', inputs: { text: '', clip: ['10', 0] } },
|
||||
'40': { class_type: 'KSampler', inputs: { model: ['10', 1], latent_image: ['10', 2],
|
||||
positive: ['20', 0], negative: kind === 'PromptManagerPair' ? ['20', 1] : ['30', 0],
|
||||
seed, steps: 1, cfg: 1, sampler_name: 'euler', scheduler: 'normal', denoise: 1 } },
|
||||
'60': { class_type: 'PMHistoryImages', inputs: { value: 64, batch: 1, latent: ['40', 0] } },
|
||||
'70': { class_type: 'SaveImage', inputs: { images: ['60', 0], filename_prefix: prefix } },
|
||||
};
|
||||
if (llm) nodes['11'] = { class_type: 'PMHistoryCached', inputs: { value: llm } };
|
||||
return nodes;
|
||||
}
|
||||
|
||||
async function saved(request, prompt, history = {}) {
|
||||
const id = randomUUID();
|
||||
const response = await request.post('/prompt', { data: { prompt, prompt_id: id } });
|
||||
expect(response.ok()).toBe(true);
|
||||
let item;
|
||||
await expect.poll(async () => {
|
||||
item = (await (await request.get(`/history/${id}`)).json())[id];
|
||||
return item?.status?.messages?.some(([kind]) => kind === 'execution_success') || false;
|
||||
}, { timeout: 15000 }).toBe(true);
|
||||
const [image] = item.outputs['70'].images;
|
||||
expect(image.type).toBe('output');
|
||||
history.item = item;
|
||||
return image.filename;
|
||||
}
|
||||
|
||||
async function linkedTo(execution, filename, expected) {
|
||||
await expect.poll(async () => (await readLinks(execution.databasePath, [filename])).links[filename],
|
||||
{ timeout: 20000, message: `${filename} is linked by the image monitor` }).toBe(expected);
|
||||
}
|
||||
|
||||
async function monitoring(execution, enabled) {
|
||||
const response = await execution.request.post('/pm_execution_fixture/monitor', { data: { enabled } });
|
||||
expect(response.ok()).toBe(true);
|
||||
expect((await response.json()).monitoring).toBe(enabled);
|
||||
}
|
||||
|
||||
async function cacheMode(execution, mode) {
|
||||
expect((await execution.request.post('/pm_execution_fixture/cache', { data: { mode } })).ok()).toBe(true);
|
||||
}
|
||||
|
||||
test.beforeEach(async ({ execution }) => {
|
||||
expect((await execution.request.post('/pm_execution_fixture/metadata', { data: { disabled: false } })).ok()).toBe(true);
|
||||
await cacheMode(execution, 'CLASSIC'); // keeps the re-run's PromptManager node cached
|
||||
await monitoring(execution, true);
|
||||
});
|
||||
test.afterEach(async ({ execution }) => {
|
||||
await monitoring(execution, false);
|
||||
await cacheMode(execution, 'RAM_PRESSURE');
|
||||
});
|
||||
|
||||
test('an output whose positive is only linked prepend text stays unlinked, never given an earlier prompt', async ({ execution }, testInfo) => {
|
||||
const tag = `${testInfo.project.name} ${randomUUID()}`;
|
||||
const earlier = `a woman holding a coffee cup ${tag}`;
|
||||
const prepended = `llm prepend landscape ${tag}`, llmText = `llm text landscape ${tag}`;
|
||||
|
||||
// A typed prompt runs first: before #342 the next unidentifiable image took it.
|
||||
const first = await saved(execution.request, graph({ text: earlier, prefix: 'misslink/earlier' }));
|
||||
await linkedTo(execution, first, earlier);
|
||||
|
||||
// The maintainer's workflow: empty editor, the generated text arrives in prepend.
|
||||
const orphan = await saved(execution.request,
|
||||
graph({ prepend: 'llm', llm: prepended, prefix: 'misslink/prepend' }));
|
||||
|
||||
// An LLM feeding the editor itself is stored and links through its run key; the
|
||||
// monitor works in creation order, so the prepend image was processed first.
|
||||
const llmGraph = graph({ kind: 'PromptManager', text: 'llm', llm: llmText, prefix: 'misslink/text' });
|
||||
const proven = await saved(execution.request, llmGraph);
|
||||
await linkedTo(execution, proven, llmText);
|
||||
const state = await readLinks(execution.databasePath, [orphan], [prepended]);
|
||||
expect(state).toEqual({ links: { [orphan]: null }, stored: { [prepended]: false } });
|
||||
|
||||
// Only the sampler seed changes: the PromptManager node is served from the cache
|
||||
// without running (the host says so), so no queue entry exists for this run and
|
||||
// only the snapshot's run match can link it.
|
||||
const rerunHistory = {};
|
||||
const rerun = await saved(execution.request, { ...llmGraph, '40': { ...llmGraph['40'],
|
||||
inputs: { ...llmGraph['40'].inputs, seed: 2 } } }, rerunHistory);
|
||||
const cached = rerunHistory.item.status.messages.find(([kind]) => kind === 'execution_cached');
|
||||
expect(cached?.[1].nodes).toEqual(expect.arrayContaining(['11', '20']));
|
||||
await linkedTo(execution, rerun, llmText);
|
||||
});
|
||||
|
||||
test('an output saved without metadata links once its recorded run names a stored prompt', async ({ execution }, testInfo) => {
|
||||
const llmText = `bare output landscape ${testInfo.project.name} ${randomUUID()}`;
|
||||
const request = execution.request;
|
||||
expect((await request.post('/pm_execution_fixture/metadata', { data: { disabled: true } })).ok()).toBe(true);
|
||||
try {
|
||||
// No embedded graph: no run key can prove the tracker entry, so only the
|
||||
// run's recorded history can, and only once the run's outputs are sealed.
|
||||
// The gate after SaveImage holds the run open while the monitor looks.
|
||||
const prompt = graph({ kind: 'PromptManager', text: 'llm', llm: llmText, prefix: 'misslink/bare' });
|
||||
prompt['80'] = { class_type: 'PMHistoryImageOutput', inputs: { images: ['70', 0], mode: 'gate' } };
|
||||
expect((await request.post('/pm_execution_fixture/gate/hold')).ok()).toBe(true);
|
||||
const id = randomUUID();
|
||||
try {
|
||||
expect((await request.post('/prompt', { data: { prompt, prompt_id: id } })).ok()).toBe(true);
|
||||
let waiting = [];
|
||||
await expect.poll(async () => {
|
||||
waiting = (await (await request.get('/pm_execution_fixture/monitor')).json()).awaiting;
|
||||
return waiting.filter(name => name.startsWith('bare_')).length;
|
||||
}, { timeout: 20000, message: 'the monitor processed the bare output while its run was open' }).toBe(1);
|
||||
const [bare] = waiting.filter(name => name.startsWith('bare_'));
|
||||
expect((await readLinks(execution.databasePath, [bare])).links[bare]).toBeNull();
|
||||
expect((await request.post('/pm_execution_fixture/gate/release')).ok()).toBe(true);
|
||||
await linkedTo(execution, bare, llmText);
|
||||
expect((await (await request.get('/pm_execution_fixture/monitor')).json()).awaiting).not.toContain(bare);
|
||||
} finally {
|
||||
expect((await request.post('/pm_execution_fixture/gate/release')).ok()).toBe(true);
|
||||
}
|
||||
} finally {
|
||||
expect((await request.post('/pm_execution_fixture/metadata', { data: { disabled: false } })).ok()).toBe(true);
|
||||
}
|
||||
});
|
||||
@@ -1,73 +0,0 @@
|
||||
// #249 S6b: a fresh install keeps its library in ComfyUI's system user folder. A pinned CPU ComfyUI starts with Prompt
|
||||
// Manager and no configuration at all (no config.json, no PROMPT_MANAGER_* paths): the library must be created in
|
||||
// user/__promptmanager, nothing may be written into the node folder, and ComfyUI's /userdata API must not serve it.
|
||||
import { execFileSync } from 'node:child_process';
|
||||
import { access, readdir } from 'node:fs/promises';
|
||||
import { createServer } from 'node:net';
|
||||
import { join } from 'node:path';
|
||||
import { fileURLToPath } from 'node:url';
|
||||
import { test, expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs';
|
||||
import { launchSharedHost } from '../../tools/vendor/comfyui-host.mjs';
|
||||
import runtime from '../../tools/vendor/comfyui-runtime.json' with { type: 'json' };
|
||||
|
||||
const repositoryRoot = fileURLToPath(new URL('../../', import.meta.url));
|
||||
|
||||
/** An unused loopback port chosen by the OS (this host shares no port with the other suites). */
|
||||
function freePort() {
|
||||
return new Promise((resolve, reject) => {
|
||||
const probe = createServer();
|
||||
probe.once('error', reject);
|
||||
probe.listen({ host: '127.0.0.1', port: 0, exclusive: true }, () => {
|
||||
const { port } = probe.address();
|
||||
probe.close(error => error ? reject(error) : resolve(port));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
test('a fresh host keeps the library in user/__promptmanager and serves none of it', async ({ playwright }, testInfo) => {
|
||||
test.skip(testInfo.project.name !== 'classic-start', 'one fresh host is enough; the renderer does not matter here');
|
||||
test.setTimeout(180000);
|
||||
const port = await freePort();
|
||||
const host = await launchSharedHost({ checkout: process.env.PM_COMFY_ROOT, repository: repositoryRoot,
|
||||
python: process.env.PM_TEST_PYTHON || 'python', port, pins: runtime, forwardOutput: false,
|
||||
logPath: join(repositoryRoot, 'tools/vendor/test-results/executions/fresh-host.log'), pluginName: 'promptmanager',
|
||||
// A fresh install: nothing written, and no inherited path variable may choose the folder.
|
||||
preparePlugin: async () => ({ PROMPT_MANAGER_CONFIG_PATH: '', PROMPT_MANAGER_DATA_DIR: '',
|
||||
PROMPT_MANAGER_EXTRA_GALLERY_ROOTS: '' }) });
|
||||
const request = await playwright.request.newContext({ baseURL: `http://127.0.0.1:${port}/` });
|
||||
try {
|
||||
await expect.poll(async () => {
|
||||
if (host.child.exitCode !== null) throw new Error('The fresh host exited before it answered');
|
||||
try { return (await request.get('/system_stats')).ok(); } catch { return false; }
|
||||
}, { timeout: 120000 }).toBe(true);
|
||||
expect((await (await request.get('/system_stats')).json()).system.comfyui_version).toBe(runtime.serverVersion);
|
||||
await expect(await request.get('/prompt_manager/stats')).toBeOK();
|
||||
|
||||
const dataFolder = join(host.stateDirectory, 'user', '__promptmanager');
|
||||
await access(join(dataFolder, 'prompts.db'));
|
||||
await access(join(dataFolder, 'logs'));
|
||||
const { location } = await (await request.get('/prompt_manager/system/library')).json();
|
||||
expect(location).toMatchObject({ path: 'user/__promptmanager/prompts.db', source: 'default', movable: false,
|
||||
offer: false, problem: null });
|
||||
|
||||
// Nothing but byte code was added to the node folder: no database, config.json, logs or caches.
|
||||
const installed = join(host.stateDirectory, 'custom_nodes', 'promptmanager');
|
||||
const tracked = new Set(execFileSync('git', ['-C', repositoryRoot, 'ls-files', '-z'], { encoding: 'utf8' })
|
||||
.split('\0').filter(Boolean).map(file => file.split('/')[0]));
|
||||
const added = (await readdir(installed)).filter(name => !tracked.has(name) && name !== '__pycache__');
|
||||
expect(added).toEqual([]);
|
||||
|
||||
// ComfyUI's /userdata serves public users only: the system user's folder is out of reach.
|
||||
for (const path of ['/userdata/..%2F__promptmanager%2Fprompts.db', '/api/userdata/..%2F__promptmanager%2Fprompts.db',
|
||||
'/userdata/%2E%2E%2F__promptmanager%2Fconfig.json', '/userdata/prompts.db']) {
|
||||
const response = await request.get(path);
|
||||
expect(response.status(), path).not.toBe(200);
|
||||
}
|
||||
const listing = await request.get('/userdata?dir=..%2F__promptmanager&recurse=true');
|
||||
expect(listing.status()).not.toBe(200);
|
||||
const asSystemUser = await request.get('/userdata/prompts.db', { headers: { 'comfy-user': '__promptmanager' } });
|
||||
expect(asSystemUser.status()).not.toBe(200);
|
||||
} finally {
|
||||
try { await request.dispose(); } finally { await host.stop(); }
|
||||
}
|
||||
});
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user