Compare commits

...
Author SHA1 Message Date
Vito Sansevero 7d28a104e5 fix(threads): refuse new previews before the monitor stops; isolate the exit child (#301)
Review findings:
- The teardown makes the API refuse new previews before it detaches the
  monitor's callbacks: an image in hand that already read the preview
  queue no longer starts a preview worker during exit (Python 3.12+
  refuses that thread and logs a traceback). New
  refuse_auto_thumbnails().
- The teardown imports its bounds inside each step and falls back to the
  standard logger, so a failing import never skips the other steps.
- The exit child keeps its library and log in its own folder
  (PROMPT_MANAGER_DATA_DIR), reports watchdog's threads too, fails with
  the child's output on a hang, and rejects "Error while stopping".
- TESTING.md: "Background threads at exit".
2026-10-27 02:15:36 -07:00
Vito Sansevero e1a717cb2b fix(threads): stop background work at exit, on cleanup and on the tests' unload (#301)
- __init__.py: _stop_background_work() detaches the monitor's callbacks,
  stops the monitor, closes the preview worker and shuts the prompt
  cleanup thread down against one 10 s deadline, once. It runs from
  atexit (ComfyUI never runs aiohttp's cleanup hooks), from on_cleanup
  where a host runs it, and from the tests' package unload; it removes
  its own atexit callback. Registered after the history hook, so at exit
  it runs before the history dispose.
- AutoThumbnailer.close/close_auto_thumbnails take a timeout and return
  whether the worker ended.
- Tests: teardown result, once-only, on_cleanup, atexit order, a failing
  step or a second Ctrl+C never skipping the tracker, one shared
  deadline, three repeated loads leaving no thread, and a child process
  that exits normally with the image worker alive.
2026-10-27 02:15:36 -07:00
Vito b1e7abd380 fix(threads): stop and join the prompt cleanup thread, the image worker and the observer (#302, #370) (#372)
* fix(threads): stop and join the prompt cleanup thread, the image worker and the observer (#302, #370)

- PromptTracker: the cleanup thread waits on a stop event instead of
  time.sleep; new shutdown(timeout=2) joins it and returns whether it ended.
- ImageGenerationHandler.shutdown(timeout=5) now joins the worker and
  returns a bool; new close() is final: schedule() refuses and a sealed
  run's history re-check is dropped, so no worker restarts afterwards.
- ImageMonitor.stop_monitoring(timeout=5) closes the handler first, then
  stops and joins the watchdog observer within one deadline.
- Every test that builds a real tracker, handler or monitor ends with its
  final stop and an assertion, before its folder and databases go.
- #370: the wrapper/session history test drains before dispose(); the
  held-write bound test uses 30 s ceilings instead of tight wall clocks.

* test(history): drain the writer before asserting a bounded close it does not test (#370)

* fix(threads): a later stop waits again for a worker or observer an earlier stop left running (#302)

Review findings on #372:
- ImageGenerationHandler keeps the workers a stop signalled that had not
  ended; every later shutdown/close joins them again, so a second call no
  longer returns True while the first worker still runs.
- ImageMonitor.stop_monitoring keeps an observer or handler it left
  running and stops and joins it again on the next call.
- A closed handler takes no awaiting history path, and a re-check that
  raced close() is not kept due.
- Tests: second stops of a busy worker and a live observer, the closed
  handler's re-check paths, the handle test asserts the worker ended and
  stops before it asserts, and a no-new-thread check instead of set
  equality.

* test(browser): the System reset restores every settings value, safety_backup included

The frontend coverage shard of #372 failed in workspace-system-tagging:
the review read "No safety backup (the policy is off)". The backups spec
saves safety_backup=false, and /__test__/system/reset restored config.json
but not PromptManagerConfig.SAFETY_BACKUP (nor INFINITE_SCROLL), so a
later spec on the same worker server saw the policy off. The reset now
captures and restores every non-tagging value POST /settings can change.
2026-10-27 02:15:23 -07:00
Vito ae755d2ce0 feat(system): Auto-tagging screen, native in System (#328, PR 3) (#368)
* feat(system): configure auto-tagging natively

* test(system): compare Auto-tagging with the concept against the real backend
2026-10-27 00:59:39 -07:00
Vito 4032f4b6ab feat(system): one tagging configuration and Tag output images as a reviewed job (#328, PR 2) (#366)
* feat(autotag): model status without paths, and a busy answer that names the operation

* feat(autotag): model downloads run as Activity jobs

* fix(autotag): keep each download route's error shape and harden the fake download

* feat(autotag): one server-side tagging configuration

* fix(autotag): the legacy batch and review work again

* feat(autotag): tag output images as a reviewed Activity job

* test(autotag): the legacy batch, review and unload page workflows

* fix(autotag): server-free tagging defaults, cancel before load, and named legacy refusals

* fix(autotag): name only the workspace batch in a legacy refusal frame
2026-10-27 00:20:13 -07:00
Vito 88fa13689a feat(system): model status, a named model gate and downloads as Activity jobs (#328, PR 1) (#365)
* feat(autotag): model status without paths, and a busy answer that names the operation

* feat(autotag): model downloads run as Activity jobs

* fix(autotag): keep each download route's error shape and harden the fake download
2026-10-26 23:49:15 -07:00
Vito cfe47d3cb0 ci: support and test Python 3.12-3.14 only (#356) (#357)
* ci: test Python 3.11-3.14 and require Python 3.11 (#356)

Python 3.10 reaches end of life in October 2026 and no PromptManager
users run it; 3.14 is current and was untested. The test matrix runs
3.11-3.14 on Linux, Windows and macOS, bandit runs on 3.11, and
requires-python is >=3.11 so the package no longer claims a version CI
does not test. The packaging test drops its 3.10-only tomllib fallback.
The required checks follow the matrix (py3.10 -> py3.14).

* ci: narrow Python support to 3.12-3.14 (#356)

ComfyUI recommends 3.13, names 3.12 as the fallback and 3.14 as working
(its README); PromptManager supports exactly those. requires-python is
>=3.12, the test matrix and the required checks drop 3.11, and bandit
runs on 3.12.
2026-10-26 22:23:06 -07:00
Vito fd949b870c test(admin): wait for the restore response, not a timer (#359) (#363)
The legacy restore specs "restore validation and cancellation preserve
live data and allow a valid retry" and "backup and restore recover from
interrupted network requests without losing data" ended with a 5 s
toBeHidden on the restore modal. The modal closes only after the server
answers, so that was a 5 s deadline on a disk-bound restore, which WebKit
CI runners sometimes exceed.

Both specs, and the downloaded-backup spec that already waited on the
response, now share confirmRestoreAndWait: it waits for the
POST /prompt_manager/restore response, asserts a 200 (showing the body on
failure) and the success body, then checks the modal closes.

TESTING.md explains how to reproduce disk-bound timeouts with
--device-write-iops in the pinned image.
2026-10-26 21:52:47 -07:00
Vito 80c6ec5340 fix(lora): find LoraManager's sidecars, previews and example images in every storage configuration (#355) (#358)
* fix(lora): read LoraManager's settings and storage layout the way it does (#355)

Add py/lora_manager_layout.py: the settings file in LoraManager's lookup
order (environment folder, portable plugin file, user config folder), only
the seven keys the import needs, the LoRA roots, the centralized sidecar
root, the example-images root and the active library, each folder through
one guard. Sidecar discovery for both storage modes and the local preview
and example-image lookup live there too.

The test runner and the browser test server point LoraManager's settings
lookup at a disposable folder, so no test reads the real settings.

* fix(lora): import the sidecars and local images LoraManager really uses (#355)

The import, the legacy scan, the trigger-word cache and the image server
now share LoraManager's layout: sidecars beside the models (alongside) or
under the sidecar root (centralized), kept only when their model exists
inside a LoRA root. Each LoRA links its LoraManager preview and example
images where they are; CivitAI downloads run only for a LoRA without any
local image. The review adds skipped, with_local_media and the layout, and
offers no token when LoraManager's settings cannot be used; the legacy scan
reports the same reason without replacing anything.

find_lora_directories becomes default_lora_roots (no scan of the
LoraManager install); the preview and example-folder helpers it replaced
are removed. The trigger cache never raises into a running node.

* fix(lora): never serve a LoraManager folder that holds the home folder (#355)

The folder guard also rejects ancestors of the home folder (/home,
C:\Users), compared in the realpath + normcase form, so a LoraManager
setting cannot widen the image server's allow-list to every user's files.

* fix(lora): show linked LoraManager videos as videos (#355)

_enrich_images marks every image without is_video by its extension, so an
.mp4 preview outside the output folders shows the Video preview label and
plays in the inspector. Tests cover serving the sidecar and example-image
folders only while the integration is on, and a filesystem root as the
sidecar folder adding nothing.

* feat(lora): the import review names the LoraManager settings it read (#355)

The review gains a first row, LoraManager ("<source> settings · <mode>
storage · library <name>", the settings file as its hover title), and its
Found row counts metadata files, LoRAs and LoRAs with local images. A plan
whose LoraManager settings cannot be used takes the close-only path:
"LoraManager settings: not usable: <reason>. Nothing was changed." The
browser fixture writes each LoRA's model file and can leave out the local
preview, so the download, gate and cancel specs still exercise CivitAI.

* docs(lora): describe how the import finds LoraManager's files (#355)

* test(lora): set the settings folder relative to the process folder (#355)

os.path.relpath fails across Windows drives (the temp folder on C:, the
checkout on D:); the test now changes into its temp folder instead.

* fix(lora): review findings for the LoraManager layout (#355)

- The example-images library subfolder goes through the folder guard like
  its root, so its canonical (normcased) form matches the media check on
  Windows for a library name with capitals.
- A failed trigger-word load is retried by a running node only after
  TRIGGER_RETRY_SECONDS, not on every prompt run.
- A review whose discovery cannot list a folder says so (folder-unreadable,
  the folder's name only) instead of a generic error.

* test(images): hold the listing clock in the deferred-walk test

test_a_deferred_walk_waits_at_least_half_a_second set the previous walk
29.99 s back against a 30 s interval and then read the real clock again:
a runner slower than 10 ms between the two lines passed the interval and
the walk started at once (failed on macOS 3.13 and Windows 3.11). The
module's clock is now held for the mark, so the deferred walk is always
the one measured.

* test(lora): write the unkept-settings fixture as raw text

CodeQL read the dict with a credential-named key flowing into write_text
as clear-text storage of a secret. The value is a marker, so the fixture
now writes the JSON text directly and names the value MARKER.

* fix(lora): link each LoRA's cover as its newest image (#355)

Workspace lists show a prompt's media newest link first, so the import
links a record's images last to first: the card shows the cover (a still
in front of a leading video) the import put first.

* fix(lora): show LoRA-linked media on the Library page and in the inspector (#355)

Library previews, the positive detail and the per-prompt images page list
a prompt's linked media under the output folders plus, while the LoRA
integration is on, LoraManager's guarded folders and the download cache
(_media_roots; output root indices keep their meaning). That media keeps
its rooted identity, so Info and the media viewer open it; the path-based
serve route answers for those roots only with a file the database records
(either path spelling), never serves or makes a preview there, and output
requests never read LoraManager's layout. Such media is marked
linked_only: no preview file, no favorite star, no Tag action. Show in
folder accepts those roots; the download cache is named, not created.

* fix(lora): expose only linked files under LoraManager's folders (#355)

Serve, inspect and Show in folder share one gate for a LoRA-linked root
(_media_roots past the output folders): a file the database does not
record as a linked image is refused with 403, checked by its given and
its real path. Before, inspect and reveal read any file under those
folders by path.

* fix(lora): check the file Show in folder would open under a LoRA root (#355)

reveal_image resolved the requested path through the client-path
resolver, which can fall back to a path anchor; under a LoRA-linked root
the file that would open now passes the same recorded-media gate, so only
a linked file is ever revealed there.
2026-10-26 21:21:33 -07:00
112 changed files with 9256 additions and 789 deletions
+3 -6
View File
@@ -2,18 +2,15 @@
"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.10 on windows-latest",
"py3.11 on windows-latest",
"py3.14 on ubuntu-latest",
"py3.12 on windows-latest",
"py3.13 on windows-latest",
"py3.10 on macos-latest",
"py3.11 on macos-latest",
"py3.14 on windows-latest",
"py3.12 on macos-latest",
"py3.13 on macos-latest",
"py3.14 on macos-latest",
"coverage gate (85%)",
"black and flake8",
"frontend modules",
+2 -2
View File
@@ -26,10 +26,10 @@ jobs:
steps:
- uses: actions/checkout@v7
- name: Set up Python 3.10
- name: Set up Python 3.12
uses: actions/setup-python@v7
with:
python-version: "3.10"
python-version: "3.12"
- name: Install bandit
run: pip install bandit[toml]
+2 -1
View File
@@ -26,6 +26,7 @@ env:
# 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,
@@ -39,7 +40,7 @@ jobs:
fail-fast: false
matrix:
os: [ubuntu-latest, windows-latest, macos-latest]
python-version: ["3.10", "3.11", "3.12", "3.13"]
python-version: ["3.12", "3.13", "3.14"]
steps:
- uses: actions/checkout@v7
+8 -4
View File
@@ -1,7 +1,7 @@
# ComfyUI Prompt Manager
![License](https://img.shields.io/badge/license-MIT-blue.svg)
![Python](https://img.shields.io/badge/python-3.10%2B-blue.svg)
![Python](https://img.shields.io/badge/python-3.12%2B-blue.svg)
![ComfyUI](https://img.shields.io/badge/ComfyUI-compatible-green.svg)
![Tests](https://github.com/ComfyAssets/ComfyUI_PromptManager/workflows/Tests/badge.svg)
![Code Quality](https://github.com/ComfyAssets/ComfyUI_PromptManager/workflows/Code%20Quality/badge.svg)
@@ -483,11 +483,15 @@ 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:
- Scans all LoRA directories (including `extra_model_paths.yaml` paths) for metadata
- Downloads example images from CivitAI as 512px thumbnails
- 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
- 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.
![LoRA Results](images/pm-lora-filtered.png)
After import, filter by the `lora-manager` category to browse your LoRA collection.
@@ -929,7 +933,7 @@ 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
```
CI runs the suite on Linux, Windows and macOS for Python 3.10 to 3.13 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.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.
### Code Style
+227 -6
View File
@@ -86,6 +86,18 @@ timer that "usually" covers a frame is a race on a loaded runner. Library
virtualization specs settle with `libraryObservation.settled()` (two consecutive
frames with no Library frame queued). To look for such races, repeat a spec under
load (`--workers 8 --repeat-each 10`, and in the image with `--cpus=1` or `2`).
CPU limits do not reproduce disk-bound timeouts (backup, restore). For those,
throttle writes in the image: with rootful Docker on a host with cgroup v2, add
`--device-write-iops /dev/<disk>:20` to the `docker run`, where `<disk>` is the
whole disk (not a partition) holding Docker's data root, because the test server
keeps its database in the container's `/tmp`. On a plain partition,
`lsblk -no pkname "$(df --output=source "$(docker info -f '{{.DockerRootDir}}')" | tail -1)"`
prints it; on LVM or LUKS, `lsblk -s` walks down to the disk. A spec that waits
on a timer then fails at its 5 s expect. A spec that waits on the response can
still fail at the 30 s test timeout when several workers share the throttled
disk; that means the disk is too slow for the test budget, not that the spec is
wrong.
On unsupported Linux
distributions, browser system dependencies may need manual installation. For a
quick browser run, use `test:workflows -- --project chromium`.
@@ -263,6 +275,72 @@ Start fixing as soon as a failed log shows the cause, but push the fix only afte
the run for that commit has finished: pull-request runs use `cancel-in-progress`,
so a push cancels the running jobs and can hide a second failure.
## Background threads at exit
PromptManager's own background threads stop through one package teardown,
`_stop_background_work()` in `__init__.py` (#301, #302). It runs once and
works against one 10 s deadline. Each step gets the smaller of its own bound
and the time left. The steps, in order:
1. Detach the monitor's callbacks (Needs review recorder, preview queue,
Images listing listener).
2. `ImageMonitor.stop_monitoring` closes the image handler, which joins the
image worker (5 s). It then calls watchdog's `observer.stop()`, which is
not bounded here, and joins the observer (5 s).
3. `close_auto_thumbnails(final=True)` stops the preview worker (30 s) and
kills ffmpeg.
4. `PromptTracker.shutdown` stops the prompt cleanup thread (2 s).
A failing step, or a second Ctrl+C, is logged with its traceback and does not
skip the later steps. A clean run logs one INFO line to PromptManager's log:
`Background work stopped in 0.4 s: monitor=ok thumbnails=ok tracker=ok`. If
any part is still running, one WARNING goes to the console:
`Background work not stopped in time (10.0 s): monitor`.
ComfyUI never runs aiohttp's `on_cleanup` hooks, because it never calls
`runner.cleanup()`. So `atexit` is the live host's path. It runs on Ctrl+C
(`main.py` catches `KeyboardInterrupt` and returns) and on a normal exit.
The teardown is registered after the execution history hook, so at exit
(newest first) it runs before the history writer's 2 s dispose.
It also runs from `on_cleanup` on hosts that do run it (our test servers) and
from the tests' package unload (`cleanup_package`), and it removes its own
`atexit` callback. Nothing runs when the process is killed: SIGTERM, a
ComfyUI-Manager restart through `os.execv`, a closed Windows console, or
ComfyUI Desktop stopping its server. SQLite's WAL journal covers those.
Stop contract: `PromptTracker.shutdown(timeout)`,
`ImageGenerationHandler.shutdown(timeout)` (restartable),
`ImageGenerationHandler.close(timeout)` (final: `schedule()` refuses and a
sealed run's re-check is dropped) and `ImageMonitor.stop_monitoring(timeout)`
each signal their thread and join it within the bound. They return True only
when it ended, and log a warning naming the part when it did not. Called on
their own thread, they only signal.
Test rules:
- Every test that constructs a real `PromptTracker`, `ImageGenerationHandler`
or `ImageMonitor` ends with its final stop and an assertion
(`tests/background_stops.py`: `close_handler`, `stop_monitor`,
`stop_tracker`). The stop comes before the folder and databases are removed.
`tearDown` runs before `addCleanup` callbacks, so a class whose `tearDown`
removes files stops the thread there.
- A test waits for a state through an `Event` set by the code under test, or
by a wrapper around it, with a 30 s ceiling. It never sleeps or polls for
that state.
- A test that asserts "still running" first waits for an `entered` event
inside the held call, then stops with `timeout=0.05`.
- A test that asserts a history writer's `dispose()`/`close()` is True right
after queuing work drains first (`drain_writer`), or holds the write itself
(#370).
- Patch the module attribute of the code under test (for example
`utils.prompt_tracker.time`, or the package's `_clock`), or patch with
`wraps=`. Never replace a `time` or `atexit` function for the whole
process.
- `tests/test_package_exit.py` runs a child process that imports the
package, leaves the image worker alive and returns. Its `atexit` reporter
checks the teardown's result. This is the interpreter path ComfyUI takes
after Ctrl+C, on all three OS.
## Performance budgets
A spec performance budget gets two tests:
@@ -1492,8 +1570,8 @@ targets of the longest legal path (1024 code points, even escaped as `\uXXXX`) f
Its fields show the server's defaults (`GET /prompt_manager/autotag/models`, the
D91 Suggestions count included) once that read lands; Start waits for it. A model
the status lists as not downloaded keeps Start `aria-disabled` and
`#tag-generation-status` says "<model> is not downloaded. Download it from the
legacy Auto-tag page." (the server's own sentence), so no job starts only to fail.
`#tag-generation-status` says "<model> is not downloaded. Download it in System ›
Auto-tagging." (the server's own sentence), so no job starts only to fail.
Start validates the fields as Generate does (native validity; nothing is sent).
The concept comparison (`images-concept.spec.mjs`) uses the concept page itself as
the oracle, with Add to collection and Remove links… removed and a copy of its Add
@@ -2575,7 +2653,7 @@ the connection's `total_changes` is 0 at close and the data reads back unchanged
`-shm` bytes are not compared: every WAL reader rewrites its read marks there, so no
reader can keep it byte-identical (`test_health_check_never_writes`).
Recorded deviations (plan D60–D75, D79–D82 renumbered D92–D107 and D111–D114; D121–D131 assigned):
Recorded deviations (plan D60–D75, D79–D82 renumbered D92–D107 and D111–D114; D121–D134 assigned):
- **D92** (shell): no "Sample configuration" chip. The Save status is truthful: "No
unsaved changes.", "Unsaved changes.", "Saved in this browser." / "Saved to the server
@@ -2609,7 +2687,8 @@ Recorded deviations (plan D60–D75, D79–D82 renumbered D92–D107 and D111–
output folder is used." when that is the case.
- **D94** (all sections): a section not delivered yet shows the hand-off panel "Your
current tools are ready" with a link to each legacy tool that still owns it (the
plan's legacy coverage table); retired section by section, gone after S12.
plan's legacy coverage table); retired section by section, gone after S12. Since S11
(#328) only Diagnostics & logs hands off (`SECTION_HANDOFFS`, to the admin page).
- **D95** (health report): the concept's `health-check` dialog titled "Health report",
with no sample note. It lists every check in the server's order: Integrity check,
Foreign keys & pair references, Canonical hashes, Pair uniqueness (added), Missing
@@ -2985,6 +3064,59 @@ Recorded deviations (plan D60–D75, D79–D82 renumbered D92–D107 and D111–
dialog on that scope. Pinned by the thumbnail dialog comparison (mapped) and its
mutation check (system-concept.spec) and `the thumbnail row counts outdated previews`
(workspace-system.spec).
- **D132** (Auto-tagging, S11, #328, OQ1/OQ5/OQ6/OQ10): "Tagging model" (`#tag-model`)
lists the four real engines in the concept's order, each `<label> · <size_label>` from
`GET /prompt_manager/autotag/models` ("WD14 SwinV2 · ~400 MB", "WD14 ViT · ~400 MB",
"GGUF JoyCaption · ~4.5 GB", "HuggingFace JoyCaption · ~16 GB"; GGUF takes the
concept's "Ollama" slot, as D61 does). Its description is the chosen engine's live
status instead of the concept's design note: "Loading model status…" (the select is
disabled until the read answers), "Model status could not be read.", "Not downloaded.
Download size <size>.", "Downloaded. Not loaded.", "Downloaded and loaded." or "Partly
downloaded. Download again to finish." (GGUF with one of its two files). The concept's
two simulated job buttons are "Download" ("Download again" for an engine on disk),
which starts the `model_download` Activity job and opens it, and "Unload", a direct
`POST /autotag/unload` with a toast ("<label> unloaded.") and a fresh status read,
enabled only while a model is loaded; a busy model's refusal reads "The tagging model
is busy: <operation>. Try again when it finishes.". "Preview tagging review" opens the
Tag image dialog for the image selected in Images with the section's current model and
values; it is `aria-disabled` without a selection (the D115 sentence "Select an image
in Images first.") or with an invalid value (the toast names the setting); a saved
engine the section does not know reads "The saved tagging model is not one of these
engines. Choose one and save preferences.". A part-downloaded GGUF offers "Download
again"; an Unload that found nothing loaded only reads the status again. At phone width
the concept's 120px select cap cuts every engine label, so `#tag-model` alone keeps its
natural width there and its row wraps the select under the text (`workspace.css`
`/* D132 */`); Existing tags stays as the concept draws it, cut at 390 as there. The
comparison maps `selectsFit` only for the engine select and measures each row's
`flexWrap` and whether its control sits under the text. Pinned by `auto-tagging matches
the concept` (both themes and sizes), `auto-tagging at 390 wraps only the engine row`,
`auto-tagging engine select fits at 320 wide` and the mutation check (it includes a
wrapped row) (system-tagging-concept.spec), and the model, Download, Unload and Preview
tests of workspace-system-tagging.spec.
- **D133** (Auto-tagging, S11): a "Suggestions per image" row (`#tag-suggestions`,
number 1–100, host key `autotag_suggestion_limit`) follows the two confidence rows:
the persisted count the Tag image dialog's Suggestions field (D91) starts from. Pinned
by the comparison's mapping (it inserts the row) and `defaults save through Save
preferences and the Tag image dialog uses them`.
- **D134** (Auto-tagging, S11, OQ2/OQ3/OQ7/OQ8): an `.operation` "Tag output images"
with "Review" follows `.row-actions`: the legacy prompt-tag batch. Review is
`aria-disabled` (the press explains) while the chosen engine is not downloaded or a
value it uses is invalid; otherwise it posts `/autotag/batch/preview` with the drafted
engine and policy (read-only) and opens "Review tag output images" in the concept's
record-changing review (its `maintenancePlan` shell): `h3` "Tag output images", the
`dl.meta-grid` Model, Existing tags, Images linked, Prompts, Already tagged, Prompts
to tag ("up to N"), Settings ("General <g> · Character <c>" or "Template of N
characters") and Backup, the upper-bound sentence in the `p.small`, the D55 backup box
and Cancel / "Start tagging", which starts the cancellable `autotag` job and opens it.
Replace after review, a model not on disk, a short disk, no linked image and nothing to
tag each show the rows with only Cancel and the reason ("Replacing tags needs a review
of each image. Use Preview tagging review, or choose Skip or Merge for a batch run.",
"Every prompt with an output image is already tagged.", "Not enough free disk space:
the safety backup needs about X MB and Y MB is free. Nothing is written until there is
room.") in the close-only review Integrations shares (`openCloseOnlyReview`). Pinned by
`tag output images review matches the concept`, its mutation check and `the
Cancel-only review matches the concept review shell` (system-tagging-concept.spec) and
the batch tests of workspace-system-tagging.spec.
- **D124** (inspector): with "Load metadata on selection" off, the Images panel and
the viewer's details show "Metadata loads when you ask for it." and a "Load metadata"
button instead of reading the file. The concept has no such state, so no concept
@@ -3410,20 +3542,109 @@ refused while a restore runs; the review is a request lease ("LoRA import review
accepts the optional `expected_revision`. LoraManager detection also searches every
`custom_nodes` folder ComfyUI registers through `folder_paths` (#144), main folder first.
**Where LoraManager keeps its files (#355).** `py/lora_manager_layout.py` repeats
LoraManager's rules without importing it (`tests/test_lora_manager_layout.py`). The
settings file is `<LORA_MANAGER_SETTINGS_DIR>/settings.json`, else the plugin's own
`settings.json` when portable mode is on (`LORA_MANAGER_PORTABLE` `1`/`0` decides, other
values are ignored; otherwise the JSON value `use_portable_settings: true`), else
`platformdirs.user_config_dir("ComfyUI-LoRA-Manager")`. It is read as UTF-8, at most 5 MB,
and only seven keys are kept (`SETTINGS_KEYS`); a missing file means LoraManager's
defaults. Discovery: in `alongside` mode every LoRA root is searched and a sidecar is kept
when `<base>.safetensors` exists beside it; in `centralized` mode only the sidecar root is
searched and a sidecar is kept when its `file_path` exists inside a LoRA root (unparseable
ones are skipped by name). `lora_files` counts the metadata files seen, `with_metadata`
the kept, readable LoRAs and `skipped` the rest. Media per LoRA, without repeats: the
preview (`preview_url`, else LoraManager's `PREVIEW_EXTENSIONS` order), then every file in
`<example-images root>[/<library>]/<sha256>/` in natural order; a leading video yields to the
first still. A file counts only with a servable extension and inside the guarded media
roots. CivitAI downloads run only for a LoRA without local media. The folder guard keeps a
settings folder only when it exists and is neither a filesystem root (`/`, `C:\`, a UNC
share root) nor the home folder or one of its ancestors (realpath + `normcase`). An
unusable layout (no install, unreadable / non-JSON / non-object / oversized settings, no
`platformdirs`, a missing or refused centralized sidecar root) gives a review without a
token ("LoraManager settings: not usable: <reason>. Nothing was changed."), a legacy scan
error with that reason, the default LoRA roots for serving, and trigger words from the
default roots in alongside mode. `tests/run_suite.py` and `tests/browser/server.py` point
`LORA_MANAGER_SETTINGS_DIR` at a disposable folder, so no test reads the developer's real
LoraManager settings; tests that call `load_layout` inject their own folder too.
`tests/browser/workspace_integrations_fixture.py` (`/__test__/integrations/*`) makes a fake
LoraManager install and LoRA folder under the server's temp root, holds or fails the
import, saves a blocking pair and answers CivitAI downloads with a stub opener, so no
test reaches the network.
```sh
python3 -m unittest tests.test_lora_import tests.test_api_lora tests.test_lora_database \
tests.test_lora_utils -v
python3 -m unittest tests.test_lora_manager_layout tests.test_lora_import tests.test_api_lora \
tests.test_lora_database tests.test_lora_utils -v
node --test tests/js/workspace-system-integrations.test.js tests/js/prompt-bridge-adapter.test.js
npm run --prefix tools/vendor test:workflows -- workspace-system-integrations.spec.mjs \
system-integrations-concept.spec.mjs \
--project chromium --project touch --project firefox --project webkit
```
### Auto-tagging contract (#328, #249 System S11)
System › Auto-tagging (`web/js/workspace/system-tagging.mjs`, D132–D134) edits one
server-side tagging configuration, `autotag.defaults` in config.json: `model_type`
(`wd14-swinv2`, `wd14-vit`, `gguf`, `hf`), `existing` (`skip`, `merge`,
`replace_after_review`), `general_threshold` and `character_threshold` (0–1),
`template` (non-blank, at most 4,000 characters and 16 KiB) and `suggestion_limit`
(1–100). They are host drafts (`autotag_model_type`, `autotag_existing`,
`autotag_general_threshold`, `autotag_character_threshold`, `autotag_template`,
`autotag_suggestion_limit`) saved by the one Save row through `POST
/prompt_manager/settings` with `expected_revision`; a bad value is refused with a 400
naming the setting by its row label ("General-tag confidence must be between 0 and
1."), nothing is written, the Save row shows that sentence (`invalid-setting`, set by
`settingRefusal` in system.mjs only for a save that carried `autotag_*` keys; every
other section's plain 400 keeps the generic status) and the control is highlighted.
The template's limit counts code points, as Python's `len` does. A hand-edited bad value
falls back to its default on load, one key at a time. `GET /autotag/models` reports the
saved values (`wd14_general_threshold`, `wd14_character_threshold`, `suggestion_limit`,
`default_prompt`, `model_type`, `existing`), so the Tag image and Auto-tag selected
dialogs start from them; a preset (the bulk review, Preview tagging review) wins.
`GET /prompt_manager/autotag/models` reads the disk in the executor, never loads or
downloads a model and carries no absolute path: per engine `name`, `label`,
`size_label`, `description`, `size`, `downloaded`, `loaded` (GGUF also `model_exists`,
`mmproj_exists`), plus `model_loaded` and `loaded_model_type`. Opening the section
makes exactly one such read per mount (`entering the section makes no model call
besides status`).
One model gate (`py/model_gate.py`) covers download, unload, both batches, the legacy
single-image tag, Generate suggestions and Auto-tag selected. A refusal is `409
model-busy` with `operation`, the holder's label: the download job's title, "Tag output
images", "Tag output images (legacy page)", "Tagging one image (legacy page)",
"Unloading the tagging model", "Auto-tag selected images", "Generate suggestions", else
"Another tagging operation". `ApiError.operation` keeps it (only on `model-busy`).
- **Download**: `POST /prompt_manager/jobs/model-download {model_type}` starts the
`model_download` job (202), title "Tagging model download · <label>", admission
`read`, not cancellable, indeterminate (Activity's generic held view: no Cancel, a
`<progress>` without a value), holding the gate until it ends. A second download is
`409 busy`; the legacy SSE route attaches to the same job and keeps its frames.
- **Unload**: `POST /prompt_manager/autotag/unload` answers `{success, message}` or
`409 model-busy`.
- **Tag output images**: `POST /prompt_manager/autotag/batch/preview {model_type,
existing}` writes nothing and opens no file; it answers `images`, `prompts`,
`already_tagged`, `to_tag` (an upper bound), `model {type, label, size_label,
downloaded}`, `existing`, `review_required`, `backup {policy}` and `disk`. `POST
/prompt_manager/jobs/autotag {model_type, existing, options}` (WD14: both thresholds;
GGUF/HF: `instructions`, the template) starts the cancellable `autotag` job (202),
title "Tag output images", admission `records`, a safety backup by the policy first.
Refusals: `400 review-required` (Replace after review never runs as a batch), `409
model-unavailable` ("<label> is not downloaded. Download it in System ›
Auto-tagging."), `507 insufficient-space`, `409 busy`, `409 model-busy`. The section
toasts `model-busy` with its operation, the other three with the server's sentence,
`busy` with the running job (D57), anything else "<what> could not start. Try again.".
Outcomes: Tagged / Unchanged / Skipped / Errors (the `image-tags` descriptor, D111).
No browser test downloads or runs a model: the injected `TestModel` answers, its
`download_model` raises unless `/__test__/tagging/download {mode: hold | succeed |
fail}` installs a fake that waits on a fixture event, and every tagging reset releases
it. Every workspace-system-tagging test asserts no `/autotag/download/` request.
Changing the routes or `py/config.py` needs a ComfyUI restart; the section itself is
HTML, CSS and JS only.
### Browser cache ownership (#222)
The legacy gallery's Clear cache used to delete every localStorage key starting `pm_` or
+153 -21
View File
@@ -22,7 +22,10 @@ Features:
- Comprehensive logging and diagnostics
"""
import atexit
import re
import threading
import time
from pathlib import Path
@@ -72,6 +75,12 @@ NODE_DISPLAY_NAME_MAPPINGS = {
# are served by our own routes instead (see py/api).
WEB_DIRECTORY = "web/comfy"
# Parts the package teardown stops (``_stop_background_work``). Each stays None
# when its start-up block below fails, and the teardown skips it.
api = None
_global_prompt_tracker = None
_global_image_monitor = None
# Add API routes (same pattern as ComfyUI_Assets)
try:
# Set extension URI
@@ -89,8 +98,6 @@ try:
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",
@@ -141,29 +148,16 @@ try:
# The Needs review pass for ambiguous new outputs runs through the API's
# capture store (the one the scan and the Tag dialog use): a linked output
# is never queued and its pending items leave the queue. It stops before
# that store is closed at shutdown.
_review_api = globals().get("api")
if _review_api is not None:
_global_image_monitor.set_review_recorder(_review_api.record_review_file)
# is never queued and its pending items leave the queue. The package
# teardown detaches it before that store is closed at shutdown.
if api is not None:
_global_image_monitor.set_review_recorder(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)
_global_image_monitor.set_thumbnail_queue(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
):
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)
_global_image_monitor.set_output_listener(api.notify_output_settled)
# Start monitoring
_global_image_monitor.start_monitoring()
@@ -233,6 +227,144 @@ except Exception:
"Failed to register execution history"
)
# ── Package teardown (#301, #302) ────────────────────────────────────────────
# ComfyUI never runs aiohttp's cleanup hooks (it never calls
# ``runner.cleanup()``), so ``atexit`` is the live host's path: Ctrl+C, and a
# normal exit. It is registered after the history hook, so at exit (newest
# first) it runs before the history dispose. Hosts that do run ``on_cleanup``
# (our test servers) run it there, and the tests' package unload calls it.
TEARDOWN_SECONDS = 10.0
_clock = time.monotonic
_background_work_lock = threading.Lock()
_background_work_result = None
def _teardown_step(name, stop, logger):
"""Run one teardown step; a failure (even Ctrl+C) is logged and is False."""
try:
return bool(stop())
except (Exception, KeyboardInterrupt):
logger.error(f"Background work: stopping {name} failed", exc_info=True)
return False
def _teardown_logger():
"""PromptManager's init logger; the standard one if that cannot load."""
try:
from .utils.logging_config import get_logger
return get_logger("prompt_manager.init")
except Exception:
import logging
return logging.getLogger("prompt_manager.init")
def _stop_background_work():
"""Stop the image monitor, the preview worker and the prompt cleanup
thread against one TEARDOWN_SECONDS deadline, once.
Returns ``{"monitor", "thumbnails", "tracker"}`` booleans (True: it
stopped), leaving out parts that never started. Later calls return a
copy of the first result.
"""
global _background_work_result
with _background_work_lock:
if _background_work_result is not None:
return dict(_background_work_result)
logger = _teardown_logger()
began = _clock()
deadline = began + TEARDOWN_SECONDS
def left(own_bound):
return max(0.0, min(own_bound, deadline - _clock()))
monitor, thumbnails, tracker = (
_global_image_monitor,
api,
_global_prompt_tracker,
)
result = {}
if monitor is not None:
def stop_monitor():
from .utils.image_monitor import OBSERVER_JOIN_SECONDS
# Callbacks first: an image in hand queues no preview into a
# closing preview worker. The worker may have read the queue
# callback already, so the API refuses new previews too (at
# exit, Python 3.12+ refuses to start their worker thread).
if thumbnails is not None:
thumbnails.refuse_auto_thumbnails()
monitor.set_review_recorder(None)
monitor.set_thumbnail_queue(None)
monitor.set_output_listener(None)
return monitor.stop_monitoring(left(OBSERVER_JOIN_SECONDS))
result["monitor"] = _teardown_step(
"the image monitor", stop_monitor, logger
)
if thumbnails is not None:
def close_previews():
from .py.api.thumbnail_auto import CLOSE_WAIT_SECONDS
return thumbnails.close_auto_thumbnails(
final=True, timeout=left(CLOSE_WAIT_SECONDS)
)
result["thumbnails"] = _teardown_step(
"the preview worker", close_previews, logger
)
if tracker is not None:
def stop_tracker():
from .utils.prompt_tracker import TRACKER_JOIN_SECONDS
return tracker.shutdown(left(TRACKER_JOIN_SECONDS))
result["tracker"] = _teardown_step(
"the prompt cleanup thread", stop_tracker, logger
)
elapsed = _clock() - began
failed = [name for name, ok in result.items() if not ok]
if failed:
logger.warning(
f"Background work not stopped in time ({elapsed:.1f} s): "
+ ", ".join(failed)
)
else:
parts = " ".join(f"{name}=ok" for name in result)
logger.info(f"Background work stopped in {elapsed:.1f} s: {parts}")
_background_work_result = result
atexit.unregister(_stop_background_work)
return dict(result)
async def _stop_background_work_on_cleanup(app):
"""aiohttp ``on_cleanup``: no request is served then, so it runs inline."""
_stop_background_work()
try:
atexit.register(_stop_background_work)
from .py.config import server_instance as _teardown_server
_teardown_host = getattr(_teardown_server, "app", None)
if _teardown_host is not None and hasattr(_teardown_host, "on_cleanup"):
_teardown_host.on_cleanup.insert(0, _stop_background_work_on_cleanup)
except Exception as e:
try:
from .utils.logging_config import get_logger
get_logger("prompt_manager.init").error(
f"Failed to register the background work teardown: {e}"
)
except Exception:
pass
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
# Print startup message with loaded tools
+4 -1
View File
@@ -192,7 +192,10 @@ class CategoryImportOperations:
outcome = "imported"
created[prompt_hash] = prompt_id
result[_COUNTERS[outcome]] += 1
for image in record.get("images") or []:
# 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})
+31
View File
@@ -1808,6 +1808,37 @@ class PromptDatabase(
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)
]
def get_all_images(
self,
limit: int = 0,
+4 -2
View File
@@ -268,8 +268,10 @@ class PromptManagerBase:
cache = get_trigger_cache()
# Lazy-load cache on first use
if not cache.is_loaded and IntegrationConfig.LORA_MANAGER_PATH:
# 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():
cache.load(IntegrationConfig.LORA_MANAGER_PATH)
modified, injected = inject_trigger_words(text, cache)
+11 -1
View File
@@ -750,7 +750,12 @@ 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."""
"""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).
"""
from urllib.parse import quote as url_quote
output_dirs = self._get_all_output_dirs()
@@ -760,6 +765,11 @@ 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
+17 -24
View File
@@ -10,12 +10,15 @@ filesystem or the database.
import math
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
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,
)
SUGGEST_FIELDS = frozenset(
{"target", "capture_token", "expected_epoch", "model_type", "options"}
@@ -24,7 +27,7 @@ _COMMON = frozenset({"suggestion_limit", "use_gpu"})
_WD14_REQUIRED = frozenset({"general_threshold", "character_threshold"})
_LLM_REQUIRED = frozenset({"instructions"})
UNAVAILABLE_HINT = "Download it from the legacy Auto-tag page."
UNAVAILABLE_HINT = "Download it in System › Auto-tagging."
class OptionError(ValueError):
@@ -32,27 +35,17 @@ class OptionError(ValueError):
def _threshold(name: str, value) -> float:
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)
try:
return validate_threshold(name, value)
except ValueError as error:
raise OptionError(str(error)) from None
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:
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
return validate_template(value)
except ValueError as error:
raise OptionError(str(error)) from None
def validate_options(model_type, options, *, default_limit: int, max_limit: int):
+632 -89
View File
@@ -1,14 +1,26 @@
"""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 .jobs import AUTOTAG_BATCH_JOB_KIND
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 .prompts import thumbnail_url_for
from . import thumbnail_settings
from .. import model_gate
@@ -16,20 +28,104 @@ 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."""
"""``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.
"""
return web.json_response(
{"success": False, "code": "model-busy", "error": MODEL_BUSY_ERROR},
{
"success": False,
"code": "model-busy",
"error": MODEL_BUSY_ERROR,
"operation": _holder_label(model_gate.holder()),
},
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"}
@@ -65,6 +161,39 @@ 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"
@@ -83,6 +212,10 @@ 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)
@@ -110,18 +243,24 @@ class AutotagRoutesMixin:
from ..config import AutotagConfig
service = get_autotag_service()
models_status = service.get_models_status()
# 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
return web.json_response(
{
"success": True,
"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,
"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,
"suggestion_limit": AutotagConfig.suggestion_limit(),
"model_type": AutotagConfig.MODEL_TYPE,
"existing": AutotagConfig.EXISTING,
}
)
@@ -159,55 +298,156 @@ class AutotagRoutesMixin:
return response
async def download_autotag_model(self, request):
"""Download an AutoTag model with streaming progress (POST only)."""
"""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
model_type = request.match_info.get("model_type")
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 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,
}
if not model_gate.try_acquire(owner):
return model_busy_response()
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:
return await self._stream_sse(request, events())
finally:
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()
)
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))
@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
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
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.
@@ -238,35 +478,48 @@ class AutotagRoutesMixin:
and not t.startswith("append:")
]
async def _apply_tags_to_prompt(self, prompt_id, tags) -> bool:
"""Merge ``tags`` into the prompt; True when something new was added."""
async def _apply_tags_to_prompt(self, prompt_id, tags) -> int:
"""Merge ``tags`` into the prompt; how many new tags were added."""
existing_prompt = await self._run_in_executor(
self.db.get_prompt_by_id, prompt_id
)
if not existing_prompt:
return False
return 0
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 = [t for t in tags if t not in existing_tags]
new_tags = list(dict.fromkeys(t for t in tags if t not in existing_tags))
if not new_tags:
return False
return 0
await self._run_in_executor(
self.db.update_prompt_metadata, prompt_id, tags=existing_tags + new_tags
)
return True
return len(new_tags)
async def _autotag_batch_events(self, service, opts: dict):
"""Yield SSE progress dicts while tagging every image linked to a prompt."""
import time as _time
"""Yield SSE progress dicts while tagging every image linked to a prompt.
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"]
status = service.get_models_status()
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)
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:
@@ -279,15 +532,21 @@ class AutotagRoutesMixin:
}
return
if opts["custom_prompt"]:
if opts.get("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.get_all_images)
images = await self._run_in_executor(self.db.autotag_batch_candidates)
total = len(images)
if total == 0:
yield {
@@ -299,7 +558,8 @@ class AutotagRoutesMixin:
"skipped_outside_roots": 0,
"status": "No images with linked prompts found",
}
service.unload_model()
if not keep_in_memory:
service.unload_model()
return
yield {
@@ -316,7 +576,7 @@ class AutotagRoutesMixin:
"skipped_outside_roots": 0,
}
tagged_prompt_ids = set()
last_update = _time.monotonic()
last_update = _monotonic()
finished = False
def progress_event(i, message):
@@ -332,6 +592,8 @@ 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
@@ -350,7 +612,7 @@ class AutotagRoutesMixin:
counts["skipped_outside_roots"] += 1
else:
skip_reason = "file missing"
elif opts["skip_tagged"] and self._real_tags(
elif opts["existing"] == "skip" and self._real_tags(
image_data.get("prompt_tags", [])
):
tagged_prompt_ids.add(prompt_id)
@@ -358,7 +620,8 @@ class AutotagRoutesMixin:
if skip_reason is not None:
counts["skipped"] += 1
now = _time.monotonic()
on_item(i, total, "skipped", image_path, SKIP_DETAILS[skip_reason])
now = _monotonic()
if (now - last_update) >= 0.5 or i == total - 1:
yield progress_event(
i, f"Skipping {i + 1}/{total} ({skip_reason})..."
@@ -368,17 +631,17 @@ class AutotagRoutesMixin:
continue
try:
tags = await run(
service.generate_tags,
allowed_path,
general_threshold=opts["general_threshold"],
character_threshold=opts["character_threshold"],
)
tags = await run(service.generate_tags, allowed_path, **tag_kwargs)
counts["processed"] += 1
if tags and await self._apply_tags_to_prompt(prompt_id, tags):
added = 0
if tags:
added = await self._apply_tags_to_prompt(prompt_id, tags)
if added:
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(
@@ -387,10 +650,11 @@ 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 = _time.monotonic()
last_update = _monotonic()
finished = True
if not keep_in_memory:
@@ -418,16 +682,22 @@ class AutotagRoutesMixin:
opts = {
"model_type": params.get("model_type", "gguf"),
"custom_prompt": params.get("prompt", ""),
"skip_tagged": str(params.get("skip_tagged", "true")).lower() == "true",
"existing": (
"skip"
if str(params.get("skip_tagged", "true")).lower() == "true"
else "merge"
),
"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:
for key in ("general_threshold", "character_threshold"):
if opts[key] is not None:
opts[key] = float(opts[key])
opts["general_threshold"], opts["character_threshold"] = (
_thresholds_or_saved(
opts["general_threshold"], opts["character_threshold"]
)
)
except (ValueError, TypeError):
return self._bad_request(
"general_threshold and character_threshold must be numeric"
@@ -457,7 +727,9 @@ 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:
if busy.kind == AUTOTAG_BATCH_JOB_KIND: # unchanged payload
# 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:
return web.json_response(
{"success": False, "error": _BUSY_ERROR}, status=409
)
@@ -467,7 +739,7 @@ class AutotagRoutesMixin:
return await self._stream_sse(request, _refusal(busy.title))
async with self._jobs.hold(
AUTOTAG_BATCH_JOB_KIND,
"Auto-tag images (legacy page)",
OWNER_LABELS["legacy-batch"],
admission="records",
) as job:
response = await self._start_autotag_batch(request, opts, events)
@@ -491,6 +763,278 @@ 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:
@@ -607,10 +1151,9 @@ class AutotagRoutesMixin:
general_threshold = data.get("general_threshold")
character_threshold = data.get("character_threshold")
try:
if general_threshold is not None:
general_threshold = float(general_threshold)
if character_threshold is not None:
character_threshold = float(character_threshold)
general_threshold, character_threshold = _thresholds_or_saved(
general_threshold, character_threshold
)
except (ValueError, TypeError):
return self._bad_request(
"general_threshold and character_threshold must be numeric"
+122 -13
View File
@@ -2295,14 +2295,43 @@ class ImageRoutesMixin:
{"success": False, "error": _safe_error(e)}, status=500
)
def _media_roots(self, outputs=None):
"""The roots a prompt's linked media is listed and inspected under.
The output folders first (``outputs``, else ``_get_all_output_dirs``),
so every output ``root_index`` keeps its meaning, then, while the
LoRA integration is on, the guarded LoraManager folders and the
CivitAI download cache (``_lora_image_dirs``, #355). Media under those
is linked-only: ``serve_output_image`` serves a file under them only
when the database records it, and no preview file is made there.
"""
roots = list(self._get_all_output_dirs() if outputs is None else outputs)
seen = {_canonical(root) for root in roots}
for extra in self._lora_image_dirs():
key = _canonical(extra)
if key not in seen:
seen.add(key)
roots.append(Path(extra))
return roots
def _lora_image_dirs(self):
"""Extra directories served when the LoRA integration is enabled."""
"""Extra directories served when the LoRA integration is enabled.
LoraManager's guarded LoRA roots, sidecar root and example-images
root (``lora_manager_layout.media_roots``, #355) and PromptManager's
CivitAI download cache. An unusable LoraManager layout gives the
default LoRA roots and the cache.
"""
try:
from ..config import IntegrationConfig
if not IntegrationConfig.LORA_MANAGER_ENABLED:
return []
from ..lora_utils import find_lora_directories, get_lora_image_cache_dir
from ..lora_utils import (
get_lora_image_cache_dir,
layout_for_install,
lora_manager_layout,
)
lora_root = IntegrationConfig.LORA_MANAGER_PATH or ""
if lora_root.strip():
@@ -2315,9 +2344,13 @@ class ImageRoutesMixin:
lora_root = resolve_lora_manager_path(lora_root)
if not lora_root:
return []
layout = layout_for_install(lora_root)
else:
layout = lora_manager_layout("")
dirs = [Path(d) for d in find_lora_directories(lora_root)]
dirs.append(Path(get_lora_image_cache_dir()))
dirs = [Path(d) for d in layout.media_roots]
# Named, not made: listing and serving never create the cache.
dirs.append(Path(get_lora_image_cache_dir(create=False)))
return dirs
except Exception:
# LoRA integration is optional; never widen access on failure.
@@ -2348,9 +2381,21 @@ class ImageRoutesMixin:
return []
async def serve_output_image(self, request):
"""Serve a media file from the ComfyUI output folder(s)."""
"""Serve a media file from the ComfyUI output folder(s).
With ``?root=N`` past the output folders (a LoRA-linked root,
``_media_roots``, #355) only a file the database records is served,
and no preview is served or made there.
"""
try:
roots = list(self._get_all_output_dirs())
outputs = list(self._get_all_output_dirs())
root_param = request.query.get("root")
roots = outputs
if root_param is not None and re.fullmatch(r"[0-9]+", root_param):
if int(root_param) >= len(outputs):
# Only a LoRA-linked request reads LoraManager's layout;
# output media (and its previews) never waits on it.
roots = await self._run_in_executor(self._media_roots, outputs)
scope_error = self._gallery_scope_error(request.query, roots)
if scope_error is not None:
return scope_error
@@ -2365,9 +2410,13 @@ class ImageRoutesMixin:
):
return _forbidden("Only media files can be served")
output_dirs = self._select_root(roots, request.query.get("root"))
output_dirs = self._select_root(roots, root_param)
if not output_dirs:
return _forbidden("No allowed image directories configured")
if len(outputs) <= int(root_param or 0):
return await self._serve_linked_only(
output_dirs[0], filepath, int(root_param), len(outputs)
)
for output_path in output_dirs:
image_path = output_path / filepath
@@ -2399,6 +2448,46 @@ class ImageRoutesMixin:
{"success": False, "error": _safe_error(e)}, status=500
)
def _is_recorded_media(self, path) -> bool:
"""Whether the database records ``path`` as a linked image.
Imported local media is stored resolved; downloaded CivitAI images
keep the cache folder's own spelling (a symlinked or junctioned data
folder), so either spelling names the record.
"""
spellings = dict.fromkeys([str(path), os.path.realpath(path)])
return any(
self.db.get_prompt_id_for_image(spelling) is not None
for spelling in spellings
)
async def _linked_only_error(self, index: int, outputs: int, path):
"""The one gate for a LoRA-linked root (``_media_roots`` past the
``outputs`` output folders, #355): only a file the database records
is served, inspected or revealed there; anything else is 403.
Output roots pass (None)."""
if index < outputs:
return None
if not await self._run_in_executor(self._is_recorded_media, path):
return _forbidden()
return None
async def _serve_linked_only(self, root, filepath, index, outputs):
"""A file under a LoRA-linked root, served only when it is recorded."""
if _is_thumbnail_cache_path(filepath):
return _forbidden()
image_path = root / filepath
if not _is_within(image_path, root):
return _forbidden()
if not image_path.is_file():
return web.json_response(
{"success": False, "error": "Image file not found"}, status=404
)
denied = await self._linked_only_error(index, outputs, image_path)
if denied is not None:
return denied
return self._file_response(image_path)
async def rescan_gallery(self, request):
"""Rewalk every output root now and answer with the new listing.
@@ -3253,7 +3342,9 @@ class ImageRoutesMixin:
if error is not None:
return error
index, scope, relative = target
roots = list(self._get_all_output_dirs())
# A LoRA-linked root too (#355), for the files the import linked only.
outputs = list(self._get_all_output_dirs())
roots = await self._run_in_executor(self._media_roots, outputs)
error = self._gallery_scope_error(
{"root_scope": scope, "root": str(index)}, roots
)
@@ -3262,6 +3353,9 @@ class ImageRoutesMixin:
root = roots[index]
if Path(relative).suffix.lower() not in SERVABLE_EXTENSIONS:
return _forbidden("Only media files can be shown")
denied = await self._linked_only_error(index, len(outputs), root / relative)
if denied is not None:
return denied
gate = self._reveal_gate
if not gate.try_enter():
return web.json_response(
@@ -3273,12 +3367,16 @@ class ImageRoutesMixin:
status=429,
)
try:
return await self._reveal(relative, root)
return await self._reveal(relative, root, linked_only=index >= len(outputs))
finally:
gate.leave()
async def _reveal(self, relative: str, root: Path):
"""Resolve and open, inside the reveal gate."""
async def _reveal(self, relative: str, root: Path, *, linked_only=False):
"""Resolve and open, inside the reveal gate.
Under a LoRA-linked root the file that would open is checked again
(it may resolve through a path anchor): only a recorded one opens.
"""
try:
path = await self._run_in_executor(self._reveal_path, relative, root)
except (OSError, ValueError) as exc:
@@ -3289,6 +3387,10 @@ class ImageRoutesMixin:
return _reveal_not_found()
if path is None:
return _reveal_not_found()
if linked_only and not await self._run_in_executor(
self._is_recorded_media, path
):
return _forbidden()
try:
await asyncio.get_running_loop().run_in_executor(
reveal_executor(), self._reveal_opener, path, root
@@ -3380,8 +3482,13 @@ class ImageRoutesMixin:
return set()
async def inspect_image(self, request):
"""Read prompt evidence for one explicitly rooted physical media file."""
roots = list(self._get_all_output_dirs())
"""Read prompt evidence for one explicitly rooted physical media file.
The roots include LoRA-linked folders (``_media_roots``, #355), so the
inspector opens a LoRA import's media (recorded files only); it only
reads them."""
outputs = list(self._get_all_output_dirs())
roots = await self._run_in_executor(self._media_roots, outputs)
scope_error = self._gallery_scope_error(request.query, roots)
if scope_error is not None:
return scope_error
@@ -3406,6 +3513,8 @@ class ImageRoutesMixin:
return _forbidden("Previews are not inspected; inspect the original")
path = Path(roots[index]) / relative
denied = _media_access_error(path, [roots[index]])
if denied is None:
denied = await self._linked_only_error(index, len(outputs), path)
if denied is not None:
return denied
if not path.is_file():
+8
View File
@@ -130,6 +130,14 @@ 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)
+4
View File
@@ -87,6 +87,8 @@ 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")
@@ -109,6 +111,8 @@ 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"
+36 -8
View File
@@ -234,9 +234,22 @@ class LibraryRoutesMixin:
)
def _library_media(
self, records: list[dict], roots: list[Path], *, linked: bool = True
self,
records: list[dict],
roots: list[Path],
*,
linked: bool = True,
outputs: int | None = None,
) -> 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"])
@@ -250,10 +263,25 @@ 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
return self._format_workspace_page(page, linked=linked)
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)
@pair_response
async def get_library_options(self, request: web.Request) -> web.Response:
@@ -294,13 +322,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
records[parent["id"]], roots, linked=False, outputs=outputs
)
await self._run_in_executor(previews)
@@ -342,9 +370,9 @@ class LibraryRoutesMixin:
result, records = await self._run_in_executor(read)
if result is None:
return _missing()
roots = list(self._get_all_output_dirs())
roots, outputs = await self._run_in_executor(self._library_roots)
result["positive"]["previews"] = await self._run_in_executor(
self._library_media, records, roots
lambda: self._library_media(records, roots, outputs=outputs)
)
return web.json_response({"success": True, **result})
@@ -358,7 +386,7 @@ class LibraryRoutesMixin:
archive = _archive_scope(request)
limit = _page_limit(request)
cursor = request.query.get("cursor")
roots = list(self._get_all_output_dirs())
roots, outputs = await self._run_in_executor(self._library_roots)
canonical_roots = [str(_canonical(root)) for root in roots]
def read() -> dict | None:
@@ -391,7 +419,7 @@ class LibraryRoutesMixin:
if result is None:
return _missing()
result["images"] = await self._run_in_executor(
self._library_media, result["images"], roots
lambda: self._library_media(result["images"], roots, outputs=outputs)
)
return web.json_response({"success": True, **result})
+47 -17
View File
@@ -5,9 +5,12 @@ 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, current_imports, blocked_by_pairs,
planned_downloads, network:{planned_downloads, api_key}}``. A plan that saved
pairs block carries no token (nothing can be applied).
``{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
@@ -31,6 +34,7 @@ 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 (
@@ -152,8 +156,9 @@ class LoraImportRoutesMixin:
return None, ("lora-path", _LORA_PATH_ERROR)
return allowed, None
async def stage_lora_import(self, files, *, on_record=None, cancelled=None):
"""Stage ``files`` off the loop: records, skips and download counts."""
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)
@@ -166,9 +171,10 @@ class LoraImportRoutesMixin:
return await self._run_in_executor(
lambda: lora_import.stage(
files,
discovery,
cache_dir,
IntegrationConfig.CIVITAI_API_KEY,
media_extensions=SERVABLE_EXTENSIONS,
on_record=report,
cancelled=cancelled or (lambda: False),
logger=self.logger,
@@ -176,8 +182,8 @@ class LoraImportRoutesMixin:
)
async def lora_files_unchanged(self, lm_path, expected):
"""Whether the metadata files under ``lm_path`` still fingerprint as
``expected`` (listed again, so an added file counts too)."""
"""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))
@@ -220,8 +226,22 @@ class LoraImportRoutesMixin:
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)
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)
@@ -248,7 +268,9 @@ class LoraImportRoutesMixin:
return _failure("The LoRA import could not be reviewed.", 500)
token = None
expires_at = None
if not blocked:
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"],
@@ -280,6 +302,9 @@ class LoraImportRoutesMixin:
"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,
@@ -342,8 +367,8 @@ class LoraImportRoutesMixin:
store.discard(token)
return _failure(INTEGRATION_CHANGED, 409, "plan-changed")
try:
files = await self._run_in_executor(lora_import.discover, lm_path)
now = await self._run_in_executor(lora_import.fingerprint, files)
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
@@ -367,9 +392,14 @@ class LoraImportRoutesMixin:
lora_import.JOB_TITLE,
admission="records",
cancellable=True,
scope={"operation": LORA_OPERATION, "lora_files": len(files)},
scope={
"operation": LORA_OPERATION,
"lora_files": len(discovery.sidecars),
},
groups=lora_import.GROUPS,
runner=lambda job: self._run_lora_import(job, db, plan, files, backup),
runner=lambda job: self._run_lora_import(
job, db, plan, discovery, backup
),
)
except JobBusy as busy:
return busy_response(busy)
@@ -378,8 +408,8 @@ class LoraImportRoutesMixin:
# ── the job ───────────────────────────────────────────────────────
async def _run_lora_import(self, job, db, plan, files, backup_on):
total = len(files)
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')}")
@@ -393,7 +423,7 @@ class LoraImportRoutesMixin:
try:
staged = await self.stage_lora_import(
files, on_record=on_record, cancelled=lambda: job.cancelled
discovery, on_record=on_record, cancelled=lambda: job.cancelled
)
except lora_import.ImportCancelled:
job.log_line(f"{CANCELLED} {CACHE_NOTE}")
+25 -7
View File
@@ -50,6 +50,15 @@ _IMPORTS_CHANGED = (
)
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."""
@@ -398,13 +407,19 @@ class LoraIntegrationMixin:
return
if isinstance(error, PlanChanged):
message = _IMPORTS_CHANGED
elif isinstance(error, LoraFilesChanged):
elif isinstance(error, (LoraFilesChanged, LoraSettingsUnusable)):
message = str(error)
else:
message = self._public_error(error)
if isinstance(
error,
(PlanChanged, LoraFilesChanged, EmptyReplacement, PairDependencyError),
(
PlanChanged,
LoraFilesChanged,
LoraSettingsUnusable,
EmptyReplacement,
PairDependencyError,
),
):
self.logger.warning(f"LoRA scan refused: {message}")
else:
@@ -432,8 +447,11 @@ class LoraIntegrationMixin:
"progress": 2,
}
)
files = await self._run_in_executor(lora_import.discover, lm_path)
total = len(files)
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)
await send_progress(
{
"type": "progress",
@@ -443,7 +461,7 @@ class LoraIntegrationMixin:
}
)
reviewed = await self._run_in_executor(lora_import.fingerprint, files)
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
@@ -467,7 +485,7 @@ class LoraIntegrationMixin:
)
if not existing:
counts["imported"] += 1
status = record["name"] if record else files[index].name
status = record["name"] if record else sidecars[index].metadata_path.name
if record and record["images"]:
status += f" ({len(record['images'])} images)"
await send_progress(
@@ -490,7 +508,7 @@ class LoraIntegrationMixin:
stop = threading.Event()
staging = asyncio.ensure_future(
self.stage_lora_import(
files,
discovery,
on_record=on_record,
cancelled=lambda: stop.is_set() or disconnected(),
)
+84 -1
View File
@@ -47,11 +47,23 @@ _SETTINGS_FILE_KEYS = {
"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"
@@ -151,6 +163,63 @@ 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
@@ -177,7 +246,7 @@ class SettingsRoutesMixin:
async def get_settings(self, request):
"""Get current settings."""
try:
from ..config import PromptManagerConfig, GalleryConfig
from ..config import AutotagConfig, PromptManagerConfig, GalleryConfig
root_paths = [
self._public_path(d) for d in GalleryConfig.MONITORING_DIRECTORIES
@@ -203,6 +272,14 @@ class SettingsRoutesMixin:
"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 "",
@@ -425,7 +502,12 @@ 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 = []
@@ -564,6 +646,7 @@ class SettingsRoutesMixin:
updates["THUMBNAIL_EDGE"] = thumbnail_settings.validate_edge(
data["thumbnail_edge"]
)
updates.update(_validate_autotag_settings(data))
return updates
def _parse_gallery_roots(self, data):
+31 -12
View File
@@ -488,10 +488,16 @@ class AutoThumbnailer:
f"Automatic preview failed for {src.name}", exc_info=True
)
def close(self):
def close(self, timeout=None):
"""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."""
dropped), kill a running ffmpeg and wait up to ``timeout`` seconds
(None: CLOSE_WAIT_SECONDS, read now); a worker still running after
that is logged.
Returns True when there was no worker or it ended; False when it is
still running, or when called on the worker itself (it only signals).
"""
wait = CLOSE_WAIT_SECONDS if timeout is None else timeout
with self._changed:
self._closed = True
self._queue.clear()
@@ -504,14 +510,18 @@ class AutoThumbnailer:
process.kill()
except OSError:
pass
if worker is None or worker is threading.current_thread():
return
worker.join(CLOSE_WAIT_SECONDS)
if worker is None:
return True
if worker is threading.current_thread():
return False
worker.join(wait)
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"
f"{wait} s; it stops at its next file"
)
return False
return True
# ── waits (tests and shutdown; never timers) ──────────────────────
@@ -609,12 +619,19 @@ class AutoThumbnailMixin:
)
return self._auto_thumbnailer
def close_auto_thumbnails(self, final=False):
def refuse_auto_thumbnails(self):
"""Refuse every later capture without stopping the worker yet (the
package teardown, before ``close_auto_thumbnails(final=True)``)."""
with AutoThumbnailMixin._auto_thumbnailer_lock:
self._auto_thumbnails_shut = True
def close_auto_thumbnails(self, final=False, timeout=None):
"""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.
Bounded by ``timeout`` (see AutoThumbnailer.close). ``final`` (the
package teardown) also keeps a late monitor event from starting a
worker again; a test reset passes False and the next use starts
fresh ones. Returns True when no capture worker is left running.
"""
with AutoThumbnailMixin._auto_thumbnailer_lock:
if final:
@@ -622,12 +639,14 @@ class AutoThumbnailMixin:
thumbnailer, self._auto_thumbnailer = self._auto_thumbnailer, None
pool, self._preview_pool_obj = self._preview_pool_obj, None
self._preview_gate_obj = None
closed = True
if thumbnailer is not None:
thumbnailer.close()
closed = thumbnailer.close(timeout)
if pool is not None:
pool.shutdown(wait=False, cancel_futures=True)
if final:
shutdown_encoders() # kill JPEG XL / ffmpeg encoders still running
return closed
def _auto_thumbnail_failures(self):
thumbnailer = self._auto_thumbnailer
+30 -28
View File
@@ -25,6 +25,23 @@ 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
@@ -45,10 +62,12 @@ except ImportError:
return logger
# Model configurations - same as standalone_tagger.py
# Model configurations - same as standalone_tagger.py. ``label`` is the short
# name System, Activity and a 409 model-busy show (#328).
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",
@@ -59,6 +78,7 @@ MODELS = {
},
"hf": {
"name": "HuggingFace",
"label": "HuggingFace JoyCaption",
"description": "Full precision model, ~16GB, higher quality",
"size": "~16GB",
"repo": "fancyfeast/llama-joycaption-beta-one-hf-llava",
@@ -66,6 +86,7 @@ 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",
@@ -75,6 +96,7 @@ MODELS = {
},
"wd14-vit": {
"name": "WD14 ViT",
"label": "WD14 ViT",
"description": "Fast classifier tagger, ~400MB, ViT architecture",
"size": "~400MB",
"repo": "SmilingWolf/wd-vit-tagger-v3",
@@ -85,17 +107,7 @@ MODELS = {
}
# Default prompts
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."
)
DEFAULT_PROMPT = DEFAULT_TAG_TEMPLATE
# Upper bound of one suggestion request: the image tag set cap.
MAX_SUGGESTION_LIMIT = MAX_IMAGE_TAGS
@@ -179,10 +191,6 @@ 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."""
@@ -205,19 +213,13 @@ class AutoTagService:
@property
def wd14_general_threshold(self) -> float:
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))
"""The saved General-tag confidence (``autotag.defaults``)."""
return _autotag_config().GENERAL_THRESHOLD
@property
def wd14_character_threshold(self) -> float:
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))
"""The saved Character-tag confidence (``autotag.defaults``)."""
return _autotag_config().CHARACTER_THRESHOLD
def get_models_status(self) -> Dict[str, Dict[str, Any]]:
"""Get availability status for all model types.
@@ -860,12 +862,12 @@ class AutoTagService:
gen_thresh = (
general_threshold
if general_threshold is not None
else self._wd14_general_threshold
else _autotag_config().GENERAL_THRESHOLD
)
char_thresh = (
character_threshold
if character_threshold is not None
else self._wd14_character_threshold
else _autotag_config().CHARACTER_THRESHOLD
)
return self._generate_wd14(
session, tags_list, image, gen_thresh, char_thresh
+96 -16
View File
@@ -563,18 +563,90 @@ 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.
``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.
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.
"""
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
@@ -587,7 +659,14 @@ class AutotagConfig:
@classmethod
def get_config(cls) -> Dict[str, Any]:
return {"defaults": {"suggestion_limit": cls.SUGGESTION_LIMIT}}
"""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}
@classmethod
def update_config(cls, new_config: Any) -> None:
@@ -598,17 +677,18 @@ class AutotagConfig:
if not isinstance(defaults, dict):
config_logger.warning("Ignoring autotag.defaults that is not an object")
return
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
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)
class PromptManagerConfig:
+103 -53
View File
@@ -1,18 +1,26 @@
"""The staged LoRA reference import (#221, #249 System S10).
"""The staged LoRA reference import (#221, #249 System S10, #355).
The import reads LoraManager's ``*.metadata.json`` files and turns each LoRA
into one prompt of the ``lora-manager`` category, tagged with its name and
trigger words and linked to its preview images. It runs in three stages, and
only the last writes to the database:
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**: list the metadata files (``discover``) and fingerprint them
(path, size and modification time; ``fingerprint``). ``preview`` adds the
counts a review shows. Nothing is written and nothing is requested.
2. **stage**: read each file and build its record; CivitAI example images are
downloaded 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 files are listed and fingerprinted once more (a file
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).
@@ -24,9 +32,15 @@ Both the System job (``POST /lora/import/apply``) and the legacy ``POST
import hashlib
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Dict, List, Optional
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
@@ -37,35 +51,44 @@ 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")
METADATA_SUFFIX = "*.metadata.json"
class ImportCancelled(Exception):
"""The import was cancelled while staging; nothing was written."""
def discover(lora_manager_path: str) -> List[Path]:
"""Every LoRA metadata file, sorted and without repeats.
def discover(lora_manager_path: Optional[str]) -> Discovery:
"""The sidecars of every LoRA whose model exists inside a LoRA root.
Raises whatever listing a folder raised: a discovery that cannot see
every folder must not replace the collection with part of it.
``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.
"""
seen = set()
files = []
for folder in lora_utils.find_lora_directories(lora_manager_path):
for path in Path(folder).rglob(METADATA_SUFFIX):
key = str(path)
if key not in seen:
seen.add(key)
files.append(path)
files.sort(key=str)
return files
return discover_sidecars(lora_utils.lora_manager_layout(lora_manager_path))
def fingerprint(files: List[Path]) -> str:
"""What a review saw of ``files``: each path, size and modification time."""
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 path in files:
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"
@@ -75,31 +98,48 @@ def fingerprint(files: List[Path]) -> str:
return "sha256:" + digest.hexdigest()
def preview(lora_manager_path: str, cache_dir: Path) -> Dict:
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, planned_downloads, fingerprint, names}``:
the metadata files found, those that parse, the example images an import
would download (not yet in the cache) and the first LoRA names.
``{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``).
"""
files = discover(lora_manager_path)
discovery = discover(lora_manager_path)
layout = discovery.layout
with_metadata = 0
with_local_media = 0
planned = 0
names = []
for path in files:
metadata = lora_utils.read_lora_metadata(path)
for sidecar in discovery.sidecars:
metadata = lora_utils.read_lora_metadata(sidecar.metadata_path)
if not metadata:
continue
with_metadata += 1
planned += lora_utils.planned_civitai_downloads(metadata, path, cache_dir)
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": len(files),
"lora_files": discovery.scanned,
"with_metadata": with_metadata,
"skipped": discovery.scanned - with_metadata,
"with_local_media": with_local_media,
"planned_downloads": planned,
"fingerprint": fingerprint(files),
"fingerprint": fingerprint(discovery),
"names": names,
"layout": layout_summary(layout, Path.home()),
}
@@ -131,35 +171,45 @@ class StagedImport:
def stage(
files: List[Path],
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 file and build its record, downloading its example images.
"""Read each sidecar and build its record with its images.
Blocking. A file that cannot be read, or whose record fails to build, is
skipped by its name. ``cancelled()`` is asked before each file (raises
ImportCancelled); ``on_record(index, record_or_None, staged)`` follows
each one.
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.
"""
staged = StagedImport(total=len(files))
for index, path in enumerate(files):
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 = lora_utils.get_preview_images_from_metadata(metadata, path)
remote = lora_utils.download_civitai_images(
metadata, path, cache_dir, api_key, stats=staged.stats
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)
+602
View File
@@ -0,0 +1,602 @@
"""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,
}
+97 -96
View File
@@ -13,6 +13,7 @@ import json
import os
import re
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path
@@ -27,6 +28,16 @@ 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 ────────────────────────────────────────────
@@ -234,23 +245,22 @@ def resolve_lora_manager_path(path) -> Optional[str]:
# ── Metadata reading ─────────────────────────────────────────────────
def find_lora_directories(lora_manager_path: str) -> List[str]:
"""Find directories that contain LoRA models (with .metadata.json files).
def default_lora_roots() -> List[str]:
"""The LoRA folders ComfyUI knows, without LoraManager's settings.
Searches: ComfyUI models/loras, extra_model_paths.yaml lora dirs,
and the LoraManager extension dir itself.
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.
"""
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()))
@@ -267,13 +277,28 @@ def find_lora_directories(lora_manager_path: str) -> 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 = []
@@ -379,46 +404,6 @@ 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
@@ -837,30 +822,16 @@ def get_lora_image_cache_dir(create: bool = True) -> Path:
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.
@@ -871,52 +842,81 @@ 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:
"""Scan LoRA metadata files and build the trigger word mapping.
"""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.
Returns:
Number of LoRAs with trigger words found.
Number of LoRAs with trigger words in the map.
"""
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
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)
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,6 +926,7 @@ class TriggerWordCache:
with self._lock:
self._cache.clear()
self._loaded = False
self._failed_at = None
# Module-level singleton
+55
View File
@@ -0,0 +1,55 @@
"""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
+1 -1
View File
@@ -3,7 +3,7 @@ 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.10"
requires-python = ">=3.12"
dependencies = ["watchdog>=2.1.0", "Pillow>=8.0.0"]
[build-system]
+27
View File
@@ -0,0 +1,27 @@
"""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))
+152
View File
@@ -0,0 +1,152 @@
// #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 -1
View File
@@ -72,6 +72,8 @@ 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 {
@@ -79,7 +81,7 @@ class TestModel:
"downloaded": name not in self.unavailable,
"loaded": self.loaded == name,
}
for name in ("gguf", "hf", "wd14-vit", "wd14-convnext", "wd14-swinv2")
for name in self.models_config
}
def is_model_loaded(self) -> bool:
@@ -166,7 +168,20 @@ 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()
@@ -194,6 +209,7 @@ 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()
@@ -270,7 +270,8 @@ function mapImportConcept(reference, plan) {
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: [['Scope', plan.scope.label], ['Found', '2 metadata files · 2 readable'],
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']] },
@@ -312,7 +313,7 @@ async function productionReview(page, request, theme) {
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 });
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`) });
@@ -324,15 +325,15 @@ for (const theme of THEMES) for (const [width, height] of SIZES) {
test('LoRA import review comparison detects mutations', async ({ page, request }) => {
await page.setViewportSize({ width: 1600, height: 1100 });
await fakeLoraManager(request, { loras: 2, civitai: true });
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')[3].textContent = 'Downloads'; });
expect(await paths()).toEqual(['grid.rows.3.0']);
await page.evaluate(() => { document.querySelectorAll('dialog#modal[open] dl.meta-grid dt')[3].textContent = 'Network'; });
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();
@@ -0,0 +1,385 @@
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([]);
});
+1 -1
View File
@@ -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 from the legacy Auto-tag page.');
await expect(status()).toHaveText('WD14 SwinV2 is not downloaded. Download it in System › Auto-tagging.');
await shot('production-unavailable');
await expect(await request.post('/__test__/tagging/models', { data: { unavailable: [] } })).toBeOK();
+15 -10
View File
@@ -257,6 +257,18 @@ 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');
@@ -287,12 +299,7 @@ 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);
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 confirmRestoreAndWait(page);
await expect(page.locator('.prompt-text')).toHaveCount(2);
await page.reload();
await expect(page.locator('.prompt-text').filter({ hasText: 'Cat portrait' })).toBeVisible();
@@ -335,8 +342,7 @@ 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 page.locator('#confirmRestore').click();
await expect(page.locator('#restoreModal')).toBeHidden();
await confirmRestoreAndWait(page);
expect(await promptSnapshot(request)).toEqual(before);
});
@@ -356,8 +362,7 @@ 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 page.locator('#confirmRestore').click();
await expect(page.locator('#restoreModal')).toBeHidden();
await confirmRestoreAndWait(page);
expect(await promptSnapshot(request)).toEqual(before);
});
+30
View File
@@ -1088,6 +1088,36 @@ 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 }) => {
@@ -322,7 +322,7 @@ test('a partial save made from Integrations keeps its outcome when Sources opens
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, previous: 2 });
await lora(request, { loras: 3, civitai: true, preview: false, previous: 2 });
const previews = [];
const applies = [];
page.on('request', item => {
@@ -339,13 +339,17 @@ test('Review import shows replacement, blocked items and network effects with th
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 readable'],
['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');
@@ -382,7 +386,7 @@ test('saved pairs block the import: the review explains it and offers only Close
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))[2]).toEqual(['Protected', '1 imported prompt used by saved pairs']);
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']);
@@ -390,6 +394,51 @@ test('saved pairs block the import: the review explains it and offers only Close
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);
@@ -408,7 +457,7 @@ test('a failed import keeps old imports and its job says so', async ({ page, req
});
test('the import can be cancelled while it downloads; then nothing is written', async ({ page, request }) => {
await lora(request, { loras: 2, civitai: true, previous: 1 });
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');
@@ -0,0 +1,356 @@
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([]);
});
+3 -3
View File
@@ -412,11 +412,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);
// Auto-tagging still hands off (S11), so its body is one known panel.
await nav.getByRole('button', { name: 'Auto-tagging', exact: true }).click();
// Diagnostics still hands off (S12), so its body is one known panel.
await nav.getByRole('button', { name: 'Diagnostics & logs', 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: Auto-tagging keeps its own body.
// Both late answers reached a section that had left: Diagnostics 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);
@@ -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 from the legacy Auto-tag page.';
const NOT_DOWNLOADED = 'WD14 SwinV2 is not downloaded. Download it in System › Auto-tagging.';
test.beforeEach(async ({ page, request, browserName }) => {
await startCoverage(page, browserName);
+1 -1
View File
@@ -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 from the legacy Auto-tag page.');
await expect(status(page)).toHaveText('WD14 SwinV2 is not downloaded. Download it in System › Auto-tagging.');
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.
+7 -5
View File
@@ -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: 'Auto-tagging', exact: true }).click();
await main.locator('.section-nav').getByRole('button', { name: 'Diagnostics & logs', exact: true }).click();
const handoff = main.locator('section.handoff');
await expect(handoff).toHaveClass(/\bsurface\b/);
expect(await handoff.evaluate(surfaceStyle)).toMatchObject(SURFACE);
@@ -465,8 +465,9 @@ test('System is native: the concept shell, section navigation, sections, hand-of
// 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) and Backup & export (S9 Task 24, workspace-system-backups.spec.mjs);
// every other section hands off.
// 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();
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}`)))
@@ -474,9 +475,10 @@ test('System is native: the concept shell, section navigation, sections, hand-of
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'] }[hash]
backups: ['div.two-col', 'section.surface', 'div.save-row'],
tagging: ['section.surface', 'div.save-row'] }[hash]
?? ['section.surface handoff', 'div.save-row']);
if (['overview', 'appearance', 'sources', 'integrations', 'maintenance', 'backups'].includes(hash)) continue;
if (['overview', 'appearance', 'sources', 'integrations', 'tagging', 'maintenance', 'backups'].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'))))
@@ -7,13 +7,18 @@ reaches a real LoraManager, ComfyUI install or the network:
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, previous}``
writes ``custom_nodes/ComfyUI-Lora-Manager`` (detected through the
``folder_paths`` stub's base path) and ``loras`` metadata files under
``models/loras``, each with a local preview PNG; ``civitai`` adds one
CivitAI example image per LoRA (served by the stub opener below);
``previous`` saves that many earlier ``lora-manager`` imports with tags and
image links;
- ``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;
@@ -30,6 +35,7 @@ test ever opens a network connection.
import io
import json
import os
import shutil
import threading
from pathlib import Path
@@ -91,7 +97,11 @@ def register_integrations_commands(app, routes, api, root: Path):
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
@@ -118,11 +128,17 @@ def register_integrations_commands(app, routes, api, root: Path):
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)
@@ -141,8 +157,10 @@ def register_integrations_commands(app, routes, api, root: Path):
"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))
(loras / f"{stem}.png").write_bytes(_png())
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(
+20 -17
View File
@@ -97,15 +97,25 @@ def register_system_commands(app, routes, api, root: Path):
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
host_values = (
PromptManagerConfig.RESULT_TIMEOUT,
PromptManagerConfig.WEBUI_DISPLAY_MODE,
PromptManagerConfig.WORKER_THREADS,
GalleryConfig.MONITORING_ENABLED,
PromptManagerConfig.THUMBNAIL_FORMAT,
PromptManagerConfig.THUMBNAIL_QUALITY,
PromptManagerConfig.THUMBNAIL_EDGE,
from py.api.settings_routes import (
_AUTOTAG_SETTINGS,
_GALLERY_SETTINGS,
_SETTINGS_FILE_KEYS,
)
# 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)
@@ -135,15 +145,8 @@ def register_system_commands(app, routes, api, root: Path):
unplugged.rmdir()
GalleryConfig.MONITORING_DIRECTORIES = list(configured)
api._health_report = None
(
PromptManagerConfig.RESULT_TIMEOUT,
PromptManagerConfig.WEBUI_DISPLAY_MODE,
PromptManagerConfig.WORKER_THREADS,
GalleryConfig.MONITORING_ENABLED,
PromptManagerConfig.THUMBNAIL_FORMAT,
PromptManagerConfig.THUMBNAIL_QUALITY,
PromptManagerConfig.THUMBNAIL_EDGE,
) = host_values
for (owner, name), value in host_values.items():
setattr(owner, name, value)
thumbnail_encode.reset_encoders()
image_monitor._monitor_instance = monitor
api.__dict__.pop("_library_resolver", None)
+63 -2
View File
@@ -12,6 +12,7 @@ import json
import os
from pathlib import Path
import shutil
import threading
import time
from typing import TYPE_CHECKING
@@ -181,6 +182,28 @@ 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:
@@ -195,6 +218,7 @@ 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.
@@ -223,6 +247,33 @@ 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).
@@ -273,18 +324,28 @@ 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
model.wd14_general_threshold = general
model.wd14_character_threshold = character
AutotagConfig.GENERAL_THRESHOLD = general
AutotagConfig.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:
+3 -1
View File
@@ -1,5 +1,6 @@
"""Deterministic nodes and queue controls installed only in an owned test host."""
import asyncio
import os
import threading
import time
@@ -291,7 +292,8 @@ async def image_monitor(request: web.Request) -> web.Response:
config.MONITORING_ENABLED = True
monitor.start_monitoring([folder_paths.get_output_directory()])
else:
monitor.stop_monitoring()
# 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})
+10
View File
@@ -388,6 +388,16 @@ test("ApiError keeps a validated running job on busy and a validated source on i
assert.deepEqual([new ApiError(409, "old caller").running, new ApiError(409, "old caller").source], [null, null]);
});
test("ApiError keeps a validated operation on model-busy only (#328)", () => {
const body = { operation: "Tagging model download · WD14 ViT" };
assert.equal(new ApiError(409, "busy", "model-busy", null, body).operation, "Tagging model download · WD14 ViT");
assert.equal(new ApiError(409, "busy", "busy", null, body).operation, null);
for (const operation of [null, 7, "", "x".repeat(201)]) {
assert.equal(new ApiError(409, "busy", "model-busy", null, { operation }).operation, null);
}
assert.equal(new ApiError(409, "busy", "model-busy").operation, null);
});
test("ApiError says whether the server itself refused (a parsed success:false body; S7 review I1)", () => {
assert.equal(new ApiError(400, "bad", null, null, { success: false, error: "bad" }).refused, true);
assert.equal(new ApiError(409, "busy", "busy", null, { success: false, code: "busy" }).refused, true);
+23
View File
@@ -225,3 +225,26 @@ test('a failed duplicate cleanup says the server reason (S7 review M3)', async (
assert.equal(finishToast(failed), `Duplicate cleanup failed · ${reason}`);
assert.equal(finishToast({ ...failed, status_line: '' }), 'Duplicate cleanup failed · 0 removed · 1 error');
});
// #328: Tag output images (`autotag`) reuses the image tagging view: Tagged / Unchanged / Skipped / Errors and Cancel.
// The legacy page's hold of the same kind counts nothing and is not cancellable, so it keeps the generic held view.
test('autotag gets Tagged / Unchanged / Skipped / Errors and Cancel; its legacy hold keeps the generic view', async () => {
const { kindOf, viewOf, canCancel } = await source;
const { startToast, finishToast } = await store;
const counters = { tagged: 3, unchanged: 1, skipped: 2, errors: 0 };
const batch = job({ kind: 'autotag', title: 'Tag output images', counters, scope: { images: 6 } });
const view = viewOf(batch);
assert.deepEqual(view.tiles.map(([group, label]) => [group, label]),
[['tagged', 'Tagged'], ['unchanged', 'Unchanged'], ['skipped', 'Skipped'], ['errors', 'Errors']]);
assert.equal(kindOf('autotag').pillVerb, 'Tagging');
assert.equal(kindOf('autotag').cancelLabel, 'Cancel tagging');
assert.equal(canCancel(batch), true);
assert.equal(startToast(batch), 'Tag output images started.');
assert.match(finishToast({ ...batch, state: 'finished', finished_at: 2 }), /^Tag output images finished/);
const held = job({ kind: 'autotag', title: 'Tag output images (legacy page)', counters: {}, cancellable: false,
progress: 0, processed: 0, total: 0 });
const legacy = viewOf(held);
assert.deepEqual([legacy.tiles, legacy.outcomes, [...legacy.openable]], [[], [], []]);
assert.equal(canCancel(held), false);
});
+47 -1
View File
@@ -197,6 +197,52 @@ test('instructions default to the server prompt', async () => {
assert.equal(c.state.options.suggestion_limit, 7, 'the suggestion count survives a model change');
});
// #328: one tagging configuration. The dialogs start on the model saved in System › Auto-tagging, unless a preset
// (a review session's choice) or the user already picked one; an unknown saved model keeps the first option.
test('dialog starts on the saved model unless a preset or the user chose one', async () => {
const saved = { [`GET ${MODELS}`]: () => response(200, models({ model_type: 'gguf' })) };
const { c } = await opened(saved);
assert.equal(c.state.model, 'gguf');
assert.equal(c.state.options.instructions, DEFAULT_PROMPT);
const { c: preset } = await opened(saved, { preset: { model: 'wd14-vit', options: {} } });
assert.equal(preset.state.model, 'wd14-vit');
const { createImageTaggingController, createTaggingOptionsController } = await load();
const h = harness(defaultRoutes(saved));
const picked = createImageTaggingController({ api: h.api, target: TARGET, source: 'images' });
picked.setModel('hf');
await picked.open();
await settle();
assert.equal(picked.state.model, 'hf', 'the user picked a model before the read landed');
// A pick made while the models read is still in flight wins over the saved model that lands after it.
const gate = deferred();
const late = harness(defaultRoutes({ [`GET ${MODELS}`]: () => gate.promise }));
const racing = createImageTaggingController({ api: late.api, target: TARGET, source: 'images' });
const opening = racing.open();
await settle();
assert.equal(racing.state.phase, 'observing');
racing.setModel('wd14-vit');
gate.resolve(response(200, models({ model_type: 'gguf' })));
await opening;
await settle();
assert.equal(racing.state.model, 'wd14-vit');
racing.dispose();
const { c: unknown } = await opened({ [`GET ${MODELS}`]: () => response(200, models({ model_type: 'wd14-convnext' })) });
assert.equal(unknown.state.model, 'wd14-swinv2');
const options = createTaggingOptionsController({ api: harness(defaultRoutes(saved)).api });
await options.open();
assert.equal(options.state.model, 'gguf');
const chosen = createTaggingOptionsController({ api: harness(defaultRoutes(saved)).api });
chosen.setModel('wd14-vit');
await chosen.open();
assert.equal(chosen.state.model, 'wd14-vit');
for (const controller of [c, preset, picked, unknown, options, chosen]) controller.dispose();
});
test('an unreadable or malformed models read falls back to the documented defaults', async () => {
for (const route of [() => fail(500, null), () => response(200, models({ default_prompt: 4, suggestion_limit: 0, wd14_general_threshold: 2, wd14_character_threshold: 'x' }))]) {
const { c } = await opened({ [`GET ${MODELS}`]: route });
@@ -542,7 +588,7 @@ test('model-unavailable and model-busy set the status line only', async () => {
await c.generate();
const before = c.state;
for (const [code, message] of [
['model-unavailable', 'WD14 SwinV2 is not downloaded. Download it from the legacy Auto-tag page.'],
['model-unavailable', 'WD14 SwinV2 is not downloaded. Download it in System › Auto-tagging.'],
['model-busy', 'A tagging model is busy with another operation. Try again when it finishes.'],
['unsupported-media', 'Images larger than 4096 px per side or 16,777,216 px are not supported.'],
]) {
@@ -129,9 +129,9 @@ test('a model that is not downloaded is named with the server sentence', () => {
const { modelUnavailable } = state;
assert.equal(modelUnavailable(null, 'wd14-vit'), '');
assert.equal(modelUnavailable({ 'wd14-vit': { downloaded: true } }, 'wd14-vit'), '');
assert.equal(modelUnavailable({ 'wd14-vit': { downloaded: false } }, 'wd14-vit'), 'WD14 ViT is not downloaded. Download it from the legacy Auto-tag page.');
assert.equal(modelUnavailable({ 'wd14-vit': { downloaded: false } }, 'wd14-vit'), 'WD14 ViT is not downloaded. Download it in System › Auto-tagging.');
assert.equal(modelUnavailable({ gguf: { downloaded: false, name: 'GGUF (Recommended)' } }, 'gguf'),
'GGUF (Recommended) is not downloaded. Download it from the legacy Auto-tag page.');
'GGUF (Recommended) is not downloaded. Download it in System › Auto-tagging.');
assert.throws(() => modelUnavailable({}, 'nope'), TypeError);
});
+31 -4
View File
@@ -108,13 +108,17 @@ test('host compatibility shows the detected renderer as status (D106)', async ()
test('the import review shows replacement, blocked items and network effects (D107)', async () => {
const { importGrid } = await integrations;
const plan = { lora_files: 3, with_metadata: 2, current_imports: 5, blocked_by_pairs: 0, planned_downloads: 4,
const plan = { lora_files: 2221, with_metadata: 615, with_local_media: 476, skipped: 1606, current_imports: 5,
blocked_by_pairs: 0, planned_downloads: 4,
layout: { source: 'portable', settings_file: '~/ComfyUI/custom_nodes/LM/settings.json', error: null, mode: 'centralized',
library: 'comfyui', example_images: true },
scope: { count: 2, label: 'Replace 5 imported prompts with 2 LoRAs' }, protected: { count: 0, reasons: [] },
backup: { policy: 'on', destination_label: 'prompts.db.backup_<date and time>' }, network: { planned_downloads: 4, api_key: true },
media: { deletes_files: false, label: 'No files are removed; image links are replaced' } };
assert.deepEqual(importGrid(plan), [
['LoraManager', 'portable settings · centralized storage · library comfyui'],
['Scope', 'Replace 5 imported prompts with 2 LoRAs'],
['Found', '3 metadata files · 2 readable'],
['Found', '2,221 metadata files · 615 LoRAs · 476 with local images'],
['Protected', 'Nothing is protected'],
['Network', 'Downloads up to 4 example images from CivitAI with your saved key'],
['Backup', 'Create and verify before changes (prompts.db.backup_<date and time>)'],
@@ -122,11 +126,34 @@ test('the import review shows replacement, blocked items and network effects (D1
]);
const blocked = importGrid({ ...plan, planned_downloads: 1, network: { api_key: false }, protected: { count: 2,
reasons: ['2 imported prompts used by saved pairs'] } });
assert.deepEqual(blocked[2], ['Protected', '2 imported prompts used by saved pairs']);
assert.deepEqual(blocked[3], ['Network', 'Downloads up to 1 example image from CivitAI']);
assert.deepEqual(blocked[3], ['Protected', '2 imported prompts used by saved pairs']);
assert.deepEqual(blocked[4], ['Network', 'Downloads up to 1 example image from CivitAI']);
// A plan without a layout (no LoraManager row) and with no counts.
assert.deepEqual(importGrid({})[0], ['Scope', '0 LoRAs']);
assert.deepEqual(importGrid({})[1], ['Found', '0 metadata files · 0 LoRAs · 0 with local images']);
assert.deepEqual(importGrid({})[3], ['Network', 'No downloads']);
assert.deepEqual(importGrid({})[5], ['Media on disk', 'No files are removed']);
assert.deepEqual(importGrid({ ...plan, lora_files: 1, with_metadata: 1, with_local_media: 0 })[2],
['Found', '1 metadata file · 1 LoRA · 0 with local images']);
});
test('the LoraManager row names what the import read, and an unusable layout closes the review (#355)', async () => {
const { layoutText, closeNote, LAYOUT_UNUSABLE, BLOCKED_NOTE, NOTHING_TO_IMPORT } = await integrations;
assert.equal(layoutText({ source: 'environment', mode: 'alongside', library: 'default', error: null }),
'environment settings · alongside storage · library default');
assert.equal(layoutText({ source: 'user config', mode: 'alongside', library: 'x', error: 'the settings file is not valid JSON' }),
'Not usable: the settings file is not valid JSON');
assert.equal(layoutText(null), null);
assert.equal(layoutText('nope'), null);
const error = 'the sidecar folder does not exist';
assert.equal(closeNote({ token: null, layout: { error } }),
'LoraManager settings: not usable: the sidecar folder does not exist. Nothing was changed.');
assert.equal(LAYOUT_UNUSABLE(error), closeNote({ token: null, layout: { error }, blocked_by_pairs: 2 }));
assert.equal(closeNote({ token: null, blocked_by_pairs: 1, layout: { error: null } }), BLOCKED_NOTE);
assert.equal(closeNote({ token: null }), NOTHING_TO_IMPORT);
assert.equal(closeNote({ token: 't', with_metadata: 0, current_imports: 0 }), NOTHING_TO_IMPORT);
assert.equal(closeNote({ token: 't', with_metadata: 1, current_imports: 0 }), null);
assert.equal(closeNote({ token: 't', with_metadata: 0, current_imports: 3 }), null);
});
test('a LoRA import job reads in Activity: tiles, outcomes and the kept collection on failure', async () => {
+16
View File
@@ -341,6 +341,22 @@ test('a host save names the settings that apply after a restart (D128)', async (
assert.equal(drafts.statusText(), 'Unsaved changes.');
});
test('a host save refused for one setting says the server reason and names no source (S11 CUJ 2a)', async () => {
const { INVALID_SETTING_CODE } = await source;
const domains = adapters({ host: { autotag_general_threshold: 0.35 } });
const drafts = await loaded(domains);
const message = 'General-tag confidence must be between 0 and 1.';
domains.outcome.host = () => { throw Object.assign(new Error(message), { code: INVALID_SETTING_CODE }); };
drafts.set('host', 'autotag_general_threshold', 1.5);
assert.deepEqual(await drafts.save(), { browser: 'clean', host: 'failed', runtime: 'clean', invalid: { source: null, message } });
assert.equal(drafts.statusText(), message);
assert.equal(drafts.invalidSource(), null);
assert.deepEqual([...drafts.failed('host')], ['autotag_general_threshold']);
assert.equal(drafts.get('host', 'autotag_general_threshold'), 1.5);
drafts.set('host', 'autotag_general_threshold', 0.4);
assert.equal(drafts.statusText(), 'Unsaved changes.');
});
test('a host save refused for one source says the server reason and names it until the next edit', async () => {
const domains = adapters({ host: { gallery_root_paths: ['output', 'output/gone'] } });
const drafts = await loaded(domains);
+583
View File
@@ -0,0 +1,583 @@
// System › Auto-tagging (#328, System S11 Tasks 6–7): the concept's rows with real data (D132 model options and live
// status, D133 suggestions per image), host drafts saved by the one Save row, and the actions: Download (an Activity
// job), Unload (direct), Preview tagging review (D115) and the reviewed Tag output images batch (D134).
const test = require('node:test');
const assert = require('node:assert/strict');
const fs = require('node:fs');
const path = require('node:path');
const tagging = import('../../web/js/workspace/system-tagging.mjs');
const stateModule = import('../../web/js/workspace/system-state.mjs');
const coreModule = import('../../web/js/shared/toolkit/core.js');
const CONCEPT = fs.readFileSync(path.join(__dirname, '../../docs/concepts/workspace/app.js'), 'utf8');
const SOURCE = fs.readFileSync(path.join(__dirname, '../../web/js/workspace/system-tagging.mjs'), 'utf8');
/** A minimal DOM whose elements count their live listeners (a listener's signal aborts it, as in a browser). */
function fakeDOM() {
const previous = global.document;
const elements = [];
class Element {
constructor(tag) {
this.tagName = tag.toUpperCase();
this.children = [];
this.attributes = {};
this.listeners = [];
this.dataset = {};
this.className = '';
this.id = '';
this.value = '';
this.disabled = false;
this._text = '';
elements.push(this);
}
get localName() { return this.tagName.toLowerCase(); }
get options() { return this.children.filter(child => child.tagName === 'OPTION'); }
get textContent() { return this._text + this.children.map(child => child.textContent).join(''); }
set textContent(text) { this.children = []; this._text = String(text); }
append(...nodes) {
for (const node of nodes) {
node.parentNode?.children && (node.parentNode.children = node.parentNode.children.filter(child => child !== node));
node.parentNode = this;
this.children.push(node);
}
}
remove() { if (this.parentNode) this.parentNode.children = this.parentNode.children.filter(child => child !== this); this.parentNode = null; }
setAttribute(name, value) { this.attributes[name] = String(value); }
getAttribute(name) { return Object.hasOwn(this.attributes, name) ? this.attributes[name] : null; }
removeAttribute(name) { delete this.attributes[name]; }
addEventListener(type, callback, options = {}) {
if (options.signal?.aborted) return;
const entry = { type, callback };
this.listeners.push(entry);
options.signal?.addEventListener('abort', () => { this.listeners = this.listeners.filter(item => item !== entry); });
}
removeEventListener(type, callback) { this.listeners = this.listeners.filter(item => item.type !== type || item.callback !== callback); }
dispatch(type) { return Promise.all(this.listeners.filter(item => item.type === type).map(({ callback }) => callback({ type, target: this }))); }
}
global.document = { createElement: tag => new Element(tag), activeElement: null };
const walk = (node, found = []) => { found.push(node); node.children.forEach(child => walk(child, found)); return found; };
return {
listenerCount: () => elements.reduce((sum, element) => sum + element.listeners.length, 0),
find: (root, predicate) => walk(root).find(predicate),
restore() { global.document = previous; },
};
}
const HOST = { autotag_model_type: 'wd14-swinv2', autotag_existing: 'skip', autotag_general_threshold: 0.35,
autotag_character_threshold: 0.85, autotag_suggestion_limit: 10, autotag_template: 'Describe the image as tags.' };
const entry = (label, size, extra = {}) => ({ name: label, label, size_label: size, downloaded: true, loaded: false, ...extra });
const MODELS = {
success: true,
models: { 'wd14-swinv2': entry('WD14 SwinV2', '~400 MB'), 'wd14-vit': entry('WD14 ViT', '~400 MB', { downloaded: false }),
gguf: entry('GGUF JoyCaption', '~4.5 GB', { downloaded: false, model_exists: true, mmproj_exists: false }),
hf: entry('HuggingFace JoyCaption', '~16 GB', { downloaded: false }) },
model_loaded: false, loaded_model_type: null,
};
const tick = () => new Promise(resolve => setImmediate(resolve));
const IMAGE = { filename: 'dunes_00001.png' };
async function mounted({ host = HOST, models = MODELS, modelsFail = false, image = IMAGE, start, post, hostFails = false } = {}) {
const { mountTagging } = await tagging;
const { createSystemDrafts } = await stateModule;
const { createScope } = await coreModule;
const dom = fakeDOM();
const drafts = createSystemDrafts({
readBrowser: () => ({}), readRuntime: () => ({}),
readHost: async () => { if (hostFails) throw new Error('offline'); return { values: host, revision: 'sha256:one' }; },
writeBrowser: async () => ({ persisted: true }), writeHost: async () => ({ success: true, persisted: true }),
writeRuntime: async () => ({}),
});
const scope = createScope();
const calls = { get: [], post: [], toasts: [], navigated: [], opened: [], shown: [], started: [] };
let answer = models;
const api = {
get: async (url, options) => {
calls.get.push(url);
assert.equal(options.cache, 'no-store');
if (modelsFail) throw new Error('offline');
return answer;
},
post: async (url, body) => { calls.post.push([url, body]); return post ? post(url, body) : { success: true }; },
};
const activity = { start: async (kind, body) => { calls.started.push([kind, body]); return start ? start(kind, body) : { job: { id: 'job1' } }; } };
const [surface] = mountTagging(null, { api, activity, scope, drafts, notify: text => calls.toasts.push(text),
navigate: hash => calls.navigated.push(hash), selected: () => image, openTagging: options => calls.opened.push(options),
show: (plan, options) => calls.shown.push([plan, options]) });
const byId = id => dom.find(surface, node => node.id === id);
const button = label => dom.find(surface, node => node.tagName === 'BUTTON' && node.textContent === label);
const description = () => byId('tag-model').parentNode.children[0].children[1].textContent;
await drafts.reloadHost();
await tick();
return { dom, drafts, scope, surface, byId, button, description, calls, setModels: value => { answer = value; } };
}
test('the section keeps the concept surface, rows, labels and actions in order', () => {
const start = CONCEPT.indexOf('case "tagging":');
const body = CONCEPT.slice(start, CONCEPT.indexOf('break;', start));
const rows = [...body.matchAll(/setting\("([^"]+)", "([^"]+)", "([^"]+)"/g)].map(([, id, label, description]) => [id, label, description]);
assert.deepEqual(rows.map(([id]) => id), ['tag-model', 'tag-existing', 'threshold-general', 'threshold-character']);
for (const [id, label] of rows) assert.ok(SOURCE.includes(`id: '${id}', label: '${label}'`), `${id} keeps "${label}"`);
// OQ6: the true descriptions are kept verbatim; the model's design note is replaced by its live status (D132).
for (const [id, , description] of rows.slice(1)) assert.ok(SOURCE.includes(`'${description}'`), `${id} description`);
assert.ok(!SOURCE.includes(rows[0][2]));
for (const option of ['Skip already tagged', 'Merge with existing', 'Replace after review']) assert.ok(body.includes(option) && SOURCE.includes(`'${option}'`));
assert.ok(body.includes('<h2>Auto-tagging</h2>') && SOURCE.includes("'h2', 'Auto-tagging'"));
assert.ok(body.includes('One tagging configuration for images and prompt-library actions. Every result can be reviewed before applying.'));
assert.ok(body.includes('LLM tag template') && SOURCE.includes("'LLM tag template'"));
assert.ok(body.includes('"Preview tagging review", "tag-review", "", "ca-primary"'));
// Prototype copy is never shipped.
for (const sample of ['Simulate download', 'Simulate unload', 'Ollama']) assert.ok(!SOURCE.includes(sample), sample);
});
test('model options are the four real engines with their size, in the concept order (D132, OQ5)', async () => {
const { modelOptions, MODEL_IDS, modelLabel, modelSize } = await tagging;
assert.deepEqual(MODEL_IDS, ['wd14-swinv2', 'wd14-vit', 'gguf', 'hf']);
assert.deepEqual(modelOptions(MODELS.models).map(([, label]) => label),
['WD14 SwinV2 · ~400 MB', 'WD14 ViT · ~400 MB', 'GGUF JoyCaption · ~4.5 GB', 'HuggingFace JoyCaption · ~16 GB']);
// Before the status answered (or when it cannot be read) the documented labels stand in.
assert.deepEqual(modelOptions(null), modelOptions(MODELS.models));
assert.equal(modelLabel('hf', { label: '' }), 'HuggingFace JoyCaption');
assert.equal(modelLabel('other', null), 'other');
assert.equal(modelSize('other', null), '');
});
test('the Tagging model description is the chosen model’s live status (D132)', async () => {
const { modelDescription, readStatus, STATUS_COPY } = await tagging;
const status = readStatus(MODELS);
assert.equal(modelDescription({ phase: 'loading' }, 'wd14-vit'), 'Loading model status…');
assert.equal(modelDescription({ phase: 'failed' }, 'wd14-vit'), 'Model status could not be read.');
assert.equal(modelDescription(status, 'wd14-swinv2'), 'Downloaded. Not loaded.');
assert.equal(modelDescription(status, 'wd14-vit'), 'Not downloaded. Download size ~400 MB.');
assert.equal(modelDescription(status, 'hf'), 'Not downloaded. Download size ~16 GB.');
assert.equal(modelDescription(status, 'gguf'), 'Partly downloaded. Download again to finish.');
const both = readStatus({ ...MODELS, models: { ...MODELS.models, gguf: entry('GGUF JoyCaption', '~4.5 GB', { downloaded: false,
model_exists: false, mmproj_exists: false }) } });
assert.equal(modelDescription(both, 'gguf'), 'Not downloaded. Download size ~4.5 GB.');
const loaded = readStatus({ ...MODELS, model_loaded: true, loaded_model_type: 'wd14-swinv2',
models: { ...MODELS.models, 'wd14-swinv2': entry('WD14 SwinV2', '~400 MB', { loaded: true }) } });
assert.equal(loaded.loadedType, 'wd14-swinv2');
assert.equal(modelDescription(loaded, 'wd14-swinv2'), STATUS_COPY.loaded);
assert.equal(modelDescription(status, 'unknown'), STATUS_COPY.failed);
assert.deepEqual(readStatus(null), { phase: 'failed', models: null, loadedType: null });
assert.equal(readStatus({ models: {}, model_loaded: true, loaded_model_type: '/abs/path' }).loadedType, null);
});
test('controls draft their host keys; numbers as numbers, other text as typed', async () => {
const { SETTING_KEYS, coerceNumber, draftValue } = await tagging;
assert.deepEqual(SETTING_KEYS, { 'tag-model': 'autotag_model_type', 'tag-existing': 'autotag_existing',
'threshold-general': 'autotag_general_threshold', 'threshold-character': 'autotag_character_threshold',
'tag-suggestions': 'autotag_suggestion_limit', 'tag-template': 'autotag_template' });
assert.deepEqual(['0.4', ' 1.5 ', '15', '', 'abc', '1e-1'].map(coerceNumber), [0.4, 1.5, 15, '', 'abc', 0.1]);
assert.equal(draftValue('autotag_suggestion_limit', '15'), 15);
assert.equal(draftValue('autotag_template', ' 12 '), ' 12 ');
assert.equal(draftValue('autotag_existing', undefined), '');
const { dom, drafts, scope, byId } = await mounted();
try {
assert.equal(byId('tag-model').value, 'wd14-swinv2');
assert.equal(byId('threshold-general').value, '0.35');
assert.equal(byId('tag-template').value, HOST.autotag_template);
assert.deepEqual([byId('threshold-general').min, byId('threshold-general').max, byId('threshold-general').step], ['0', '1', '0.05']);
assert.deepEqual([byId('tag-suggestions').min, byId('tag-suggestions').max], ['1', '100']);
byId('tag-existing').value = 'merge'; await byId('tag-existing').dispatch('change');
byId('threshold-general').value = '0.4'; await byId('threshold-general').dispatch('change');
byId('tag-suggestions').value = '15'; await byId('tag-suggestions').dispatch('change');
byId('tag-template').value = 'Tags only.'; await byId('tag-template').dispatch('input');
assert.deepEqual(drafts.changes('host'), { autotag_existing: 'merge', autotag_general_threshold: 0.4,
autotag_suggestion_limit: 15, autotag_template: 'Tags only.' });
// Every control names its stored key, so a refused save highlights it (D92).
for (const id of Object.keys(SETTING_KEYS)) assert.equal(byId(id).dataset.key, SETTING_KEYS[id]);
assert.equal(byId('tag-template').dataset.setting, 'tag-template');
} finally { scope.dispose(); assert.equal(dom.listenerCount(), 0); dom.restore(); }
});
test('opening the section reads the model status once; the select waits for it', async () => {
const { dom, scope, byId, description, calls, surface } = await mounted();
try {
assert.deepEqual(calls.get, ['/prompt_manager/autotag/models']);
assert.deepEqual(calls.post, []);
assert.equal(byId('tag-model').disabled, false);
assert.equal(description(), 'Downloaded. Not loaded.');
assert.deepEqual(byId('tag-model').options.map(option => option.textContent),
['WD14 SwinV2 · ~400 MB', 'WD14 ViT · ~400 MB', 'GGUF JoyCaption · ~4.5 GB', 'HuggingFace JoyCaption · ~16 GB']);
assert.deepEqual(surface.children.map(node => node.className),
['', '', 'setting-row', 'setting-row', 'setting-row', 'setting-row', 'setting-row', 'field', 'row-actions', 'operation']);
} finally { scope.dispose(); dom.restore(); }
const failed = await mounted({ modelsFail: true });
try {
assert.equal(failed.description(), 'Model status could not be read.');
assert.equal(failed.button('Download').getAttribute('aria-disabled'), 'true');
assert.equal(failed.button('Unload').disabled, true);
await failed.button('Download').dispatch('click');
assert.deepEqual(failed.calls.toasts, ['Model status could not be read.']);
assert.deepEqual(failed.calls.started, []);
} finally { failed.scope.dispose(); failed.dom.restore(); }
});
test('server settings that cannot be read leave the rows disabled and say so', async () => {
const { dom, scope, byId, button, description, calls } = await mounted({ hostFails: true });
try {
assert.equal(byId('tag-existing').disabled, true);
assert.equal(byId('tag-existing').parentNode.children[0].children.at(-1).textContent, 'Server settings could not be read.');
assert.equal(description(), 'Model status could not be read.');
await button('Preview tagging review').dispatch('click');
assert.deepEqual(calls.toasts, ['Server settings could not be read.']);
} finally { scope.dispose(); dom.restore(); }
});
test('Download starts the model download job and opens it; the label follows the disk (OQ10)', async () => {
const { dom, drafts, scope, button, calls, byId } = await mounted();
try {
assert.ok(button('Download again'));
byId('tag-model').value = 'wd14-vit'; await byId('tag-model').dispatch('change');
assert.equal(drafts.get('host', 'autotag_model_type'), 'wd14-vit');
await button('Download').dispatch('click');
assert.deepEqual(calls.started, [['model-download', { model_type: 'wd14-vit' }]]);
assert.deepEqual(calls.navigated, ['#/activity/job1']);
} finally { scope.dispose(); dom.restore(); }
});
test('startJob toasts by code and never navigates on a refusal', async () => {
const { startJob, startRefusal } = await tagging;
const { createScope } = await coreModule;
const error = (code, message = 'Server says.', extra = {}) => Object.assign(new Error(message), { code, ...extra });
assert.equal(startRefusal(error('model-busy', 'x', { operation: 'Tagging model download · WD14 ViT' }), 'X'),
'The tagging model is busy: Tagging model download · WD14 ViT. Try again when it finishes.');
assert.equal(startRefusal(error('model-busy', 'x'), 'X'),
'The tagging model is busy: Another tagging operation. Try again when it finishes.');
for (const code of ['model-unavailable', 'review-required', 'insufficient-space']) assert.equal(startRefusal(error(code), 'X'), 'Server says.');
assert.equal(startRefusal(error('weird'), 'Tag output images'), 'Tag output images could not start. Try again.');
assert.equal(startRefusal(new Error('offline'), 'Tag output images'), 'Tag output images could not start. Try again.');
const scope = createScope();
const toasts = [];
const navigated = [];
const run = start => startJob({ activity: { start }, scope, kind: 'autotag', body: {}, what: 'Tag output images',
notify: text => toasts.push(text), navigate: hash => navigated.push(hash) });
assert.equal(await run(async () => { throw error('model-busy', 'busy', { operation: 'Generate suggestions' }); }), null);
assert.equal(await run(async () => ({ busy: { id: 'a1', title: 'Scan for metadata' } })), null);
assert.equal(await run(async () => ({ busy: null })), null);
assert.deepEqual(toasts, ['The tagging model is busy: Generate suggestions. Try again when it finishes.',
['Another job is running: ', { text: 'Scan for metadata', href: '#/activity/a1' }, '.'],
'Another job is running. Try again when it ends.']);
assert.deepEqual(await run(async () => ({ job: { id: 'j2' } })), { id: 'j2' });
assert.deepEqual(navigated, ['#/activity/j2']);
scope.dispose();
assert.equal(await run(async () => { throw error('weird'); }), null);
assert.deepEqual(await run(async () => ({ job: { id: 'j3' } })), { id: 'j3' });
assert.deepEqual(navigated, ['#/activity/j2']);
assert.equal(toasts.length, 3);
});
test('Unload is enabled only while a model is loaded; a busy refusal names the operation', async () => {
const loadedModels = { ...MODELS, model_loaded: true, loaded_model_type: 'wd14-swinv2' };
let refuse = null;
const post = url => {
if (refuse) throw refuse;
return url.endsWith('/unload') ? { success: true, message: 'WD14SWINV2 model unloaded successfully', model_loaded: false } : {};
};
const { dom, scope, button, calls, setModels } = await mounted({ models: loadedModels, post });
try {
assert.equal(button('Unload').disabled, false);
refuse = Object.assign(new Error('A tagging model is busy'), { code: 'model-busy', operation: 'Tagging model download · WD14 ViT' });
await button('Unload').dispatch('click');
assert.deepEqual(calls.toasts, ['The tagging model is busy: Tagging model download · WD14 ViT. Try again when it finishes.']);
refuse = new Error('offline');
await button('Unload').dispatch('click');
assert.equal(calls.toasts.at(-1), 'The tagging model could not be unloaded. Try again.');
refuse = null;
setModels(MODELS);
await button('Unload').dispatch('click');
assert.equal(calls.toasts.at(-1), 'WD14 SwinV2 unloaded.');
assert.deepEqual(calls.post.map(([url]) => url), Array(3).fill('/prompt_manager/autotag/unload'));
// Each attempt reads the status again; nothing is loaded now.
assert.equal(calls.get.length, 4);
assert.equal(button('Unload').disabled, true);
} finally { scope.dispose(); dom.restore(); }
});
test('Preview tagging review opens the Tag image dialog with the section’s values, else says why (D115)', async () => {
const { NO_SELECTION, taggingPreset, draftProblem } = await tagging.then(async module => ({ ...module,
NO_SELECTION: (await import('../../web/js/workspace/vocabulary.mjs')).NO_SELECTION }));
assert.equal(NO_SELECTION, 'Select an image in Images first.');
assert.deepEqual(taggingPreset(HOST), { model: 'wd14-swinv2', options: { general_threshold: 0.35, character_threshold: 0.85,
suggestion_limit: 10 } });
assert.deepEqual(taggingPreset({ ...HOST, autotag_model_type: 'gguf' }), { model: 'gguf',
options: { instructions: HOST.autotag_template, suggestion_limit: 10 } });
assert.equal(draftProblem(HOST), null);
assert.deepEqual(draftProblem({ ...HOST, autotag_general_threshold: 1.5 }),
{ key: 'autotag_general_threshold', message: 'General-tag confidence must be between 0 and 1.' });
assert.equal(draftProblem({ ...HOST, autotag_character_threshold: '' }).key, 'autotag_character_threshold');
assert.equal(draftProblem({ ...HOST, autotag_suggestion_limit: 2.5 }).key, 'autotag_suggestion_limit');
assert.equal(draftProblem({ ...HOST, autotag_suggestion_limit: 0 }, { limit: false }), null);
assert.deepEqual(draftProblem({ ...HOST, autotag_model_type: 'hf', autotag_template: ' ' }),
{ key: 'autotag_template', message: 'LLM tag template must be non-empty text.' });
assert.deepEqual(draftProblem({ ...HOST, autotag_model_type: 'hf', autotag_template: 'x'.repeat(4001) }),
{ key: 'autotag_template', message: 'LLM tag template must be at most 4000 characters.' });
// Characters are code points, as Python's len counts them: 4000 emoji (8000 UTF-16 units) are allowed.
assert.equal(draftProblem({ ...HOST, autotag_model_type: 'hf', autotag_template: '🏜'.repeat(4000) }), null);
assert.equal(draftProblem({ ...HOST, autotag_model_type: 'hf', autotag_template: '🏜'.repeat(4001) }).key, 'autotag_template');
// A WD14 run never reads the template, nor a language model the thresholds.
assert.equal(draftProblem({ ...HOST, autotag_template: '' }), null);
assert.equal(draftProblem({ ...HOST, autotag_model_type: 'gguf', autotag_general_threshold: 7 }), null);
const none = await mounted({ image: null });
try {
const preview = none.button('Preview tagging review');
assert.equal(preview.getAttribute('aria-disabled'), 'true');
assert.equal(preview.className, 'ca-primary');
await preview.dispatch('click');
assert.deepEqual(none.calls.toasts, ['Select an image in Images first.']);
assert.deepEqual(none.calls.opened, []);
} finally { none.scope.dispose(); none.dom.restore(); }
const { dom, scope, button, byId, calls } = await mounted();
try {
await button('Preview tagging review').dispatch('click');
assert.equal(calls.opened.length, 1);
assert.equal(calls.opened[0].image, IMAGE);
assert.equal(calls.opened[0].source, 'tags');
assert.deepEqual(calls.opened[0].preset, taggingPreset(HOST));
byId('threshold-general').value = '1.5'; await byId('threshold-general').dispatch('change');
assert.equal(button('Preview tagging review').getAttribute('aria-disabled'), 'true');
await button('Preview tagging review').dispatch('click');
assert.deepEqual(calls.toasts, ['General-tag confidence must be between 0 and 1.']);
assert.equal(calls.opened.length, 1);
} finally { scope.dispose(); dom.restore(); }
});
const PLAN = { images: 3, prompts: 2, already_tagged: 1, to_tag: 1,
model: { type: 'wd14-swinv2', label: 'WD14 SwinV2', size_label: '~400 MB', downloaded: true }, existing: 'skip',
review_required: false, backup: { policy: 'on' }, disk: { free_bytes: 10 * 2 ** 30, required_bytes: 2 ** 20, sufficient: true } };
test('the batch review grid, body and Cancel-only reasons (D134, CUJ 4–5)', async () => {
const { batchGrid, batchBody, batchCloseNote, existingLabel, REPLACE_NOTE, NOTHING_TO_TAG, NO_LINKED_IMAGES } = await tagging;
assert.deepEqual(batchGrid(PLAN, HOST), [['Model', 'WD14 SwinV2 · ~400 MB'], ['Existing tags', 'Skip already tagged'],
['Images linked', '3'], ['Prompts', '2'], ['Already tagged', '1'], ['Prompts to tag', 'up to 1'],
['Settings', 'General 0.35 · Character 0.85'], ['Backup', 'Create and verify before changes']]);
const llm = { ...HOST, autotag_model_type: 'gguf' };
assert.deepEqual(batchGrid({ ...PLAN, model: { type: 'gguf' }, backup: { policy: 'off' } }, llm).slice(-2),
[['Settings', 'Template of 27 characters'], ['Backup', 'No safety backup (the policy is off)']]);
assert.equal(batchGrid({}, HOST)[2][1], '0');
assert.equal(existingLabel('odd'), 'odd');
assert.deepEqual(batchBody(HOST), { model_type: 'wd14-swinv2', existing: 'skip',
options: { general_threshold: 0.35, character_threshold: 0.85 } });
assert.deepEqual(batchBody(llm), { model_type: 'gguf', existing: 'skip', options: { instructions: HOST.autotag_template } });
assert.equal(batchCloseNote(PLAN), null);
assert.equal(batchCloseNote({ ...PLAN, review_required: true }), REPLACE_NOTE);
assert.equal(REPLACE_NOTE, 'Replacing tags needs a review of each image. Use Preview tagging review, or choose Skip or Merge for a batch run.');
assert.equal(batchCloseNote({ ...PLAN, model: { ...PLAN.model, downloaded: false } }), 'WD14 SwinV2 is not downloaded. Download it first.');
assert.equal(batchCloseNote({ ...PLAN, disk: { free_bytes: 2 ** 20, required_bytes: 2 ** 21, sufficient: false } }),
'Not enough free disk space: the safety backup needs about 2.0 MB and 1.0 MB is free. Nothing is written until there is room.');
assert.equal(batchCloseNote({ ...PLAN, prompts: 0, to_tag: 0 }), NO_LINKED_IMAGES);
assert.equal(batchCloseNote({ ...PLAN, to_tag: 0 }), NOTHING_TO_TAG);
assert.equal(NOTHING_TO_TAG, 'Every prompt with an output image is already tagged.');
});
test('showBatchPlan opens the review with the backup box, or the Cancel-only review', async () => {
const { showBatchPlan, REVIEW_TITLE, BATCH_SENTENCE } = await tagging;
const { createScope } = await coreModule;
const scope = createScope();
const reviews = [];
const closed = [];
const started = [];
const options = { values: HOST, scope, activity: { start: async (kind, body) => { started.push([kind, body]); return { job: { id: 'b1' } }; } },
notify: () => {}, navigate: () => {}, openReview: review => reviews.push(review), openCloseOnly: review => closed.push(review) };
showBatchPlan(PLAN, options);
assert.equal(reviews.length, 1);
const [review] = reviews;
assert.equal(review.title, REVIEW_TITLE);
assert.equal(review.title, 'Review tag output images');
assert.equal(review.heading, 'Tag output images');
assert.equal(review.sentence, BATCH_SENTENCE);
assert.equal(review.primary.label, 'Start tagging');
assert.equal(review.cancelLabel, 'Cancel');
assert.equal(review.backup.copy.heading, 'Back up the database first (recommended)');
await review.onStart();
assert.deepEqual(started, [['autotag', { model_type: 'wd14-swinv2', existing: 'skip',
options: { general_threshold: 0.35, character_threshold: 0.85 } }]]);
showBatchPlan({ ...PLAN, backup: { policy: 'off' } }, options);
assert.equal(reviews[1].backup, null);
showBatchPlan({ ...PLAN, review_required: true }, options);
assert.equal(closed.length, 1);
assert.equal(closed[0].note.startsWith('Replacing tags needs a review'), true);
assert.equal(closed[0].grid[0][0], 'Model');
scope.dispose();
});
test('Tag output images › Review previews the drafts, else says why; a refused preview is toasted', async () => {
let refuse = null;
const post = url => {
if (refuse) throw refuse;
assert.equal(url, '/prompt_manager/autotag/batch/preview');
return PLAN;
};
const { dom, scope, button, byId, calls } = await mounted({ post });
try {
const review = button('Review');
assert.equal(review.getAttribute('aria-disabled'), null);
await review.dispatch('click');
await tick();
assert.deepEqual(calls.post, [['/prompt_manager/autotag/batch/preview', { model_type: 'wd14-swinv2', existing: 'skip' }]]);
assert.equal(calls.shown.length, 1);
assert.equal(calls.shown[0][0], PLAN);
assert.deepEqual(calls.shown[0][1].values, HOST);
assert.equal(review.textContent, 'Review');
refuse = Object.assign(new Error('busy'), { code: 'busy', running: { id: 'r1', title: 'Optimize database' } });
await review.dispatch('click');
await tick();
assert.deepEqual(calls.toasts.at(-1), ['Another job is running: ', { text: 'Optimize database', href: '#/activity/r1' }, '.']);
refuse = new Error('offline');
await review.dispatch('click');
await tick();
assert.equal(calls.toasts.at(-1), 'The plan could not be reviewed. Nothing was changed.');
// A model that is not on disk cannot be reviewed: the button explains instead.
byId('tag-model').value = 'wd14-vit'; await byId('tag-model').dispatch('change');
assert.equal(review.getAttribute('aria-disabled'), 'true');
await review.dispatch('click');
assert.equal(calls.toasts.at(-1), 'WD14 ViT is not downloaded. Download it first.');
assert.equal(calls.post.length, 3);
} finally { scope.dispose(); dom.restore(); }
});
test('settingRefusal marks only an Auto-tagging save\'s plain 400 as invalid-setting (CUJ 2a)', async () => {
const { settingRefusal } = await import('../../web/js/workspace/system.mjs');
const refusal = (extra = {}) => Object.assign(new Error('General-tag confidence must be between 0 and 1.'),
{ status: 400, refused: true, code: null, ...extra });
const tagging = { autotag_general_threshold: 1.5, expected_revision: 'sha256:one' };
assert.equal(settingRefusal(refusal(), tagging).code, 'invalid-setting');
// Another section's refusal keeps the generic copy, even with a plain 400.
const other = refusal({ message: 'Result timeout must be between 0 and 300.' });
assert.equal(settingRefusal(other, { result_timeout: 999, expected_revision: 'sha256:one' }).code, null);
assert.equal(settingRefusal(refusal(), undefined).code, null);
for (const extra of [{ status: 409 }, { refused: false }, { code: 'jxl-unavailable' }]) {
assert.notEqual(settingRefusal(refusal(extra), tagging).code, 'invalid-setting');
}
assert.equal(settingRefusal(undefined, tagging), undefined);
});
test('a part-downloaded GGUF offers Download again', async () => {
const { downloadLabel, partlyDownloaded } = await tagging;
const partial = MODELS.models.gguf;
assert.equal(partlyDownloaded('gguf', partial), true);
assert.equal(downloadLabel('gguf', partial), 'Download again');
assert.equal(downloadLabel('gguf', { ...partial, model_exists: false }), 'Download');
assert.equal(downloadLabel('hf', MODELS.models.hf), 'Download');
assert.equal(downloadLabel('wd14-swinv2', MODELS.models['wd14-swinv2']), 'Download again');
assert.equal(partlyDownloaded('hf', { model_exists: true, mmproj_exists: false }), false);
const { dom, scope, byId, button, description } = await mounted();
try {
byId('tag-model').value = 'gguf'; await byId('tag-model').dispatch('change');
assert.equal(description(), 'Partly downloaded. Download again to finish.');
assert.ok(button('Download again'));
} finally { scope.dispose(); dom.restore(); }
});
test('the disk note says what is short in the server\'s words, and guesses no unknown byte counts', async () => {
const { diskShortNote, DISK_SHORT_UNKNOWN } = await tagging;
assert.equal(diskShortNote({ required_bytes: 3 * 2 ** 20, free_bytes: 2 ** 19, sufficient: false }),
'Not enough free disk space: the safety backup needs about 3.0 MB and 0.5 MB is free. Nothing is written until there is room.');
for (const disk of [{ required_bytes: NaN, free_bytes: 1 }, { required_bytes: 1, free_bytes: Infinity }, { required_bytes: -1, free_bytes: 1 },
{ required_bytes: '1', free_bytes: 1 }, null]) {
assert.equal(diskShortNote(disk), DISK_SHORT_UNKNOWN);
}
assert.equal(DISK_SHORT_UNKNOWN, 'Not enough free disk space for the safety backup. Nothing is written until there is room.');
assert.ok(!diskShortNote({ required_bytes: 1, free_bytes: 0 }).includes('apply'));
});
test('a drafts repaint never re-enables Download while its start is in flight', async () => {
let answer;
const start = () => new Promise(resolve => { answer = resolve; });
const { dom, scope, byId, button, calls } = await mounted({ start });
try {
const download = button('Download again');
const pressed = download.dispatch('click');
await tick();
assert.equal(download.getAttribute('aria-disabled'), 'true');
byId('tag-existing').value = 'merge'; await byId('tag-existing').dispatch('change');
assert.equal(download.getAttribute('aria-disabled'), 'true');
// A second press while it starts is ignored.
await download.dispatch('click');
assert.equal(calls.started.length, 1);
answer({ job: { id: 'd1' } });
await pressed;
assert.equal(download.getAttribute('aria-disabled'), null);
assert.deepEqual(calls.navigated, ['#/activity/d1']);
} finally { scope.dispose(); dom.restore(); }
});
test('an Unload that found nothing loaded says nothing and reads the status again', async () => {
const { dom, scope, button, calls } = await mounted({ models: { ...MODELS, model_loaded: true, loaded_model_type: 'wd14-swinv2' },
post: () => ({ success: true, message: 'No model was loaded' }) });
try {
await button('Unload').dispatch('click');
assert.deepEqual(calls.toasts, []);
assert.equal(calls.get.length, 2);
} finally { scope.dispose(); dom.restore(); }
});
test('an unknown saved engine has its own reason', async () => {
const { UNKNOWN_MODEL } = await tagging;
const { dom, scope, button, calls } = await mounted({ host: { ...HOST, autotag_model_type: 'wd14-convnext' } });
try {
await button('Preview tagging review').dispatch('click');
await button('Download').dispatch('click');
await button('Review').dispatch('click');
assert.deepEqual(calls.toasts, [UNKNOWN_MODEL, UNKNOWN_MODEL, UNKNOWN_MODEL]);
assert.equal(UNKNOWN_MODEL, 'The saved tagging model is not one of these engines. Choose one and save preferences.');
assert.deepEqual(calls.started, []);
} finally { scope.dispose(); dom.restore(); }
});
test('a review whose dialog fails to open is toasted, not lost', async () => {
const { dom, scope, button, calls } = await mounted({ post: () => PLAN });
try {
calls.shown.push = () => { throw new Error('dialog failed'); };
await button('Review').dispatch('click');
await tick();
assert.equal(calls.toasts.at(-1), 'The plan could not be reviewed. Nothing was changed.');
assert.equal(button('Review').getAttribute('aria-disabled'), null);
} finally { scope.dispose(); dom.restore(); }
});
test('a status read, an unload and a preview that finish after the section left do nothing', async () => {
const { mountTagging } = await tagging;
const { createSystemDrafts } = await stateModule;
const { createScope } = await coreModule;
const dom = fakeDOM();
const pending = [];
const later = value => new Promise(resolve => pending.push(() => resolve(value)));
const drafts = createSystemDrafts({ readBrowser: () => ({}), readRuntime: () => ({}),
readHost: async () => ({ values: HOST, revision: 'sha256:one' }), writeBrowser: async () => ({}), writeHost: async () => ({}),
writeRuntime: async () => ({}) });
await drafts.load('host');
const scope = createScope();
const toasts = [];
const shown = [];
let gets = 0;
const loaded = { ...MODELS, model_loaded: true, loaded_model_type: 'wd14-swinv2',
models: { ...MODELS.models, 'wd14-swinv2': entry('WD14 SwinV2', '~400 MB', { loaded: true }) } };
const api = { get: () => { gets += 1; return gets === 1 ? Promise.resolve(loaded) : later(MODELS); }, post: url => later(url.endsWith('/unload') ? { success: true, model_loaded: false } : PLAN) };
const [surface] = mountTagging(null, { api, activity: {}, scope, drafts, notify: text => toasts.push(text), navigate: () => {},
selected: () => IMAGE, openTagging: () => {}, show: plan => shown.push(plan) });
try {
await tick();
const find = label => dom.find(surface, node => node.tagName === 'BUTTON' && node.textContent === label);
const description = () => dom.find(surface, node => node.id === 'tag-model').parentNode.children[0].children[1].textContent;
assert.equal(description(), 'Downloaded and loaded.');
const unloading = find('Unload').dispatch('click');
const reviewing = find('Review').dispatch('click');
await tick();
scope.dispose();
for (const settle of pending.splice(0)) settle();
await unloading;
await reviewing;
await tick();
for (const settle of pending.splice(0)) settle();
await tick();
assert.deepEqual(toasts, []);
assert.deepEqual(shown, []);
// The status read after the unload never reached the page either.
assert.equal(description(), 'Downloaded and loaded.');
assert.equal(dom.listenerCount(), 0);
} finally { dom.restore(); }
});
+5
View File
@@ -841,6 +841,11 @@ test('openReviewDialog lays out heading, grid, options, sentence, backup box and
[['dt', 'dd'], 'Scope', '2 stale image links'], [['dt', 'dd'], 'Media on disk', 'No source files removed']]);
review.setValue('Scope', '3 stale image links');
assert.equal(grid.children[0].children[1].textContent, '3 stale image links');
// #355: a row's hover title (the LoRA review's settings file); an unknown term changes nothing.
review.setTitle('Media on disk', '~/settings.json');
review.setTitle('Absent', 'x');
review.setValue('Absent', 'x');
assert.equal(grid.children[1].children[1].title, '~/settings.json');
const actions = byId(dom, 'modal-body').children.at(-1);
assert.deepEqual(actions.children.map(node => [node.textContent, node.className]),
[['Cancel', ''], ['Apply reviewed operation', 'ca-primary']]);
+29 -10
View File
@@ -33,7 +33,9 @@ written next to it (TESTING.md "Time limits and hangs").
import argparse
import faulthandler
import os
import shutil
import sys
import tempfile
import unittest
from pathlib import Path
@@ -41,6 +43,9 @@ DEFAULT_PER_TEST_LIMIT = 120
DEFAULT_SHUTDOWN_LIMIT = 120
NO_TESTS_EXIT = 5
USAGE_EXIT = 2
# LoraManager's settings-folder override (#355): every test reads an empty,
# disposable LoraManager settings folder, never the developer's real one.
LORA_MANAGER_SETTINGS_ENV = "LORA_MANAGER_SETTINGS_DIR"
class LimitError(ValueError):
@@ -115,6 +120,16 @@ def parse_args(argv):
return parser.parse_args(argv)
def isolate_lora_manager_settings():
"""Point LoraManager's settings lookup at a fresh temporary folder.
Returns the folder; the caller removes it when the run ends.
"""
folder = tempfile.mkdtemp(prefix="pm-lora-settings-")
os.environ[LORA_MANAGER_SETTINGS_ENV] = folder
return folder
def main(argv=None):
args = parse_args(sys.argv[1:] if argv is None else argv)
# Keep the real stderr: a test may replace sys.stderr, and the watchdog
@@ -128,16 +143,20 @@ def main(argv=None):
return USAGE_EXIT
faulthandler.dump_traceback_later(per_test_limit, exit=True, file=stream)
start = os.path.abspath(args.start_directory)
suite = unittest.TestLoader().discover(start, args.pattern, top_level_dir=start)
runner = unittest.TextTestRunner(
stream=stream,
verbosity=args.verbosity,
# As unittest.main does: show warnings unless -W/PYTHONWARNINGS set them.
warnings=None if sys.warnoptions else "default",
resultclass=make_result_class(stream, per_test_limit, shutdown_limit),
)
result = runner.run(suite)
settings_folder = isolate_lora_manager_settings()
try:
start = os.path.abspath(args.start_directory)
suite = unittest.TestLoader().discover(start, args.pattern, top_level_dir=start)
runner = unittest.TextTestRunner(
stream=stream,
verbosity=args.verbosity,
# As unittest.main does: show warnings unless -W/PYTHONWARNINGS set them.
warnings=None if sys.warnoptions else "default",
resultclass=make_result_class(stream, per_test_limit, shutdown_limit),
)
result = runner.run(suite)
finally:
shutil.rmtree(settings_folder, ignore_errors=True)
if result.testsRun == 0 and not result.skipped:
return NO_TESTS_EXIT
return 0 if result.wasSuccessful() else 1
+1 -1
View File
@@ -36,7 +36,7 @@ SUGGEST = "/prompt_manager/image-annotations/suggest"
WAIT = 30.0
WD14 = {"general_threshold": 0.35, "character_threshold": 0.85}
LLM = {"instructions": "Generate tags."}
UNAVAILABLE = "Download it from the legacy Auto-tag page."
UNAVAILABLE = "Download it in System › Auto-tagging."
SCORED = [("sandstone", 0.98), ("desert", 0.96)] + [
(f"tag_{i:03d}", round(0.9 - i * 0.005, 4)) for i in range(150)
+2
View File
@@ -1145,8 +1145,10 @@ class TestEnvelopes(AnnotationAPITestCase):
sorted(await self.json_of(resp)),
[
"default_prompt",
"existing", # #328 adds the saved batch policy and model
"loaded_model_type",
"model_loaded",
"model_type",
"models",
"success",
"suggestion_limit", # B2 adds it (D68); nothing else changes
+1066 -3
View File
File diff suppressed because it is too large Load Diff
+127 -5
View File
@@ -12,6 +12,7 @@ import pathlib
import sys
import tempfile
import threading
import types
import unittest
from contextlib import ExitStack
from pathlib import Path, PurePath
@@ -1317,7 +1318,8 @@ class TestServeErrorBranches(ImageRouteCoverageCase):
with (
patch.object(IntegrationConfig, "LORA_MANAGER_ENABLED", True),
patch.object(IntegrationConfig, "LORA_MANAGER_PATH", ""),
patch("py.lora_utils.find_lora_directories", return_value=[lora]),
patch("py.lora_utils.detect_lora_manager", return_value=None),
patch("py.lora_utils.default_lora_roots", return_value=[lora]),
patch("py.lora_utils.get_lora_image_cache_dir", return_value=cache),
):
resp = await self.client.request(
@@ -1328,7 +1330,8 @@ class TestServeErrorBranches(ImageRouteCoverageCase):
with (
patch.object(IntegrationConfig, "LORA_MANAGER_ENABLED", True),
patch.object(IntegrationConfig, "LORA_MANAGER_PATH", ""),
patch("py.lora_utils.find_lora_directories", side_effect=RuntimeError),
patch("py.lora_utils.detect_lora_manager", return_value=None),
patch("py.lora_utils.default_lora_roots", side_effect=RuntimeError),
):
resp = await self.client.request(
"GET", f"/prompt_manager/images/{image_id}/file"
@@ -1347,7 +1350,12 @@ class TestLoraManagerPathAllowList(ImageRouteCoverageCase):
)
stack.enter_context(patch.object(IntegrationConfig, "LORA_MANAGER_PATH", path))
stack.enter_context(
patch("py.lora_utils.find_lora_directories", return_value=[lora])
patch("py.lora_utils.default_lora_roots", return_value=[lora])
)
# Never the developer's LoraManager settings (#355).
settings = stack.enter_context(tempfile.TemporaryDirectory())
stack.enter_context(
patch.dict(os.environ, {"LORA_MANAGER_SETTINGS_DIR": settings})
)
stack.enter_context(
patch("py.lora_utils.get_lora_image_cache_dir", return_value=cache)
@@ -1389,6 +1397,115 @@ class TestLoraManagerPathAllowList(ImageRouteCoverageCase):
resolve.assert_called_once_with(configured)
class TestLoraManagerMediaServing(ImageRouteCoverageCase):
"""#355: LoraManager's sidecar and example-image folders are served while
the integration is on, only through the folder guard."""
async def setUpAsync(self):
await super().setUpAsync()
from py.config import IntegrationConfig
self._lm_tmp = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
self.addCleanup(self._lm_tmp.cleanup)
root = Path(os.path.realpath(self._lm_tmp.name))
custom_nodes = root / "custom_nodes"
self.lm = custom_nodes / "ComfyUI-Lora-Manager"
(self.lm / "py").mkdir(parents=True)
(self.lm / "__init__.py").write_text("")
self.settings = root / "settings"
self.settings.mkdir()
self.sidecars = root / "sidecars"
self.examples = root / "examples"
self.cache = root / "cache"
for folder in (self.sidecars, self.examples, self.cache):
folder.mkdir()
self.write_settings(str(self.sidecars))
self.enabled = patch.object(IntegrationConfig, "LORA_MANAGER_ENABLED", True)
for item in (
self.enabled,
patch.object(IntegrationConfig, "LORA_MANAGER_PATH", str(self.lm)),
patch(
"py.lora_utils.custom_nodes_directories", return_value=[custom_nodes]
),
patch("py.lora_utils.default_lora_roots", return_value=[]),
patch("py.lora_utils.get_lora_image_cache_dir", return_value=self.cache),
patch.dict(os.environ, {"LORA_MANAGER_SETTINGS_DIR": str(self.settings)}),
):
item.start()
self.addCleanup(item.stop)
def write_settings(self, sidecar_root):
(self.settings / "settings.json").write_text(
json.dumps(
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": sidecar_root,
"example_images_path": str(self.examples),
"civitai_api_key": "never-served",
}
)
)
async def _status(self, image_id):
resp = await self.client.request(
"GET", f"/prompt_manager/images/{image_id}/file"
)
await resp.read() # release the served file before the folder goes
return resp.status
async def test_sidecar_and_example_media_are_served_while_enabled(self):
preview = self.sidecars / "loras-0a1b2c3d" / "neon.webp"
preview.parent.mkdir()
preview.write_bytes(b"RIFF")
example = self.examples / "ab12" / "1.mp4"
example.parent.mkdir()
example.write_bytes(b"video")
ids = [self._link_image(preview, "a"), self._link_image(example, "b")]
for image_id in ids:
self.assertEqual(await self._status(image_id), 200)
self.enabled.stop()
for image_id in ids:
self.assertEqual(await self._status(image_id), 403)
self.enabled.start()
async def test_a_filesystem_root_as_sidecar_folder_adds_nothing(self):
self.write_settings(os.path.abspath(os.sep))
dirs = self.api._lora_image_dirs()
# An unusable layout: the default LoRA roots (none here) and the cache.
self.assertEqual(dirs, [self.cache])
async def test_the_layout_folders_are_the_served_ones(self):
dirs = self.api._lora_image_dirs()
self.assertEqual(
dirs,
[
Path(os.path.normcase(str(self.sidecars))),
Path(os.path.normcase(str(self.examples))),
self.cache,
],
)
class TestEnrichMarksVideos(ImageAPITestCase):
"""#355: a linked video outside the output folders is marked as video."""
async def test_enrich_sets_is_video_and_media_type(self):
elsewhere = Path(tempfile.gettempdir()) / "lora-media"
images = [
{"id": 1, "image_path": str(elsewhere / "a.MP4")},
{"id": 2, "image_path": str(elsewhere / "b.png")},
{"id": 3, "image_path": str(elsewhere / "c.mp4"), "is_video": False},
{"id": 4, "image_path": ""},
]
self.api._enrich_images(images)
self.assertEqual(
[(i.get("is_video"), i.get("media_type")) for i in images],
[(True, "video"), (False, "image"), (False, None), (None, None)],
)
self.assertEqual(images[0]["url"], "/prompt_manager/images/1/file")
class TestLinkImageRelativePaths(ImageAPITestCase):
"""POST /images/link accepts the path forms the API hands out."""
@@ -3343,8 +3460,13 @@ class TestSharedListing(ImageAPITestCase):
record = await self._listed()
self._interval(30)
# The interval ends in a few milliseconds: still never a zero delay.
record.last_started_at = images_module._time.monotonic() - 29.99
self.api.mark_root_dirty(self._output("b.png"))
# The module's clock is held, so a slow runner cannot pass the end
# of the interval between these lines (the walk would start at once).
now = images_module._time.monotonic()
record.last_started_at = now - 29.99
held = types.SimpleNamespace(monotonic=lambda: now)
with patch.object(images_module, "_time", held):
self.api.mark_root_dirty(self._output("b.png"))
handle = record.rewalk_handle
self.assertIsNotNone(handle)
self.assertGreaterEqual(self._remaining(handle), 0.45)
+183
View File
@@ -1,9 +1,12 @@
"""Native library endpoints use bounded cursors and rooted media identities."""
import json
import os
import unittest
from unittest.mock import patch
from utils.hashing import generate_prompt_hash
try:
from test_api_images import ImageAPITestCase, make_png
except ImportError:
@@ -978,6 +981,186 @@ class TestLibraryAPI(ImageAPITestCase):
self.assertNotIn("literal", json.dumps(result))
class TestLibraryLoraMedia(ImageAPITestCase):
"""#355 (G6): media a LoRA import linked outside the output folders shows
on the Library page and opens in the inspector while the integration is on."""
async def setUpAsync(self):
await super().setUpAsync()
import tempfile
from pathlib import Path
self._lora_tmp = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
self.addCleanup(self._lora_tmp.cleanup)
self.lora_dir = Path(self._lora_tmp.name).resolve() / "sidecars"
self.lora_dir.mkdir()
self.lora_on = True
self.api._lora_image_dirs = lambda: [self.lora_dir] if self.lora_on else []
async def get(self, path="", **params):
response = await self.client.get(
"/prompt_manager/library" + path, params=params
)
return response, await response.json()
def _imported(self):
"""A LoRA prompt linked like the import links it: the cover first."""
still = make_png(self.lora_dir / "abc" / "cover.png")
video = self.lora_dir / "abc" / "clip.mp4"
video.write_bytes(b"video")
self.api.db.replace_category_prompts(
"lora-manager",
[
{
"name": "neon",
"text": "neon street",
"hash": generate_prompt_hash("neon street"),
"tags": ["lora-manager"],
"images": [str(still), str(video)],
}
],
)
positive = self.api.db.search_prompts(category="lora-manager")[0]["id"]
ids = {
os.path.basename(row["image_path"]): row["id"]
for row in self.api.db.get_prompt_images(positive)
}
return positive, ids
async def test_list_previews_show_linked_lora_media_cover_first(self):
positive, _ids = self._imported()
response, result = await self.get()
self.assertEqual(response.status, 200, result)
item = next(i for i in result["items"] if i["id"] == positive)
previews = item["previews"]
self.assertEqual([p["filename"] for p in previews], ["cover.png", "clip.mp4"])
# The media keeps the rooted identity every workspace view checks; no
# preview is made inside LoraManager's folders, so the card loads the
# original and a video shows its label.
self.assertEqual({p["root_index"] for p in previews}, {1})
self.assertTrue(
all(
p["url"].startswith("/prompt_manager/images/serve/abc/")
for p in previews
)
)
self.assertEqual([p["thumbnail_url"] for p in previews], [None, None])
self.assertEqual([p["is_video"] for p in previews], [False, True])
self.assertEqual([p["linked_only"] for p in previews], [True, True])
self.assertNotIn(str(self.lora_dir), json.dumps(result))
file = await self.client.get(previews[0]["url"])
self.assertEqual(file.status, 200)
await file.read()
self.assertFalse((self.lora_dir / "thumbnails").exists())
async def test_detail_and_images_page_list_lora_media(self):
positive, _ids = self._imported()
make_png(self.output_dir / "out.png")
_, detail = await self.get(f"/{positive}")
self.assertEqual(
[p["filename"] for p in detail["positive"]["previews"]],
["cover.png", "clip.mp4"],
)
_, page = await self.get(f"/{positive}/images")
self.assertEqual(
[i["filename"] for i in page["images"]], ["cover.png", "clip.mp4"]
)
# No favorite star: favorites belong to output listings.
self.assertNotIn("favorite", page["images"][0])
self.assertTrue(page["images"][0]["linked_only"])
async def test_inspector_reads_lora_media_by_its_root(self):
positive, _ids = self._imported()
_, page = await self.get(f"/{positive}/images")
image = page["images"][0]
response = await self.client.get(
"/prompt_manager/images/inspect/" + image["path"],
params={
"root": str(image["root_index"]),
"root_scope": image["root_scope"],
},
)
self.assertEqual(response.status, 200, await response.text())
async def test_the_inspector_reads_only_linked_files_under_a_lora_root(self):
# Security review: a LoRA root exposes only the files the import linked.
self._imported()
make_png(self.lora_dir / "abc" / "unlinked.png")
scope = self.api._gallery_root_scope(self.lora_dir, 1)
for name, status in (("unlinked.png", 403), ("cover.png", 200)):
with self.subTest(name=name):
response = await self.client.get(
"/prompt_manager/images/inspect/abc/" + name,
params={"root": "1", "root_scope": scope},
)
self.assertEqual(response.status, status, await response.text())
async def test_only_recorded_lora_media_is_served_by_path(self):
self._imported()
make_png(self.lora_dir / "abc" / "unlinked.png")
served = await self.client.get(
"/prompt_manager/images/serve/abc/cover.png", params={"root": "1"}
)
self.assertEqual(served.status, 200)
await served.read()
unlinked = await self.client.get(
"/prompt_manager/images/serve/abc/unlinked.png", params={"root": "1"}
)
self.assertEqual(unlinked.status, 403)
preview = await self.client.get(
"/prompt_manager/images/serve/thumbnails/abc/cover_thumb.webp",
params={"root": "1"},
)
self.assertEqual(preview.status, 403)
self.assertFalse((self.lora_dir / "thumbnails").exists())
# Without a root, only the output folders are searched.
unrooted = await self.client.get("/prompt_manager/images/serve/abc/cover.png")
self.assertEqual(unrooted.status, 404)
async def test_a_download_recorded_through_a_symlinked_cache_is_served(self):
from pathlib import Path
real = self.lora_dir.parent / "real-cache"
real.mkdir()
link = self.lora_dir.parent / "cache-link"
try:
os.symlink(real, link, target_is_directory=True)
except (OSError, NotImplementedError):
self.skipTest("symlinks are not available here")
self.api._lora_image_dirs = lambda: [link]
image = make_png(real / "neon" / "a1b2.jpg")
stored = str(link / "neon" / "a1b2.jpg")
positive = self._save_prompt("downloaded")
self.api.db.link_image_to_prompt(positive, stored)
_, page = await self.get(f"/{positive}/images")
self.assertEqual([i["filename"] for i in page["images"]], [image.name])
served = await self.client.get(page["images"][0]["url"])
self.assertEqual(served.status, 200)
await served.read()
self.assertTrue(Path(stored).is_file())
async def test_a_malformed_root_is_refused_not_an_error(self):
response = await self.client.get(
"/prompt_manager/images/serve/abc/cover.png", params={"root": "\u00b2"}
)
self.assertEqual(response.status, 403)
async def test_lora_media_is_hidden_while_the_integration_is_off(self):
positive, _ids = self._imported()
self.lora_on = False
_, result = await self.get()
item = next(i for i in result["items"] if i["id"] == positive)
self.assertEqual(item["previews"], [])
response = await self.client.get(
"/prompt_manager/images/inspect/abc/cover.png", params={"root": "1"}
)
self.assertEqual(response.status, 400)
served = await self.client.get(
"/prompt_manager/images/serve/abc/cover.png", params={"root": "1"}
)
self.assertEqual(served.status, 403)
class TestNameCursorBudget(unittest.TestCase):
def test_lru_byte_budget_scope_and_oversize_keys_are_bounded(self):
from py.api.library import _NameCursorStore
+71 -18
View File
@@ -70,8 +70,16 @@ class LoraAPITestCase(AioHTTPTestCase):
# Never the repository's config.json (live user data): every save in
# these tests goes to a file in this test's temp directory.
self.config_file = self.tmp_root / "config.json"
# Never the developer's LoraManager settings either (#355): an empty
# settings folder of this test's own.
self.lm_settings = self.tmp_root / "lora-manager-settings"
self.lm_settings.mkdir()
self._env_patch = patch.dict(
os.environ, {"PROMPT_MANAGER_CONFIG_PATH": str(self.config_file)}
os.environ,
{
"PROMPT_MANAGER_CONFIG_PATH": str(self.config_file),
"LORA_MANAGER_SETTINGS_DIR": str(self.lm_settings),
},
)
self._env_patch.start()
@@ -194,6 +202,8 @@ def _make_lora_manager(where):
def _write_metadata(directory, stem, words=None, images=None, prompt=None):
"""A LoraManager sidecar beside its model, as LoraManager names them:
``<stem>.metadata.json`` next to ``<stem>.safetensors`` (#355)."""
meta = {"file_name": f"{stem}.safetensors", "model_name": stem, "civitai": {}}
if words is not None:
meta["civitai"]["trainedWords"] = words
@@ -201,7 +211,8 @@ def _write_metadata(directory, stem, words=None, images=None, prompt=None):
meta["civitai"]["images"] = [
{"url": url, "meta": {"prompt": prompt} if prompt else {}} for url in images
]
path = directory / f"{stem}.safetensors.metadata.json"
(directory / f"{stem}.safetensors").write_bytes(b"model")
path = directory / f"{stem}.metadata.json"
path.write_text(json.dumps(meta))
return path
@@ -300,9 +311,7 @@ class TestLoraEnableRoute(LoraAPITestCase):
_write_metadata(loras, "neon", words=["glow"])
self.detect_mock.return_value = str(lm)
with patch.object(
lora_utils, "find_lora_directories", return_value=[str(loras)]
):
with patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]):
resp = await self._post_json(
"/prompt_manager/lora/enable",
{
@@ -848,9 +857,7 @@ class TestLoraScanRoute(LoraAPITestCase):
cache_dir = self.tmp_root / "img_cache"
with (
patch.object(
lora_utils, "find_lora_directories", return_value=[str(loras)]
),
patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]),
patch.object(
lora_utils, "get_lora_image_cache_dir", return_value=cache_dir
),
@@ -880,9 +887,10 @@ class TestLoraScanRoute(LoraAPITestCase):
images=["https://civitai.com/a.png", "https://evil.example/b.png"],
prompt="neon street",
)
(loras / "neon.png").write_bytes(_png_bytes())
# No local image: its CivitAI images are downloaded (#355).
_write_metadata(loras, "plain")
(loras / "broken.metadata.json").write_text("{{{")
(loras / "broken.safetensors").write_bytes(b"model")
existing_id = self.api.db.save_prompt(
text="plain",
category=None,
@@ -910,8 +918,56 @@ class TestLoraScanRoute(LoraAPITestCase):
self.assertEqual(imported[0]["text"], "neon street")
self.assertIn("lora:neon", imported[0]["tags"])
self.assertIn("glow", imported[0]["tags"])
self.assertEqual(len(self.api.db.get_prompt_images(imported[0]["id"])), 2)
self.assertEqual(len(self.api.db.get_prompt_images(imported[0]["id"])), 1)
self.assertEqual(len(self.api.db.get_prompt_images(existing_id)), 0)
# The unreadable sidecar is reported by its file name.
statuses = [e.get("status") for e in events if e.get("processed")]
self.assertIn("broken.metadata.json", statuses)
async def test_scan_links_local_media_without_downloading(self):
self._enable()
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(
loras, "neon", images=["https://civitai.com/a.png"], prompt="neon street"
)
(loras / "neon.preview.webp").write_bytes(b"x")
resp, events, opener_calls = await self._scan(loras)
self.assertEqual(events[-1]["imported"], 1)
self.assertEqual(opener_calls, [])
imported = self.api.db.search_prompts(category="lora-manager")
images = self.api.db.get_prompt_images(imported[0]["id"])
self.assertEqual(
[os.path.basename(i["image_path"]) for i in images], ["neon.preview.webp"]
)
async def test_scan_reports_unusable_settings_without_replacing(self):
self._enable()
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(loras, "neon")
(self.lm_settings / "settings.json").write_text("{not json")
old_id = self.api.db.save_prompt(
text="kept",
category="lora-manager",
tags=[],
rating=None,
prompt_hash=generate_prompt_hash("kept"),
)
resp, events, _ = await self._scan(loras)
self.assertEqual(resp.status, 200)
self.assertEqual(events[-1]["type"], "error")
self.assertEqual(
events[-1]["message"],
"LoraManager settings: not usable: the settings file is not valid "
"JSON. Nothing was changed.",
)
kept = self.api.db.search_prompts(category="lora-manager")
self.assertEqual([p["id"] for p in kept], [old_id])
async def test_save_failure_is_counted_as_skipped(self):
self._enable()
@@ -949,9 +1005,7 @@ class TestLoraTriggerWordRoutes(LoraAPITestCase):
loras.mkdir()
_write_metadata(loras, "neon", words=["glow", "bright"])
with patch.object(
lora_utils, "find_lora_directories", return_value=[str(loras)]
):
with patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]):
resp = await self.client.request(
"GET", "/prompt_manager/lora/trigger-words?name=NEON"
)
@@ -985,17 +1039,16 @@ class TestLoraTriggerWordRoutes(LoraAPITestCase):
loras.mkdir()
_write_metadata(loras, "neon", words=["glow"])
with patch.object(
lora_utils, "find_lora_directories", return_value=[str(loras)]
):
with patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]):
resp = await self.client.request(
"POST", "/prompt_manager/lora/refresh-cache"
)
data = await resp.json()
self.assertTrue(data["success"])
# keyed by both the file_name stem and the metadata-file stem
self.assertEqual(data["loras_with_trigger_words"], 2)
# Keyed by the model name and the sidecar name, which are the same
# for LoraManager's alongside layout.
self.assertEqual(data["loras_with_trigger_words"], 1)
async def test_refresh_error_is_500(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
+29 -12
View File
@@ -19,8 +19,15 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from PIL import Image
from py import autotag
from py.autotag import MODELS, AutoTagService
# py.autotag reads its saved defaults from py.config, which imports ComfyUI's
# ``server`` at load: stand in for it when this module runs alone.
_mock_server = MagicMock()
_mock_server.PromptServer.instance.routes = MagicMock()
sys.modules.setdefault("server", _mock_server)
from py import autotag # noqa: E402
from py.autotag import MODELS, AutoTagService # noqa: E402
from py.config import AutotagConfig # noqa: E402
@contextlib.contextmanager
@@ -254,12 +261,19 @@ class TestWd14Thresholding(EngineTestCase):
self.assertEqual(list(feeds.keys()), ["input"])
self.assertEqual(tuple(feeds["input"].shape), (1, 8, 8, 3))
def test_service_threshold_setters_clamp(self):
service = self._loaded([0.0] * 6)
service.wd14_general_threshold = 1.7
service.wd14_character_threshold = -2
self.assertEqual(service.wd14_general_threshold, 1.0)
self.assertEqual(service.wd14_character_threshold, 0.0)
def test_service_thresholds_are_the_saved_defaults(self):
# #328: one configuration; the service has no thresholds of its own.
service = self._loaded([0.9, 0.3, 0.8, 0.99, 0.99, 0.5])
with (
patch.object(AutotagConfig, "GENERAL_THRESHOLD", 0.25),
patch.object(AutotagConfig, "CHARACTER_THRESHOLD", 0.75),
):
self.assertEqual(
(service.wd14_general_threshold, service.wd14_character_threshold),
(0.25, 0.75),
)
tags = service.generate_tags(str(self.image_path))
self.assertEqual(tags, ["1girl", "hatsune_miku", "long_hair", "smile"])
class TestGgufNormalisation(EngineTestCase):
@@ -369,10 +383,13 @@ class TestTagSuggestions(EngineTestCase):
def test_thresholds_default_to_the_service_values(self):
service = self._wd14()
service.wd14_general_threshold = 0.95
service.wd14_character_threshold = 0.95
suggestions = service.generate_tag_suggestions(str(self.image_path), limit=10)
with (
patch.object(AutotagConfig, "GENERAL_THRESHOLD", 0.95),
patch.object(AutotagConfig, "CHARACTER_THRESHOLD", 0.95),
):
suggestions = service.generate_tag_suggestions(
str(self.image_path), limit=10
)
self.assertEqual(suggestions, [("sandstone", 0.98), ("desert", 0.96)])
+16
View File
@@ -79,5 +79,21 @@ class TestBrowserServerLifetime(unittest.TestCase):
process.wait(timeout=1)
class TestBrowserServerLoraIsolation(unittest.TestCase):
"""#355: page tests never read the developer's LoraManager settings."""
def test_settings_folder_is_inside_the_server_root(self):
from unittest.mock import patch
from tests.browser.server import isolate_lora_manager_settings
with tempfile.TemporaryDirectory() as root:
with patch.dict(os.environ, {"LORA_MANAGER_SETTINGS_DIR": "/real/one"}):
folder = isolate_lora_manager_settings(Path(root))
self.assertEqual(os.environ["LORA_MANAGER_SETTINGS_DIR"], str(folder))
self.assertEqual(folder.parent, Path(root))
self.assertTrue(folder.is_dir())
if __name__ == "__main__":
unittest.main()
+103 -2
View File
@@ -413,9 +413,21 @@ class TestAutotagConfig(unittest.TestCase):
self.config_class = getattr(config_module, "AutotagConfig", None)
self.assertIsNotNone(self.config_class, "AutotagConfig is not implemented")
self._orig = PromptManagerConfig.get_config()
self._defaults = {
attr: getattr(self.config_class, attr, None)
for attr in (
"MODEL_TYPE",
"EXISTING",
"GENERAL_THRESHOLD",
"CHARACTER_THRESHOLD",
"TEMPLATE",
"SUGGESTION_LIMIT",
)
}
def tearDown(self):
self.config_class.SUGGESTION_LIMIT = 10
for attr, value in self._defaults.items():
setattr(self.config_class, attr, value)
PromptManagerConfig.update_config(self._orig)
def _load(self, value):
@@ -446,8 +458,23 @@ class TestAutotagConfig(unittest.TestCase):
self.assertEqual(self.config_class.suggestion_limit(), 10)
self.assertEqual(
PromptManagerConfig.get_config()["autotag"],
{"defaults": {"suggestion_limit": 10}},
{
"defaults": {
"model_type": "wd14-swinv2",
"existing": "skip",
"general_threshold": 0.35,
"character_threshold": 0.85,
"suggestion_limit": 10,
}
},
)
# The built-in template is not written; a saved one is.
self.config_class.TEMPLATE = "Saved."
self.assertEqual(
PromptManagerConfig.get_config()["autotag"]["defaults"]["template"],
"Saved.",
)
self.config_class.TEMPLATE = config_module.DEFAULT_TAG_TEMPLATE
with tempfile.TemporaryDirectory() as tmpdir:
path = os.path.join(tmpdir, "config.json")
PromptManagerConfig.update_config(
@@ -466,6 +493,80 @@ class TestAutotagConfig(unittest.TestCase):
)
self.assertEqual(self.config_class.suggestion_limit(), value)
def test_autotag_defaults_fall_back_per_key_on_load(self):
"""#328: each saved tagging default is checked alone on load; a bad
hand-edited value falls back to its default with one warning and
leaves the other keys as saved."""
good = {
"model_type": "gguf",
"existing": "merge",
"general_threshold": 0.4,
"character_threshold": 0.7,
"template": "Tags only.",
"suggestion_limit": 15,
}
fallback = {
"model_type": "wd14-swinv2",
"existing": "skip",
"general_threshold": 0.35,
"character_threshold": 0.85,
"template": config_module.DEFAULT_TAG_TEMPLATE,
"suggestion_limit": 10,
}
bad_values = {
"model_type": ["wd14-convnext", 3, None],
"existing": ["replace", "", True],
"general_threshold": [1.5, -0.1, float("nan"), "0.4", True],
"character_threshold": [2, float("inf"), None],
# 4000 characters are at most 16,000 UTF-8 bytes, so the 16 KiB
# byte cap can only be met past the character cap.
"template": ["", " ", "x" * 4001, "\U0001f600" * 4097, "\ud800", 7],
}
for key, values in bad_values.items():
for value in values:
with self.subTest(key=key, value=value):
saved = dict(good, **{key: value})
with tempfile.TemporaryDirectory() as tmpdir:
path = os.path.join(tmpdir, "config.json")
with open(path, "w") as f:
json.dump({"autotag": {"defaults": saved}}, f)
with self.assertLogs(
"prompt_manager.config", level="WARNING"
) as logs:
PromptManagerConfig.load_from_file(path)
self.assertEqual(len(logs.records), 1, logs.output)
self.assertIn(f"autotag.defaults.{key}", logs.output[0])
loaded = {
k: getattr(self.config_class, attr)
for k, (attr, _v, _d) in self.config_class._KEYS.items()
}
self.assertEqual(loaded, dict(good, **{key: fallback[key]}))
def test_tagging_validators_moved_with_their_messages(self):
validate_threshold = config_module.validate_threshold
validate_template = config_module.validate_template
self.assertEqual(validate_threshold("general_threshold", 1), 1.0)
for value, message in (
(float("nan"), "general_threshold must be a finite number"),
("0.4", "general_threshold must be a finite number"),
(1.01, "general_threshold must be between 0 and 1"),
):
with self.assertRaises(ValueError) as caught:
validate_threshold("general_threshold", value)
self.assertEqual(str(caught.exception), message)
self.assertEqual(validate_template("Tags."), "Tags.")
for value, message in (
(" ", "instructions must be non-empty text"),
("x" * 4001, "instructions must be at most 4000 characters"),
("\ud800", "instructions must be valid Unicode text"),
):
with self.assertRaises(ValueError) as caught:
validate_template(value)
self.assertEqual(str(caught.exception), message)
self.assertEqual(
config_module.MODEL_TYPES, {"gguf", "hf", "wd14-swinv2", "wd14-vit"}
)
def test_suggestion_limit_setting_validated_on_load(self):
for value in (0, 101, "10", float("nan"), 2.5, True, None, [10]):
with self.subTest(value=value):
+44
View File
@@ -1398,6 +1398,50 @@ class TestModelHousekeeping(DatabaseTestCase):
self.assertIsNot(self.db.model.get_connection(), failing)
class TestAutotagBatchCandidates(DatabaseTestCase):
"""#328: the batch's candidate read never loads workflow documents."""
def test_autotag_candidates_skip_workflow_data(self):
from utils.hashing import generate_prompt_hash
first = self.db.save_prompt(
text="a", tags=["one", "two"], prompt_hash=generate_prompt_hash("a")
)
second = self.db.save_prompt(text="b", prompt_hash=generate_prompt_hash("b"))
old = self.db.link_image_to_prompt(
first, "/out/a.png", {"workflow": {"big": "x" * 1000}}
)
new = self.db.link_image_to_prompt(second, "/out/b.png")
statements = []
conn = self.db.model.get_connection() # this thread's own connection
conn.set_trace_callback(statements.append)
try:
rows = self.db.autotag_batch_candidates()
finally:
conn.set_trace_callback(None)
self.assertEqual(
rows[0],
{
"id": new,
"image_path": "/out/b.png",
"prompt_id": second,
"prompt_tags": [],
},
)
self.assertEqual(
{**rows[1], "prompt_tags": sorted(rows[1]["prompt_tags"])},
{
"id": old,
"image_path": "/out/a.png",
"prompt_id": first,
"prompt_tags": ["one", "two"],
},
)
self.assertEqual(len(rows), 2)
self.assertTrue(statements)
self.assertFalse(any("workflow_data" in sql for sql in statements))
if __name__ == "__main__":
unittest.main()
+10 -3
View File
@@ -28,6 +28,8 @@ except ImportError:
Status = namedtuple("ExecutionStatus", "status_str completed messages")
JOB = "a0c26da3-eb26-405c-b143-c662a84c3112"
# A cross-thread wait or a duration check fails only past this (#370).
CEILING_SECONDS = 30.0
class HostQueue:
@@ -209,6 +211,9 @@ class HistoryHooksTests(unittest.TestCase):
self.queue.get = unrelated_after
self.dequeue(0)
self.queue.task_done(0, {}, status())
# The writer drains first: this test checks wrappers and the session,
# not the 2 s bound, which a slow runner can spend on the write (#370).
self.drain(first)
self.assertTrue(first.dispose())
self.assertTrue(first.dispose())
self.assertIs(self.queue.get, unrelated_after)
@@ -339,15 +344,17 @@ class HistoryHooksTests(unittest.TestCase):
with patch.object(self.db, "write_execution_start", blocked):
registration = self.register()
self.dequeue(0)
self.assertTrue(entered.wait(1))
# Ceilings, not tight wall clocks (#370): the return values show
# the bound; a ceiling still fails a call that waits unbounded.
self.assertTrue(entered.wait(CEILING_SECONDS))
began = time.monotonic()
self.queue.task_done(0, {}, status())
self.dequeue(1)
self.assertLess(time.monotonic() - began, 0.2)
self.assertLess(time.monotonic() - began, CEILING_SECONDS)
began = time.monotonic()
self.assertFalse(registration.dispose())
self.assertGreaterEqual(time.monotonic() - began, 1.8)
self.assertLess(time.monotonic() - began, 2.3)
self.assertLess(time.monotonic() - began, CEILING_SECONDS)
release.set()
self.assertTrue(registration.dispose(2))
+2
View File
@@ -642,6 +642,8 @@ class HistoryWriterTests(unittest.TestCase):
== "succeeded"
)
self.assertTrue(writer.submit(self.start_command(2)))
# Drain first: the 2 s bound is not under test here (#370).
drain_writer(writer)
self.assertTrue(writer.close(2))
self.assertEqual(len(set(threads)), 1)
self.assertNotEqual(threads[0], threading.get_ident())
+2
View File
@@ -245,6 +245,8 @@ class OutputHooksTests(fixtures.unittest.TestCase):
registration = self.register()
self.dequeue(0)
late = registration._task_done
# Drain first: the 2 s bound is not under test here (#370).
self.drain(registration)
self.assertTrue(registration.dispose())
thread = threading.Thread(
target=lambda: late((0, history(), fixtures.status()), {})
+3
View File
@@ -17,6 +17,7 @@ from utils.execution_contexts import MAX_ATTEMPT_BYTES
from utils.execution_output_capture import seal_outputs
from utils.execution_outputs import snapshot_output_roots
from utils.execution_history_writer import HistoryWriter
from tests.history_waits import drain_writer
try:
from tests import test_execution_history_writer as fixtures
@@ -161,6 +162,8 @@ class OutputWriterTests(fixtures.unittest.TestCase):
start = replace(self.start_command(), run_id=str(fixtures.uuid4()))
self.assertTrue(writer.submit(start))
self.finalization(writer, start)
# Drain first: the 2 s bound is not under test here (#370).
drain_writer(writer)
self.assertTrue(writer.close(2))
self.assertEqual(
self.db.get_execution_run(start.run_id)["outcome"], "succeeded"
+2
View File
@@ -12,6 +12,7 @@ import unittest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from tests.background_stops import stop_tracker # noqa: E402
from utils.hashing import (
generate_prompt_hash,
is_duplicate_prompt,
@@ -49,6 +50,7 @@ class TestSingleHashImplementation(unittest.TestCase):
db = mock.Mock()
db.get_prompt_by_hash.return_value = {"id": 42}
tracker = PromptTracker(db)
self.addCleanup(stop_tracker, self, tracker)
tracker.set_current_prompt(" Hello World ")
db.get_prompt_by_hash.assert_called_once_with(
generate_prompt_hash(" Hello World ")
+2
View File
@@ -8,6 +8,7 @@ import unittest.mock as mock
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from tests.background_stops import close_handler # noqa: E402
from database.operations import PromptDatabase
from utils.hashing import generate_prompt_hash
from utils.image_monitor import ImageGenerationHandler
@@ -52,6 +53,7 @@ class LinkingTestCase(unittest.TestCase):
)
def tearDown(self):
close_handler(self, self.handler) # and the tracker's cleanup thread
self.db.close_all() # Windows cannot unlink an open database file
for path in (
self.image,
+231 -12
View File
@@ -136,13 +136,21 @@ class TestMonitorSingletonLifecycle(unittest.TestCase):
class FakeObserver:
"""Records what the monitor does with the watchdog observer."""
"""Records what the monitor does with the watchdog observer.
def __init__(self):
``calls`` lists ``stop`` and ``join`` in order (shared with a test's own
recorders when given); ``stay_alive`` keeps it alive after ``join``;
``stop_error`` makes ``stop`` raise.
"""
def __init__(self, calls=None, stay_alive=False, stop_error=None):
self.scheduled = []
self.started = False
self.stopped = False
self.joined = False
self.calls = calls if calls is not None else []
self.stay_alive = stay_alive
self.stop_error = stop_error
def schedule(self, handler, path, recursive=False):
self.scheduled.append((handler, path, recursive))
@@ -151,12 +159,18 @@ class FakeObserver:
self.started = True
def stop(self):
self.calls.append(("stop",))
if self.stop_error is not None:
raise self.stop_error
self.stopped = True
def join(self, timeout=None):
self.calls.append(("join", timeout))
self.joined = True
def is_alive(self):
if self.stay_alive:
return self.started
return self.started and not self.stopped
@@ -171,6 +185,8 @@ class TestStartStop(unittest.TestCase):
self.missing = os.path.join(self.tmp.name, "does-not-exist")
def tearDown(self):
# Before the folder is removed (addCleanup runs after tearDown).
self.assertTrue(self.monitor.stop_monitoring())
im_mod._monitor_instance = None
def test_only_nonexistent_directories_leaves_monitoring_off(self):
@@ -225,9 +241,8 @@ class TestStartStop(unittest.TestCase):
worker = handler._worker
self.assertTrue(worker.is_alive())
self.monitor.stop_monitoring()
self.assertTrue(self.monitor.stop_monitoring())
worker.join(timeout=2)
self.assertFalse(worker.is_alive())
self.assertEqual(handler.pending_count(), 0)
# A restart gets a fresh handler and observer
@@ -303,6 +318,61 @@ class TestStartStop(unittest.TestCase):
self.assertEqual(self.monitor.monitored_directories, [self.tmp.name])
self.monitor.stop_monitoring()
def _start_with(self, fake):
with patch.object(im_mod, "Observer", return_value=fake):
self.monitor.start_monitoring([self.tmp.name])
handler = self.monitor.handler
real_close = handler.close
def close(timeout=im_mod.WORKER_JOIN_SECONDS):
fake.calls.append(("close",))
return real_close(timeout)
handler.close = close
return handler
def test_stop_monitoring_joins_observer_and_worker(self):
fake = FakeObserver()
handler = self._start_with(fake)
self.assertTrue(self.monitor.stop_monitoring(timeout=3.0))
self.assertEqual([c[0] for c in fake.calls], ["close", "stop", "join"])
join_timeout = fake.calls[2][1]
self.assertIsNotNone(join_timeout)
self.assertLessEqual(join_timeout, 3.0)
self.assertTrue(handler._closed)
def test_stop_monitoring_reports_an_observer_still_running(self):
fake = FakeObserver(stay_alive=True)
self._start_with(fake)
with self.assertLogs("prompt_manager.image_monitor", "WARNING") as logs:
self.assertFalse(self.monitor.stop_monitoring(timeout=0.05))
self.assertIn("watchdog observer still running", "\n".join(logs.output))
self.assertFalse(self.monitor.is_monitoring)
fake.stay_alive = False # it ends; tearDown's stop waits for it
def test_a_second_stop_monitoring_waits_for_the_observer(self):
fake = FakeObserver(stay_alive=True)
self._start_with(fake)
with self.assertLogs("prompt_manager.image_monitor", "WARNING"):
self.assertFalse(self.monitor.stop_monitoring(timeout=0.05))
self.assertFalse(self.monitor.stop_monitoring(timeout=0.05))
fake.stay_alive = False
self.assertTrue(self.monitor.stop_monitoring())
self.assertFalse(self.monitor._unfinished)
def test_stop_monitoring_reports_an_observer_stop_that_raised(self):
fake = FakeObserver(stop_error=OSError("descriptor gone"))
self._start_with(fake)
with self.assertLogs("prompt_manager.image_monitor", "WARNING") as logs:
self.assertFalse(self.monitor.stop_monitoring())
self.assertIn("Error while stopping image monitoring", "\n".join(logs.output))
fake.stop_error = None # tearDown's stop calls stop() again
def test_stop_monitoring_without_an_observer_returns_true(self):
self.assertTrue(self.monitor.stop_monitoring())
class FakeEvent:
def __init__(self, path, is_directory=False):
@@ -335,6 +405,10 @@ class TestHandlerScheduling(unittest.TestCase):
self.processed.append((path, prompt_snapshot))
)
def tearDown(self):
# Before the folder is removed (addCleanup runs after tearDown).
self.assertTrue(self.handler.close())
def _image(self, name, content=b"x" * 16):
path = os.path.join(self.tmp.name, name)
with open(path, "wb") as f:
@@ -399,23 +473,127 @@ class TestHandlerScheduling(unittest.TestCase):
empty = self._image("empty.png", content=b"")
self.assertFalse(self.handler.wait_until_settled(empty))
def test_shutdown_stops_the_worker_and_clears_pending(self):
def test_shutdown_joins_the_worker_within_the_bound(self):
self.handler.processing_delay = 10 # keep the first image waiting
self.handler.schedule(self._image("s1.png"))
self.handler.schedule(self._image("s2.png"))
worker = self.handler._worker
self.assertTrue(worker.is_alive())
self.handler.shutdown()
self.assertTrue(self.handler.shutdown())
worker.join(timeout=2)
self.assertFalse(worker.is_alive())
self.assertEqual(self.handler.pending_count(), 0)
def test_shutdown_reports_a_worker_still_busy(self):
entered, release = threading.Event(), threading.Event()
def busy(path, prompt_snapshot=None):
entered.set()
release.wait(30)
self.handler.process_new_image = busy
self.handler.schedule(self._image("busy.png"))
worker = self.handler._worker
self.addCleanup(worker.join, 30)
self.addCleanup(release.set)
self.assertTrue(entered.wait(30), "the worker never took the image")
with self.assertLogs("prompt_manager.image_monitor", "WARNING") as logs:
self.assertFalse(self.handler.shutdown(timeout=0.05))
self.assertIn("image worker still running", "\n".join(logs.output))
release.set()
worker.join(30)
self.assertFalse(worker.is_alive())
def test_a_second_stop_waits_for_a_worker_still_finishing(self):
entered, release = threading.Event(), threading.Event()
def busy(path, prompt_snapshot=None):
entered.set()
release.wait(30)
self.handler.process_new_image = busy
self.handler.schedule(self._image("twice.png"))
worker = self.handler._worker
self.addCleanup(worker.join, 30)
self.addCleanup(release.set)
self.assertTrue(entered.wait(30), "the worker never took the image")
with self.assertLogs("prompt_manager.image_monitor", "WARNING"):
self.assertFalse(self.handler.close(timeout=0.05))
self.assertFalse(self.handler.close(timeout=0.05))
release.set()
self.assertTrue(self.handler.close(timeout=30))
self.assertFalse(worker.is_alive())
def test_shutdown_on_the_worker_itself_only_signals(self):
entered, result = threading.Event(), {}
def stop_from_inside(path, prompt_snapshot=None):
result["value"] = self.handler.shutdown()
entered.set()
self.handler.process_new_image = stop_from_inside
self.handler.schedule(self._image("self.png"))
worker = self.handler._worker
self.assertTrue(entered.wait(30), "the worker never took the image")
worker.join(30)
self.assertFalse(result["value"])
self.assertFalse(worker.is_alive())
def _live_workers(self):
return {
t for t in threading.enumerate() if t.name == "PromptManagerImageWorker"
}
def _seal(self, handler, name):
handler._await_history(os.path.join(self.tmp.name, name))
handler._on_outputs_sealed(
"run-1",
[{"folder_type": "output", "subfolder": "", "filename": name}],
)
def test_close_refuses_new_work(self):
path = self._image("late.png")
with patch.object(im_mod, "host_output_directory", return_value=self.tmp.name):
self.assertTrue(self.handler.close())
before = self._live_workers()
self.assertFalse(self.handler.schedule(path))
# The image in hand can still ask for a history re-check and its
# run can be sealed afterwards, on the history writer's thread.
self.handler._await_history(path)
self.assertEqual(self.handler.awaiting_history(), [])
self._seal(self.handler, "late.png")
self.handler._schedule_recheck(path)
self.assertIsNone(self.handler._worker)
self.assertEqual(list(self.handler._history_due), [])
self.assertEqual(self._live_workers() - before, set())
def test_a_sealed_run_starts_a_worker_on_an_open_handler(self):
self._image("open.png")
with patch.object(im_mod, "host_output_directory", return_value=self.tmp.name):
self._seal(self.handler, "open.png")
self.assertIsNotNone(self.handler._worker)
self.assertTrue(self.handler.close())
def test_close_is_idempotent(self):
self.handler.schedule(self._image("c1.png"))
self.assertTrue(self.handler.close())
self.assertTrue(self.handler.close())
def test_on_created_after_close_takes_no_lock(self):
self.assertTrue(self.handler.close())
self.handler.on_created(FakeEvent(self._image("after.png")))
self.tracker.keyed_snapshot.assert_not_called()
self.assertEqual(self.handler.pending_count(), 0)
def test_shutdown_is_idempotent_and_safe_before_start(self):
fresh = im_mod.ImageGenerationHandler(MagicMock(), self.tracker)
fresh.shutdown()
fresh.shutdown()
self.assertTrue(fresh.shutdown())
self.assertTrue(fresh.shutdown())
self.assertTrue(fresh.close())
self.handler.schedule(self._image("s3.png"))
self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0))
self.handler.shutdown()
@@ -434,7 +612,7 @@ class TestHandlerScheduling(unittest.TestCase):
self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0))
self.assertIn(path, [p for p, _ in self.processed])
self.assertIsNot(self.handler._worker, old_worker)
self.handler.shutdown()
self.assertTrue(self.handler.close())
def test_processing_errors_do_not_kill_the_worker(self):
def explode(path, prompt_snapshot=None):
@@ -449,6 +627,47 @@ class TestHandlerScheduling(unittest.TestCase):
self.assertEqual(self.processed, [first, second])
class TestCloseReleasesHandles(unittest.TestCase):
"""``close`` joins the worker, so the test folder can be removed (#302)."""
def test_close_releases_handles(self):
from PIL import Image
from database.operations import PromptDatabase
from tests.open_handles import assert_closed
from utils.prompt_tracker import PromptTracker
with tempfile.TemporaryDirectory() as folder:
db = PromptDatabase(os.path.join(folder, "prompts.db"))
tracker = PromptTracker(db)
handler = im_mod.ImageGenerationHandler(db, tracker)
handler.processing_delay = 0
handler.settle_interval = 0.01
done = threading.Event()
real_process = handler.process_new_image
def process(path, prompt_snapshot=None):
try:
real_process(path, prompt_snapshot=prompt_snapshot)
finally:
done.set()
handler.process_new_image = process
path = os.path.join(folder, "real.png")
Image.new("RGB", (8, 8), "red").save(path)
try:
self.assertTrue(handler.schedule(path))
worker = handler._worker
self.assertTrue(done.wait(30), "the worker never processed the image")
finally:
closed, shut = handler.close(), tracker.shutdown()
db.model.close_all()
self.assertTrue(closed)
self.assertTrue(shut)
self.assertFalse(worker.is_alive())
assert_closed(self, folder)
class TestOutputListener(unittest.TestCase):
"""Every settled output is reported to the Images listing (#348, PR 2)."""
@@ -456,7 +675,7 @@ class TestOutputListener(unittest.TestCase):
self.tmp = tempfile.TemporaryDirectory()
self.addCleanup(self.tmp.cleanup)
self.handler = im_mod.ImageGenerationHandler(MagicMock(), MagicMock())
self.addCleanup(self.handler.shutdown)
self.addCleanup(lambda: self.assertTrue(self.handler.close()))
self.handler.processing_delay = 0
self.handler.settle_interval = 0.01
self.handler.process_new_image = lambda path, prompt_snapshot=None: None
@@ -549,7 +768,7 @@ class TestOutputListener(unittest.TestCase):
self.assertIsNone(monitor.output_listener)
self.assertIsNone(monitor.handler.output_listener)
finally:
monitor.stop_monitoring()
self.assertTrue(monitor.stop_monitoring())
class TestModuleIsolation(unittest.TestCase):
+35
View File
@@ -8,6 +8,7 @@ import ctypes
import json
import logging
import os
import shutil
import subprocess
import sys
import tempfile
@@ -694,6 +695,40 @@ class TestRevealRoute(ImageAPITestCase):
# Its own worker, never ComfyUI's shared default executor (L1).
self.assertTrue(self.threads[0].startswith("pm-reveal"), self.threads)
async def test_a_lora_root_reveals_only_linked_files(self):
# #355 security review: a LoRA root exposes only the files the import
# linked; any other file under it is refused before the opener runs.
lora = Path(tempfile.mkdtemp()).resolve()
self.addCleanup(shutil.rmtree, lora, True)
self.api._lora_image_dirs = lambda: [lora]
linked = make_png(lora / "abc" / "cover.png")
make_png(lora / "abc" / "unlinked.png")
self.api.db.link_image_to_prompt(self._save_prompt("lora"), str(linked))
scope = self.api._gallery_root_scope(lora, 1)
status, _data = await self.post(
{"root_index": 1, "root_scope": scope, "path": "abc/unlinked.png"}
)
self.assertEqual(status, 403)
self.assertEqual(self.opened, [])
status, data = await self.post(
{"root_index": 1, "root_scope": scope, "path": "abc/cover.png"}
)
self.assertEqual((status, data), (200, {"success": True}))
self.assertEqual(self.opened, [Path(os.path.realpath(linked))])
# The file that would open is checked too, whatever the request named
# (a path anchor could resolve it elsewhere in the root).
other = make_png(lora / "abc" / "other.png")
from py.api.images import ImageRoutesMixin
with patch.object(
ImageRoutesMixin, "_resolve_client_media_path", return_value=other
):
status, _data = await self.post(
{"root_index": 1, "root_scope": scope, "path": "abc/cover.png"}
)
self.assertEqual(status, 403)
self.assertEqual(len(self.opened), 1)
async def test_target_changed_before_launch_is_404(self):
make_png(self.output_dir / "sub" / "a.png")
self.outcome = file_reveal.RevealTargetChanged()
+1 -1
View File
@@ -310,7 +310,7 @@ class TestAutoTag(ImageTagJobTestCase):
self.assertTrue(
any(
line.endswith(
"is not downloaded. Download it from the legacy Auto-tag page."
"is not downloaded. Download it in System › Auto-tagging."
)
for line in job.log
),
+274 -12
View File
@@ -4,6 +4,8 @@ The package is loaded the way ComfyUI loads it (``__init__`` with relative
imports) with ComfyUI's ``server`` and ``folder_paths`` modules stubbed.
"""
import asyncio
import atexit
import contextlib
import importlib.util
import io
@@ -16,6 +18,7 @@ import threading
import types
import unittest
import unittest.mock as mock
from pathlib import Path
from aiohttp import web
@@ -40,13 +43,20 @@ class FakePromptQueue:
return self.results.pop(0) if self.results else None
def fake_server_module():
"""A ``server`` module whose PromptServer.instance accepts routes and hooks."""
def fake_server_module(app=None, queue_task_done=False):
"""A ``server`` module whose PromptServer.instance accepts routes and hooks.
``app`` replaces the default app (one without aiohttp hooks);
``queue_task_done`` gives the queue a ``task_done``, so history registers.
"""
queue = FakePromptQueue()
if queue_task_done:
queue.task_done = lambda *a, **k: None
instance = types.SimpleNamespace(
routes=web.RouteTableDef(),
prompt_queue=FakePromptQueue(),
prompt_queue=queue,
handlers=[],
app=types.SimpleNamespace(middlewares=[]),
app=app if app is not None else types.SimpleNamespace(middlewares=[]),
)
instance.add_on_prompt_handler = instance.handlers.append
module = types.ModuleType("server")
@@ -92,7 +102,7 @@ def _import_package(name, db_path, output_dir, server=None):
try:
spec.loader.exec_module(package)
except BaseException:
cleanup_package(package)
cleanup_package(package, require_stopped=False)
raise
finally:
for key, module in saved.items():
@@ -103,20 +113,34 @@ def _import_package(name, db_path, output_dir, server=None):
return package, server.PromptServer.instance
def cleanup_package(package):
"""Stop fixture-owned workers before releasing its other resources."""
def cleanup_package(package, require_stopped=True):
"""Stop fixture-owned workers before releasing its other resources.
The package's own teardown runs first (the reload case of #301); with
``require_stopped`` every part it reports must have stopped. The direct
stops after it also cover an import that failed before the teardown
existed, and a test that made a teardown step fail on purpose.
"""
teardown = getattr(package, "_stop_background_work", None)
if callable(teardown):
result = teardown()
if require_stopped:
assert all(result.values()), f"package teardown left work running: {result}"
monitor = getattr(package, "_global_image_monitor", None)
if monitor is not None:
assert monitor.stop_monitoring(), "package image monitor still running"
tracker = getattr(package, "_global_prompt_tracker", None)
if tracker is not None:
assert tracker.shutdown(), "package prompt cleanup thread still running"
registration = getattr(package, "_history_registration", None)
if registration is not None:
assert registration.dispose(3), "package history worker leaked"
monitor = getattr(package, "_global_image_monitor", None)
if monitor is not None:
monitor.stop_monitoring()
close_every_database()
sys.modules.pop(package.__name__, None)
@contextlib.contextmanager
def load_package(name, db_path, output_dir, server=None):
def load_package(name, db_path, output_dir, server=None, require_stopped=True):
"""Dispose package workers before closing its isolated logging handlers."""
root_logger = logging.getLogger("prompt_manager")
loggers = [root_logger] + [
@@ -144,7 +168,7 @@ def load_package(name, db_path, output_dir, server=None):
try:
yield package, instance
finally:
cleanup_package(package)
cleanup_package(package, require_stopped=require_stopped)
finally:
current_loggers = [root_logger] + [
logger
@@ -556,3 +580,241 @@ class TestPackageLoggerIsolation(unittest.TestCase):
self.assert_logging_restored(
captured["handlers"], captured["package"]._history_registration
)
PACKAGE_THREAD_NAMES = ("PromptManagerPromptCleanup", "PromptManagerImageWorker")
def package_threads():
"""Live threads this design owns: the tracker, the image worker, and the
watchdog observer and emitters (which have no name of their own)."""
from watchdog.utils import BaseThread
return {
thread
for thread in threading.enumerate()
if thread.name in PACKAGE_THREAD_NAMES or isinstance(thread, BaseThread)
}
def hooked_app():
return types.SimpleNamespace(
middlewares=[], on_cleanup=[], on_startup=[], on_shutdown=[]
)
class PackageTeardownCase(unittest.TestCase):
"""Each test loads its own package (#301): the teardown runs once."""
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
self.addCleanup(self.tmp.cleanup)
self.addCleanup(lambda: self.assertEqual(open_handles_under(self.tmp.name), []))
self.output_dir = os.path.join(self.tmp.name, "output")
os.makedirs(self.output_dir)
self.db_path = os.path.join(self.tmp.name, "prompts.db")
def load(self, server=None, require_stopped=True, suffix=""):
name = f"pm_teardown_{self._testMethodName}{suffix}"
stack = contextlib.ExitStack()
self.addCleanup(stack.close)
stack.enter_context(contextlib.redirect_stdout(io.StringIO()))
package, _ = stack.enter_context(
load_package(
name,
self.db_path,
self.output_dir,
server,
require_stopped=require_stopped,
)
)
return package
def busy_worker(self, package, name="busy.png"):
"""The image worker alive in its processing delay."""
handler = package._global_image_monitor.handler
self.assertIsNotNone(handler, "the package's monitor did not start")
handler.processing_delay = 3600
# Outside the watched folder, so the observer does not schedule it first.
path = os.path.join(self.tmp.name, name)
with open(path, "wb") as image:
image.write(b"x" * 16)
self.assertTrue(handler.schedule(path))
self.assertTrue(handler._worker.is_alive())
return handler._worker
class TestPackageTeardown(PackageTeardownCase):
def test_teardown_stops_every_part_once(self):
package = self.load()
monitor = package._global_image_monitor
observer = monitor.observer
tracker_thread = package._global_prompt_tracker.cleanup_thread
worker = self.busy_worker(package)
with self.assertLogs("prompt_manager.init", "INFO") as logs:
result = package._stop_background_work()
self.assertEqual(result, {"monitor": True, "thumbnails": True, "tracker": True})
self.assertEqual(list(result), ["monitor", "thumbnails", "tracker"])
self.assertIn("Background work stopped in", "\n".join(logs.output))
self.assertIn("monitor=ok thumbnails=ok tracker=ok", "\n".join(logs.output))
self.assertFalse(worker.is_alive())
self.assertFalse(observer.is_alive())
self.assertFalse(tracker_thread.is_alive())
result["monitor"] = False # a copy: the stored result is unchanged
with mock.patch.object(monitor, "stop_monitoring") as again:
self.assertEqual(
package._stop_background_work(),
{"monitor": True, "thumbnails": True, "tracker": True},
)
again.assert_not_called()
def test_teardown_detaches_the_monitor_callbacks(self):
package = self.load()
monitor = package._global_image_monitor
self.assertIsNotNone(monitor.review_recorder)
package._stop_background_work()
self.assertIsNone(monitor.review_recorder)
self.assertIsNone(monitor.thumbnail_queue)
self.assertIsNone(monitor.output_listener)
def test_an_image_in_hand_queues_no_preview_during_teardown(self):
package = self.load()
monitor = package._global_image_monitor
real_stop = monitor.stop_monitoring
image = os.path.join(self.output_dir, "in_hand.png")
from PIL import Image
Image.new("RGB", (8, 8), "red").save(image)
seen = {}
def stop_with_an_image_in_hand(timeout):
# The worker read ``thumbnail_queue`` before the detach and now
# hands its finished image over (at exit, Python 3.12+ refuses a
# new thread here and logs a traceback).
seen["queued"] = package.api.queue_capture_thumbnail(image)
seen["worker"] = package.api._auto_thumbnailer
return real_stop(timeout)
with (
mock.patch.object(
package.api,
"_get_all_output_dirs",
return_value=[Path(os.path.realpath(self.output_dir))],
),
mock.patch.object(
monitor, "stop_monitoring", side_effect=stop_with_an_image_in_hand
),
):
result = package._stop_background_work()
self.assertEqual(seen["queued"], "skipped")
self.assertIsNone(seen["worker"])
self.assertTrue(all(result.values()), result)
def test_teardown_removes_its_atexit_callback(self):
package = self.load()
with mock.patch("atexit.unregister", wraps=atexit.unregister) as unregister:
package._stop_background_work()
unregister.assert_called_once_with(package._stop_background_work)
def test_on_cleanup_runs_the_teardown(self):
app = hooked_app()
package = self.load(fake_server_module(app=app))
self.assertIs(app.on_cleanup[0], package._stop_background_work_on_cleanup)
self.assertIsNone(package._background_work_result)
async def run_cleanup():
for callback in list(app.on_cleanup):
await callback(app)
asyncio.run(run_cleanup())
self.assertEqual(
package._background_work_result,
{"monitor": True, "thumbnails": True, "tracker": True},
)
def test_teardown_runs_before_history_dispose_at_exit(self):
with mock.patch("atexit.register", wraps=atexit.register) as register:
package = self.load(fake_server_module(queue_task_done=True))
registered = [call.args[0] for call in register.call_args_list]
registration = package._history_registration
self.assertIsNotNone(registration, "history did not register")
self.assertIn(registration.dispose, registered)
self.assertIn(package._stop_background_work, registered)
# atexit runs newest first: the teardown before the history dispose.
self.assertGreater(
registered.index(package._stop_background_work),
registered.index(registration.dispose),
)
class TestPackageTeardownFailures(PackageTeardownCase):
"""A failing step never skips the later ones (design D14, R2)."""
def assert_tracker_stops_after(self, error):
package = self.load(require_stopped=False)
monitor = package._global_image_monitor
tracker_thread = package._global_prompt_tracker.cleanup_thread
with (
mock.patch.object(monitor, "stop_monitoring", side_effect=error),
self.assertLogs("prompt_manager.init", "WARNING") as logs,
):
result = package._stop_background_work()
self.assertFalse(result["monitor"])
self.assertTrue(result["tracker"])
self.assertFalse(tracker_thread.is_alive())
text = "\n".join(logs.output)
self.assertIn("Background work not stopped in time", text)
self.assertRegex(text, r"not stopped in time \([0-9.]+ s\): monitor$")
self.assertIn("Traceback", text)
def test_teardown_step_failure_does_not_skip_later_steps(self):
self.assert_tracker_stops_after(RuntimeError("stop failed"))
def test_second_interrupt_does_not_skip_the_tracker(self):
self.assert_tracker_stops_after(KeyboardInterrupt())
def test_teardown_shares_one_deadline(self):
package = self.load(require_stopped=False)
now = [100.0]
seen = []
def step(timeout, **_):
seen.append(timeout)
now[0] += 6.0
return True
with (
mock.patch.object(package, "_clock", lambda: now[0]),
mock.patch.object(
package._global_image_monitor, "stop_monitoring", side_effect=step
),
mock.patch.object(
package.api,
"close_auto_thumbnails",
side_effect=lambda final=False, timeout=None: step(timeout),
),
mock.patch.object(
package._global_prompt_tracker, "shutdown", side_effect=step
),
):
package._stop_background_work()
self.assertEqual(seen, [5.0, 4.0, 0.0])
class TestRepeatedLoads(PackageTeardownCase):
def test_repeated_loads_leave_no_threads(self):
for round_ in range(3):
before = package_threads()
with contextlib.redirect_stdout(io.StringIO()):
with load_package(
f"pm_teardown_repeat_{round_}",
self.db_path,
self.output_dir,
) as (package, _):
self.busy_worker(package, f"busy_{round_}.png")
started = package_threads() - before
self.assertTrue(started, "the package started no thread")
left = {t for t in package_threads() - before if t.is_alive()}
self.assertEqual(left, set(), f"round {round_}")
+1 -1
View File
@@ -460,7 +460,7 @@ class TestLegacyRoutesHoldTheRegistry(_LegacyHoldCase):
async def test_second_legacy_autotag_batch_keeps_its_json_409(self):
from py.api.autotag_routes import _BUSY_ERROR
self._run_job("autotag", "Auto-tag images (legacy page)")
self._run_job("autotag", "Tag output images (legacy page)")
resp = await self.client.request(
"POST", "/prompt_manager/autotag/start", json={}
)
+3
View File
@@ -29,6 +29,7 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database import library_move as lm # noqa: E402
from database import models # noqa: E402
from database.operations import PromptDatabase # noqa: E402
from tests.background_stops import close_handler # noqa: E402
from tests.open_handles import assert_closed # noqa: E402
from utils import library_location as ll # noqa: E402
from utils import usage_tracking # noqa: E402
@@ -204,6 +205,7 @@ class TestBackgroundWritesSurviveTheSwitch(_MoveCase):
image = os.path.join(self.comfy, "output", "late.png")
writer_db = PromptDatabase(self.old_db)
handler = ImageGenerationHandler(writer_db, mock.MagicMock())
self.addCleanup(close_handler, self, handler)
started = threading.Event()
probe = _CrossThreadCloseProbe()
outcome = {}
@@ -256,6 +258,7 @@ class TestBackgroundWritesSurviveTheSwitch(_MoveCase):
db = PromptDatabase(self.old_db)
handler = ImageGenerationHandler(db, mock.MagicMock())
self.addCleanup(close_handler, self, handler)
image = os.path.join(self.comfy, "output", "orphan.png")
with mock.patch.object(handler, "logger") as log:
handler.link_image_to_prompt(image, {"id": 99999}, {})
+312 -28
View File
@@ -252,7 +252,7 @@ class LoraImportAPITestCase(LoraAPITestCase):
self._patches = [
patch.object(
lora_utils, "find_lora_directories", return_value=[str(self.loras)]
lora_utils, "default_lora_roots", return_value=[str(self.loras)]
),
patch.object(
lora_utils, "get_lora_image_cache_dir", return_value=self.cache_dir
@@ -276,6 +276,8 @@ class LoraImportAPITestCase(LoraAPITestCase):
await super().tearDownAsync()
def _two_loras(self):
"""``neon`` has only a CivitAI image (downloaded); ``plain`` has a
local preview (linked where it is, nothing downloaded)."""
_write_metadata(
self.loras,
"neon",
@@ -283,8 +285,8 @@ class LoraImportAPITestCase(LoraAPITestCase):
images=["https://civitai.com/a.png"],
prompt="neon street",
)
(self.loras / "neon.png").write_bytes(_png_bytes())
_write_metadata(self.loras, "plain")
_write_metadata(self.loras, "plain", images=["https://civitai.com/p.png"])
(self.loras / "plain.png").write_bytes(_png_bytes())
async def _post(self, url, body=None):
resp = await self._post_json(url, {} if body is None else body)
@@ -321,8 +323,34 @@ class LoraImportRouteTests(LoraImportAPITestCase):
self.assertEqual(resp.status, 200, plan)
self.assertEqual(
{k: plan[k] for k in ("lora_files", "with_metadata", "current_imports")},
{"lora_files": 3, "with_metadata": 2, "current_imports": 2},
{
k: plan[k]
for k in (
"lora_files",
"with_metadata",
"skipped",
"with_local_media",
"current_imports",
)
},
{
"lora_files": 3,
"with_metadata": 2,
"skipped": 1,
"with_local_media": 1,
"current_imports": 2,
},
)
self.assertEqual(
plan["layout"],
{
"source": "environment",
"settings_file": str(self.lm_settings / "settings.json"),
"error": None,
"mode": "alongside",
"library": "default",
"example_images": False,
},
)
self.assertEqual(plan["blocked_by_pairs"], 0)
self.assertEqual(plan["planned_downloads"], 1)
@@ -336,7 +364,95 @@ class LoraImportRouteTests(LoraImportAPITestCase):
self.assertEqual(_size(db_path + "-shm"), shm)
self.assertEqual(_snapshot(self.api.db), rows)
self.assertNotIn(SECRET_KEY, json.dumps(plan))
self.assertNotIn(str(self.tmp_root), json.dumps(plan))
# The settings file is the only absolute path the review shows.
shown = dict(plan, layout=dict(plan["layout"], settings_file=None))
self.assertNotIn(str(self.tmp_root), json.dumps(shown))
async def test_unusable_settings_review_offers_no_token(self):
self._two_loras()
(self.lm_settings / "settings.json").write_text("[1, 2]")
resp, plan = await self._post(PREVIEW)
self.assertEqual(resp.status, 200, plan)
self.assertIsNone(plan["token"])
self.assertEqual(
plan["layout"]["error"], "the settings file is not a JSON object"
)
self.assertEqual((plan["lora_files"], plan["with_metadata"]), (0, 0))
self.assertEqual(self.requests, [])
async def test_centralized_review_counts_what_it_skips(self):
sidecars = self.tmp_root / "sidecars"
examples = self.tmp_root / "examples"
(self.lm_settings / "settings.json").write_text(
json.dumps(
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": str(sidecars),
"example_images_path": str(examples),
"libraries": {"comfyui": {}},
"active_library": "comfyui",
}
)
)
component = sidecars / "loras-0a1b2c3d"
component.mkdir(parents=True)
model = self.loras / "neon.safetensors"
model.write_bytes(b"m")
meta = {"model_name": "neon", "file_path": str(model), "sha256": "AB12"}
(component / "neon.metadata.json").write_text(json.dumps(meta))
(examples / "ab12").mkdir(parents=True)
(examples / "ab12" / "image_1.png").write_bytes(_png_bytes())
gone = dict(meta, file_path=str(self.loras / "gone.safetensors"))
(component / "gone.metadata.json").write_text(json.dumps(gone))
(component / "broken.metadata.json").write_text("{{{")
resp, plan = await self._post(PREVIEW)
self.assertEqual(resp.status, 200, plan)
self.assertEqual(
(
plan["lora_files"],
plan["with_metadata"],
plan["skipped"],
plan["with_local_media"],
plan["planned_downloads"],
),
(3, 1, 2, 1, 0),
)
self.assertEqual(
(plan["layout"]["mode"], plan["layout"]["library"]),
("centralized", "comfyui"),
)
self.assertTrue(plan["layout"]["example_images"])
resp, body = await self._post(APPLY, {"token": plan["token"]})
job = self.api._jobs.get(body["job"]["id"])
await job.task
self.assertEqual(job.state, "finished", job.log)
self.assertEqual(self.requests, [])
self.assertEqual(
[row["detail"] for row in job.outcomes["skipped"]],
["broken.metadata.json"],
)
imported = self.api.db.search_prompts(category=CATEGORY)
images = self.api.db.get_prompt_images(imported[0]["id"])
self.assertEqual(
[os.path.basename(i["image_path"]) for i in images], ["image_1.png"]
)
async def test_an_unlistable_folder_says_so_in_the_review(self):
self._two_loras()
error = PermissionError(13, "Permission denied", str(self.loras / "lost+found"))
with patch.object(lora_import, "discover_sidecars", side_effect=error):
resp, body = await self._post(PREVIEW)
self.assertEqual((resp.status, body["code"]), (500, "folder-unreadable"))
self.assertEqual(
body["error"],
"A LoRA folder could not be listed (Permission denied: lost+found). "
"Nothing was changed.",
)
self.assertNotIn(str(self.tmp_root), json.dumps(body))
async def test_preview_requires_the_integration(self):
IntegrationConfig.LORA_MANAGER_ENABLED = False
@@ -358,7 +474,7 @@ class LoraImportRouteTests(LoraImportAPITestCase):
self._two_loras()
resp, plan = await self._post(PREVIEW)
rows = _snapshot(self.api.db)
meta = self.loras / "plain.safetensors.metadata.json"
meta = self.loras / "plain.metadata.json"
meta.write_text(meta.read_text() + " ")
resp, body = await self._post(APPLY, {"token": plan["token"]})
@@ -397,7 +513,17 @@ class LoraImportRouteTests(LoraImportAPITestCase):
imported = self.api.db.search_prompts(category=CATEGORY)
self.assertEqual(sorted(p["text"] for p in imported), ["neon street", "plain"])
neon = next(p for p in imported if p["text"] == "neon street")
self.assertEqual(len(self.api.db.get_prompt_images(neon["id"])), 2)
self.assertEqual(len(self.api.db.get_prompt_images(neon["id"])), 1)
# plain's local preview is linked where it is; its CivitAI image is not
# downloaded (#355).
plain = next(p for p in imported if p["text"] == "plain")
self.assertEqual(
[i["image_path"] for i in self.api.db.get_prompt_images(plain["id"])],
[os.path.realpath(self.loras / "plain.png")],
)
self.assertEqual(
[r.full_url for r in self.requests], ["https://civitai.com/a.png"]
)
self.assertEqual(len(self.api.db.search_prompts(category="portraits")), 1)
# The token is single-use.
resp, body = await self._post(APPLY, {"token": plan["token"]})
@@ -409,7 +535,7 @@ class LoraImportRouteTests(LoraImportAPITestCase):
# Legacy scan: discovery raises once the stream is open.
with patch.object(
lora_utils, "find_lora_directories", side_effect=OSError("share gone")
lora_import, "discover_sidecars", side_effect=OSError("share gone")
):
resp = await self.client.request("POST", SCAN)
events = self._sse(await resp.text())
@@ -565,6 +691,7 @@ class LoraImportRouteTests(LoraImportAPITestCase):
async def test_legacy_scan_events_unchanged_on_success(self):
self._two_loras()
(self.loras / "broken.metadata.json").write_text("{{{")
(self.loras / "broken.safetensors").write_bytes(b"model")
resp = await self.client.request("POST", SCAN)
events = self._sse(await resp.text())
@@ -853,13 +980,13 @@ class LoraImportRefusalTests(LoraImportAPITestCase):
async def test_files_changed_while_staging_fail_the_job(self):
# I1: LoraManager rewrites a metadata file while the downloads run.
meta = self.loras / "plain.safetensors.metadata.json"
meta = self.loras / "plain.metadata.json"
await self._change_during_download(
lambda: meta.write_text(meta.read_text() + " ")
)
async def test_a_file_removed_while_staging_fails_the_job(self):
meta = self.loras / "plain.safetensors.metadata.json"
meta = self.loras / "plain.metadata.json"
await self._change_during_download(meta.unlink)
async def test_a_file_added_while_staging_fails_the_job(self):
@@ -869,7 +996,7 @@ class LoraImportRefusalTests(LoraImportAPITestCase):
self._two_loras()
before = _snapshot(self.api.db)
self.gate = threading.Event()
meta = self.loras / "plain.safetensors.metadata.json"
meta = self.loras / "plain.metadata.json"
resp = await self.client.request("POST", SCAN)
loop = asyncio.get_running_loop()
self.assertTrue(await loop.run_in_executor(None, self.downloading.wait, 10))
@@ -986,37 +1113,194 @@ class LoraImportRestoreTests(LoraImportAPITestCase):
class LoraImportStageTests(unittest.TestCase):
"""The stages over a disposable LoraManager install (#355)."""
def setUp(self):
self.tmp = Path(tempfile.mkdtemp())
self.tmp = Path(tempfile.mkdtemp()).resolve()
self.addCleanup(lambda: __import__("shutil").rmtree(self.tmp, True))
self.lm = _make_lora_manager(self.tmp / "custom_nodes")
self.loras = self.tmp / "loras"
self.loras.mkdir()
self.settings = self.tmp / "settings"
self.settings.mkdir()
self.roots = [str(self.loras)]
patchers = [
patch.object(
lora_utils,
"custom_nodes_directories",
return_value=[self.tmp / "custom_nodes"],
),
patch.object(lora_utils, "default_lora_roots", lambda: list(self.roots)),
patch.dict(os.environ, {"LORA_MANAGER_SETTINGS_DIR": str(self.settings)}),
]
for patcher in patchers:
patcher.start()
self.addCleanup(patcher.stop)
def settings_file(self, data):
(self.settings / "settings.json").write_text(json.dumps(data))
def discover(self):
return lora_import.discover(str(self.lm))
def test_discover_sorts_and_drops_repeats(self):
a = self.tmp / "a"
b = self.tmp / "a" / "b"
a = self.loras / "a"
b = self.loras / "a" / "b"
b.mkdir(parents=True)
_write_metadata(b, "zeta")
_write_metadata(a, "alpha")
with patch.object(
lora_utils, "find_lora_directories", return_value=[str(a), str(b)]
):
files = lora_import.discover(str(self.tmp))
self.assertEqual([f.name for f in files][0], "alpha.safetensors.metadata.json")
self.assertEqual(len(files), 2)
self.roots = [str(a), str(b)]
discovery = self.discover()
names = [s.metadata_path.name for s in discovery.sidecars]
self.assertEqual(names, ["alpha.metadata.json", "zeta.metadata.json"])
self.assertEqual(discovery.scanned, 2)
def test_discover_without_an_install_finds_nothing(self):
_write_metadata(self.loras, "alpha")
discovery = lora_import.discover(str(self.tmp))
self.assertEqual(discovery.layout.error, "no LoraManager install found")
self.assertEqual(discovery.sidecars, [])
def test_fingerprint_follows_size_and_mtime(self):
meta = _write_metadata(self.tmp, "one")
first = lora_import.fingerprint([meta])
self.assertEqual(lora_import.fingerprint([meta]), first)
meta = _write_metadata(self.loras, "one")
first = lora_import.fingerprint(self.discover())
self.assertEqual(lora_import.fingerprint(self.discover()), first)
meta.write_text(meta.read_text() + " ")
self.assertNotEqual(lora_import.fingerprint([meta]), first)
self.assertNotEqual(lora_import.fingerprint(self.discover()), first)
changed = lora_import.fingerprint(self.discover())
meta.unlink()
self.assertNotEqual(lora_import.fingerprint([meta]), first)
self.assertNotEqual(lora_import.fingerprint(self.discover()), changed)
def test_fingerprint_ignores_unrelated_settings(self):
_write_metadata(self.loras, "one")
self.settings_file({"civitai_api_key": "a"})
first = lora_import.fingerprint(self.discover())
self.settings_file({"civitai_api_key": "b", "theme": "dark"})
self.assertEqual(lora_import.fingerprint(self.discover()), first)
sidecars = self.tmp / "sidecars"
sidecars.mkdir()
self.settings_file(
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": str(sidecars),
}
)
self.assertNotEqual(lora_import.fingerprint(self.discover()), first)
def test_fingerprint_follows_a_missing_file_between_stat_calls(self):
meta = _write_metadata(self.loras, "one")
discovery = self.discover()
first = lora_import.fingerprint(discovery)
meta.unlink()
self.assertNotEqual(lora_import.fingerprint(discovery), first)
def test_stage_cancel_before_a_file_raises(self):
meta = _write_metadata(self.tmp, "one")
_write_metadata(self.loras, "one")
with self.assertRaises(lora_import.ImportCancelled):
lora_import.stage([meta], self.tmp, cancelled=lambda: True)
lora_import.stage(
self.discover(),
self.tmp,
media_extensions=MEDIA,
cancelled=lambda: True,
)
def test_local_media_means_no_download(self):
_write_metadata(self.loras, "neon", images=["https://civitai.com/a.png"])
(self.loras / "neon.mp4").write_bytes(b"video")
def refuse(*args, **kwargs):
raise AssertionError("a LoRA with local media downloads nothing")
with patch.object(lora_utils, "download_civitai_images", refuse):
staged = lora_import.stage(
self.discover(), self.tmp, media_extensions=MEDIA
)
self.assertEqual(staged.skipped, [])
self.assertEqual(
staged.records[0]["images"], [os.path.realpath(self.loras / "neon.mp4")]
)
def test_a_lora_without_local_media_downloads(self):
_write_metadata(self.loras, "neon", images=["https://civitai.com/a.png"])
calls = []
def download(metadata, path, cache_dir, api_key, stats=None):
calls.append(path.name)
return ["/cache/a.jpg"]
with patch.object(lora_utils, "download_civitai_images", download):
staged = lora_import.stage(
self.discover(), self.tmp, media_extensions=MEDIA
)
self.assertEqual(calls, ["neon.metadata.json"])
self.assertEqual(staged.records[0]["images"], ["/cache/a.jpg"])
def test_alongside_records_match_the_previous_import(self):
_write_metadata(
self.loras,
"neon",
words=["glow", "glow"],
images=["https://civitai.com/a.png"],
prompt="neon street",
)
(self.loras / "neon.png").write_bytes(b"x")
staged = lora_import.stage(self.discover(), self.tmp, media_extensions=MEDIA)
self.assertEqual(
staged.records,
[
{
"name": "neon",
"text": "neon street",
"hash": generate_prompt_hash("neon street"),
"tags": ["lora-manager", "lora:neon", "glow"],
"images": [os.path.realpath(self.loras / "neon.png")],
}
],
)
def test_preview_counts_local_media_and_downloads(self):
_write_metadata(self.loras, "local", images=["https://civitai.com/l.png"])
(self.loras / "local.webp").write_bytes(b"x")
_write_metadata(
self.loras,
"remote",
images=["https://civitai.com/r1.png", "https://civitai.com/r2.png"],
)
(self.loras / "orphan.metadata.json").write_text("{}")
found = lora_import.preview(str(self.lm), self.tmp / "cache", MEDIA)
self.assertEqual(
{
k: found[k]
for k in (
"lora_files",
"with_metadata",
"skipped",
"with_local_media",
"planned_downloads",
)
},
{
"lora_files": 3,
"with_metadata": 2,
"skipped": 1,
"with_local_media": 1,
"planned_downloads": 2,
},
)
self.assertEqual(found["names"], ["local", "remote"])
self.assertEqual(found["layout"]["mode"], "alongside")
def test_preview_of_an_unusable_layout(self):
(self.settings / "settings.json").write_text("nope")
_write_metadata(self.loras, "one")
found = lora_import.preview(str(self.lm), self.tmp / "cache", MEDIA)
self.assertEqual(
found["layout"]["error"], "the settings file is not valid JSON"
)
self.assertEqual((found["lora_files"], found["with_metadata"]), (0, 0))
MEDIA = frozenset({".png", ".jpg", ".jpeg", ".webp", ".mp4", ".webm"})
if __name__ == "__main__":
+858
View File
@@ -0,0 +1,858 @@
"""Where LoraManager keeps its files (#355): ``py/lora_manager_layout.py``.
Every test builds its own LoraManager install, settings file, LoRA roots,
sidecar root and example-images folder in a temporary folder, and injects the
environment and the user config folder: no test reads the developer's real
LoraManager settings or opens the network.
"""
import json
import ntpath
import os
import shutil
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from py import lora_manager_layout as layout_module
from py.lora_manager_layout import (
METADATA_SUFFIX,
PREVIEW_EXTENSIONS,
Sidecar,
discover_sidecars,
layout_summary,
load_layout,
local_media,
)
MEDIA = frozenset(
{".png", ".jpg", ".jpeg", ".webp", ".gif", ".mp4", ".webm", ".mov", ".bmp"}
)
MARKER = "marker-from-unkept-settings-0123"
class LayoutTestCase(unittest.TestCase):
"""A disposable LoraManager install and the folders its settings name."""
def setUp(self):
self.tmp = Path(tempfile.mkdtemp()).resolve()
self.addCleanup(shutil.rmtree, self.tmp, True)
self.plugin = self.tmp / "custom_nodes" / "ComfyUI-Lora-Manager"
(self.plugin / "py").mkdir(parents=True)
(self.plugin / "__init__.py").write_text("")
self.user_config = self.tmp / "user-config" / "ComfyUI-LoRA-Manager"
self.env_dir = self.tmp / "env-settings"
self.default_loras = self.tmp / "models" / "loras"
self.default_loras.mkdir(parents=True)
self.environ = {}
self.config_calls = []
def user_config_dir(self):
self.config_calls.append(True)
return str(self.user_config)
def write_settings(self, folder, data):
folder.mkdir(parents=True, exist_ok=True)
path = folder / "settings.json"
path.write_text(json.dumps(data), encoding="utf-8")
return path
def load(self, path="__plugin__", **kwargs):
lora_manager = str(self.plugin) if path == "__plugin__" else path
kwargs.setdefault("default_roots", lambda: [str(self.default_loras)])
kwargs.setdefault("environ", self.environ)
kwargs.setdefault("user_config_dir", self.user_config_dir)
return load_layout(lora_manager, **kwargs)
def model(self, folder, stem):
folder.mkdir(parents=True, exist_ok=True)
path = folder / f"{stem}.safetensors"
path.write_bytes(b"model")
return path
def sidecar(self, folder, stem, data=None):
folder.mkdir(parents=True, exist_ok=True)
path = folder / f"{stem}{METADATA_SUFFIX}"
path.write_text(json.dumps(data if data is not None else {}), "utf-8")
return path
class SettingsLookupTests(LayoutTestCase):
def test_environment_folder_wins_over_portable_and_user_config(self):
env_file = self.write_settings(self.env_dir, {"sidecar_storage_mode": "x"})
self.write_settings(self.plugin, {"use_portable_settings": True})
self.write_settings(self.user_config, {})
self.environ["LORA_MANAGER_SETTINGS_DIR"] = str(self.env_dir)
layout = self.load()
self.assertEqual(layout.source, "environment")
self.assertEqual(layout.settings_path, env_file)
self.assertIsNone(layout.error)
self.assertEqual(self.config_calls, [])
def test_an_empty_environment_folder_is_ignored(self):
self.environ["LORA_MANAGER_SETTINGS_DIR"] = ""
self.assertEqual(self.load().source, "user config")
def test_portable_flag_in_the_plugin_settings(self):
plugin_file = self.write_settings(self.plugin, {"use_portable_settings": True})
layout = self.load()
self.assertEqual(layout.source, "portable")
self.assertEqual(layout.settings_path, plugin_file)
self.assertEqual(self.config_calls, [])
def test_portable_environment_one_turns_it_on(self):
self.environ["LORA_MANAGER_PORTABLE"] = "1"
layout = self.load()
self.assertEqual(layout.source, "portable")
self.assertEqual(layout.settings_path, self.plugin / "settings.json")
def test_portable_environment_zero_turns_it_off_despite_the_file(self):
self.write_settings(self.plugin, {"use_portable_settings": True})
self.environ["LORA_MANAGER_PORTABLE"] = "0"
self.assertEqual(self.load().source, "user config")
def test_other_portable_environment_values_are_ignored(self):
for value in ("true", "yes"):
with self.subTest(value=value):
self.environ["LORA_MANAGER_PORTABLE"] = value
self.assertEqual(self.load().source, "user config")
self.write_settings(self.plugin, {"use_portable_settings": True})
self.assertEqual(self.load().source, "portable")
(self.plugin / "settings.json").unlink()
def test_a_string_portable_flag_is_ignored(self):
self.write_settings(self.plugin, {"use_portable_settings": "true"})
self.assertEqual(self.load().source, "user config")
def test_an_unreadable_plugin_file_means_not_portable(self):
(self.plugin / "settings.json").write_text("{not json")
self.assertEqual(self.load().source, "user config")
def test_user_config_folder_is_the_injected_one(self):
user_file = self.write_settings(self.user_config, {})
layout = self.load()
self.assertEqual(layout.source, "user config")
self.assertEqual(layout.settings_path, user_file)
self.assertEqual(self.config_calls, [True])
def test_the_real_user_config_folder_is_never_opened(self):
opened = []
real_open = open
def spy(path, *args, **kwargs):
opened.append(os.path.realpath(str(path)))
return real_open(path, *args, **kwargs)
with patch("builtins.open", spy):
layout = self.load()
self.assertEqual(layout.settings_path, self.user_config / "settings.json")
for path in opened:
self.assertTrue(path.startswith(str(self.tmp)), path)
def test_a_missing_settings_file_means_defaults(self):
layout = self.load()
self.assertIsNone(layout.error)
self.assertEqual(layout.mode, "alongside")
self.assertIsNone(layout.example_images_root)
self.assertIsNone(layout.sidecar_root)
self.assertEqual(layout.library, "default")
self.assertEqual(
layout.lora_roots, (os.path.normcase(str(self.default_loras)),)
)
def test_settings_dir_is_expanded_and_made_absolute(self):
self.write_settings(self.env_dir, {})
# Relative to the process folder (relpath fails across Windows drives).
# The cleanup returns before the temp folder is removed (LIFO).
self.addCleanup(os.chdir, os.getcwd())
os.chdir(self.tmp)
self.environ["LORA_MANAGER_SETTINGS_DIR"] = self.env_dir.name
self.assertEqual(self.load().settings_path, self.env_dir / "settings.json")
class UnusableLayoutTests(LayoutTestCase):
def test_no_install(self):
layout = self.load(path=None)
self.assertEqual(layout.error, "no LoraManager install found")
self.assertIsNone(layout.settings_path)
self.assertIsNone(layout.source)
# The default LoRA roots stay known for serving and trigger words.
self.assertEqual(
layout.lora_roots, (os.path.normcase(str(self.default_loras)),)
)
self.assertEqual(layout.media_roots, layout.lora_roots)
self.assertEqual(discover_sidecars(layout).sidecars, [])
def test_unreadable_settings(self):
self.write_settings(self.user_config, {})
with patch.object(layout_module, "_read_bytes", side_effect=PermissionError):
layout = self.load()
self.assertEqual(layout.error, "the settings file could not be read")
def test_settings_that_are_not_json(self):
self.user_config.mkdir(parents=True)
(self.user_config / "settings.json").write_text("{nope")
self.assertEqual(self.load().error, "the settings file is not valid JSON")
def test_settings_that_are_not_utf8(self):
self.user_config.mkdir(parents=True)
(self.user_config / "settings.json").write_bytes(b'{"a": "\xff"}')
self.assertEqual(self.load().error, "the settings file is not valid JSON")
def test_settings_that_are_not_an_object(self):
self.write_settings(self.user_config, ["a"])
self.assertEqual(self.load().error, "the settings file is not a JSON object")
def test_oversized_settings(self):
self.user_config.mkdir(parents=True)
with open(self.user_config / "settings.json", "wb") as handle:
handle.write(b" " * (layout_module.MAX_SETTINGS_BYTES + 1))
self.assertEqual(self.load().error, "the settings file is larger than 5 MB")
def test_platformdirs_import_failure(self):
def missing():
raise ImportError("no platformdirs")
layout = self.load(user_config_dir=missing)
self.assertEqual(layout.error, "platformdirs is not available")
self.assertEqual(layout.mode, "alongside")
def test_default_user_config_dir_needs_platformdirs(self):
with patch.dict(sys.modules, {"platformdirs": None}):
with self.assertRaises(ImportError):
layout_module._platform_user_config_dir()
def test_default_user_config_dir_uses_the_app_name(self):
fake = type(sys)("platformdirs")
calls = []
fake.user_config_dir = lambda *a, **k: calls.append((a, k)) or "/x"
with patch.dict(sys.modules, {"platformdirs": fake}):
self.assertEqual(layout_module._platform_user_config_dir(), "/x")
self.assertEqual(calls, [(("ComfyUI-LoRA-Manager",), {"appauthor": False})])
def test_missing_centralized_sidecar_root(self):
self.write_settings(
self.user_config,
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": str(self.tmp / "gone"),
},
)
layout = self.load()
self.assertEqual(layout.error, "the sidecar folder does not exist")
self.assertEqual(layout.mode, "centralized")
self.assertEqual(
layout.lora_roots, (os.path.normcase(str(self.default_loras)),)
)
self.assertEqual(layout.media_roots, layout.lora_roots)
self.assertEqual(discover_sidecars(layout).scanned, 0)
def test_a_centralized_sidecar_root_that_is_a_filesystem_root(self):
self.write_settings(
self.user_config,
{"sidecar_storage_mode": "centralized", "sidecar_storage_path": "/"},
)
layout = self.load()
self.assertEqual(
layout.error, "the sidecar folder is a filesystem root or the home folder"
)
self.assertNotIn(os.path.normcase(os.path.realpath("/")), layout.media_roots)
class SecretTests(LayoutTestCase):
def test_other_settings_keys_never_leave_the_module(self):
# Raw text keeps credential-named keys out of test variables.
self.user_config.mkdir(parents=True)
(self.user_config / "settings.json").write_text(
'{"civitai_api_key": "%s", "nested": {"token": "%s"}}' % (MARKER, MARKER),
encoding="utf-8",
)
layout = self.load()
self.assertNotIn(MARKER, repr(layout))
self.assertNotIn(MARKER, json.dumps(layout_summary(layout, Path.home())))
def test_only_the_seven_keys_are_kept(self):
self.assertEqual(
layout_module.SETTINGS_KEYS,
frozenset(
{
"sidecar_storage_mode",
"sidecar_storage_path",
"example_images_path",
"folder_paths",
"extra_folder_paths",
"libraries",
"active_library",
}
),
)
kept = layout_module._settings_subset(
json.loads('{"civitai_api_key": "%s", "libraries": {}}' % MARKER)
)
self.assertEqual(kept, {"libraries": {}})
class ActiveLibraryTests(LayoutTestCase):
def setUp(self):
super().setUp()
self.top = self.tmp / "top-loras"
self.lib_a = self.tmp / "a-loras"
self.lib_b = self.tmp / "b-loras"
self.extra = self.tmp / "b-extra"
for folder in (self.top, self.lib_a, self.lib_b, self.extra):
folder.mkdir()
self.examples = self.tmp / "examples"
self.examples.mkdir()
def settings(self, **extra):
data = {
"folder_paths": {"loras": [str(self.top)]},
"libraries": {
"first lib": {"folder_paths": {"loras": [str(self.lib_a)]}},
"second": {
"folder_paths": {"loras": [str(self.lib_b)]},
"extra_folder_paths": {"loras": [str(self.extra)]},
},
},
"example_images_path": str(self.examples),
}
data.update(extra)
self.write_settings(self.user_config, data)
def test_the_active_library_lists_are_used(self):
self.settings(active_library="second")
(self.examples / "second").mkdir()
layout = self.load()
self.assertEqual(layout.library, "second")
self.assertIn(os.path.normcase(str(self.lib_b)), layout.lora_roots)
self.assertIn(os.path.normcase(str(self.extra)), layout.lora_roots)
self.assertNotIn(os.path.normcase(str(self.top)), layout.lora_roots)
self.assertNotIn(os.path.normcase(str(self.lib_a)), layout.lora_roots)
self.assertIn(os.path.normcase(str(self.default_loras)), layout.lora_roots)
self.assertEqual(
layout.example_images_root, os.path.normcase(str(self.examples / "second"))
)
def test_an_unknown_active_library_falls_back_to_the_first(self):
self.settings(active_library="missing")
(self.examples / "first_lib").mkdir()
layout = self.load()
self.assertEqual(layout.library, "first lib")
self.assertIn(os.path.normcase(str(self.lib_a)), layout.lora_roots)
# The sanitized library name is the example-images subfolder.
self.assertEqual(
layout.example_images_root,
os.path.normcase(str(self.examples / "first_lib")),
)
def test_no_libraries_means_default_and_the_top_level_lists(self):
self.write_settings(
self.user_config,
{
"folder_paths": {"loras": str(self.top)},
"extra_folder_paths": {"loras": [str(self.extra), 7]},
"example_images_path": str(self.examples),
},
)
layout = self.load()
self.assertEqual(layout.library, "default")
self.assertIn(os.path.normcase(str(self.top)), layout.lora_roots)
self.assertIn(os.path.normcase(str(self.extra)), layout.lora_roots)
self.assertEqual(
layout.example_images_root, os.path.normcase(str(self.examples))
)
def test_one_library_has_no_library_subfolder(self):
self.write_settings(
self.user_config,
{
"libraries": {
"comfyui": {"folder_paths": {"loras": [str(self.lib_a)]}}
},
"active_library": "comfyui",
"example_images_path": str(self.examples),
},
)
layout = self.load()
self.assertEqual(layout.library, "comfyui")
self.assertEqual(
layout.example_images_root, os.path.normcase(str(self.examples))
)
def test_a_mixed_case_library_folder_keeps_its_example_images(self):
# Review finding: the library subfolder must be canonical like its
# root, or every example image fails the media check on Windows
# (normcase lowercases there; native proof is the Windows CI job).
self.write_settings(
self.user_config,
{
"libraries": {"Studio A": {}, "other": {}},
"active_library": "Studio A",
"example_images_path": str(self.examples),
},
)
image = self.examples / "Studio_A" / "abc" / "1.png"
image.parent.mkdir(parents=True)
image.write_bytes(b"x")
model = self.default_loras / "m.safetensors"
model.write_bytes(b"m")
(self.default_loras / "m.metadata.json").write_text("{}")
layout = self.load()
self.assertEqual(
layout.example_images_root,
os.path.normcase(os.path.realpath(self.examples / "Studio_A")),
)
entry = discover_sidecars(layout).sidecars[0]
self.assertEqual(
local_media(entry, {"sha256": "ABC"}, layout, MEDIA),
[os.path.realpath(image)],
)
def test_a_missing_library_subfolder_means_no_example_images(self):
self.settings(active_library="second")
self.assertIsNone(self.load().example_images_root)
def test_an_unset_or_missing_example_folder_means_none(self):
self.write_settings(self.user_config, {"example_images_path": " "})
self.assertIsNone(self.load().example_images_root)
self.write_settings(
self.user_config, {"example_images_path": str(self.tmp / "gone")}
)
self.assertIsNone(self.load().example_images_root)
def test_sanitize_library_name(self):
sanitize = layout_module._sanitize_library_name
self.assertEqual(sanitize("My Lib/2"), "My_Lib_2")
self.assertEqual(sanitize("ok-name_1.0"), "ok-name_1.0")
self.assertEqual(sanitize(""), "default")
def test_odd_library_shapes_are_ignored(self):
self.write_settings(
self.user_config,
{"libraries": {"x": "not a dict"}, "folder_paths": ["not", "a dict"]},
)
layout = self.load()
self.assertEqual(layout.library, "x")
self.assertEqual(
layout.lora_roots, (os.path.normcase(str(self.default_loras)),)
)
class CentralizedDiscoveryTests(LayoutTestCase):
def setUp(self):
super().setUp()
self.loras = self.tmp / "loras"
self.checkpoints = self.tmp / "checkpoints"
self.checkpoints.mkdir()
self.sidecars = self.tmp / "sidecars"
self.sidecars.mkdir()
self.write_settings(
self.user_config,
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": str(self.sidecars),
"folder_paths": {"loras": [str(self.loras)]},
},
)
self.loras.mkdir()
def test_component_folders_and_the_legacy_nested_form(self):
a = self.model(self.loras, "alpha")
b = self.model(self.loras / "sub", "beta")
self.sidecar(self.sidecars / "loras-1a2b3c4d", "alpha", {"file_path": str(a)})
self.sidecar(
self.sidecars / "comfyui" / "loras-1a2b3c4d" / "sub",
"beta",
{"file_path": b.as_posix()},
)
discovery = discover_sidecars(self.load())
self.assertEqual(discovery.scanned, 2)
self.assertEqual(
sorted(s.model_path.name for s in discovery.sidecars),
["alpha.safetensors", "beta.safetensors"],
)
self.assertEqual(discovery.unreadable, [])
self.assertEqual(
[str(s.metadata_path) for s in discovery.sidecars],
sorted(str(s.metadata_path) for s in discovery.sidecars),
)
def test_missing_models_and_checkpoints_are_skipped(self):
ckpt = self.model(self.checkpoints, "base")
self.sidecar(self.sidecars / "c", "base", {"file_path": str(ckpt)})
self.sidecar(
self.sidecars / "l",
"gone",
{"file_path": str(self.loras / "gone.safetensors")},
)
self.sidecar(self.sidecars / "l", "nopath", {"model_name": "x"})
discovery = discover_sidecars(self.load())
self.assertEqual(discovery.scanned, 3)
self.assertEqual(discovery.sidecars, [])
def test_leftover_sidecars_beside_models_are_not_scanned(self):
model = self.model(self.loras, "alpha")
self.sidecar(self.loras, "alpha", {"file_path": str(model)})
discovery = discover_sidecars(self.load())
self.assertEqual((discovery.scanned, discovery.sidecars), (0, []))
def test_unparseable_sidecars_are_recorded(self):
(self.sidecars / f"broken{METADATA_SUFFIX}").write_text("{{{")
(self.sidecars / f"list{METADATA_SUFFIX}").write_text("[1]")
discovery = discover_sidecars(self.load())
self.assertEqual(discovery.scanned, 2)
self.assertEqual(
sorted(discovery.unreadable),
[f"broken{METADATA_SUFFIX}", f"list{METADATA_SUFFIX}"],
)
def test_default_sidecar_root_is_beside_the_settings(self):
self.write_settings(
self.user_config,
{
"sidecar_storage_mode": "centralized",
"folder_paths": {"loras": [str(self.loras)]},
},
)
(self.user_config / "sidecars").mkdir()
layout = self.load()
self.assertEqual(
layout.sidecar_root, os.path.normcase(str(self.user_config / "sidecars"))
)
self.assertIn(layout.sidecar_root, layout.media_roots)
def test_an_unknown_mode_is_alongside(self):
self.write_settings(self.user_config, {"sidecar_storage_mode": "elsewhere"})
self.assertEqual(self.load().mode, "alongside")
def test_an_unlistable_folder_raises(self):
layout = self.load()
def refuse(error):
raise error
def failing_walk(top, onerror=None, followlinks=False):
onerror(PermissionError(13, "denied", top))
return iter(())
with patch.object(layout_module.os, "walk", failing_walk):
with self.assertRaises(OSError):
discover_sidecars(layout)
class AlongsideDiscoveryTests(LayoutTestCase):
def test_a_sidecar_with_its_model_is_kept(self):
model = self.model(self.default_loras, "alpha")
meta = self.sidecar(self.default_loras, "alpha")
self.sidecar(self.default_loras, "orphan")
discovery = discover_sidecars(self.load())
self.assertEqual(discovery.scanned, 2)
self.assertEqual(
discovery.sidecars,
[Sidecar(Path(os.path.realpath(meta)), Path(os.path.realpath(model)))],
)
def test_a_folder_reached_through_two_nested_roots_is_listed_once(self):
nested = self.default_loras / "nested"
self.model(nested, "alpha")
self.sidecar(nested, "alpha")
layout = self.load(default_roots=lambda: [str(self.default_loras), str(nested)])
discovery = discover_sidecars(layout)
self.assertEqual((discovery.scanned, len(discovery.sidecars)), (1, 1))
@unittest.skipUnless(hasattr(os, "symlink"), "symlinks")
def test_a_symlink_loop_is_walked_once(self):
self.model(self.default_loras, "alpha")
self.sidecar(self.default_loras, "alpha")
try:
os.symlink(self.default_loras, self.default_loras / "loop")
except (OSError, NotImplementedError):
self.skipTest("symlinks need privileges here")
discovery = discover_sidecars(self.load())
self.assertEqual((discovery.scanned, len(discovery.sidecars)), (1, 1))
class LocalMediaTests(LayoutTestCase):
def setUp(self):
super().setUp()
self.examples = self.tmp / "examples"
self.examples.mkdir()
self.write_settings(
self.user_config, {"example_images_path": str(self.examples)}
)
self.model(self.default_loras, "alpha")
self.meta = self.sidecar(self.default_loras, "alpha")
self.layout = self.load()
self.entry = discover_sidecars(self.layout).sidecars[0]
def media(self, metadata=None):
return local_media(self.entry, metadata or {}, self.layout, MEDIA)
def touch(self, path):
path.parent.mkdir(parents=True, exist_ok=True)
path.write_bytes(b"x")
return os.path.realpath(path)
def test_preview_url_comes_first(self):
preview = self.touch(self.default_loras / "custom.jpeg")
self.touch(self.default_loras / "alpha.webp")
self.assertEqual(self.media({"preview_url": preview}), [preview])
def test_preview_url_outside_media_roots_or_not_servable_is_ignored(self):
outside = self.touch(self.tmp / "elsewhere" / "p.png")
avif = self.touch(self.default_loras / "p.avif")
fallback = self.touch(self.default_loras / "alpha.png")
for url in (
outside,
avif,
"relative.png",
7,
str(self.default_loras / "x.png"),
):
with self.subTest(url=url):
self.assertEqual(self.media({"preview_url": url}), [fallback])
def test_suffix_order_follows_lora_manager(self):
self.assertEqual(
PREVIEW_EXTENSIONS,
(
".webp",
".preview.webp",
".preview.png",
".preview.jpeg",
".preview.jpg",
".preview.mp4",
".png",
".jpeg",
".jpg",
".mp4",
".gif",
".webm",
".avif",
".jxl",
),
)
mp4 = self.touch(self.default_loras / "alpha.mp4")
self.assertEqual(self.media(), [mp4])
webp = self.touch(self.default_loras / "alpha.preview.webp")
self.assertEqual(self.media(), [webp])
def test_an_avif_only_preview_gives_nothing(self):
self.touch(self.default_loras / "alpha.avif")
self.assertEqual(self.media(), [])
def test_example_images_in_numeric_order(self):
folder = self.examples / "abcdef"
names = ["image_10.png", "image_2.png", "image_1.jpg", "notes.txt"]
for name in names:
self.touch(folder / name)
(folder / "sub").mkdir()
preview = self.touch(self.default_loras / "alpha.png")
found = self.media({"sha256": "ABCDEF"})
self.assertEqual(
[os.path.basename(p) for p in found],
["alpha.png", "image_1.jpg", "image_2.png", "image_10.png"],
)
self.assertEqual(found[0], preview)
def test_missing_sha256_or_folder_gives_no_example_images(self):
self.touch(self.examples / "abc" / "1.png")
self.assertEqual(self.media(), [])
self.assertEqual(self.media({"sha256": "zzz"}), [])
self.assertEqual(self.media({"sha256": 5}), [])
def test_an_unlistable_example_folder_gives_what_was_found(self):
self.touch(self.examples / "abc" / "1.png")
with patch.object(layout_module.os, "scandir", side_effect=PermissionError):
self.assertEqual(self.media({"sha256": "abc"}), [])
def test_cover_moves_the_first_still_in_front_of_a_video(self):
video = self.touch(self.default_loras / "alpha.mp4")
still = self.touch(self.examples / "abc" / "2.png")
other = self.touch(self.examples / "abc" / "1.webm")
self.assertEqual(self.media({"sha256": "abc"}), [still, video, other])
def test_videos_only_keep_their_order(self):
video = self.touch(self.default_loras / "alpha.mp4")
self.assertEqual(self.media(), [video])
def test_repeats_are_dropped(self):
preview = self.touch(self.default_loras / "alpha.png")
self.assertEqual(self.media({"preview_url": preview}), [preview])
def test_video_suffixes_match_the_video_module(self):
from utils.video_metadata import VIDEO_EXTENSIONS
self.assertEqual(layout_module.VIDEO_SUFFIXES, frozenset(VIDEO_EXTENSIONS))
class FolderGuardTests(unittest.TestCase):
"""Portable path-flavor controls (TESTING.md Windows portability preflight)."""
def setUp(self):
self.tmp = Path(tempfile.mkdtemp()).resolve()
self.addCleanup(shutil.rmtree, self.tmp, True)
def guard_nt(self, path, home="C:\\Users\\Vito"):
return layout_module._guard_folder(
path, home, pathmod=ntpath, isdir=lambda _p: True
)
def test_filesystem_roots_and_home_are_rejected(self):
self.assertIsNone(layout_module._guard_folder("/", str(self.tmp)))
self.assertIsNone(self.guard_nt("C:\\"))
self.assertIsNone(self.guard_nt("c:/"))
self.assertIsNone(self.guard_nt("\\\\server\\share\\"))
self.assertIsNone(self.guard_nt("c:\\USERS\\vito"))
self.assertIsNone(self.guard_nt("C:\\Users\\Vito\\"))
self.assertEqual(
self.guard_nt("C:\\Users\\Vito\\Loras"), "c:\\users\\vito\\loras"
)
self.assertIsNone(layout_module._guard_folder(str(self.tmp), str(self.tmp)))
def test_ancestors_of_home_are_rejected(self):
# Security review (#355): a settings folder that contains the home
# folder would serve every user's files.
home = self.tmp / "home" / "me"
inside = home / "loras"
inside.mkdir(parents=True)
guard = layout_module._guard_folder
self.assertIsNone(guard(str(self.tmp / "home"), str(home)))
self.assertIsNone(guard(str(self.tmp), str(home)))
self.assertIsNone(guard(str(home), str(home)))
self.assertEqual(guard(str(inside), str(home)), os.path.normcase(str(inside)))
# A sibling sharing the name prefix is not an ancestor.
sibling = self.tmp / "home" / "me2"
sibling.mkdir()
self.assertEqual(guard(str(sibling), str(home)), os.path.normcase(str(sibling)))
def test_windows_ancestors_of_home_are_rejected(self):
self.assertIsNone(self.guard_nt("C:\\Users", home="C:\\Users\\me"))
self.assertIsNone(self.guard_nt("c:/users/", home="C:\\Users\\me"))
self.assertEqual(
self.guard_nt("C:\\Users\\me\\Loras", home="C:\\Users\\me"),
"c:\\users\\me\\loras",
)
self.assertEqual(
self.guard_nt("C:\\Users\\meta", home="C:\\Users\\me"),
"c:\\users\\meta",
)
def test_settings_folders_that_contain_home_are_not_served(self):
home = self.tmp / "home" / "me"
home.mkdir(parents=True)
settings = self.tmp / "settings"
settings.mkdir()
(settings / "settings.json").write_text(
json.dumps(
{
"folder_paths": {"loras": [str(self.tmp / "home")]},
"example_images_path": str(self.tmp),
}
)
)
with patch.object(layout_module.Path, "home", return_value=home):
layout = load_layout(
str(self.tmp),
default_roots=lambda: [],
environ={"LORA_MANAGER_SETTINGS_DIR": str(settings)},
user_config_dir=lambda: str(self.tmp / "unused"),
)
self.assertIsNone(layout.error)
self.assertEqual(layout.lora_roots, ())
self.assertIsNone(layout.example_images_root)
self.assertEqual(layout.media_roots, ())
def test_missing_folders_add_nothing(self):
self.assertIsNone(layout_module._guard_folder(str(self.tmp / "x"), "/home"))
self.assertIsNone(
layout_module._guard_folder("\\\\server\\share\\loras", str(self.tmp))
)
self.assertEqual(
layout_module._guard_folder(str(self.tmp), "/nonexistent-home"),
os.path.normcase(str(self.tmp)),
)
def test_windows_model_path_matches_a_lower_case_root(self):
self.assertTrue(
layout_module._inside(
"C:/Models/Loras/A.safetensors", ("c:\\models\\loras",), ntpath
)
)
self.assertFalse(
layout_module._inside(
"C:/Models/LorasX/A.safetensors", ("c:\\models\\loras",), ntpath
)
)
self.assertFalse(
layout_module._inside(
"\\\\other\\share\\A.safetensors", ("c:\\models\\loras",), ntpath
)
)
def test_mixed_case_native_containment(self):
inside = str(self.tmp / "Loras" / "a.png")
root = os.path.normcase(str(self.tmp / "Loras"))
self.assertTrue(layout_module._inside(inside, (root,)))
self.assertFalse(layout_module._inside(str(self.tmp / "other.png"), (root,)))
class SummaryTests(LayoutTestCase):
def test_summary_names_what_the_import_read(self):
examples = self.tmp / "ex"
examples.mkdir()
sidecars = self.plugin / "sidecars"
sidecars.mkdir()
self.write_settings(
self.plugin,
{
"use_portable_settings": True,
"sidecar_storage_mode": "centralized",
"example_images_path": str(examples),
"libraries": {"comfyui": {}},
},
)
layout = self.load()
summary = layout_summary(layout, self.tmp)
self.assertEqual(
summary,
{
"source": "portable",
"settings_file": os.path.join(
"~", "custom_nodes", "ComfyUI-Lora-Manager", "settings.json"
),
"error": None,
"mode": "centralized",
"library": "comfyui",
"example_images": True,
},
)
def test_summary_outside_home_and_without_settings(self):
summary = layout_summary(self.load(path=None), self.tmp / "home")
self.assertEqual(summary["settings_file"], None)
self.assertEqual(summary["error"], "no LoraManager install found")
self.assertFalse(summary["example_images"])
layout = self.load()
self.assertEqual(
layout_summary(layout, self.tmp / "home")["settings_file"],
str(self.user_config / "settings.json"),
)
if __name__ == "__main__":
unittest.main()
+173 -52
View File
@@ -1263,11 +1263,12 @@ class TestTriggerWordCache(unittest.TestCase):
"file_name": "my_lora.safetensors",
"civitai": {"trainedWords": ["trigger1", "trigger2"]},
}
meta_path = Path(tmpdir) / "my_lora.safetensors.metadata.json"
meta_path = Path(tmpdir) / "my_lora.metadata.json"
meta_path.write_text(json.dumps(meta))
(Path(tmpdir) / "my_lora.safetensors").write_bytes(b"m")
# Patch find_lora_directories to return our temp dir
with patch("py.lora_utils.find_lora_directories", return_value=[tmpdir]):
# No LoraManager install: the default LoRA folders, alongside.
with patch("py.lora_utils.default_lora_roots", return_value=[tmpdir]):
count = self.cache.load(tmpdir)
self.assertTrue(self.cache.is_loaded)
@@ -1283,11 +1284,10 @@ class TestTriggerWordCache(unittest.TestCase):
"file_name": "MyLoRA.safetensors",
"civitai": {"trainedWords": ["word1"]},
}
(Path(tmpdir) / "MyLoRA.safetensors.metadata.json").write_text(
json.dumps(meta)
)
(Path(tmpdir) / "MyLoRA.metadata.json").write_text(json.dumps(meta))
(Path(tmpdir) / "MyLoRA.safetensors").write_bytes(b"m")
with patch("py.lora_utils.find_lora_directories", return_value=[tmpdir]):
with patch("py.lora_utils.default_lora_roots", return_value=[tmpdir]):
self.cache.load(tmpdir)
self.assertEqual(self.cache.get_trigger_words("mylora"), ["word1"])
@@ -1512,7 +1512,7 @@ class TestDetectLoraManager(unittest.TestCase):
self.assertEqual(lora_utils.detect_lora_manager(), str(lm.resolve()))
class TestFindLoraDirectories(unittest.TestCase):
class TestDefaultLoraRoots(unittest.TestCase):
def setUp(self):
self.tmp = Path(tempfile.mkdtemp())
@@ -1520,7 +1520,7 @@ class TestFindLoraDirectories(unittest.TestCase):
self.root = _make_comfy_root(self.tmp)
self.lm = _make_lora_manager(self.root / "custom_nodes")
def test_collects_default_extra_runtime_and_metadata_dirs(self):
def test_collects_default_extra_and_runtime_dirs_only(self):
extra_abs = self.tmp / "extra_abs"
extra_abs.mkdir()
(self.root / "rel_loras").mkdir()
@@ -1542,14 +1542,14 @@ class TestFindLoraDirectories(unittest.TestCase):
{"folder_paths": _FakeFolderPaths(str(self.root), [str(runtime)])},
),
):
dirs = lora_utils.find_lora_directories(str(self.lm))
dirs = lora_utils.default_lora_roots()
# #355: folders inside the LoraManager install are no longer roots.
expected = {
str((self.root / "models" / "loras").resolve()),
str((self.root / "rel_loras").resolve()),
str(extra_abs.resolve()),
str(runtime.resolve()),
str(nested.resolve()),
}
self.assertEqual(set(dirs), expected)
self.assertEqual(dirs, sorted(dirs))
@@ -1560,12 +1560,21 @@ class TestFindLoraDirectories(unittest.TestCase):
(self.root / "extra_model_paths.yml").write_text("a: [unclosed\n")
self.assertEqual(lora_utils._get_extra_lora_paths(self.root), [])
def test_without_root_only_metadata_dirs_are_found(self):
def test_without_root_nothing_is_found(self):
with (
patch.object(lora_utils, "find_comfyui_root", return_value=None),
patch.dict(sys.modules, {"folder_paths": None}),
):
self.assertEqual(lora_utils.find_lora_directories(str(self.lm)), [])
self.assertEqual(lora_utils.default_lora_roots(), [])
def test_the_removed_helpers_are_gone(self):
for name in (
"find_lora_directories",
"get_preview_images_from_metadata",
"get_preview_image_from_metadata",
"get_example_images_dir",
):
self.assertFalse(hasattr(lora_utils, name), name)
class TestModelNameAndPreviews(unittest.TestCase):
@@ -1581,44 +1590,6 @@ class TestModelNameAndPreviews(unittest.TestCase):
self.assertEqual(get_name({"file_name": "f.safetensors"}), "f.safetensors")
self.assertEqual(get_name({}), "unknown")
def test_preview_images_found_by_stem(self):
meta_path = self.tmp / "mylora.safetensors.metadata.json"
meta_path.write_text("{}")
(self.tmp / "mylora.png").write_bytes(b"x")
(self.tmp / "mylora.preview.jpg").write_bytes(b"x")
found = lora_utils.get_preview_images_from_metadata({}, meta_path)
self.assertEqual(
sorted(os.path.basename(f) for f in found),
["mylora.png", "mylora.preview.jpg"],
)
self.assertEqual(
lora_utils.get_preview_image_from_metadata({}, meta_path), found[0]
)
def test_preview_uses_file_name_from_metadata(self):
meta_path = self.tmp / "other.metadata.json"
(self.tmp / "named.webp").write_bytes(b"x")
found = lora_utils.get_preview_images_from_metadata(
{"file_name": "named.safetensors"}, meta_path
)
self.assertEqual([os.path.basename(f) for f in found], ["named.webp"])
self.assertIsNone(lora_utils.get_preview_image_from_metadata({}, meta_path))
def test_example_images_dir(self):
self.assertIsNone(lora_utils.get_example_images_dir(str(self.tmp)))
nested = self.tmp / "user" / "example_images"
nested.mkdir(parents=True)
self.assertEqual(
lora_utils.get_example_images_dir(str(self.tmp)), str(nested.resolve())
)
direct = self.tmp / "example_images"
direct.mkdir()
self.assertEqual(
lora_utils.get_example_images_dir(str(self.tmp)), str(direct.resolve())
)
class TestDownloadEdgeCases(unittest.TestCase):
@@ -1703,11 +1674,161 @@ class TestInjectTriggerWords(unittest.TestCase):
(tmp / "named.metadata.json").write_text(
json.dumps({"civitai": {"trainedWords": ["w"]}})
)
for stem in ("bad", "nowords", "named"):
(tmp / f"{stem}.safetensors").write_bytes(b"m")
cache = TriggerWordCache()
with patch.object(lora_utils, "find_lora_directories", return_value=[str(tmp)]):
with patch.object(lora_utils, "default_lora_roots", return_value=[str(tmp)]):
self.assertEqual(cache.load(str(tmp)), 1)
self.assertEqual(cache.get_trigger_words("named"), ["w"])
class TestTriggerWordCacheLayout(unittest.TestCase):
"""#355: trigger words come from the sidecars LoraManager really uses."""
def setUp(self):
self.tmp = Path(tempfile.mkdtemp()).resolve()
self.addCleanup(shutil.rmtree, self.tmp, True)
self.custom_nodes = self.tmp / "custom_nodes"
self.lm = _make_lora_manager(self.custom_nodes)
self.loras = self.tmp / "loras"
self.loras.mkdir()
self.sidecars = self.tmp / "sidecars"
self.settings = self.tmp / "settings"
self.settings.mkdir()
patchers = [
patch.object(
lora_utils, "custom_nodes_directories", return_value=[self.custom_nodes]
),
patch.object(lora_utils, "default_lora_roots", return_value=[]),
patch.dict(os.environ, {"LORA_MANAGER_SETTINGS_DIR": str(self.settings)}),
]
for patcher in patchers:
patcher.start()
self.addCleanup(patcher.stop)
def centralized(self):
self.sidecars.mkdir(exist_ok=True)
(self.settings / "settings.json").write_text(
json.dumps(
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": str(self.sidecars),
"folder_paths": {"loras": [str(self.loras)]},
}
)
)
def sidecar(self, folder, stem, model, words):
folder.mkdir(parents=True, exist_ok=True)
meta = {
"file_name": stem,
"file_path": str(model),
"civitai": {"trainedWords": words},
}
(folder / f"{stem}.metadata.json").write_text(json.dumps(meta))
def test_centralized_sidecars_are_read(self):
self.centralized()
model = self.loras / "neon.safetensors"
model.write_bytes(b"m")
self.sidecar(self.sidecars / "loras-0a1b2c3d", "neon", model, ["glow"])
# A leftover sidecar beside the model is not what LoraManager reads.
self.sidecar(self.loras, "neon", model, ["stale"])
cache = TriggerWordCache()
self.assertEqual(cache.load(str(self.lm)), 1)
self.assertEqual(cache.get_trigger_words("neon"), ["glow"])
def test_both_the_model_and_the_sidecar_name_are_keys(self):
self.centralized()
model = self.loras / "model_name.safetensors"
model.write_bytes(b"m")
self.sidecar(self.sidecars / "c", "sidecar_name", model, ["w"])
cache = TriggerWordCache()
self.assertEqual(cache.load(str(self.lm)), 2)
self.assertEqual(cache.get_trigger_words("model_name"), ["w"])
self.assertEqual(cache.get_trigger_words("sidecar_name"), ["w"])
def test_an_empty_path_means_auto_detection(self):
self.centralized()
model = self.loras / "neon.safetensors"
model.write_bytes(b"m")
self.sidecar(self.sidecars / "c", "neon", model, ["glow"])
cache = TriggerWordCache()
with patch.object(lora_utils, "detect_lora_manager", return_value=str(self.lm)):
self.assertEqual(cache.load(""), 1)
def test_an_unlistable_folder_keeps_the_previous_map(self):
cache = TriggerWordCache()
cache._cache = {"old": ["word"]}
cache._loaded = True
with (
patch.object(
lora_utils, "discover_sidecars", side_effect=OSError("share gone")
),
self.assertLogs("prompt_manager.lora_utils", level="WARNING") as logs,
):
self.assertEqual(cache.load(str(self.lm)), 1)
self.assertEqual(cache.get_trigger_words("old"), ["word"])
self.assertIn("share gone", "\n".join(logs.output))
def test_a_failed_load_waits_before_a_node_retries(self):
# Review finding: a node calls load while the cache is not loaded; a
# failing discovery must not rescan on every prompt run.
cache = TriggerWordCache()
self.assertTrue(cache.retry_due())
clock = [100.0]
with (
patch.object(lora_utils.time, "monotonic", lambda: clock[0]),
patch.object(lora_utils, "discover_sidecars", side_effect=OSError("x")),
self.assertLogs("prompt_manager.lora_utils", level="WARNING"),
):
cache.load(str(self.lm))
self.assertFalse(cache.is_loaded)
self.assertFalse(cache.retry_due())
clock[0] += lora_utils.TRIGGER_RETRY_SECONDS
self.assertTrue(cache.retry_due())
cache.load(str(self.lm))
self.assertTrue(cache.is_loaded)
self.assertTrue(cache.retry_due())
def test_any_other_failure_never_raises(self):
cache = TriggerWordCache()
with (
patch.object(lora_utils, "load_layout", side_effect=RuntimeError("bad")),
self.assertLogs("prompt_manager.lora_utils", level="WARNING"),
):
self.assertEqual(cache.load(str(self.lm)), 0)
self.assertFalse(cache.is_loaded)
def test_an_unusable_layout_falls_back_to_the_default_roots(self):
(self.settings / "settings.json").write_text("{not json")
model = self.loras / "neon.safetensors"
model.write_bytes(b"m")
self.sidecar(self.loras, "neon", model, ["glow"])
cache = TriggerWordCache()
with (
patch.object(
lora_utils, "default_lora_roots", return_value=[str(self.loras)]
),
self.assertLogs("prompt_manager.lora_utils", level="WARNING") as logs,
):
self.assertEqual(cache.load(str(self.lm)), 1)
self.assertEqual(cache.get_trigger_words("neon"), ["glow"])
self.assertIn("not valid JSON", "\n".join(logs.output))
def test_layout_wraps_the_install_and_the_default_roots(self):
with patch.object(
lora_utils, "default_lora_roots", return_value=[str(self.loras)]
):
layout = lora_utils.lora_manager_layout(str(self.lm))
self.assertIsNone(layout.error)
self.assertEqual(layout.source, "environment")
self.assertIn(os.path.normcase(str(self.loras)), layout.lora_roots)
self.assertEqual(
lora_utils.lora_manager_layout(str(self.tmp)).error,
"no LoraManager install found",
)
if __name__ == "__main__":
unittest.main()
+37
View File
@@ -315,6 +315,43 @@ class TestLoraTriggerWords(NodeTestCase):
)
cache.load.assert_called_once_with("/loras")
def test_a_failing_discovery_logs_and_the_node_still_runs(self):
# #355: the trigger cache reads LoraManager's folders while a node
# runs; a folder that cannot be listed must never stop the node.
from py import lora_utils
cache = lora_utils.get_trigger_cache()
cache.clear()
self.addCleanup(cache.clear)
with (
mock.patch.dict(sys.modules, {"py.config": fake_integration_config(True)}),
mock.patch.object(
lora_utils, "discover_sidecars", side_effect=OSError("share gone")
),
self.assertLogs("prompt_manager.lora_utils", level="WARNING") as logs,
):
text = self.node._inject_lora_trigger_words("a cat <lora:neon:1>")
self.assertEqual(text, "a cat <lora:neon:1>")
self.assertIn("share gone", "\n".join(logs.output))
def test_a_failed_load_is_not_retried_on_every_run(self):
from py import lora_utils
cache = lora_utils.get_trigger_cache()
cache.clear()
self.addCleanup(cache.clear)
with (
mock.patch.dict(sys.modules, {"py.config": fake_integration_config(True)}),
mock.patch.object(
lora_utils, "discover_sidecars", side_effect=OSError("share gone")
) as discover,
mock.patch.object(lora_utils, "default_lora_roots", return_value=[]),
self.assertLogs("prompt_manager.lora_utils", level="WARNING"),
):
for _ in range(3):
self.node._inject_lora_trigger_words("a cat <lora:neon:1>")
self.assertEqual(discover.call_count, 1)
class TestPromptManagerEncoding(NodeTestCase):
node_cls = PromptManager
+2
View File
@@ -24,6 +24,7 @@ from unittest.mock import patch
from PIL import Image, PngImagePlugin
from database.operations import PromptDatabase
from tests.background_stops import close_handler
from tests.history_waits import drain_writer
from tests.test_api_admin import AdminAPITestCase
from tests.test_api_images import ImageAPITestCase
@@ -224,6 +225,7 @@ class TestMonitorIngestsNonFiniteMetadata(StrictRows, unittest.TestCase):
self.handler.logger.disabled = True
def tearDown(self):
close_handler(self, self.handler)
self.db.close_all()
assert_closed(self, self.temp.name)
self.temp.cleanup()
+142
View File
@@ -0,0 +1,142 @@
"""The package teardown runs at a normal interpreter exit (#301, #302).
ComfyUI catches Ctrl+C in ``main.py`` and returns, so its exit is a normal
return: ``atexit`` callbacks then run, newest first. A child process imports
the package with ComfyUI stubbed, leaves the image worker alive in its
processing delay and returns. Its own ``atexit`` reporter, registered before
the package, runs after the teardown and prints the result as one JSON line.
"""
import json
import os
import subprocess
import sys
import tempfile
import textwrap
import unittest
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
# The parent never waits longer than this for the child (the hang watchdog
# allows 120 s per test).
CHILD_CEILING_SECONDS = 60
REPORT_PREFIX = "PM_EXIT_REPORT "
CHILD = textwrap.dedent("""
import atexit
import contextlib
import io
import json
import os
import sys
import threading
root, tmp = sys.argv[1], sys.argv[2]
sys.path.insert(0, root)
state = {}
def report():
from watchdog.utils import BaseThread
package = state.get("package")
result = getattr(package, "_background_work_result", None)
# The image worker, the tracker, and watchdog's observer and emitters.
alive = sorted(
t.name
for t in threading.enumerate()
if t.name in ("PromptManagerImageWorker", "PromptManagerPromptCleanup")
or isinstance(t, BaseThread)
)
print("PM_EXIT_REPORT " + json.dumps({"result": result, "alive": alive}))
sys.stdout.flush()
# First registered, so it runs last: after the package's teardown and
# after the history dispose.
atexit.register(report)
from tests.test_init_package import _import_package, fake_server_module
output = os.path.join(tmp, "output")
with contextlib.redirect_stdout(io.StringIO()):
package, _ = _import_package(
"pm_exit_child",
os.path.join(tmp, "prompts.db"),
output,
fake_server_module(queue_task_done=True),
)
state["package"] = package
monitor = package._global_image_monitor
handler = monitor.handler if monitor is not None else None
if handler is None:
sys.exit(3)
handler.processing_delay = 3600
path = os.path.join(tmp, "busy.png") # outside the watched folder
with open(path, "wb") as image:
image.write(b"x" * 16)
if not handler.schedule(path) or not handler._worker.is_alive():
sys.exit(4)
# Return normally with the worker alive, as ComfyUI does after Ctrl+C.
""")
class TestPackageExit(unittest.TestCase):
def test_teardown_runs_at_a_normal_exit_with_the_worker_alive(self):
with tempfile.TemporaryDirectory() as tmp:
os.makedirs(os.path.join(tmp, "output"))
script = os.path.join(tmp, "child.py")
with open(script, "w", encoding="utf-8") as handle:
handle.write(CHILD)
env = dict(os.environ)
env.pop("PROMPT_MANAGER_CONFIG_PATH", None)
# The child's library and log file live in its own folder, never
# in the checkout's logs/ (which this runner may hold open) or a
# data folder the caller set.
env["PROMPT_MANAGER_DATA_DIR"] = tmp
env["PYTHONIOENCODING"] = "utf-8"
try:
completed = subprocess.run(
[sys.executable, script, ROOT, tmp],
cwd=ROOT,
env=env,
capture_output=True,
timeout=CHILD_CEILING_SECONDS,
encoding="utf-8",
errors="replace",
)
except subprocess.TimeoutExpired as hung:
self.fail(
f"the child did not exit within {CHILD_CEILING_SECONDS} s\n"
f"stdout:\n{hung.stdout}\nstderr:\n{hung.stderr}"
)
# Leaving the block removes the folder: on Windows that fails
# if the child left a handle open, which it cannot once it exited,
# so the check here is that the child exited at all.
detail = f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}"
self.assertEqual(completed.returncode, 0, detail)
reports = [
line[len(REPORT_PREFIX) :]
for line in completed.stdout.splitlines()
if line.startswith(REPORT_PREFIX)
]
self.assertEqual(len(reports), 1, detail)
report = json.loads(reports[0])
# "thumbnails" is True here because no preview worker started; its
# close result is tested in test_thumbnail_auto and test_init_package.
self.assertEqual(
report["result"],
{"monitor": True, "thumbnails": True, "tracker": True},
detail,
)
self.assertEqual(report["alive"], [], detail)
# The teardown's own warnings. Import-time warnings (the API's
# output-folder search finds no ComfyUI here) are not about exit.
for text in (
"Traceback",
"Background work not stopped in time",
"still running after",
"Error while stopping",
):
self.assertNotIn(text, completed.stderr, detail)
if __name__ == "__main__":
unittest.main()
+5 -23
View File
@@ -26,27 +26,9 @@ def _read(name):
def _pyproject():
try:
import tomllib # Python 3.11+
except ImportError: # pragma: no cover - exercised on 3.10 only
tomllib = None
text = _read("pyproject.toml")
if tomllib is not None:
return tomllib.loads(text)
# Minimal fallback for 3.10: pull the dependency list out with a regex.
match = re.search(r"dependencies\s*=\s*\[(.*?)\]", text, re.S)
deps = re.findall(r'"([^"]*)"', match.group(1)) if match else []
return {
"project": {
"dependencies": deps,
"requires-python": (
re.search(r'requires-python\s*=\s*"([^"]+)"', text).group(1)
if "requires-python" in text
else None
),
},
"build-system": {} if "[build-system]" in text else None,
}
import tomllib
return tomllib.loads(_read("pyproject.toml"))
class PyprojectContract(unittest.TestCase):
@@ -65,8 +47,8 @@ class PyprojectContract(unittest.TestCase):
}
self.assertEqual(deps, wanted)
def test_requires_python_floor_is_3_10(self):
self.assertEqual(_pyproject()["project"]["requires-python"], ">=3.10")
def test_requires_python_floor_is_3_12(self):
self.assertEqual(_pyproject()["project"]["requires-python"], ">=3.12")
def test_build_system_is_declared(self):
self.assertIsNotNone(_pyproject().get("build-system"))
+2
View File
@@ -23,6 +23,7 @@ try:
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed
from tests.background_stops import close_handler # noqa: E402
from database.operations import PromptDatabase
from prompt_manager import PromptManager
from prompt_manager_pair import PromptManagerPair
@@ -278,6 +279,7 @@ class MonitorCase(DatabaseCase):
fd, self.image = tempfile.mkstemp(suffix=".png", dir=self.tmp.name)
os.close(fd)
self.handler = ImageGenerationHandler(self.db, self.tracker)
self.addCleanup(close_handler, self, self.handler)
self.linked = []
self.handler.link_image_to_prompt = lambda path, prompt, meta: (
self.linked.append(prompt["id"])
+90
View File
@@ -5,10 +5,12 @@ Tests PromptTracker, PromptExecutionContext, and the singleton
get_prompt_tracker. Uses a mock db_manager to avoid database deps.
"""
import logging
import os
import sys
import time
import threading
import types
import unittest
from unittest.mock import MagicMock, patch
@@ -312,5 +314,93 @@ class TestPromptQueue(unittest.TestCase):
self.assertEqual(popped_ids, ids)
def _real_tracker():
"""A PromptTracker built by its own ``__init__``, so its cleanup thread runs."""
mock_db = MagicMock()
mock_db.get_prompt_by_hash.return_value = None
return PromptTracker(mock_db)
class TestTrackerShutdown(unittest.TestCase):
"""The cleanup thread stops on a signal and is joined within a bound (#301)."""
def test_shutdown_stops_the_cleanup_thread_within_the_bound(self):
tracker = _real_tracker()
thread = tracker.cleanup_thread
self.assertTrue(thread.is_alive())
self.assertEqual(thread.name, "PromptManagerPromptCleanup")
self.assertTrue(tracker.shutdown())
self.assertFalse(thread.is_alive())
def test_shutdown_is_idempotent(self):
tracker = _real_tracker()
self.assertTrue(tracker.shutdown())
self.assertTrue(tracker.shutdown())
def test_cleanup_error_wait_is_interruptible(self):
import utils.prompt_tracker as tracker_mod
real_time = time.time
calls = {"n": 0}
def failing_once():
calls["n"] += 1
if calls["n"] == 1:
raise RuntimeError("clock unavailable")
return real_time()
stub = types.SimpleNamespace(
time=failing_once, sleep=time.sleep, monotonic=time.monotonic
)
logged = threading.Event()
logger = logging.getLogger("prompt_manager.prompt_tracker")
real_error = logger.error
def error(*args, **kwargs):
real_error(*args, **kwargs)
logged.set()
with (
patch.object(tracker_mod, "time", stub),
patch.object(logger, "error", side_effect=error),
):
tracker = _real_tracker()
self.addCleanup(tracker.shutdown)
self.assertTrue(logged.wait(30), "the cleanup error was never logged")
# The thread is now in its 60 s error wait.
self.assertTrue(tracker.shutdown())
self.assertFalse(tracker.cleanup_thread.is_alive())
def test_tracker_works_after_shutdown(self):
tracker = _real_tracker()
self.assertTrue(tracker.shutdown())
tracker.set_current_prompt("still tracked", {"prompt_id": 9})
self.assertEqual(tracker.pop_next_prompt()["id"], 9)
def test_shutdown_of_a_tracker_built_without_init(self):
self.assertTrue(_make_tracker().shutdown())
def test_shutdown_on_the_cleanup_thread_only_signals(self):
tracker = _real_tracker()
self.addCleanup(lambda: self.assertTrue(tracker.shutdown()))
result = {}
with patch.object(tracker, "cleanup_thread", threading.current_thread()):
result["value"] = tracker.shutdown()
self.assertFalse(result["value"])
self.assertTrue(tracker._stop.is_set())
def test_shutdown_reports_a_thread_still_running(self):
tracker = _real_tracker()
self.assertTrue(tracker.shutdown())
busy = MagicMock()
busy.is_alive.return_value = True
tracker.cleanup_thread = busy
with self.assertLogs("prompt_manager.prompt_tracker", "WARNING") as logs:
self.assertFalse(tracker.shutdown(timeout=0.05))
busy.join.assert_called_once_with(0.05)
self.assertIn("prompt cleanup thread still running", "\n".join(logs.output))
if __name__ == "__main__":
unittest.main()
+4 -1
View File
@@ -27,6 +27,7 @@ except ImportError: # discovered with tests/ as the top-level directory
from PIL import Image, PngImagePlugin
from tests.background_stops import close_handler, stop_monitor
from database.operations import PromptDatabase
from database.review_schema import canonical_evidence, evidence_digest
from py.api import PromptManagerAPI
@@ -790,6 +791,7 @@ class TestReviewEdges(ReviewTestCase):
class TestMonitorRecording(ReviewTestCase):
def _handler(self):
handler = ImageGenerationHandler(self.api.db, PromptTracker(self.api.db))
self.addCleanup(close_handler, self, handler)
return handler
def _linked_prompt(self, path):
@@ -953,7 +955,7 @@ class TestMonitorRecording(ReviewTestCase):
monitor.set_review_recorder(None)
self.assertIsNone(monitor.handler.review_recorder)
finally:
monitor.stop_monitoring()
stop_monitor(self, monitor)
if __name__ == "__main__":
@@ -1092,6 +1094,7 @@ class TestCaptureStoreShutdown(ReviewTestCase):
self._scan_record(path)
# The monitor's pass records nothing, so it needs no store.
handler = ImageGenerationHandler(self.api.db, PromptTracker(self.api.db))
self.addCleanup(close_handler, self, handler)
handler.review_recorder = self.api.record_review_file
handler._record_review(str(path))
self.assertEqual(self.api.record_review_file(str(path)), "skipped")
+3 -1
View File
@@ -17,6 +17,7 @@ except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed
from database.operations import PromptDatabase
from tests.background_stops import close_handler
from tests.test_api_images import ImageAPITestCase
from tests.test_prompt_attribution import issue_graph, pair_evidence
from tools import review_image_links as tool
@@ -158,8 +159,8 @@ class ReviewToolTests(ToolFixture):
def test_conflicting_sources_are_left_to_the_maintainer(self):
db = PromptDatabase(self.db_path)
handler = tool.ImageGenerationHandler(db, None)
try:
handler = tool.ImageGenerationHandler(db, None)
row = {
"id": 1,
"prompt_id": self.coffee,
@@ -170,6 +171,7 @@ class ReviewToolTests(ToolFixture):
):
verdict = tool.judge(db, handler, row, self.output)
finally:
close_handler(self, handler)
db.close_all()
self.assertEqual(verdict["verdict"], "conflict")
+3
View File
@@ -28,6 +28,7 @@ except ImportError: # discovered with tests/ as the top-level directory
from test_api_review import ReviewRoutesTestCase
from test_review import unlinkable_graph, write_png
from tests.background_stops import close_handler
from utils import inspection_semantics, review_capture
from utils.hashing import generate_prompt_hash
from utils.image_monitor import ImageGenerationHandler
@@ -437,6 +438,7 @@ class TestApplyInspectionCandidate(InspectionTestCase):
class TestMonitorNeverQueues(InspectionTestCase):
def _handler(self):
handler = ImageGenerationHandler(self.api.db, PromptTracker(self.api.db))
self.addCleanup(close_handler, self, handler)
handler.review_recorder = self.api.record_review_file
return handler
@@ -760,6 +762,7 @@ class TestRecordingEdges(InspectionTestCase):
text="latest", prompt_hash=generate_prompt_hash("latest")
)
handler = ImageGenerationHandler(self.api.db, PromptTracker(self.api.db))
self.addCleanup(close_handler, self, handler)
handler.link_image_to_prompt(
str(path),
{"id": prompt_id, "text": "latest", "fallback": True},
+2
View File
@@ -38,6 +38,7 @@ except ImportError: # discovered with tests/ as the top-level directory
write_png,
)
from tests.background_stops import close_handler
from utils.hashing import generate_prompt_hash
from utils.image_monitor import ImageGenerationHandler
from utils.prompt_tracker import PromptTracker
@@ -299,6 +300,7 @@ class TestQueueSelfHeals(OutcomeTestCase):
class TestMonitorRule(OutcomeTestCase):
def _handler(self):
handler = ImageGenerationHandler(self.api.db, PromptTracker(self.api.db))
self.addCleanup(close_handler, self, handler)
handler.review_recorder = self.api.record_review_file
return handler
+4 -2
View File
@@ -11,6 +11,7 @@ from unittest import mock
from database.run_attribution import stored_run_prompt
from tests import test_execution_history_database as history_tests
from tests.background_stops import close_handler
from tests.test_api_admin import AdminAPITestCase
from tests.test_execution_context_database import snapshot
from tests.test_prompt_attribution import issue_graph, llm_text_graph, make_tracker
@@ -151,6 +152,7 @@ class TestStoredRunPrompt(RecordedRuns, unittest.TestCase):
prompt_id = self.save(LLM_TEXT)
self.record_run(self.db, self.output, [("sub", "img_00001_.png")], [LLM_TEXT])
handler = ImageGenerationHandler(self.db, make_tracker(self.db))
self.addCleanup(close_handler, self, handler)
linked = []
handler.link_image_to_prompt = lambda path, prompt, meta: linked.append(
prompt["id"]
@@ -332,7 +334,7 @@ class TestMonitorHistoryRecheck(RecordedRuns, unittest.TestCase):
self.handler = ImageGenerationHandler(self.db, make_tracker(self.db))
self.handler.processing_delay = 0
self.handler.settle_interval = 0.01
self.addCleanup(self.handler.shutdown)
self.addCleanup(close_handler, self, self.handler)
self.linked = []
self.handler.link_image_to_prompt = lambda path, prompt, meta: (
self.linked.append((os.path.basename(path), prompt["id"]))
@@ -650,7 +652,7 @@ class TestLinkedOutputRoot(RecordedRuns, unittest.TestCase):
handler = ImageGenerationHandler(db, make_tracker(db))
handler.processing_delay = 0
handler.settle_interval = 0.01
self.addCleanup(handler.shutdown)
self.addCleanup(close_handler, self, handler)
linked = []
handler.link_image_to_prompt = lambda path, prompt, meta: linked.append(
prompt["id"]
+24
View File
@@ -152,6 +152,30 @@ class RunSuiteTests(unittest.TestCase):
result = run_suite(self.folder, None, "-p", "check_*.py")
self.assertEqual(result.returncode, 1, result.stderr)
def test_lora_manager_settings_read_a_disposable_folder(self):
# #355: no test reads the developer's real LoraManager settings.
record = self.folder / "seen.txt"
self.write(
"test_probe_lora.py",
f"""
import os
import unittest
class Probe(unittest.TestCase):
def test_isolated(self):
folder = os.environ["LORA_MANAGER_SETTINGS_DIR"]
with open({str(record)!r}, "w", encoding="utf-8") as handle:
handle.write(folder)
self.assertTrue(os.path.isdir(folder))
self.assertEqual(os.listdir(folder), [])
""",
)
result = run_suite(self.folder, {"LORA_MANAGER_SETTINGS_DIR": "/real/one"})
self.assertEqual(result.returncode, 0, result.stderr)
folder = record.read_text(encoding="utf-8")
self.assertNotEqual(folder, "/real/one")
self.assertFalse(os.path.exists(folder), "removed when the run ends")
def test_empty_folder_exits_five(self):
result = run_suite(self.folder)
self.assertEqual(result.returncode, 5, result.stderr)
+126
View File
@@ -976,3 +976,129 @@ class TestSourceFormError(unittest.TestCase):
if __name__ == "__main__":
unittest.main()
class TestAutotagDefaultsSettings(SourcesTestCase):
"""#328: the tagging defaults save through POST /settings like the others."""
ATTRS = (
"MODEL_TYPE",
"EXISTING",
"GENERAL_THRESHOLD",
"CHARACTER_THRESHOLD",
"TEMPLATE",
"SUGGESTION_LIMIT",
)
SAVED = {
"autotag_model_type": "gguf",
"autotag_existing": "merge",
"autotag_general_threshold": 0.4,
"autotag_character_threshold": 0.8,
"autotag_template": "Tags only.",
"autotag_suggestion_limit": 15,
}
async def asyncSetUp(self):
await super().asyncSetUp()
from py.config import AutotagConfig
self.autotag = AutotagConfig
for attr in self.ATTRS:
self.addCleanup(setattr, AutotagConfig, attr, getattr(AutotagConfig, attr))
def _runtime(self):
return {attr: getattr(self.autotag, attr) for attr in self.ATTRS}
async def test_autotag_defaults_save_round_trip(self):
self._write_config({"custom": {"keep": True}})
revision = (await self._get())["revision"]
status, data = await self._post(
"/prompt_manager/settings", {"expected_revision": revision, **self.SAVED}
)
self.assertEqual(status, 200, data)
self.assertEqual(data["restart_required_fields"], [])
saved = self._read_config()
self.assertEqual(
saved["autotag"]["defaults"],
{key[len("autotag_") :]: value for key, value in self.SAVED.items()},
)
self.assertEqual(saved["custom"], {"keep": True})
self.assertEqual(
self._runtime(),
{
"MODEL_TYPE": "gguf",
"EXISTING": "merge",
"GENERAL_THRESHOLD": 0.4,
"CHARACTER_THRESHOLD": 0.8,
"TEMPLATE": "Tags only.",
"SUGGESTION_LIMIT": 15,
},
)
settings = (await self._get())["settings"]
self.assertEqual({key: settings[key] for key in self.SAVED}, self.SAVED)
async def test_stale_revision_changes_no_default(self):
revision = (await self._get())["revision"]
self._write_config({"custom": {"edited": "elsewhere"}})
before = (self._config_bytes(), self._runtime())
status, data = await self._post(
"/prompt_manager/settings", {"expected_revision": revision, **self.SAVED}
)
self.assertEqual(status, 409, data)
self.assertEqual(data["code"], "settings-changed")
self.assertEqual((self._config_bytes(), self._runtime()), before)
async def test_bad_default_message_names_the_setting(self):
cases = (
(
"autotag_general_threshold",
1.5,
"General-tag confidence must be between 0 and 1.",
),
(
"autotag_character_threshold",
"0.8",
"Character-tag confidence must be a finite number.",
),
(
"autotag_character_threshold",
True,
"Character-tag confidence must be a finite number.",
),
(
"autotag_model_type",
"wd14-convnext",
"Tagging model must be one of gguf, hf, wd14-swinv2, wd14-vit.",
),
(
"autotag_existing",
"replace",
"Existing tags must be skip, merge or replace_after_review.",
),
("autotag_template", " ", "LLM tag template must be non-empty text."),
(
"autotag_template",
"x" * 4001,
"LLM tag template must be at most 4000 characters.",
),
(
"autotag_suggestion_limit",
0,
"Suggestions per image must be a whole number from 1 to 100.",
),
(
"autotag_suggestion_limit",
2.5,
"Suggestions per image must be a whole number from 1 to 100.",
),
)
self._write_config({"custom": {"keep": True}})
before = (self._config_bytes(), self._runtime())
for key, value, message in cases:
with self.subTest(key=key, value=value):
status, data = await self._post(
"/prompt_manager/settings", {**self.SAVED, key: value}
)
self.assertEqual(status, 400, data)
self.assertEqual(data["error"], message)
self.assertEqual((self._config_bytes(), self._runtime()), before)
+20 -2
View File
@@ -32,6 +32,7 @@ from PIL import Image # noqa: E402
import py.api.images as images_module # noqa: E402
import py.api.thumbnail_auto as auto # noqa: E402
from tests.background_stops import close_handler, stop_monitor # noqa: E402
from py.api.jobs import THUMBNAILS_JOB_KIND # noqa: E402
SERVE = "/prompt_manager/images/serve/"
@@ -737,11 +738,27 @@ class TestShutdown(AutoThumbnailCase):
patch.object(auto, "CLOSE_WAIT_SECONDS", 0),
self.assertLogs("prompt_manager.api", "WARNING") as logs,
):
self.api.close_auto_thumbnails()
self.assertFalse(self.api.close_auto_thumbnails())
process.kill.assert_called_once()
self.assertIn("still running", "\n".join(logs.output))
release.set()
async def test_close_reports_an_ended_worker(self):
path = make_png(self.output_dir / "a.png", size=(64, 64))
self.api.queue_capture_thumbnail(str(path))
worker = self.api._auto_thumbnails()._worker
await self._idle()
loop = asyncio.get_running_loop()
closed = await loop.run_in_executor(
None, lambda: self.api.close_auto_thumbnails(timeout=WAIT)
)
self.assertTrue(closed)
self.assertFalse(worker.is_alive())
def test_close_without_a_thumbnailer_returns_true(self):
self.assertIsNone(self.api._auto_thumbnailer)
self.assertTrue(self.api.close_auto_thumbnails())
class TestHistorySharesTheAdmission(AutoThumbnailCase):
async def _shared_flight(self, name, is_video, first):
@@ -980,6 +997,7 @@ class TestMonitorQueuesThumbnails(AutoThumbnailCase):
from utils.image_monitor import ImageGenerationHandler
handler = ImageGenerationHandler(self.api.db, unittest.mock.MagicMock())
self.addCleanup(close_handler, self, handler)
handler.processing_delay = 0
handler.settle_interval = 0
return handler
@@ -1022,7 +1040,7 @@ class TestMonitorQueuesThumbnails(AutoThumbnailCase):
monitor.set_thumbnail_queue(None)
self.assertIsNone(monitor.handler.thumbnail_queue)
finally:
monitor.stop_monitoring()
stop_monitor(self, monitor)
class TestScanHandsOffThumbnails(AutoThumbnailCase):
+2
View File
@@ -16,6 +16,7 @@ import uuid
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from tests.background_stops import stop_tracker # noqa: E402
from database.operations import PromptDatabase
from prompt_manager_base import PromptManagerBase
from utils.hashing import generate_prompt_hash
@@ -596,6 +597,7 @@ class TestNodeRoles(DbTestCase):
class TestTrackerQueueFlag(unittest.TestCase):
def setUp(self):
self.tracker = PromptTracker(mock.Mock())
self.addCleanup(stop_tracker, self, self.tracker)
def test_push_to_queue_false_keeps_the_batch_queue_empty(self):
self.tracker.set_current_prompt("a cat", {"prompt_id": 1}, push_to_queue=False)
+6
View File
@@ -545,6 +545,12 @@
"functions": 98,
"statements": 98
},
"workspace/system-tagging.mjs": {
"lines": 80,
"branches": 70,
"functions": 75,
"statements": 80
},
"workspace/system-thumbnails.mjs": {
"lines": 98,
"branches": 95,
+133 -21
View File
@@ -70,6 +70,11 @@ MAX_AWAITING_HISTORY = 1000
AWAIT_HISTORY_SECONDS = 3600
# ``schedule`` snapshot meaning "history-only re-check" (not a new image)
HISTORY_ONLY = object()
# How long a stop waits for the image worker, and for the watchdog observer
# after its stop (#302). Every wait before the image in hand is interrupted
# at once, so a bound is reached only while one image is still being linked.
WORKER_JOIN_SECONDS = 5.0
OBSERVER_JOIN_SECONDS = 5.0
class ImageGenerationHandler(FileSystemEventHandler):
@@ -143,6 +148,11 @@ class ImageGenerationHandler(FileSystemEventHandler):
self._history_due: "OrderedDict[str, None]" = OrderedDict()
self._awaiting_lock = threading.Lock()
self._sealed_subscribed = False
# Set by ``close``: no new work, no new worker, ever (#302).
self._closed = False
# Workers a stop signalled that had not ended yet: every later stop
# waits for them again, so it never reports True too early.
self._finishing: List[threading.Thread] = []
def on_created(self, event):
"""Handle filesystem creation events.
@@ -157,6 +167,10 @@ class ImageGenerationHandler(FileSystemEventHandler):
Args:
event: FileSystemEvent object containing event details
"""
# Closed: return before any log or tracker lock, so watchdog's
# ``observer.stop()`` never waits on this dispatch.
if self._closed:
return
if event.is_directory or not self.is_image_file(event.src_path):
return
self.logger.info(f"New image detected: {event.src_path}")
@@ -180,9 +194,12 @@ class ImageGenerationHandler(FileSystemEventHandler):
and once its size has stopped changing.
Returns:
False when the path is already pending or the pending set is full
False when the path is already pending, the pending set is full or
the handler is closed
"""
with self._pending_lock:
if self._closed:
return False
if image_path in self._pending:
self.logger.debug(f"Already pending, not rescheduled: {image_path}")
return False
@@ -203,12 +220,19 @@ class ImageGenerationHandler(FileSystemEventHandler):
with self._pending_lock:
return len(self._pending)
def shutdown(self) -> None:
"""Stop the worker thread and drop every image still waiting.
def shutdown(self, timeout: float = WORKER_JOIN_SECONDS) -> bool:
"""Stop the worker thread, drop every image still waiting, and wait
for the worker within ``timeout`` seconds.
The worker wakes from any delay or settle wait, finishes nothing more
and exits; the next ``schedule`` starts a fresh worker with its own
queue. Safe to call repeatedly and before anything was scheduled.
The worker wakes from any delay or settle wait, finishes at most the
image in hand and exits; the next ``schedule`` starts a fresh worker
with its own queue (use ``close`` for a final stop). Safe to call
repeatedly and before anything was scheduled.
Returns:
True when no worker is left running. False when the worker is
still finishing its image after ``timeout`` (a warning is
logged), or when called on the worker itself (it only signals).
"""
with self._awaiting_lock:
self._awaiting.clear()
@@ -218,9 +242,40 @@ class ImageGenerationHandler(FileSystemEventHandler):
worker, self._worker = self._worker, None
work, self._work = self._work, queue.Queue()
stop, self._stop = self._stop, threading.Event()
if worker is not None:
self._finishing.append(worker)
finishing = [w for w in self._finishing if w.is_alive()]
self._finishing = finishing
stop.set()
if worker is not None and worker.is_alive():
work.put(None) # sentinel: wake a worker idle on an empty queue
if threading.current_thread() in finishing:
self.logger.debug("shutdown called on the image worker itself")
return False
deadline = time.monotonic() + timeout
for thread in finishing:
thread.join(max(0.0, deadline - time.monotonic()))
with self._pending_lock:
self._finishing = [w for w in self._finishing if w.is_alive()]
still_running = bool(self._finishing)
if still_running:
self.logger.warning(
f"image worker still running after {timeout} s; "
"it stops after the image in hand"
)
return False
return True
def close(self, timeout: float = WORKER_JOIN_SECONDS) -> bool:
"""Final stop: ``shutdown`` and refuse every later ``schedule``.
After ``close`` no worker starts again, not even from a sealed run's
history re-check on the history writer's thread. Returns what
``shutdown`` returns.
"""
with self._pending_lock:
self._closed = True
return self.shutdown(timeout)
def _ensure_worker(self) -> None:
"""Start the single daemon worker thread if it isn't running.
@@ -521,7 +576,12 @@ class ImageGenerationHandler(FileSystemEventHandler):
self._awaiting.popitem(last=False)
def _await_history(self, image_path: str) -> None:
"""Re-check ``image_path`` once, history only, when its run is sealed."""
"""Re-check ``image_path`` once, history only, when its run is sealed.
A closed handler re-checks nothing.
"""
if self._closed:
return
if not self._sealed_subscribed:
subscribe_outputs_sealed(self._on_outputs_sealed)
self._sealed_subscribed = True
@@ -570,10 +630,19 @@ class ImageGenerationHandler(FileSystemEventHandler):
self._schedule_recheck(image_path)
def _schedule_recheck(self, image_path: str) -> None:
"""Schedule the history-only re-check, or keep it due for later."""
"""Schedule the history-only re-check, or keep it due for later.
A closed handler drops it.
"""
if self._closed:
return
if self.schedule(image_path, HISTORY_ONLY):
return
with self._awaiting_lock:
# ``close`` sets the flag before ``shutdown`` clears this set, so
# a close that raced the refusal above is seen here.
if self._closed:
return
self._history_due.pop(image_path, None)
self._history_due[image_path] = None
while len(self._history_due) > MAX_AWAITING_HISTORY:
@@ -872,6 +941,9 @@ class ImageMonitor:
self.observer = None
self.handler = None
self.monitored_directories = []
# (observer, handler) pairs a stop left running: the next stop waits
# for them again instead of reporting True.
self._unfinished: List[Tuple[Any, Any]] = []
self.logger = get_logger("prompt_manager.image_monitor")
self.review_recorder: Optional[Callable[[str], str]] = None
self.thumbnail_queue: Optional[Callable[[str], str]] = None
@@ -999,29 +1071,69 @@ class ImageMonitor:
"""Whether a filesystem observer is running."""
return self.observer is not None
def stop_monitoring(self):
"""Stop the image monitoring system.
def stop_monitoring(self, timeout: float = OBSERVER_JOIN_SECONDS) -> bool:
"""Stop the image monitoring system within one ``timeout`` deadline.
Cleanly shuts down the filesystem watcher and clears all monitoring state.
This method should be called before program exit to ensure proper cleanup.
Safe to call when monitoring never started.
Closes the handler first (its worker is joined and its dispatch
returns at once), then stops the watchdog observer and joins it.
watchdog's ``observer.stop()`` joins its emitters without a limit;
it counts against the deadline but is not bounded here. A part an
earlier call left running is waited for again. Safe to call when
monitoring never started and more than once.
Returns:
True when every worker and observer it stopped has ended.
"""
deadline = time.monotonic() + timeout
observer, self.observer = self.observer, None
handler, self.handler = self.handler, None
self.monitored_directories = []
pairs, self._unfinished = self._unfinished, []
if observer is not None:
pairs.append((observer, handler))
if not pairs:
return True
def remaining() -> float:
return max(0.0, deadline - time.monotonic())
for pair_observer, pair_handler in pairs:
handler_ok = self._close_handler(pair_handler, remaining())
observer_ok = self._stop_observer(pair_observer, remaining())
if not observer_ok:
self.logger.warning(
f"watchdog observer still running after {timeout} s"
)
if not (handler_ok and observer_ok):
self._unfinished.append(
(
None if observer_ok else pair_observer,
None if handler_ok else pair_handler,
)
)
self.logger.debug("Image monitoring stopped")
return not self._unfinished
def _close_handler(self, handler, timeout: float) -> bool:
if handler is None:
return True
try:
return handler.close(timeout)
except Exception as e:
self.logger.warning(f"Error while stopping the image worker: {e}")
return False
def _stop_observer(self, observer, timeout: float) -> bool:
"""``observer.stop()`` (safe to repeat), then a bounded join."""
if observer is None:
return
return True
try:
observer.stop()
observer.join()
observer.join(timeout)
except Exception as e:
self.logger.warning(f"Error while stopping image monitoring: {e}")
if handler is not None:
try:
handler.shutdown()
except Exception as e:
self.logger.warning(f"Error while stopping the image worker: {e}")
self.logger.debug("Image monitoring stopped")
return False
return not observer.is_alive()
def detect_comfyui_output_dirs(self) -> list:
"""Auto-detect ComfyUI output directories.
+39 -6
View File
@@ -44,6 +44,9 @@ except ImportError:
from utils.hashing import generate_prompt_hash
from utils.logging_config import get_logger
# How long ``PromptTracker.shutdown`` waits for the cleanup thread (#301).
TRACKER_JOIN_SECONDS = 2.0
class PromptTracker:
"""Thread-safe tracking of current prompt executions.
@@ -92,9 +95,12 @@ class PromptTracker:
self.cleanup_interval = 300 # 5 minutes
self.prompt_timeout = 600 # 10 minutes
# Start cleanup thread
# Start cleanup thread; ``shutdown`` sets ``_stop`` to end it.
self._stop = threading.Event()
self.cleanup_thread = threading.Thread(
target=self._cleanup_expired_prompts, daemon=True
target=self._cleanup_expired_prompts,
name="PromptManagerPromptCleanup",
daemon=True,
)
self.cleanup_thread.start()
self.logger.debug("PromptTracker initialization completed")
@@ -464,9 +470,10 @@ class PromptTracker:
Runs continuously as a daemon thread, periodically removing expired
prompt contexts from the global tracking dictionary. This prevents
memory leaks from accumulating old prompt data.
memory leaks from accumulating old prompt data. Every wait ends at
once when ``shutdown`` sets the stop event.
"""
while True:
while not self._stop.is_set():
try:
current_time = time.time()
expired_ids = []
@@ -497,11 +504,37 @@ class PromptTracker:
if pruned:
self.logger.debug(f"Pruned {pruned} expired queue entries")
time.sleep(self.cleanup_interval)
if self._stop.wait(self.cleanup_interval):
break
except Exception as e:
self.logger.error(f"Error in cleanup thread: {e}")
time.sleep(60) # Wait a minute before retrying
if self._stop.wait(60): # Wait a minute before retrying
break
def shutdown(self, timeout: float = TRACKER_JOIN_SECONDS) -> bool:
"""Stop the cleanup thread and wait for it within ``timeout`` seconds.
Returns True when the thread has ended (or never existed). The tracker
keeps working afterwards; only the background pruning of expired
entries stops. Safe to call more than once. Called on the cleanup
thread itself it only signals and returns False.
"""
stop = getattr(self, "_stop", None)
thread = getattr(self, "cleanup_thread", None)
if stop is None or thread is None:
return True # built without __init__: no thread was started
stop.set()
if thread is threading.current_thread():
self.logger.debug("shutdown called on the prompt cleanup thread")
return False
thread.join(timeout)
if thread.is_alive():
self.logger.warning(
f"prompt cleanup thread still running after {timeout} s"
)
return False
return True
def get_active_prompts(self) -> Dict[str, Dict[str, Any]]:
"""
+2 -1
View File
@@ -1428,12 +1428,13 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr
<div class="w-full bg-pm-input rounded-full h-1.5">
<div id="downloadProgressBar" class="bg-pm-accent h-1.5 rounded-full transition-all duration-300" style="width: 0%"></div>
</div>
<p class="text-pm-secondary text-xs mt-2">The download continues in Activity.</p>
</div>
<div class="flex justify-center">
<button id="cancelDownloadBtn"
class="px-3 py-1.5 text-pm-error bg-transparent hover:bg-pm-error-tint border border-pm text-xs font-medium rounded-pm-sm transition-colors">
Cancel
Close
</button>
</div>
</div>
+2 -1
View File
@@ -730,13 +730,14 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr
<div class="w-full bg-pm-input rounded-full h-2">
<div id="downloadProgressBar" class="bg-pm-accent h-2 rounded-full transition-all duration-300" style="width: 0%"></div>
</div>
<p class="text-pm-secondary text-xs mt-2">The download continues in Activity.</p>
</div>
<!-- Cancel Button -->
<div class="flex justify-center">
<button id="cancelDownloadBtn"
class="px-3 py-1 bg-pm-error hover:bg-pm-error text-pm text-xs font-medium rounded-pm-sm transition-colors">
Cancel
Close
</button>
</div>
</div>
+13 -7
View File
@@ -3689,6 +3689,7 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
document.getElementById('downloadStatus').textContent = 'Preparing...';
document.getElementById('downloadProgressPercent').textContent = '0%';
document.getElementById('downloadProgressBar').style.width = '0%';
document.getElementById('downloadProgressBar').classList.remove('animate-pulse');
this.hideModal("autoTagModal");
this.showModal("autoTagDownloadModal");
@@ -3699,18 +3700,21 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
this.autoTagState.downloadEventSource.onmessage = (event) => {
const data = JSON.parse(event.data);
if (data.status === 'downloading') {
const percent = Math.round(data.progress * 100);
document.getElementById('downloadStatus').textContent = data.message || 'Downloading...';
document.getElementById('downloadProgressPercent').textContent = `${percent}%`;
document.getElementById('downloadProgressBar').style.width = `${percent}%`;
} else if (data.status === 'complete') {
// The server sends `type` (progress / complete / error); the download is an
// Activity job with no byte progress, so the bar is indeterminate (#328).
if (data.type === 'progress') {
document.getElementById('downloadStatus').textContent = 'Downloading\u2026';
document.getElementById('downloadProgressPercent').textContent = '';
const bar = document.getElementById('downloadProgressBar');
bar.style.width = '100%';
bar.classList.add('animate-pulse');
} else if (data.type === 'complete') {
this.autoTagState.downloadEventSource.close();
this.hideModal("autoTagDownloadModal");
this.showModal("autoTagModal");
this.checkAutoTagModels();
this.showNotification('Model downloaded successfully!', 'success');
} else if (data.status === 'error') {
} else if (data.type === 'error') {
this.autoTagState.downloadEventSource.close();
this.hideModal("autoTagDownloadModal");
this.showModal("autoTagModal");
@@ -3966,6 +3970,8 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`;
}
document.getElementById('reviewImage').src = imageUrl;
// The record's file name, else the last segment of its path (#328: `filename` was never declared).
const filename = image.filename || String(image.image_path || image.path || '').split(/[\\/]/).pop();
document.getElementById('reviewImageName').textContent = filename || 'Unknown';
document.getElementById('reviewTagsVisible').innerHTML = '<div class="text-pm-secondary">Generating tags...</div>';
document.getElementById('reviewTagsHidden').innerHTML = '';

Some files were not shown because too many files have changed in this diff Show More