Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5a9971784c | ||
|
|
3dd45c0efd | ||
|
|
a320acc41b | ||
|
|
b7b99004c7 | ||
|
|
168af0d868 | ||
|
|
19097ed6e1 | ||
|
|
476b1f73dc | ||
|
|
7dc8d99363 | ||
|
|
a478ae0a1d | ||
|
|
02be260a14 | ||
|
|
2dbf79c9ae |
@@ -1,16 +1,19 @@
|
||||
{
|
||||
"required_status_checks": {
|
||||
"strict": false,
|
||||
"strict": true,
|
||||
"contexts": [
|
||||
"py3.10 on ubuntu-latest",
|
||||
"py3.11 on ubuntu-latest",
|
||||
"py3.12 on ubuntu-latest",
|
||||
"py3.13 on ubuntu-latest",
|
||||
"py3.14 on ubuntu-latest",
|
||||
"py3.10 on windows-latest",
|
||||
"py3.11 on windows-latest",
|
||||
"py3.12 on windows-latest",
|
||||
"py3.13 on windows-latest",
|
||||
"py3.14 on windows-latest",
|
||||
"py3.10 on macos-latest",
|
||||
"py3.11 on macos-latest",
|
||||
"py3.12 on macos-latest",
|
||||
"py3.13 on macos-latest",
|
||||
"py3.14 on macos-latest",
|
||||
"coverage gate (85%)",
|
||||
"black and flake8",
|
||||
"frontend modules",
|
||||
|
||||
@@ -4,7 +4,6 @@ permissions:
|
||||
contents: read
|
||||
|
||||
on:
|
||||
merge_group:
|
||||
push:
|
||||
branches: [main]
|
||||
pull_request:
|
||||
@@ -27,10 +26,10 @@ jobs:
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
|
||||
- name: Set up Python 3.12
|
||||
- name: Set up Python 3.10
|
||||
uses: actions/setup-python@v7
|
||||
with:
|
||||
python-version: "3.12"
|
||||
python-version: "3.10"
|
||||
|
||||
- name: Install bandit
|
||||
run: pip install bandit[toml]
|
||||
|
||||
+15
-40
@@ -4,7 +4,6 @@ permissions:
|
||||
contents: read
|
||||
|
||||
on:
|
||||
merge_group:
|
||||
pull_request:
|
||||
branches: [main]
|
||||
push:
|
||||
@@ -27,7 +26,6 @@ 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,
|
||||
@@ -41,7 +39,7 @@ jobs:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os: [ubuntu-latest, windows-latest, macos-latest]
|
||||
python-version: ["3.12", "3.13", "3.14"]
|
||||
python-version: ["3.10", "3.11", "3.12", "3.13"]
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
@@ -64,18 +62,16 @@ jobs:
|
||||
comfyui:
|
||||
name: real ComfyUI compatibility
|
||||
runs-on: ubuntu-latest
|
||||
# Chromium and its system libraries come from the pinned Playwright image (tag = @playwright/test in
|
||||
# tools/vendor/package-lock.json; tests/vendor/browser-image.test.mjs checks it). The runner install
|
||||
# pulled them from Ubuntu's apt mirror on every run, and three runs (#349, #378, #382) crawled at
|
||||
# ~26 kB/s for 11-12 minutes until the job hit its limit after every test had passed.
|
||||
container:
|
||||
image: mcr.microsoft.com/playwright:v1.64.1-noble
|
||||
options: --user 1001
|
||||
timeout-minutes: 24
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
with:
|
||||
persist-credentials: false
|
||||
- uses: actions/setup-node@v7
|
||||
with:
|
||||
node-version: "24"
|
||||
cache: npm
|
||||
cache-dependency-path: tools/vendor/package-lock.json
|
||||
- name: Read the reviewed runtime pins
|
||||
id: runtime
|
||||
run: |
|
||||
@@ -100,42 +96,21 @@ jobs:
|
||||
ref: ${{ steps.runtime.outputs.kikoRevision }}
|
||||
path: .ci-kikotools
|
||||
persist-credentials: false
|
||||
# The image's own Python and Node 24: setup-python/setup-node install the runner's toolcache builds,
|
||||
# whose entry points name host paths that do not exist in the container. The image has no ensurepip
|
||||
# and the job has no root, so a checksum-pinned pip wheel bootstraps a user-space venv. The image
|
||||
# Python must be the reviewed host Python, so a pin or image bump that disagrees fails here.
|
||||
- name: Python environment on the reviewed host Python
|
||||
env:
|
||||
PIP_WHEEL: https://files.pythonhosted.org/packages/f3/6e/1736e5b4ae2b778ef2f81c47d797de9f891d4d8acb047a24ca37a60294dd/pip-26.2.1-py3-none-any.whl
|
||||
PIP_SHA256: 71138adf1f4ca900cdb7d289c21b7494329f2332b6d85f0e1c42108c0384ed3e
|
||||
PINNED_PYTHON: ${{ steps.runtime.outputs.pythonVersion }}
|
||||
run: |
|
||||
image="$(python3 -c 'import sys; print("%d.%d" % sys.version_info[:2])')"
|
||||
[ "$image" = "$PINNED_PYTHON" ] || { echo "::error::image Python $image is not the pinned $PINNED_PYTHON"; exit 1; }
|
||||
venv="$RUNNER_TEMP/venv" wheel="$RUNNER_TEMP/pip-26.2.1-py3-none-any.whl"
|
||||
python3 -m venv --without-pip "$venv"
|
||||
curl -fsSL --retry 3 -o "$wheel" "$PIP_WHEEL"
|
||||
echo "$PIP_SHA256 $wheel" | sha256sum -c -
|
||||
"$venv/bin/python" "$wheel/pip" install --no-index --no-cache-dir "$wheel"
|
||||
echo "$venv/bin" >> "$GITHUB_PATH"
|
||||
- run: python --version && node --version && node --version | grep -q '^v24\.'
|
||||
# pip's download cache (torch alone is ~200 MB); inside the container ~ is the job's HOME (/github/home).
|
||||
- name: Cache pip downloads
|
||||
uses: actions/cache@v6
|
||||
- uses: actions/setup-python@v7
|
||||
with:
|
||||
path: ~/.cache/pip
|
||||
key: pip-${{ runner.os }}-playwright-v1.64.1-${{ hashFiles('requirements-dev.txt', '.ci-comfyui/requirements.txt', 'tools/vendor/comfyui-runtime.json') }}
|
||||
python-version: ${{ steps.runtime.outputs.pythonVersion }}
|
||||
cache: pip
|
||||
cache-dependency-path: |
|
||||
requirements-dev.txt
|
||||
.ci-comfyui/requirements.txt
|
||||
tools/vendor/comfyui-runtime.json
|
||||
- name: Install CPU runtime and test dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install "torch==${{ steps.runtime.outputs.torchVersion }}" "torchvision==${{ steps.runtime.outputs.torchvisionVersion }}" --index-url https://download.pytorch.org/whl/cpu
|
||||
pip install -r .ci-comfyui/requirements.txt -r requirements-dev.txt "torch==${{ steps.runtime.outputs.torchVersion }}" "torchvision==${{ steps.runtime.outputs.torchvisionVersion }}"
|
||||
# npm's download cache, keyed on the lockfile. node_modules itself is still built fresh by npm ci.
|
||||
- name: Cache npm downloads
|
||||
uses: actions/cache@v6
|
||||
with:
|
||||
path: ~/.npm
|
||||
key: npm-${{ runner.os }}-playwright-v1.64.1-${{ hashFiles('tools/vendor/package-lock.json') }}
|
||||
- run: npm ci --prefix tools/vendor --ignore-scripts
|
||||
- run: npx --prefix tools/vendor playwright install --with-deps chromium
|
||||
- name: Classic and Nodes 2.0 startup and compatibility
|
||||
env:
|
||||
PM_COMFY_ROOT: ${{ github.workspace }}/.ci-comfyui
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
name: Viewer.js vendor and compatibility
|
||||
|
||||
on:
|
||||
merge_group:
|
||||
pull_request:
|
||||
branches: [main]
|
||||
push:
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# ComfyUI Prompt Manager
|
||||
|
||||

|
||||

|
||||

|
||||

|
||||

|
||||

|
||||
@@ -483,15 +483,11 @@ If you use [ComfyUI-Lora-Manager](https://github.com/willmiao/ComfyUI-Lora-Manag
|
||||
|
||||
Click **Import LoRA Data** in the settings panel to start the import. A progress popup shows real-time status as each LoRA is processed:
|
||||
|
||||
- Reads LoRA Manager's own settings the way LoRA Manager does: from the folder in `LORA_MANAGER_SETTINGS_DIR`, from the plugin's `settings.json` in portable mode, or from your user config folder, and only the keys that say where its files are
|
||||
- Finds each LoRA's metadata in both storage modes: beside the model (the default) or in the central sidecar folder, and imports only LoRAs whose model file still exists in a LoRA folder (the review says how many metadata files it skipped)
|
||||
- Links LoRA Manager's preview (including `.preview.webp` and video previews) and every downloaded example image and video where they are, without copying them
|
||||
- Downloads example images from CivitAI as 512px thumbnails only for LoRAs that have no local image
|
||||
- Scans all LoRA directories (including `extra_model_paths.yaml` paths) for metadata
|
||||
- Downloads example images from CivitAI as 512px thumbnails
|
||||
- Creates a prompt entry for each LoRA with trigger words, example prompts, and preview images
|
||||
- All imported LoRAs are tagged with `lora-manager` category for easy filtering
|
||||
|
||||
The System › Integrations review names the settings it read (source, storage mode and LoRA Manager library) and counts metadata files, LoRAs and LoRAs with local images. If LoRA Manager's settings cannot be used, it says why and nothing is changed.
|
||||
|
||||

|
||||
|
||||
After import, filter by the `lora-manager` category to browse your LoRA collection.
|
||||
@@ -933,7 +929,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.12 to 3.14 with an 85% coverage gate. `python -m unittest discover -s tests` still works; see `TESTING.md` ("Time limits and hangs") for the watchdog. Tests that delete a database file first check that nothing in the process still holds it open (`tests/open_handles.py`), so a connection leak fails on Linux the way it would on Windows.
|
||||
CI runs the suite on Linux, Windows and macOS for Python 3.10 to 3.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.
|
||||
|
||||
### Code Style
|
||||
|
||||
|
||||
+7
-353
@@ -86,18 +86,6 @@ 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`.
|
||||
@@ -275,77 +263,6 @@ 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. Refuse new previews (`refuse_auto_thumbnails()`), then detach the
|
||||
monitor's callbacks (Needs review recorder, preview queue, Images listing
|
||||
listener).
|
||||
2. `ImageMonitor.stop_monitoring` works against one 5 s deadline. Closing the
|
||||
image handler joins the image worker; watchdog's `observer.stop()` (not
|
||||
bounded here) and the observer join get what is left, which can be 0 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`. A Ctrl+C that lands
|
||||
between two steps (each step catches its own) is re-raised after the result is
|
||||
stored, with the parts not reached False, and the `atexit` callback is
|
||||
removed: the later steps are skipped, no summary line is logged, and a later
|
||||
call returns that result and runs nothing.
|
||||
|
||||
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:
|
||||
@@ -416,75 +333,6 @@ hold a walk open with a gated scan instead of sleeping.
|
||||
listing gets a 500. Rescan waits for a walk that started after the click
|
||||
(by walk number, not by clock), Rescans clicked during one walk share the
|
||||
next, and a failed walk is a 500 even when a previous listing exists.
|
||||
- **New outputs (added without a walk, PR 2).** The image monitor reports
|
||||
every settled output it accepts (by default `.png .jpg .jpeg .webp .gif`;
|
||||
videos only when their extension is in `monitoring.extensions`), linked or
|
||||
not, with its `realpath`, through `set_output_listener`. On its worker
|
||||
thread the API's `notify_output_settled` only does
|
||||
`loop.call_soon_threadsafe(output_settled, path)`, and does nothing before
|
||||
the loop is known or after `cleanup_gallery_listings`. `output_settled`
|
||||
matches the path against each walk's `resolved_root` (`normcase(realpath)`,
|
||||
computed in the executor; `root_contains` ends the match at a separator
|
||||
and ignores a Windows `\\?\` or `\\?\UNC\` prefix on either side).
|
||||
- **The add.** A root with a stored listing and no walk running gets the
|
||||
file added without a walk. `relative_part` slices the path below the
|
||||
root in its own case (Windows compares folded, slices the original);
|
||||
`_run_adds` places the pending files of one root in one executor batch
|
||||
(`_patch_listing`): `_settled_entry` applies the walk's filters
|
||||
(not a symlink or folder by `lstat`, `_name_suffix` in `MEDIA_EXTENSIONS`,
|
||||
no `thumbnails` folder, at most `MAX_SCAN_DEPTH` folders, every folder on
|
||||
the way listable by `os.access`, which on Windows ignores ACLs: a folder
|
||||
whose ACL denies listing is added there and corrected by the next walk). Where names fold case the stored
|
||||
spelling must match first, as the filters compare it: Windows through
|
||||
`nt._getfinalpathname` (its `realpath` keeps the given case when it cannot
|
||||
open the file), macOS through `F_GETPATH` on a non-blocking open (a FIFO
|
||||
never holds the executor); a case-insensitive mount on Linux is not
|
||||
checked and is corrected by the next walk. `insert_entry` builds a new list with
|
||||
the `(Path, mtime)` entry where the walk's newest-first sort puts it,
|
||||
replacing the same path's entry. `add_folder` copies both folder sets. The
|
||||
batch is stored only when no walk started and no drop happened while it
|
||||
ran (a walk started after the report lists the file itself); `walked_at`
|
||||
is untouched, so the 30 s refresh still walks and corrects any drift, and
|
||||
`walk_root` keeps the patched list object when that walk finds the same.
|
||||
Tests compare every added listing with a real walk (`_walked`).
|
||||
- **When it walks instead.** The root is marked dirty (`mark_root_dirty`,
|
||||
`_mark_dirty`) when the add cannot be exact: a walk is running (it may have
|
||||
passed the folder: the trailing walk), the path cannot be sliced, the
|
||||
`lstat` or the spelling read fails, the stored spelling differs, another
|
||||
file has the same mtime (a walk orders those by its
|
||||
visit order), the listing is at `MAX_GALLERY_FILES`, or the batch raised.
|
||||
A file the walk's filters exclude is neither added nor walked. A root
|
||||
without a listing, or a dropped one, is walked by its next read. A dirty
|
||||
root gets one background walk, no sooner than `rewalk_interval` =
|
||||
`min(15, max(5, 5 × last walk's duration))` seconds after the previous walk
|
||||
started. `_schedule_walk` never arms a timer while a walk runs (the walk's
|
||||
end re-checks `dirty`: the trailing walk), arms at most one per root and
|
||||
never sooner than 0.5 s, and keeps the earlier deadline when the stale
|
||||
refresh and a dirty root both want one; every walk start (read, Rescan,
|
||||
warm-up or timer) cancels it and forgets the pending adds. A failed walk
|
||||
clears `dirty`; a cancelled walk starts no trailing walk. When the output
|
||||
roots change (`reset_gallery_roots_cache`) every record is `retired`: its
|
||||
timer is cancelled and no add, mark, timer or trailing walk touches it
|
||||
until a reader or the warm-up asks for that root again (the monitor keeps
|
||||
watching a removed root until restart). Tests patch
|
||||
`REWALK_FLOOR`/`REWALK_CAP` and run the pending timer's callback by hand
|
||||
(`_fire`) instead of waiting for it, and hold an add batch in the executor
|
||||
(`_hold_adds`) to race it with a walk, a drop or the cleanup.
|
||||
- **Workspace caches after an add.** A patched listing shares every
|
||||
unchanged entry object with the one before it, so `listing_delta` finds
|
||||
the change by identity (a walk builds new tuples and is rebuilt in full,
|
||||
as is any change of more than `MAX_DELTA` entries). The per-file values
|
||||
build only the added files (`_workspace_listed_step`), the newest order
|
||||
places them by bisection (`_newest_step`, only when the per-file cache
|
||||
still holds exactly the values this call built: a page read and the warm
|
||||
build run in executor threads at once, and an order labelled with older
|
||||
values than it holds would add a file twice), the favorites lookup is copied
|
||||
and extended (`_favorite_listing_step`, rebuilt when a key collides), and
|
||||
the other orders, wanted paths and locators rebuild lazily. On a synthetic
|
||||
20,000-file listing a one-file add costs about 3 ms in the executor and
|
||||
15 ms for the warm build, against about 155 ms + 80 ms for a walk and a
|
||||
full build (`test_a_warm_build_after_an_add_builds_only_the_added_values`
|
||||
counts the values built and the sort keys read).
|
||||
- **Measuring it.** `tools/vendor/perf/images_load_probe.mjs` is a maintainer
|
||||
tool, never run in CI. After `npm ci --prefix tools/vendor` and
|
||||
`npx --prefix tools/vendor playwright install chromium`, run
|
||||
@@ -1575,8 +1423,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 in System ›
|
||||
Auto-tagging." (the server's own sentence), so no job starts only to fail.
|
||||
`#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.
|
||||
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
|
||||
@@ -2658,7 +2506,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–D134 assigned):
|
||||
Recorded deviations (plan D60–D75, D79–D82 renumbered D92–D107 and D111–D114; D121–D131 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
|
||||
@@ -2692,8 +2540,7 @@ 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. Since S11
|
||||
(#328) only Diagnostics & logs hands off (`SECTION_HANDOFFS`, to the admin page).
|
||||
plan's legacy coverage table); retired section by section, gone after S12.
|
||||
- **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
|
||||
@@ -3069,59 +2916,6 @@ 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
|
||||
@@ -3496,9 +3290,7 @@ workflow (README changelog).
|
||||
### LoRA reference import contract (#221, #249 System S10)
|
||||
|
||||
The LoRA import is staged (`py/lora_import.py`): discover the LoraManager
|
||||
`*.metadata.json` files and fingerprint every scanned one, kept or skipped (path, size,
|
||||
modification time and whether it was kept, #362; a skipped sidecar added, removed or edited
|
||||
after the review refuses Start with `plan-changed`, `test_start_refused_when_a_skipped_sidecar_changed`); stage each
|
||||
`*.metadata.json` files and fingerprint them (path, size, modification time); stage each
|
||||
record (CivitAI example images go into the image cache first, best effort, counted in
|
||||
`downloaded` / `download_errors`; each image is written to a temp file beside its cache
|
||||
entry and moved in place, so a half-written image is never seen as cached); the files are
|
||||
@@ -3549,158 +3341,20 @@ 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. Only an absolute `preview_url` is used; a relative one is skipped with a DEBUG line
|
||||
naming the sidecar. A sidecar that is not valid UTF-8 JSON is skipped like any unreadable
|
||||
one (#362). CivitAI downloads run only for a LoRA without local media; a record's
|
||||
images are linked in reverse order, so its first image (the cover) is the newest link and
|
||||
the card image. The folder guard keeps a folder from the settings, and each of ComfyUI's
|
||||
default LoRA folders (`models/loras`, `extra_model_paths.yaml`, `folder_paths`), only when
|
||||
it exists and is neither a filesystem root (`/`, `C:\`, a UNC share root) nor the home
|
||||
folder or a folder that contains it (realpath + `normcase`); a refused default folder that
|
||||
exists is listed in the layout's `refused_defaults` and logged as one warning each time the
|
||||
refused set changes (`RefusedDefaultsTests`). An unusable layout (no install, unreadable /
|
||||
non-JSON / non-object / oversized settings, no `platformdirs`, a missing or refused
|
||||
centralized sidecar root, "the sidecar folder is a filesystem root, the home folder or a
|
||||
folder that contains it") 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 (no example-images root), and trigger words from the default roots
|
||||
in alongside mode. A folder discovery cannot list answers the review with 500, code
|
||||
`folder-unreadable` ("A LoRA folder could not be listed (...). Nothing was changed.").
|
||||
`TriggerWordCache.load` never raises: an `OSError` is a warning, anything else is logged
|
||||
with its traceback (#362), both keep the previous map, and a running node retries a failed
|
||||
load only after 60 s (`TRIGGER_RETRY_SECONDS`). `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.
|
||||
|
||||
**Linked-only gate and the layout cache (#362).** While the integration is on, the LoRA
|
||||
allow-list (`ImagesMixin._lora_image_dirs`) is the layout's `media_roots` plus the CivitAI
|
||||
download cache, and `_media_roots(outputs)` appends it after the output folders for Library
|
||||
previews, the Library prompt detail, the per-prompt images page, `serve_output_image` for a
|
||||
root past the outputs, `inspect_image` and `reveal_image`. Under such a root only a file the
|
||||
database records as a linked image is served, inspected or revealed (`_linked_only_error`,
|
||||
by its given and its real path); anything else answers 403. Library responses mark that
|
||||
media `linked_only`: no preview file, no favorite, no Tag button. The allow-list's layout
|
||||
comes from `py/lora_layout_cache.py` (`tests/test_lora_layout_cache.py`, which injects the
|
||||
clocks, `stat` and `ntpath` and never sleeps): one entry keyed by the stripped configured
|
||||
path, reused while the plugin settings file's and the settings file's `(st_mtime_ns,
|
||||
st_size)` (or "missing") are unchanged, the install and every media root still resolve to
|
||||
themselves (a folder swapped for a symlink reloads and the guard runs again), and the entry
|
||||
is under `MAX_AGE_SECONDS` (30 s, monotonic clock) old. A load is stored only when no
|
||||
watched file was modified within `RACY_SECONDS` (3 s) of the load (git's racy rule: a write
|
||||
during the load, LoraManager's in-place write, FAT/HFS+ coarse times), every watched file
|
||||
could be checked (any `stat` error but not-found, and an offline Windows share, is "cannot
|
||||
validate"), the settings were readable, and no `clear()` ran meanwhile (generation
|
||||
counter). Hits take no lock; concurrent misses each load. A finished import (System job or
|
||||
legacy scan, both through `replace_lora_imports`) clears it. The integration-off check runs
|
||||
before it, and the configured-path rule (a configured path that is not an install under
|
||||
`custom_nodes` adds no folders) on its result; a hit also needs the same home folder, which
|
||||
the guard compares against. `serve_image` builds the allow-list in the
|
||||
executor. The cache writes one INFO line, "LoraManager layout changed: <source> settings,
|
||||
<mode> storage, <n> media folders", when a stored layout differs from the last one logged.
|
||||
The review, Start, the legacy scan and the trigger cache read the layout fresh. Image route
|
||||
tests clear the per-process cache in `ImageRouteCoverageCase.setUp`; `write_settings` in
|
||||
`TestLoraManagerMediaServing` back-dates its file 10 s so the first load is stored.
|
||||
|
||||
`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_layout_cache tests.test_lora_manager_layout \
|
||||
tests.test_lora_import tests.test_api_lora tests.test_lora_database tests.test_lora_utils \
|
||||
tests.test_api_images -v
|
||||
python3 -m unittest 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 in
|
||||
steps of 0.05, #369),
|
||||
`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. Every page shows the thresholds in
|
||||
steps of 0.05 (the System rows, the Tag image dialogs, the legacy sliders), so a saved
|
||||
threshold off that grid is refused ("General-tag confidence must be a multiple of
|
||||
0.05.", `validate_saved_threshold`; `draftProblem` holds back Preview tagging review and
|
||||
the batch Review with the same sentence) and falls back on load; an accepted one is
|
||||
stored exactly on the grid (#369). A per-request threshold keeps its own check: `/image-annotations/suggest` and
|
||||
`/jobs/autotag` take 0–1; the legacy `/autotag/single` and `/autotag/start` only require a number. `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
|
||||
|
||||
+19
-166
@@ -22,10 +22,7 @@ Features:
|
||||
- Comprehensive logging and diagnostics
|
||||
"""
|
||||
|
||||
import atexit
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
@@ -75,12 +72,6 @@ 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
|
||||
@@ -98,6 +89,8 @@ 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",
|
||||
@@ -148,16 +141,25 @@ 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. 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)
|
||||
# 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)
|
||||
# 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(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(api.notify_output_settled)
|
||||
_global_image_monitor.set_thumbnail_queue(_review_api.queue_capture_thumbnail)
|
||||
_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)
|
||||
api.close_auto_thumbnails(final=True)
|
||||
|
||||
_review_host.on_cleanup.insert(0, _stop_review_recording)
|
||||
|
||||
# Start monitoring
|
||||
_global_image_monitor.start_monitoring()
|
||||
@@ -227,155 +229,6 @@ 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. A Ctrl+C inside a step is logged and the
|
||||
step is False; one between steps is re-raised after the result (parts
|
||||
not reached are False) is stored and the ``atexit`` entry removed.
|
||||
"""
|
||||
global _background_work_result
|
||||
with _background_work_lock:
|
||||
if _background_work_result is not None:
|
||||
return dict(_background_work_result)
|
||||
monitor, thumbnails, tracker = (
|
||||
_global_image_monitor,
|
||||
api,
|
||||
_global_prompt_tracker,
|
||||
)
|
||||
# A part whose global is set has started; False until its step says
|
||||
# otherwise, so an interrupt between steps leaves an honest result.
|
||||
started = (
|
||||
("monitor", monitor),
|
||||
("thumbnails", thumbnails),
|
||||
("tracker", tracker),
|
||||
)
|
||||
result = {name: False for name, part in started if part is not None}
|
||||
try:
|
||||
logger = _teardown_logger()
|
||||
began = _clock()
|
||||
deadline = began + TEARDOWN_SECONDS
|
||||
|
||||
def left(own_bound):
|
||||
return max(0.0, min(own_bound, deadline - _clock()))
|
||||
|
||||
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}")
|
||||
finally:
|
||||
_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
|
||||
|
||||
@@ -192,10 +192,7 @@ class CategoryImportOperations:
|
||||
outcome = "imported"
|
||||
created[prompt_hash] = prompt_id
|
||||
result[_COUNTERS[outcome]] += 1
|
||||
# Linked last to first: every workspace list shows a prompt's media
|
||||
# newest link first, so the record's first image (the LoRA cover,
|
||||
# #355) is the one a Library card shows.
|
||||
for image in reversed(record.get("images") or []):
|
||||
for image in record.get("images") or []:
|
||||
if self._insert_image_link(conn, prompt_id, image):
|
||||
result["images"] += 1
|
||||
result["rows"].append({"name": name, "outcome": outcome})
|
||||
|
||||
@@ -1808,37 +1808,6 @@ 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,
|
||||
|
||||
@@ -268,10 +268,8 @@ class PromptManagerBase:
|
||||
|
||||
cache = get_trigger_cache()
|
||||
|
||||
# Lazy-load cache on first use; a failed load waits before the next try
|
||||
# (#355: never a rescan of LoraManager's folders on every run).
|
||||
retry_due = getattr(cache, "retry_due", lambda: True)
|
||||
if not cache.is_loaded and IntegrationConfig.LORA_MANAGER_PATH and retry_due():
|
||||
# Lazy-load cache on first use
|
||||
if not cache.is_loaded and IntegrationConfig.LORA_MANAGER_PATH:
|
||||
cache.load(IntegrationConfig.LORA_MANAGER_PATH)
|
||||
|
||||
modified, injected = inject_trigger_words(text, cache)
|
||||
|
||||
+7
-52
@@ -362,11 +362,8 @@ class PromptManagerAPI(
|
||||
# background walks refreshing them (py/api/gallery_listing.py, #348).
|
||||
self._root_listings: dict[str, RootListing] = {}
|
||||
self._listing_tasks: set = set()
|
||||
# The loop the walks run on, set by the startup warm-up or the first
|
||||
# walk; the monitor's thread hops onto it (notify_output_settled).
|
||||
# The loop the walks run on, set by the first walk.
|
||||
self._listing_loop = None
|
||||
# Set by cleanup_gallery_listings: no background walk starts after it.
|
||||
self._listings_closed = False
|
||||
# The workspace caches rebuilt behind a stored walk that changed the
|
||||
# listing, once the host's startup warm-up has run (#348).
|
||||
self._workspace_warm_enabled = False
|
||||
@@ -495,7 +492,7 @@ class PromptManagerAPI(
|
||||
"""The shared listing record of one output root (created empty)."""
|
||||
return self._root_listings.setdefault(str(root), RootListing())
|
||||
|
||||
def invalidate_gallery_cache(self, retire=False):
|
||||
def invalidate_gallery_cache(self):
|
||||
"""Drop every root's listing, so the next read waits for a new walk.
|
||||
|
||||
Called after files are deleted through the app (duplicate cleanup and
|
||||
@@ -506,9 +503,6 @@ class PromptManagerAPI(
|
||||
_workspace_listed and _card_facts); its sort orders are dropped, as
|
||||
the walk replaces them anyway, and so is the favorites lookup, which
|
||||
holds the old list.
|
||||
|
||||
``retire`` (the roots changed): no background walk of a root starts
|
||||
again until a reader or the warm-up asks for it.
|
||||
"""
|
||||
loop = self._listing_loop
|
||||
if loop is not None and loop.is_running() and _running_loop() is not loop:
|
||||
@@ -519,7 +513,7 @@ class PromptManagerAPI(
|
||||
|
||||
def drop():
|
||||
try:
|
||||
self._drop_listings(retire)
|
||||
self._drop_listings()
|
||||
finally:
|
||||
dropped.set()
|
||||
|
||||
@@ -527,19 +521,11 @@ class PromptManagerAPI(
|
||||
if not dropped.wait(_OFF_LOOP_DROP_WAIT):
|
||||
self.logger.warning("The Images listing drop is still queued")
|
||||
return
|
||||
self._drop_listings(retire)
|
||||
self._drop_listings()
|
||||
|
||||
def _drop_listings(self, retire=False):
|
||||
def _drop_listings(self):
|
||||
for record in list(self._root_listings.values()):
|
||||
record.drop()
|
||||
if retire:
|
||||
# The roots changed: a removed root must not come back through
|
||||
# a pending timer, a trailing walk or a monitor mark.
|
||||
record.retired = True
|
||||
record.dirty = False
|
||||
handle, record.rewalk_handle = record.rewalk_handle, None
|
||||
if handle is not None:
|
||||
handle.cancel()
|
||||
self._workspace_order_cache = {}
|
||||
self._favorite_listing_cache = {}
|
||||
|
||||
@@ -553,36 +539,15 @@ class PromptManagerAPI(
|
||||
# Each stored walk that changes a listing then rebuilds the workspace's
|
||||
# per-file values, orders and favorites behind it (_warm_workspace).
|
||||
self._workspace_warm_enabled = True
|
||||
self._listing_loop = asyncio.get_running_loop()
|
||||
self._listings_closed = False
|
||||
try:
|
||||
for root in self._get_all_output_dirs():
|
||||
self._root_listing(root).retired = False
|
||||
self._schedule_walk(root)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Could not warm the Images listing: {e}")
|
||||
|
||||
def notify_output_settled(self, path):
|
||||
"""The image monitor settled an output at *path* (any thread).
|
||||
|
||||
Called on the monitor's worker thread with the resolved path. It only
|
||||
hands the path to the event loop (``output_settled``) and never
|
||||
touches a listing from this thread; before the loop is known, or once
|
||||
the listings are cleaned up, it does nothing.
|
||||
"""
|
||||
loop = self._listing_loop
|
||||
if loop is None or self._listings_closed or not loop.is_running():
|
||||
return
|
||||
try:
|
||||
loop.call_soon_threadsafe(self.output_settled, path)
|
||||
except RuntimeError:
|
||||
pass # the loop closed meanwhile (shutdown)
|
||||
|
||||
async def cleanup_gallery_listings(self, app=None):
|
||||
"""Stop the background listing walks (aiohttp on_cleanup)."""
|
||||
self._workspace_warm_enabled = False
|
||||
# No mark, deferred walk or trailing walk starts a walk from now on.
|
||||
self._listings_closed = True
|
||||
# Before any await, so no delayed refresh can start a walk meanwhile.
|
||||
for record in self._root_listings.values():
|
||||
handle, record.rewalk_handle = record.rewalk_handle, None
|
||||
@@ -606,7 +571,7 @@ class PromptManagerAPI(
|
||||
|
||||
def reset_gallery_roots_cache(self):
|
||||
"""Drop every cached listing and derived workspace value (roots changed)."""
|
||||
self.invalidate_gallery_cache(retire=True)
|
||||
self.invalidate_gallery_cache()
|
||||
self._workspace_listed_cache = {}
|
||||
self._workspace_card_cache = {}
|
||||
|
||||
@@ -750,12 +715,7 @@ class PromptManagerAPI(
|
||||
# ── Shared utilities used by multiple mixins ──────────────────────
|
||||
|
||||
def _enrich_images(self, images):
|
||||
"""Add url and thumbnail_url to a flat list of image dicts.
|
||||
|
||||
Each dict without ``is_video`` also gets ``is_video`` and
|
||||
``media_type`` from its extension (#355: a linked LoraManager video
|
||||
outside the output folders is shown as a video, not a picture).
|
||||
"""
|
||||
"""Add url and thumbnail_url to a flat list of image dicts."""
|
||||
from urllib.parse import quote as url_quote
|
||||
|
||||
output_dirs = self._get_all_output_dirs()
|
||||
@@ -765,11 +725,6 @@ class PromptManagerAPI(
|
||||
if not image_path_str:
|
||||
continue
|
||||
|
||||
if "is_video" not in image:
|
||||
video = is_video_path(image_path_str)
|
||||
image["is_video"] = video
|
||||
image["media_type"] = "video" if video else "image"
|
||||
|
||||
img_path = Path(image_path_str)
|
||||
|
||||
# Set fallback url via image ID
|
||||
|
||||
@@ -10,15 +10,12 @@ filesystem or the database.
|
||||
|
||||
import math
|
||||
|
||||
from ..tagging_defaults import ( # noqa: F401 - re-exported for the route
|
||||
LLM_MODEL_TYPES as LLM_MODELS,
|
||||
MAX_INSTRUCTION_BYTES,
|
||||
MAX_INSTRUCTION_CHARS,
|
||||
MODEL_TYPES,
|
||||
WD14_MODEL_TYPES as WD14_MODELS,
|
||||
validate_template,
|
||||
validate_threshold,
|
||||
)
|
||||
WD14_MODELS = frozenset({"wd14-swinv2", "wd14-vit"})
|
||||
LLM_MODELS = frozenset({"gguf", "hf"})
|
||||
MODEL_TYPES = WD14_MODELS | LLM_MODELS
|
||||
|
||||
MAX_INSTRUCTION_CHARS = 4000
|
||||
MAX_INSTRUCTION_BYTES = 16 * 1024
|
||||
|
||||
SUGGEST_FIELDS = frozenset(
|
||||
{"target", "capture_token", "expected_epoch", "model_type", "options"}
|
||||
@@ -27,7 +24,7 @@ _COMMON = frozenset({"suggestion_limit", "use_gpu"})
|
||||
_WD14_REQUIRED = frozenset({"general_threshold", "character_threshold"})
|
||||
_LLM_REQUIRED = frozenset({"instructions"})
|
||||
|
||||
UNAVAILABLE_HINT = "Download it in System › Auto-tagging."
|
||||
UNAVAILABLE_HINT = "Download it from the legacy Auto-tag page."
|
||||
|
||||
|
||||
class OptionError(ValueError):
|
||||
@@ -35,17 +32,27 @@ class OptionError(ValueError):
|
||||
|
||||
|
||||
def _threshold(name: str, value) -> float:
|
||||
try:
|
||||
return validate_threshold(name, value)
|
||||
except ValueError as error:
|
||||
raise OptionError(str(error)) from None
|
||||
if type(value) not in (int, float) or not math.isfinite(value):
|
||||
raise OptionError(f"{name} must be a finite number")
|
||||
if not 0 <= value <= 1:
|
||||
raise OptionError(f"{name} must be between 0 and 1")
|
||||
return float(value)
|
||||
|
||||
|
||||
def _instructions(value) -> str:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise OptionError("instructions must be non-empty text")
|
||||
if len(value) > MAX_INSTRUCTION_CHARS:
|
||||
raise OptionError(
|
||||
f"instructions must be at most {MAX_INSTRUCTION_CHARS} characters"
|
||||
)
|
||||
try:
|
||||
return validate_template(value)
|
||||
except ValueError as error:
|
||||
raise OptionError(str(error)) from None
|
||||
encoded = value.encode("utf-8")
|
||||
except UnicodeEncodeError:
|
||||
raise OptionError("instructions must be valid Unicode text") from None
|
||||
if len(encoded) > MAX_INSTRUCTION_BYTES:
|
||||
raise OptionError("instructions are too long")
|
||||
return value
|
||||
|
||||
|
||||
def validate_options(model_type, options, *, default_limit: int, max_limit: int):
|
||||
|
||||
+85
-628
@@ -1,26 +1,14 @@
|
||||
"""AutoTag API routes for PromptManager."""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import functools
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
from urllib.parse import quote
|
||||
|
||||
from .job_routes import busy_response
|
||||
from .jobs import (
|
||||
AUTOTAG_BATCH_JOB_KIND,
|
||||
AUTOTAG_GROUP_NAMES,
|
||||
MODEL_DOWNLOAD_JOB_KIND,
|
||||
JobBusy,
|
||||
)
|
||||
from .admin import _plural, _scan_location, _scan_roots
|
||||
from .maintenance import disk_plan, short_space_message
|
||||
from .jobs import AUTOTAG_BATCH_JOB_KIND
|
||||
from .prompts import thumbnail_url_for
|
||||
from . import thumbnail_settings
|
||||
from .. import model_gate
|
||||
@@ -28,104 +16,20 @@ from .. import model_gate
|
||||
from aiohttp import web
|
||||
|
||||
_INTERNAL_ERROR = "An internal error occurred. Check server logs for details."
|
||||
|
||||
# Tag output images (#328): the native prompt-tag batch.
|
||||
AUTOTAG_JOB_TITLE = "Tag output images"
|
||||
AUTOTAG_PREVIEW_LEASE = "Tag output images review"
|
||||
AUTOTAG_OWNER_KIND = "autotag-job"
|
||||
BATCH_PREVIEW_BODY_MAX = 1024
|
||||
BATCH_JOB_BODY_MAX = 32 * 1024
|
||||
REVIEW_REQUIRED = (
|
||||
"Replacing tags needs a review of each image. Use Preview tagging review, "
|
||||
"or choose Skip or Merge for a batch run."
|
||||
)
|
||||
BATCH_CANCELLED = "Cancelled. Tags already added were kept."
|
||||
_BUSY_ERROR = "Auto-tagging is already running"
|
||||
MODEL_BUSY_ERROR = (
|
||||
"A tagging model is busy with another operation. Try again when it finishes."
|
||||
)
|
||||
|
||||
|
||||
# Who holds the model gate, as a 409 model-busy names it (``operation``). An
|
||||
# owner is ``"<kind>:<uuid>"`` (``model_gate.new_owner``); a download names
|
||||
# its own job, whose title says which model.
|
||||
MODEL_DOWNLOAD_OWNER_KIND = "model-download"
|
||||
OWNER_LABELS = {
|
||||
"autotag-job": "Tag output images",
|
||||
"legacy-batch": "Tag output images (legacy page)",
|
||||
"legacy-single": "Tagging one image (legacy page)",
|
||||
"legacy-unload": "Unloading the tagging model",
|
||||
"image-tag-job": "Auto-tag selected images",
|
||||
"image-annotations": "Generate suggestions",
|
||||
}
|
||||
ANOTHER_OPERATION = "Another tagging operation"
|
||||
_DOWNLOAD_FAILED = "The download failed. Check server logs for details."
|
||||
# POST /jobs/model-download carries only {model_type}.
|
||||
MODEL_DOWNLOAD_BODY_MAX = 1024
|
||||
# The running download's gate owner and job title. Only one download runs at a
|
||||
# time and each needs the gate, so a new start replaces the entry.
|
||||
_download_titles = {}
|
||||
|
||||
|
||||
def _holder_label(holder: Optional[str]) -> str:
|
||||
if not holder:
|
||||
return ANOTHER_OPERATION
|
||||
kind = holder.split(":", 1)[0]
|
||||
if kind == MODEL_DOWNLOAD_OWNER_KIND:
|
||||
return _download_titles.get(holder, ANOTHER_OPERATION)
|
||||
return OWNER_LABELS.get(kind, ANOTHER_OPERATION)
|
||||
|
||||
|
||||
def model_busy_response() -> web.Response:
|
||||
"""``409 model-busy``: another operation holds the shared model gate.
|
||||
|
||||
``operation`` names the holder (``OWNER_LABELS``), or "Another tagging
|
||||
operation" when the gate was freed meanwhile or the kind is unknown.
|
||||
"""
|
||||
"""``409 model-busy``: another operation holds the shared model gate."""
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "model-busy",
|
||||
"error": MODEL_BUSY_ERROR,
|
||||
"operation": _holder_label(model_gate.holder()),
|
||||
},
|
||||
{"success": False, "code": "model-busy", "error": MODEL_BUSY_ERROR},
|
||||
status=409,
|
||||
)
|
||||
|
||||
|
||||
_SIZE = re.compile(r"^(~?[0-9.]+)\s*([KMGT]?B)$")
|
||||
|
||||
|
||||
def size_label(size: str) -> str:
|
||||
"""The engine table's size with a space before the unit (``~400 MB``)."""
|
||||
match = _SIZE.match(str(size or "").strip())
|
||||
return f"{match.group(1)} {match.group(2)}" if match else str(size or "")
|
||||
|
||||
|
||||
def public_models(status: dict, loaded_type: Optional[str]) -> dict:
|
||||
"""Each engine model's public entry: names and sizes from ``MODELS``, the
|
||||
flags from ``status``; never a path (``model_path`` is dropped)."""
|
||||
from ..autotag import MODELS
|
||||
|
||||
models = {}
|
||||
for model_type, config in MODELS.items():
|
||||
flags = status.get(model_type) or {}
|
||||
entry = {
|
||||
"name": config["name"],
|
||||
"label": config["label"],
|
||||
"description": config["description"],
|
||||
"size": config["size"],
|
||||
"size_label": size_label(config["size"]),
|
||||
"downloaded": bool(flags.get("downloaded")),
|
||||
"loaded": loaded_type == model_type,
|
||||
}
|
||||
for key in ("model_exists", "mmproj_exists"):
|
||||
if key in flags:
|
||||
entry[key] = bool(flags[key])
|
||||
models[model_type] = entry
|
||||
return models
|
||||
|
||||
|
||||
# Only these may be handed to the tagging engine or listed by the scanner.
|
||||
_MEDIA_EXTENSIONS = frozenset(
|
||||
{".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".tif"}
|
||||
@@ -161,39 +65,6 @@ def path_is_within(candidate, root) -> bool:
|
||||
return child != parent and child.is_relative_to(parent)
|
||||
|
||||
|
||||
def _thresholds_or_saved(general, character):
|
||||
"""The request's WD14 thresholds as floats; one it omits is the saved
|
||||
default (``autotag.defaults``, OQ4). Raises ValueError/TypeError."""
|
||||
from ..config import AutotagConfig
|
||||
|
||||
return (
|
||||
AutotagConfig.GENERAL_THRESHOLD if general is None else float(general),
|
||||
AutotagConfig.CHARACTER_THRESHOLD if character is None else float(character),
|
||||
)
|
||||
|
||||
|
||||
async def _single(event):
|
||||
"""A stream of one event."""
|
||||
yield event
|
||||
|
||||
|
||||
_monotonic = time.monotonic # the legacy frames' throttle clock (patched in tests)
|
||||
|
||||
# What one image came to in a batch (``on_item``): fixed strings, never an
|
||||
# exception's text.
|
||||
SKIP_DETAILS = {
|
||||
"unlinked": "File not found",
|
||||
"prompt already processed": "Prompt handled by another image",
|
||||
"outside allowed directories": "Outside the output folders",
|
||||
"file missing": "File not found",
|
||||
"already tagged": "Prompt already tagged",
|
||||
}
|
||||
|
||||
|
||||
def _added(count):
|
||||
return f"Added {count} tag" if count == 1 else f"Added {count} tags"
|
||||
|
||||
|
||||
async def _refusal(title):
|
||||
"""The single error event of an auto-tag batch refused by another job."""
|
||||
message = f"{title} already running"
|
||||
@@ -212,10 +83,6 @@ class AutotagRoutesMixin:
|
||||
async def download_autotag_model_route(request):
|
||||
return await self.download_autotag_model(request)
|
||||
|
||||
@routes.post("/prompt_manager/autotag/batch/preview")
|
||||
async def preview_autotag_batch_route(request):
|
||||
return await self.preview_autotag_batch(request)
|
||||
|
||||
@routes.post("/prompt_manager/autotag/start")
|
||||
async def start_autotag_route(request):
|
||||
return await self.start_autotag(request)
|
||||
@@ -243,24 +110,18 @@ class AutotagRoutesMixin:
|
||||
from ..config import AutotagConfig
|
||||
|
||||
service = get_autotag_service()
|
||||
# Off the loop: the HuggingFace check walks the whole HF cache.
|
||||
status = await self._run_in_executor(service.get_models_status)
|
||||
loaded = service.is_model_loaded()
|
||||
loaded_type = service.get_loaded_model_type() if loaded else None
|
||||
models_status = service.get_models_status()
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"models": public_models(status, loaded_type),
|
||||
# The saved tagging defaults (#328, System › Auto-tagging).
|
||||
"default_prompt": AutotagConfig.TEMPLATE,
|
||||
"model_loaded": loaded,
|
||||
"loaded_model_type": loaded_type,
|
||||
"wd14_general_threshold": AutotagConfig.GENERAL_THRESHOLD,
|
||||
"wd14_character_threshold": AutotagConfig.CHARACTER_THRESHOLD,
|
||||
"models": models_status,
|
||||
"default_prompt": service.default_prompt,
|
||||
"model_loaded": service.is_model_loaded(),
|
||||
"loaded_model_type": service.get_loaded_model_type(),
|
||||
"wd14_general_threshold": service.wd14_general_threshold,
|
||||
"wd14_character_threshold": service.wd14_character_threshold,
|
||||
"suggestion_limit": AutotagConfig.suggestion_limit(),
|
||||
"model_type": AutotagConfig.MODEL_TYPE,
|
||||
"existing": AutotagConfig.EXISTING,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -298,156 +159,55 @@ class AutotagRoutesMixin:
|
||||
return response
|
||||
|
||||
async def download_autotag_model(self, request):
|
||||
"""Download a tagging model, streaming the legacy events (POST only).
|
||||
|
||||
The download is the model_download Activity job: this request
|
||||
attaches to a running download of the same model or starts one, and
|
||||
streams the legacy frames until the job ends. A client disconnect
|
||||
ends only the stream; the download goes on. A running download of
|
||||
another model answers 409 model-busy (it holds the gate).
|
||||
"""
|
||||
from ..autotag import MODELS, get_autotag_service
|
||||
|
||||
"""Download an AutoTag model with streaming progress (POST only)."""
|
||||
model_type = request.match_info.get("model_type")
|
||||
if model_type not in MODELS:
|
||||
invalid = {"type": "error", "message": f"Invalid model type: {model_type}"}
|
||||
return await self._stream_sse(request, _single(invalid))
|
||||
try:
|
||||
job = self._jobs.active(MODEL_DOWNLOAD_JOB_KIND)
|
||||
if job is not None and job.scope.get("model_type") != model_type:
|
||||
return model_busy_response()
|
||||
if job is None:
|
||||
job, refusal = self._start_model_download(
|
||||
model_type, get_autotag_service()
|
||||
owner = model_gate.new_owner("legacy-download")
|
||||
|
||||
async def events():
|
||||
try:
|
||||
from ..autotag import get_autotag_service
|
||||
|
||||
service = get_autotag_service()
|
||||
if model_type not in service.models_config:
|
||||
yield {
|
||||
"type": "error",
|
||||
"message": f"Invalid model type: {model_type}",
|
||||
}
|
||||
return
|
||||
|
||||
yield {
|
||||
"type": "progress",
|
||||
"progress": 0,
|
||||
"status": "Starting download...",
|
||||
}
|
||||
|
||||
def progress_callback(status: str, progress: float):
|
||||
pass
|
||||
|
||||
success = await model_gate.run(
|
||||
owner, service.download_model, model_type, progress_callback
|
||||
)
|
||||
if refusal is not None:
|
||||
return refusal
|
||||
except Exception:
|
||||
self.logger.exception("Download model error")
|
||||
failed = {"type": "error", "message": _INTERNAL_ERROR}
|
||||
return await self._stream_sse(request, _single(failed))
|
||||
return await self._stream_sse(request, self._download_frames(job))
|
||||
if success:
|
||||
yield {
|
||||
"type": "complete",
|
||||
"progress": 100,
|
||||
"status": "Download complete",
|
||||
}
|
||||
else:
|
||||
yield {"type": "error", "message": "Download failed"}
|
||||
except Exception:
|
||||
self.logger.exception("Download model error")
|
||||
yield {
|
||||
"type": "error",
|
||||
"message": _INTERNAL_ERROR,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
async def _download_frames(job):
|
||||
"""The legacy download frames for job: start, then its end."""
|
||||
yield {"type": "progress", "progress": 0, "status": "Starting download..."}
|
||||
async with contextlib.aclosing(job.events()) as events:
|
||||
async for _payload in events:
|
||||
pass
|
||||
if job.state == "finished":
|
||||
yield {"type": "complete", "progress": 100, "status": "Download complete"}
|
||||
else:
|
||||
yield {"type": "error", "message": "Download failed"}
|
||||
|
||||
# ── model download job (#328) ───────────────────────────────────────
|
||||
|
||||
async def start_model_download_job(self, request):
|
||||
"""``POST /jobs/model-download {model_type}``: 202 with the job.
|
||||
|
||||
400 for a bad body, 409 ``busy`` while a download (or an exclusive
|
||||
job) runs, 409 ``model-busy`` while another operation holds the gate.
|
||||
"""
|
||||
data, error = await self._read_json_body(
|
||||
request, max_bytes=MODEL_DOWNLOAD_BODY_MAX
|
||||
)
|
||||
if error is not None:
|
||||
return error
|
||||
if set(data) != {"model_type"}:
|
||||
return self._bad_request("The body must be exactly {model_type}")
|
||||
from ..autotag import MODELS, get_autotag_service
|
||||
|
||||
model_type = data["model_type"]
|
||||
if not isinstance(model_type, str) or model_type not in MODELS:
|
||||
return self._bad_request(
|
||||
"model_type must be one of " + ", ".join(sorted(MODELS))
|
||||
)
|
||||
try:
|
||||
job, refusal = self._start_model_download(model_type, get_autotag_service())
|
||||
except Exception as e:
|
||||
self.logger.exception("Model download start error")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e)}, status=500
|
||||
)
|
||||
if refusal is not None:
|
||||
return refusal
|
||||
return web.json_response({"success": True, "job": job.summary()}, status=202)
|
||||
|
||||
def _start_model_download(self, model_type, service):
|
||||
"""Register the download job: ``(job, None)`` or ``(None, refusal)``.
|
||||
|
||||
Synchronous (no await between the checks and the start). The gate is
|
||||
taken before the job registers and released when the runner ends,
|
||||
also when the task is cancelled before the runner ever starts.
|
||||
"""
|
||||
from ..autotag import MODELS
|
||||
|
||||
busy = self._jobs.blocker(MODEL_DOWNLOAD_JOB_KIND, admission="read")
|
||||
if busy is not None:
|
||||
return None, busy_response(JobBusy(busy))
|
||||
owner = model_gate.new_owner(MODEL_DOWNLOAD_OWNER_KIND)
|
||||
if not model_gate.try_acquire(owner):
|
||||
return None, model_busy_response()
|
||||
config = MODELS[model_type]
|
||||
label = config["label"]
|
||||
title = f"Tagging model download · {label}"
|
||||
_download_titles.clear()
|
||||
_download_titles[owner] = title
|
||||
return model_busy_response()
|
||||
try:
|
||||
job = self._jobs.start(
|
||||
MODEL_DOWNLOAD_JOB_KIND,
|
||||
title,
|
||||
admission="read",
|
||||
cancellable=False,
|
||||
scope={
|
||||
"model_type": model_type,
|
||||
"label": label,
|
||||
"size_label": size_label(config["size"]),
|
||||
},
|
||||
runner=functools.partial(
|
||||
self._run_model_download, service, model_type, owner
|
||||
),
|
||||
)
|
||||
except BaseException as refused:
|
||||
# JobBusy, or anything else the registry raised: give the gate back.
|
||||
model_gate.release(owner)
|
||||
_download_titles.pop(owner, None)
|
||||
if isinstance(refused, JobBusy):
|
||||
return None, busy_response(refused)
|
||||
raise
|
||||
|
||||
def ended(_task):
|
||||
model_gate.release(owner)
|
||||
if model_gate.holder() != owner: # not while a thread still runs
|
||||
_download_titles.pop(owner, None)
|
||||
|
||||
job.task.add_done_callback(ended)
|
||||
return job, None
|
||||
|
||||
async def _run_model_download(self, service, model_type, owner, job):
|
||||
label = job.scope["label"]
|
||||
job.update(status_line=f"Downloading {label} ({job.scope['size_label']})")
|
||||
job.log_line(f"Download started: {label}")
|
||||
job.publish({"type": "status", "status": job.status_line})
|
||||
try:
|
||||
ok = await model_gate.run(owner, service.download_model, model_type, None)
|
||||
except Exception:
|
||||
self.logger.exception(f"Tagging model download failed: {model_type}")
|
||||
ok = False
|
||||
return await self._stream_sse(request, events())
|
||||
finally:
|
||||
# Before the job ends, so a reader that sees the end finds the
|
||||
# gate free (a running download thread defers the release).
|
||||
model_gate.release(owner)
|
||||
if ok:
|
||||
job.update(status_line=f"Downloaded {label}")
|
||||
job.log_line(f"Download finished: {label}")
|
||||
job.publish({"type": "status", "status": job.status_line})
|
||||
job.finish("finished")
|
||||
return
|
||||
job.update(status_line=_DOWNLOAD_FAILED)
|
||||
job.log_line(f"Download failed: {label}")
|
||||
job.publish({"type": "status", "status": job.status_line})
|
||||
job.finish("failed")
|
||||
|
||||
async def _read_start_params(self, request) -> Optional[dict]:
|
||||
"""Merge query-string and optional JSON-body parameters for autotag/start.
|
||||
@@ -478,48 +238,35 @@ class AutotagRoutesMixin:
|
||||
and not t.startswith("append:")
|
||||
]
|
||||
|
||||
async def _apply_tags_to_prompt(self, prompt_id, tags) -> int:
|
||||
"""Merge ``tags`` into the prompt; how many new tags were added."""
|
||||
async def _apply_tags_to_prompt(self, prompt_id, tags) -> bool:
|
||||
"""Merge ``tags`` into the prompt; True when something new was added."""
|
||||
existing_prompt = await self._run_in_executor(
|
||||
self.db.get_prompt_by_id, prompt_id
|
||||
)
|
||||
if not existing_prompt:
|
||||
return 0
|
||||
return False
|
||||
existing_tags = existing_prompt.get("tags", [])
|
||||
if isinstance(existing_tags, str):
|
||||
existing_tags = [t.strip() for t in existing_tags.split(",") if t.strip()]
|
||||
new_tags = list(dict.fromkeys(t for t in tags if t not in existing_tags))
|
||||
new_tags = [t for t in tags if t not in existing_tags]
|
||||
if not new_tags:
|
||||
return 0
|
||||
return False
|
||||
await self._run_in_executor(
|
||||
self.db.update_prompt_metadata, prompt_id, tags=existing_tags + new_tags
|
||||
)
|
||||
return len(new_tags)
|
||||
return True
|
||||
|
||||
async def _autotag_batch_events(self, service, opts: dict):
|
||||
"""Yield SSE progress dicts while tagging every image linked to a prompt.
|
||||
"""Yield SSE progress dicts while tagging every image linked to a prompt."""
|
||||
import time as _time
|
||||
|
||||
Both batches run here: the legacy page's (its frames unchanged) and
|
||||
the ``autotag`` job's. ``opts``: ``model_type``, ``keep_in_memory``,
|
||||
``existing`` (``skip`` leaves prompts with real tags alone, ``merge``
|
||||
adds to them), the two thresholds, and optionally ``gate_owner``,
|
||||
``prompt`` (the template, passed to each call; the legacy page sets
|
||||
``custom_prompt`` instead), ``cancelled`` (checked before each image:
|
||||
a cancelled run returns without a ``complete`` frame) and
|
||||
``on_item(index, total, group, image_path, detail)`` (one outcome per
|
||||
image, ``SKIP_DETAILS`` for the skips).
|
||||
"""
|
||||
model_type = opts["model_type"]
|
||||
keep_in_memory = opts["keep_in_memory"]
|
||||
cancelled = opts.get("cancelled") or (lambda: False)
|
||||
on_item = opts.get("on_item") or (lambda *_args: None)
|
||||
status = await self._run_in_executor(service.get_models_status)
|
||||
status = service.get_models_status()
|
||||
if not status.get(model_type, {}).get("downloaded"):
|
||||
yield {"type": "error", "message": f"Model {model_type} not downloaded"}
|
||||
return
|
||||
|
||||
if cancelled():
|
||||
return
|
||||
yield {"type": "progress", "progress": 0, "status": "Loading model..."}
|
||||
run = self._gated_runner(opts.get("gate_owner"))
|
||||
try:
|
||||
@@ -532,21 +279,15 @@ class AutotagRoutesMixin:
|
||||
}
|
||||
return
|
||||
|
||||
if opts.get("custom_prompt"):
|
||||
if opts["custom_prompt"]:
|
||||
service.custom_prompt = opts["custom_prompt"]
|
||||
tag_kwargs = {
|
||||
"general_threshold": opts["general_threshold"],
|
||||
"character_threshold": opts["character_threshold"],
|
||||
}
|
||||
if opts.get("prompt"):
|
||||
tag_kwargs["prompt"] = opts["prompt"]
|
||||
|
||||
yield {
|
||||
"type": "progress",
|
||||
"progress": 5,
|
||||
"status": "Model loaded. Fetching all images from database...",
|
||||
}
|
||||
images = await self._run_in_executor(self.db.autotag_batch_candidates)
|
||||
images = await self._run_in_executor(self.db.get_all_images)
|
||||
total = len(images)
|
||||
if total == 0:
|
||||
yield {
|
||||
@@ -558,8 +299,7 @@ class AutotagRoutesMixin:
|
||||
"skipped_outside_roots": 0,
|
||||
"status": "No images with linked prompts found",
|
||||
}
|
||||
if not keep_in_memory:
|
||||
service.unload_model()
|
||||
service.unload_model()
|
||||
return
|
||||
|
||||
yield {
|
||||
@@ -576,7 +316,7 @@ class AutotagRoutesMixin:
|
||||
"skipped_outside_roots": 0,
|
||||
}
|
||||
tagged_prompt_ids = set()
|
||||
last_update = _monotonic()
|
||||
last_update = _time.monotonic()
|
||||
finished = False
|
||||
|
||||
def progress_event(i, message):
|
||||
@@ -592,8 +332,6 @@ class AutotagRoutesMixin:
|
||||
|
||||
try:
|
||||
for i, image_data in enumerate(images):
|
||||
if cancelled():
|
||||
return
|
||||
image_path = image_data.get("image_path")
|
||||
prompt_id = image_data.get("prompt_id")
|
||||
skip_reason = None
|
||||
@@ -612,7 +350,7 @@ class AutotagRoutesMixin:
|
||||
counts["skipped_outside_roots"] += 1
|
||||
else:
|
||||
skip_reason = "file missing"
|
||||
elif opts["existing"] == "skip" and self._real_tags(
|
||||
elif opts["skip_tagged"] and self._real_tags(
|
||||
image_data.get("prompt_tags", [])
|
||||
):
|
||||
tagged_prompt_ids.add(prompt_id)
|
||||
@@ -620,8 +358,7 @@ class AutotagRoutesMixin:
|
||||
|
||||
if skip_reason is not None:
|
||||
counts["skipped"] += 1
|
||||
on_item(i, total, "skipped", image_path, SKIP_DETAILS[skip_reason])
|
||||
now = _monotonic()
|
||||
now = _time.monotonic()
|
||||
if (now - last_update) >= 0.5 or i == total - 1:
|
||||
yield progress_event(
|
||||
i, f"Skipping {i + 1}/{total} ({skip_reason})..."
|
||||
@@ -631,17 +368,17 @@ class AutotagRoutesMixin:
|
||||
continue
|
||||
|
||||
try:
|
||||
tags = await run(service.generate_tags, allowed_path, **tag_kwargs)
|
||||
tags = await run(
|
||||
service.generate_tags,
|
||||
allowed_path,
|
||||
general_threshold=opts["general_threshold"],
|
||||
character_threshold=opts["character_threshold"],
|
||||
)
|
||||
counts["processed"] += 1
|
||||
added = 0
|
||||
if tags:
|
||||
added = await self._apply_tags_to_prompt(prompt_id, tags)
|
||||
if added:
|
||||
if tags and await self._apply_tags_to_prompt(prompt_id, tags):
|
||||
counts["tagged"] += 1
|
||||
on_item(i, total, "tagged", image_path, _added(added))
|
||||
else:
|
||||
counts["skipped"] += 1
|
||||
on_item(i, total, "unchanged", image_path, "No new tags")
|
||||
tagged_prompt_ids.add(prompt_id)
|
||||
except Exception as img_err:
|
||||
self.logger.error(
|
||||
@@ -650,11 +387,10 @@ class AutotagRoutesMixin:
|
||||
)
|
||||
counts["errors"] += 1
|
||||
counts["processed"] += 1
|
||||
on_item(i, total, "errors", image_path, "Tagging failed")
|
||||
|
||||
yield progress_event(i, f"Processing {i + 1}/{total}...")
|
||||
await asyncio.sleep(0.01)
|
||||
last_update = _monotonic()
|
||||
last_update = _time.monotonic()
|
||||
|
||||
finished = True
|
||||
if not keep_in_memory:
|
||||
@@ -682,22 +418,16 @@ class AutotagRoutesMixin:
|
||||
opts = {
|
||||
"model_type": params.get("model_type", "gguf"),
|
||||
"custom_prompt": params.get("prompt", ""),
|
||||
"existing": (
|
||||
"skip"
|
||||
if str(params.get("skip_tagged", "true")).lower() == "true"
|
||||
else "merge"
|
||||
),
|
||||
"skip_tagged": str(params.get("skip_tagged", "true")).lower() == "true",
|
||||
"keep_in_memory": str(params.get("keep_in_memory", "true")).lower()
|
||||
== "true",
|
||||
"general_threshold": params.get("general_threshold"),
|
||||
"character_threshold": params.get("character_threshold"),
|
||||
}
|
||||
try:
|
||||
opts["general_threshold"], opts["character_threshold"] = (
|
||||
_thresholds_or_saved(
|
||||
opts["general_threshold"], opts["character_threshold"]
|
||||
)
|
||||
)
|
||||
for key in ("general_threshold", "character_threshold"):
|
||||
if opts[key] is not None:
|
||||
opts[key] = float(opts[key])
|
||||
except (ValueError, TypeError):
|
||||
return self._bad_request(
|
||||
"general_threshold and character_threshold must be numeric"
|
||||
@@ -727,9 +457,7 @@ class AutotagRoutesMixin:
|
||||
# never runs beside a scan, restore or the workspace tagging job.
|
||||
busy = self._jobs.blocker(AUTOTAG_BATCH_JOB_KIND, admission="records")
|
||||
if busy is not None:
|
||||
# A second legacy batch keeps its JSON 409; the workspace's Tag
|
||||
# output images job is named in a stream frame.
|
||||
if busy.kind == AUTOTAG_BATCH_JOB_KIND and busy.title != AUTOTAG_JOB_TITLE:
|
||||
if busy.kind == AUTOTAG_BATCH_JOB_KIND: # unchanged payload
|
||||
return web.json_response(
|
||||
{"success": False, "error": _BUSY_ERROR}, status=409
|
||||
)
|
||||
@@ -739,7 +467,7 @@ class AutotagRoutesMixin:
|
||||
return await self._stream_sse(request, _refusal(busy.title))
|
||||
async with self._jobs.hold(
|
||||
AUTOTAG_BATCH_JOB_KIND,
|
||||
OWNER_LABELS["legacy-batch"],
|
||||
"Auto-tag images (legacy page)",
|
||||
admission="records",
|
||||
) as job:
|
||||
response = await self._start_autotag_batch(request, opts, events)
|
||||
@@ -763,278 +491,6 @@ class AutotagRoutesMixin:
|
||||
finally:
|
||||
_autotag_batch_lock.release()
|
||||
|
||||
# ── Tag output images (#328) ────────────────────────────────────────
|
||||
|
||||
async def _batch_body(self, request, keys, max_bytes):
|
||||
"""``(data, None)`` for a JSON object with exactly ``keys``."""
|
||||
data, error = await self._read_json_body(request, max_bytes=max_bytes)
|
||||
if error is not None:
|
||||
return None, error
|
||||
if set(data) != set(keys):
|
||||
return None, self._bad_request(
|
||||
"The body must be exactly {" + ", ".join(keys) + "}"
|
||||
)
|
||||
return data, None
|
||||
|
||||
@staticmethod
|
||||
def _batch_choice(data):
|
||||
"""The model's engine entry and the policy, or a 400 message."""
|
||||
from ..autotag import MODELS
|
||||
from ..config import EXISTING_POLICIES
|
||||
|
||||
model_type, existing = data.get("model_type"), data.get("existing")
|
||||
if not isinstance(model_type, str) or model_type not in MODELS:
|
||||
return None, "model_type must be one of " + ", ".join(sorted(MODELS))
|
||||
if not isinstance(existing, str) or existing not in EXISTING_POLICIES:
|
||||
return None, "existing must be one of " + ", ".join(EXISTING_POLICIES)
|
||||
return MODELS[model_type], None
|
||||
|
||||
async def preview_autotag_batch(self, request):
|
||||
"""Review Tag output images: counts, model, backup and disk; no write.
|
||||
|
||||
Opens no image file. Refused (409 ``busy``) while an exclusive job runs.
|
||||
"""
|
||||
try:
|
||||
async with self._jobs.lease(AUTOTAG_PREVIEW_LEASE):
|
||||
return await self._preview_autotag_batch(request)
|
||||
except JobBusy as busy:
|
||||
return busy_response(busy)
|
||||
|
||||
async def _preview_autotag_batch(self, request):
|
||||
from ..autotag import get_autotag_service
|
||||
|
||||
data, error = await self._batch_body(
|
||||
request, ("model_type", "existing"), BATCH_PREVIEW_BODY_MAX
|
||||
)
|
||||
if error is not None:
|
||||
return error
|
||||
config, message = self._batch_choice(data)
|
||||
if message:
|
||||
return self._bad_request(message)
|
||||
model_type, existing = data["model_type"], data["existing"]
|
||||
policy = self._safety_backup_policy()
|
||||
try:
|
||||
candidates = await self._run_in_executor(self.db.autotag_batch_candidates)
|
||||
status = await self._run_in_executor(
|
||||
get_autotag_service().get_models_status
|
||||
)
|
||||
disk, _short, _size = await self._run_in_executor(
|
||||
functools.partial(
|
||||
disk_plan,
|
||||
self.db.model.db_path,
|
||||
backup=policy != "off",
|
||||
vacuum=False,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.exception("Tag output images could not be reviewed")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e)}, status=500
|
||||
)
|
||||
prompts = {row["prompt_id"]: row["prompt_tags"] for row in candidates}
|
||||
tagged = sum(1 for tags in prompts.values() if self._real_tags(tags))
|
||||
to_tag = len(prompts) - tagged if existing == "skip" else len(prompts)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"images": len(candidates),
|
||||
"prompts": len(prompts),
|
||||
"already_tagged": tagged,
|
||||
"to_tag": to_tag,
|
||||
"model": {
|
||||
"type": model_type,
|
||||
"label": config["label"],
|
||||
"size_label": size_label(config["size"]),
|
||||
"downloaded": bool(
|
||||
(status.get(model_type) or {}).get("downloaded")
|
||||
),
|
||||
},
|
||||
"existing": existing,
|
||||
"review_required": existing == "replace_after_review",
|
||||
"backup": {"policy": policy},
|
||||
"disk": disk,
|
||||
}
|
||||
)
|
||||
|
||||
async def start_autotag_job(self, request):
|
||||
"""``POST /jobs/autotag``: start Tag output images (202), or refuse.
|
||||
|
||||
In order: 400 body or options, 400 ``review-required``; 409
|
||||
``model-unavailable``; 507 ``insufficient-space``; 409 ``busy``
|
||||
(another records job, the legacy batch's hold included); 409
|
||||
``model-busy``. No await follows the busy check.
|
||||
"""
|
||||
from ..autotag import get_autotag_service
|
||||
from ..config import MAX_SUGGESTION_LIMIT, AutotagConfig
|
||||
from .annotation_suggest import OptionError, validate_options
|
||||
|
||||
data, error = await self._batch_body(
|
||||
request, ("model_type", "existing", "options"), BATCH_JOB_BODY_MAX
|
||||
)
|
||||
if error is not None:
|
||||
return error
|
||||
config, message = self._batch_choice(data)
|
||||
if message:
|
||||
return self._bad_request(message)
|
||||
model_type, existing, options = (
|
||||
data["model_type"],
|
||||
data["existing"],
|
||||
data["options"],
|
||||
)
|
||||
if isinstance(options, dict) and {"suggestion_limit", "use_gpu"} & set(options):
|
||||
return self._bad_request("options take no suggestion_limit or use_gpu")
|
||||
try:
|
||||
_model, options = validate_options(
|
||||
model_type,
|
||||
options,
|
||||
default_limit=AutotagConfig.suggestion_limit(),
|
||||
max_limit=MAX_SUGGESTION_LIMIT,
|
||||
)
|
||||
except OptionError as error:
|
||||
return self._bad_request(str(error))
|
||||
if existing == "replace_after_review":
|
||||
return web.json_response(
|
||||
{"success": False, "code": "review-required", "error": REVIEW_REQUIRED},
|
||||
status=400,
|
||||
)
|
||||
service = get_autotag_service()
|
||||
backup_on = self._safety_backup_policy() == "on"
|
||||
try:
|
||||
status = await self._run_in_executor(service.get_models_status)
|
||||
candidates = await self._run_in_executor(self.db.autotag_batch_candidates)
|
||||
_disk, short, _size = await self._run_in_executor(
|
||||
functools.partial(
|
||||
disk_plan, self.db.model.db_path, backup=backup_on, vacuum=False
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.exception("Tag output images could not start")
|
||||
return web.json_response(
|
||||
{"success": False, "error": self._public_error(e)}, status=500
|
||||
)
|
||||
if not (status.get(model_type) or {}).get("downloaded"):
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "model-unavailable",
|
||||
"error": f"{config['label']} is not downloaded. "
|
||||
"Download it in System › Auto-tagging.",
|
||||
},
|
||||
status=409,
|
||||
)
|
||||
if backup_on and short is not None:
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "insufficient-space",
|
||||
"error": short_space_message(short),
|
||||
},
|
||||
status=507,
|
||||
)
|
||||
busy = self._jobs.blocker(AUTOTAG_BATCH_JOB_KIND, admission="records")
|
||||
if busy is not None:
|
||||
return busy_response(JobBusy(busy))
|
||||
owner = model_gate.new_owner(AUTOTAG_OWNER_KIND)
|
||||
if not model_gate.try_acquire(owner):
|
||||
return model_busy_response()
|
||||
opts = {
|
||||
"model_type": model_type,
|
||||
"existing": existing,
|
||||
"keep_in_memory": True,
|
||||
"gate_owner": owner,
|
||||
"general_threshold": options.get("general_threshold"),
|
||||
"character_threshold": options.get("character_threshold"),
|
||||
"prompt": options.get("instructions"),
|
||||
}
|
||||
try:
|
||||
job = self._jobs.start(
|
||||
AUTOTAG_BATCH_JOB_KIND,
|
||||
AUTOTAG_JOB_TITLE,
|
||||
admission="records",
|
||||
cancellable=True,
|
||||
scope={
|
||||
"model_type": model_type,
|
||||
"existing": existing,
|
||||
"images": len(candidates),
|
||||
"label": config["label"],
|
||||
},
|
||||
groups=AUTOTAG_GROUP_NAMES,
|
||||
runner=functools.partial(
|
||||
self._run_autotag_job, service, opts, backup_on
|
||||
),
|
||||
)
|
||||
except BaseException as refused:
|
||||
model_gate.release(owner)
|
||||
if isinstance(refused, JobBusy):
|
||||
return busy_response(refused)
|
||||
raise
|
||||
job.task.add_done_callback(lambda _task: model_gate.release(owner))
|
||||
return web.json_response({"success": True, "job": job.summary()}, status=202)
|
||||
|
||||
async def _run_autotag_job(self, service, opts, backup_on, job):
|
||||
ok, _backup = await self._safety_backup_step(job, self.db, backup_on)
|
||||
if not ok:
|
||||
return
|
||||
try:
|
||||
await self._tag_output_images(service, opts, job)
|
||||
except Exception:
|
||||
self.logger.exception("Tag output images failed")
|
||||
job.update(status_line=_INTERNAL_ERROR)
|
||||
job.log_line(f"Tagging failed · {_INTERNAL_ERROR}")
|
||||
job.finish("failed")
|
||||
finally:
|
||||
if backup_on:
|
||||
await self._prune_backup_pool(job, self.db)
|
||||
|
||||
async def _tag_output_images(self, service, opts, job):
|
||||
roots = _scan_roots(self._get_all_output_dirs())
|
||||
job.log_line(
|
||||
f"Tagging started · {job.scope['label']} · "
|
||||
+ ("skip prompts already tagged" if opts["existing"] == "skip" else "merge")
|
||||
)
|
||||
job.update(status_line="Loading the tagging model")
|
||||
|
||||
def on_item(index, total, group, image_path, detail):
|
||||
path, root, root_index = _scan_location(image_path or "", roots)
|
||||
job.update(
|
||||
processed=index + 1,
|
||||
total=total,
|
||||
progress=100 * (index + 1) // total,
|
||||
status_line=f"Tagging {index + 1} of {total}",
|
||||
)
|
||||
job.record(
|
||||
group,
|
||||
{
|
||||
"path": path,
|
||||
"root": root,
|
||||
"root_index": root_index,
|
||||
"detail": detail,
|
||||
},
|
||||
f"{path}: {detail}",
|
||||
)
|
||||
job.publish({"type": "progress", "progress": job.progress})
|
||||
|
||||
final = None
|
||||
async for event in self._autotag_batch_events(
|
||||
service,
|
||||
{**opts, "cancelled": lambda: job.cancelled, "on_item": on_item},
|
||||
):
|
||||
if event.get("type") in ("complete", "error"):
|
||||
final = event
|
||||
if final is not None and final["type"] == "error":
|
||||
job.update(status_line=final["message"])
|
||||
job.log_line(f"Tagging failed · {final['message']}")
|
||||
job.finish("failed")
|
||||
elif job.cancelled: # a cancel wins over a run that then completed
|
||||
job.update(status_line=BATCH_CANCELLED)
|
||||
job.log_line(BATCH_CANCELLED)
|
||||
job.finish("cancelled")
|
||||
else:
|
||||
done = f"Tagged {_plural(job.counters.get('tagged', 0), 'prompt')}"
|
||||
job.update(status_line=done, progress=100)
|
||||
job.log_line(done)
|
||||
job.finish("finished")
|
||||
|
||||
def _gated_runner(self, owner):
|
||||
"""Executor calls under the model gate *owner* holds (plain when None)."""
|
||||
if owner is None:
|
||||
@@ -1151,9 +607,10 @@ class AutotagRoutesMixin:
|
||||
general_threshold = data.get("general_threshold")
|
||||
character_threshold = data.get("character_threshold")
|
||||
try:
|
||||
general_threshold, character_threshold = _thresholds_or_saved(
|
||||
general_threshold, character_threshold
|
||||
)
|
||||
if general_threshold is not None:
|
||||
general_threshold = float(general_threshold)
|
||||
if character_threshold is not None:
|
||||
character_threshold = float(character_threshold)
|
||||
except (ValueError, TypeError):
|
||||
return self._bad_request(
|
||||
"general_threshold and character_threshold must be numeric"
|
||||
|
||||
+20
-211
@@ -6,20 +6,13 @@ loop; the walk itself, the comparison with the previous listing and both
|
||||
folder sets run in the executor through ``walk_root``. Routes live in
|
||||
``py/api/images.py`` (``_listing``); this module has no aiohttp import so
|
||||
tests can exercise it without a server.
|
||||
|
||||
A settled output the image monitor reports is added to its root's stored
|
||||
listing without a walk (``output_settled``, ``insert_entry``), in the place
|
||||
a walk would list it. When that cannot be exact, the root is marked dirty
|
||||
(``mark_root_dirty``) instead; a dirty root gets one background walk, no
|
||||
sooner than ``rewalk_interval`` after its previous walk started (PR 2).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from bisect import bisect_left
|
||||
from dataclasses import dataclass, field
|
||||
from dataclasses import dataclass
|
||||
from pathlib import PurePath
|
||||
from typing import Callable, NamedTuple
|
||||
|
||||
@@ -33,19 +26,6 @@ REFRESH_AFTER = 30
|
||||
# (#348 live test: reads during the walk took 2.2 s instead of 0.07 s).
|
||||
REFRESH_DELAY = 1.0
|
||||
|
||||
# No background walk of a root starts sooner than rewalk_interval(last walk's
|
||||
# duration) after its previous walk started: FACTOR times that duration, at
|
||||
# least FLOOR and at most CAP seconds. Walks up to CAP / FACTOR (3 s) then
|
||||
# hold the GIL at most ~20% of the time; the cap keeps a new output listed
|
||||
# within ~20 s, at the cost of a higher share (back to back from 15 s) on
|
||||
# slower walks. Read through this module so tests can patch them.
|
||||
REWALK_FLOOR = 5
|
||||
REWALK_FACTOR = 5
|
||||
REWALK_CAP = 15
|
||||
|
||||
# A deferred walk is never armed sooner than this: no zero-delay loop.
|
||||
MIN_REWALK_DELAY = 0.5
|
||||
|
||||
# A reader with no listing gives up (500) after this many walks in a row were
|
||||
# discarded by invalidations made while they ran.
|
||||
MAX_DISCARDED_WALKS = 3
|
||||
@@ -100,35 +80,14 @@ class RootListing:
|
||||
flight: asyncio.Future | None = None
|
||||
# True while walks keep failing: one warning per streak.
|
||||
failing: bool = False
|
||||
# The one deferred background walk not started yet (loop.call_later), or
|
||||
# None: a stale refresh or a dirty root's rewalk, whichever is due first.
|
||||
# The delayed refresh not started yet (loop.call_later), or None.
|
||||
rewalk_handle: asyncio.TimerHandle | None = None
|
||||
# normcase(realpath(root)), computed by the walk in the executor: monitor
|
||||
# paths are matched against it without file I/O on the loop.
|
||||
resolved_root: str | None = None
|
||||
# Monotonic start and duration of the last walk (rewalk_interval).
|
||||
last_started_at: float | None = None
|
||||
last_duration: float | None = None
|
||||
# A settled output was reported under this root; cleared when a walk starts.
|
||||
dirty: bool = False
|
||||
# Dropped because the output roots changed: no background walk until a
|
||||
# reader (or the warm-up) asks for this root again. The monitor keeps
|
||||
# watching a removed root until ComfyUI restarts.
|
||||
retired: bool = False
|
||||
# Relative paths of settled outputs waiting to be added to ``files``
|
||||
# without a walk, and the one task adding them (``output_settled``).
|
||||
pending_adds: list = field(default_factory=list)
|
||||
add_task: asyncio.Task | None = None
|
||||
# Settled outputs added to the stored listing since the last walk stored.
|
||||
patched: int = 0
|
||||
|
||||
def drop(self) -> None:
|
||||
self.files = None
|
||||
self.folders = None
|
||||
self.folders_with_ancestors = None
|
||||
self.walked_at = None
|
||||
self.pending_adds = []
|
||||
self.patched = 0
|
||||
self.generation += 1
|
||||
|
||||
|
||||
@@ -140,174 +99,26 @@ def folder_sets(root, files) -> tuple[frozenset, frozenset]:
|
||||
files, plus every ancestor joined by ``/``.
|
||||
"""
|
||||
# From the path strings: relative_to per file was half of a large walk.
|
||||
sep, prefix = _root_prefix(root)
|
||||
folders = set()
|
||||
for path, _mtime in files:
|
||||
rel_dir = _folder_of(root, path, sep, prefix)
|
||||
if rel_dir:
|
||||
folders.add(rel_dir)
|
||||
with_ancestors = set(folders)
|
||||
for folder in folders:
|
||||
with_ancestors.update(_ancestors(folder))
|
||||
return frozenset(folders), frozenset(with_ancestors)
|
||||
|
||||
|
||||
def add_folder(root, path, folders, with_ancestors) -> tuple[frozenset, frozenset]:
|
||||
"""Both folder sets once *path* is listed too (the same sets when it adds none)."""
|
||||
sep, prefix = _root_prefix(root)
|
||||
rel_dir = _folder_of(root, path, sep, prefix)
|
||||
if not rel_dir or rel_dir in folders:
|
||||
return folders, with_ancestors
|
||||
return folders | {rel_dir}, with_ancestors | {rel_dir, *_ancestors(rel_dir)}
|
||||
|
||||
|
||||
def _root_prefix(root) -> tuple[str, str]:
|
||||
# The separator is the path flavour's own (tests use PureWindowsPath).
|
||||
sep = str(type(root)("a", "b"))[1]
|
||||
prefix = str(root)
|
||||
if not prefix.endswith(sep):
|
||||
prefix += sep
|
||||
return sep, prefix
|
||||
|
||||
|
||||
def _folder_of(root, path, sep, prefix) -> str:
|
||||
"""The folder ``/gallery/subfolders`` names for *path*; "" for the root."""
|
||||
text = str(path)
|
||||
if text.startswith(prefix):
|
||||
rel_dir = text[len(prefix) :].rpartition(sep)[0]
|
||||
else:
|
||||
rel_dir = str(path.relative_to(root).parent)
|
||||
return "" if rel_dir == "." else rel_dir
|
||||
|
||||
|
||||
def _ancestors(folder: str) -> list[str]:
|
||||
parts = folder.replace("\\", "/").split("/")
|
||||
return ["/".join(parts[:i]) for i in range(1, len(parts))]
|
||||
|
||||
|
||||
def relative_part(resolved_root: str, path, pathmod=os.path) -> str | None:
|
||||
"""The part of the resolved *path* below *resolved_root*, as a walk spells it.
|
||||
|
||||
Kept in *path*'s own case (``realpath`` gives a file's on-disk case on
|
||||
Windows) with the native separator. None when *path* is not below the
|
||||
root, or when the case folding changed the length so the slice would
|
||||
not be exact.
|
||||
"""
|
||||
sep = pathmod.sep
|
||||
text = str(path)
|
||||
if sep == "\\":
|
||||
text = _plain(text.replace("/", sep), pathmod)
|
||||
folded = pathmod.normcase(text)
|
||||
prefix = _plain(resolved_root, pathmod)
|
||||
if not prefix.endswith(sep):
|
||||
prefix += sep
|
||||
if len(folded) != len(text) or not folded.startswith(prefix):
|
||||
return None
|
||||
relative = text[len(prefix) :]
|
||||
if any(part in ("", ".", "..") for part in relative.split(sep)):
|
||||
return None
|
||||
return relative
|
||||
|
||||
|
||||
def insert_entry(files: list, entry: tuple, max_files: int) -> list | None:
|
||||
"""A new listing with *entry* ``(path, mtime)`` where a walk would list it.
|
||||
|
||||
*files* is newest first; an entry for the same path is replaced. Returns
|
||||
*files* itself when it already holds *entry*, and None when the place is
|
||||
not exact: another file has the same mtime (a walk orders those by its
|
||||
own visit order), or the listing is at the walk's *max_files* cap.
|
||||
Blocking (one pass over *files*): run it in the executor.
|
||||
"""
|
||||
path, mtime = entry
|
||||
text = str(path)
|
||||
kept = files
|
||||
for index, (listed, listed_mtime) in enumerate(files):
|
||||
if str(listed) == text:
|
||||
if listed_mtime == mtime:
|
||||
return files
|
||||
kept = files[:index] + files[index + 1 :]
|
||||
break
|
||||
else:
|
||||
if len(files) >= max_files:
|
||||
return None
|
||||
index = bisect_left(kept, -mtime, key=lambda item: -item[1])
|
||||
if index < len(kept) and kept[index][1] == mtime:
|
||||
return None
|
||||
return [*kept[:index], entry, *kept[index:]]
|
||||
|
||||
|
||||
# A listing that differs from the cached one by more entries than this is
|
||||
# rebuilt in full by the workspace caches (listing_delta).
|
||||
MAX_DELTA = 256
|
||||
|
||||
|
||||
def listing_delta(old: list, new: list, limit: int = MAX_DELTA):
|
||||
"""``(added, removed)`` entries from *old* to *new*, or None past *limit*.
|
||||
|
||||
A listing patched by ``insert_entry`` shares every unchanged entry object
|
||||
with the one it came from, so entries are told apart by identity: a walk
|
||||
builds new tuples, shares none and is rebuilt in full. *old* must still
|
||||
be alive (the
|
||||
caller's cache holds it), so no identity was reused.
|
||||
"""
|
||||
if abs(len(new) - len(old)) > limit:
|
||||
return None
|
||||
old_ids = set(map(id, old))
|
||||
added = []
|
||||
for entry in new:
|
||||
if id(entry) not in old_ids:
|
||||
added.append(entry)
|
||||
if len(added) > limit:
|
||||
return None
|
||||
if new and len(added) == len(new):
|
||||
return None # nothing shared: a walk
|
||||
if len(old) + len(added) == len(new):
|
||||
return added, []
|
||||
new_ids = set(map(id, new))
|
||||
return added, [entry for entry in old if id(entry) not in new_ids]
|
||||
|
||||
|
||||
def rewalk_interval(last_duration) -> float:
|
||||
"""Seconds from a walk's start to the next background walk of its root."""
|
||||
scaled = REWALK_FACTOR * (last_duration or 0.0)
|
||||
return min(REWALK_CAP, max(REWALK_FLOOR, scaled))
|
||||
|
||||
|
||||
def root_contains(resolved_root: str, path, pathmod=os.path) -> bool:
|
||||
"""True when the resolved *path* lies under *resolved_root* (no file I/O).
|
||||
|
||||
*resolved_root* is already ``normcase``d; *path* is folded the same way,
|
||||
so Windows matches case-insensitively and with either separator. The
|
||||
match ends at a separator: ``/out2/x.png`` is not under ``/out``.
|
||||
"""
|
||||
target = _plain(pathmod.normcase(str(path)), pathmod)
|
||||
prefix = _plain(resolved_root, pathmod)
|
||||
if not prefix.endswith(pathmod.sep):
|
||||
prefix += pathmod.sep
|
||||
return target.startswith(prefix)
|
||||
|
||||
|
||||
def _plain(path: str, pathmod) -> str:
|
||||
"""*path* (normcased or not) without a Windows extended-length prefix.
|
||||
|
||||
``realpath`` may keep ``\\\\?\\`` on a file past MAX_PATH while its shorter
|
||||
root resolves without it.
|
||||
"""
|
||||
if pathmod.sep != "\\":
|
||||
return path
|
||||
if path[:8].lower() == "\\\\?\\unc\\":
|
||||
return "\\\\" + path[8:]
|
||||
if path.startswith("\\\\?\\"):
|
||||
return path[4:]
|
||||
return path
|
||||
|
||||
|
||||
def resolve_root(root) -> str:
|
||||
"""``normcase(realpath(root))``, the spelling monitor paths are matched in.
|
||||
|
||||
File I/O (realpath follows links): call it in the executor.
|
||||
"""
|
||||
return os.path.normcase(os.path.realpath(str(root)))
|
||||
folders = set()
|
||||
for path, _mtime in files:
|
||||
text = str(path)
|
||||
if text.startswith(prefix):
|
||||
rel_dir = text[len(prefix) :].rpartition(sep)[0]
|
||||
else:
|
||||
rel_dir = str(path.relative_to(root).parent)
|
||||
if rel_dir and rel_dir != ".":
|
||||
folders.add(rel_dir)
|
||||
with_ancestors = set(folders)
|
||||
for folder in folders:
|
||||
parts = folder.replace("\\", "/").split("/")
|
||||
for i in range(1, len(parts)):
|
||||
with_ancestors.add("/".join(parts[:i]))
|
||||
return frozenset(folders), frozenset(with_ancestors)
|
||||
|
||||
|
||||
def same_listing(old, new) -> bool:
|
||||
@@ -319,10 +130,8 @@ def same_listing(old, new) -> bool:
|
||||
|
||||
def walk_root(
|
||||
scan_fn: Callable[[PurePath], list], root, previous_files
|
||||
) -> tuple[list, frozenset, frozenset, str]:
|
||||
"""The single executor job of a walk.
|
||||
|
||||
Returns ``(files, folders, with_ancestors, resolved_root)``.
|
||||
) -> tuple[list, frozenset, frozenset]:
|
||||
"""The single executor job of a walk: ``(files, folders, with_ancestors)``.
|
||||
|
||||
Opens the root first, so a missing or unreadable root raises ``OSError``
|
||||
instead of listing nothing. Returns *previous_files* itself when nothing
|
||||
@@ -334,4 +143,4 @@ def walk_root(
|
||||
if same_listing(previous_files, files):
|
||||
files = previous_files
|
||||
folders, with_ancestors = folder_sets(root, files)
|
||||
return files, folders, with_ancestors, resolve_root(root)
|
||||
return files, folders, with_ancestors
|
||||
|
||||
+54
-519
@@ -10,12 +10,10 @@ import io
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import stat
|
||||
import subprocess # nosec B404 - ffmpeg frames through run_encoder, fixed argv
|
||||
import sys
|
||||
import time as _time
|
||||
import urllib.parse
|
||||
from bisect import bisect_left, bisect_right, insort
|
||||
from bisect import bisect_left, bisect_right
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Callable, NamedTuple
|
||||
@@ -591,10 +589,6 @@ def _workspace_key(entry: tuple, order: str) -> tuple:
|
||||
return (mtime, entry.name, entry.tie, entry.index)
|
||||
|
||||
|
||||
def _newest_key(entry: _Listed) -> tuple:
|
||||
return _workspace_key(entry, "newest")
|
||||
|
||||
|
||||
def _rated_key(entry: _Listed, rank: int) -> tuple:
|
||||
"""Highest rated: rating 5..1, linked unrated (0), unlinked (-1); then newest."""
|
||||
return (-rank, *_workspace_key(entry, "rating"))
|
||||
@@ -770,34 +764,6 @@ def _favorite_relative_key(relative: str) -> str:
|
||||
return _path_flavor.normcase(relative).replace(_path_flavor.sep, "/")
|
||||
|
||||
|
||||
def _listing_posix(root, prefix, path) -> str:
|
||||
"""*path*'s POSIX spelling below *root* (*prefix*: ``join(root, "")``)."""
|
||||
text = str(path)
|
||||
if text.startswith(prefix):
|
||||
relative = text[len(prefix) :]
|
||||
return relative if os.sep == "/" else relative.replace(os.sep, "/")
|
||||
return path.relative_to(root).as_posix()
|
||||
|
||||
|
||||
def _favorite_listing_step(root, listing, added, removed) -> dict | None:
|
||||
"""*listing* after a patched listing's change, or None to rebuild it.
|
||||
|
||||
Only when every removed file is added again under the same spelling (a
|
||||
rewritten output) and no added file's key is already taken by another
|
||||
spelling: the full build keeps the first spelling in listing order.
|
||||
"""
|
||||
prefix = os.path.join(str(root), "")
|
||||
spellings = {_listing_posix(root, prefix, path) for path, _mtime in added}
|
||||
if any(_listing_posix(root, prefix, path) not in spellings for path, _ in removed):
|
||||
return None
|
||||
listing = dict(listing)
|
||||
for posix in spellings:
|
||||
known = listing.setdefault(_favorite_relative_key(posix), posix)
|
||||
if known != posix:
|
||||
return None
|
||||
return listing
|
||||
|
||||
|
||||
def _favorite_root_key(root) -> str:
|
||||
"""Server-side identity of a gallery root: its real, case-normalized path.
|
||||
|
||||
@@ -861,91 +827,6 @@ def _walk_media(output_path, max_files):
|
||||
stack.append((entry.path, rel_parts + (entry.name,)))
|
||||
|
||||
|
||||
def _settled_entry(root, relative):
|
||||
"""``(Path, mtime)`` of a settled output as a walk would list it.
|
||||
|
||||
*relative* is its part below *root* (``gallery_listing.relative_part``).
|
||||
None when a walk would not list it (``_walk_media``'s filters, and every
|
||||
folder on the way must be listable); raises ``OSError`` when that cannot
|
||||
be told (the file is gone or unreadable, or its on-disk spelling differs
|
||||
or cannot be read). Blocking: run it in the executor.
|
||||
"""
|
||||
path = Path(os.path.join(os.fspath(root), relative))
|
||||
info = os.lstat(path)
|
||||
if stat.S_ISLNK(info.st_mode) or stat.S_ISDIR(info.st_mode):
|
||||
return None
|
||||
# Before the name filters: they compare the walk's (stored) spelling.
|
||||
on_disk = _on_disk_path(path)
|
||||
if on_disk is not None and not on_disk.endswith(os.sep + relative):
|
||||
raise OSError(f"{relative} is spelled differently on disk")
|
||||
*folders, name = relative.split(os.sep)
|
||||
if _name_suffix(name).lower() not in MEDIA_EXTENSIONS:
|
||||
return None
|
||||
if "thumbnails" in folders or len(folders) > MAX_SCAN_DEPTH:
|
||||
return None
|
||||
folder = os.fspath(root)
|
||||
for part in ("", *folders):
|
||||
folder = os.path.join(folder, part) if part else folder
|
||||
if not os.access(folder, os.R_OK | os.X_OK):
|
||||
return None # the walk cannot list this folder
|
||||
return path, info.st_mtime
|
||||
|
||||
|
||||
def _on_disk_path(path, platform=None):
|
||||
"""The file's path as its folders store the names, where that can differ.
|
||||
|
||||
Windows and macOS file systems ignore case. Windows' ``realpath`` gives
|
||||
the stored case, but keeps the given one when it cannot open the file;
|
||||
macOS' ``realpath`` always keeps the given one. A walk lists the names as
|
||||
stored, so an add checks them: Windows through ``GetFinalPathNameByHandle``
|
||||
(``nt._getfinalpathname``), macOS through ``F_GETPATH``. Raises OSError
|
||||
when it cannot be read; None elsewhere (case-sensitive file systems).
|
||||
"""
|
||||
platform = platform or sys.platform
|
||||
if platform == "win32":
|
||||
import nt
|
||||
|
||||
return nt._getfinalpathname(os.fspath(path))
|
||||
if platform != "darwin":
|
||||
return None
|
||||
import fcntl
|
||||
|
||||
# Non-blocking: a FIFO named like an image must not hold the executor.
|
||||
descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NONBLOCK", 0))
|
||||
try:
|
||||
raw = fcntl.fcntl(descriptor, fcntl.F_GETPATH, bytes(1024))
|
||||
finally:
|
||||
os.close(descriptor)
|
||||
return os.fsdecode(raw.split(b"\0", 1)[0])
|
||||
|
||||
|
||||
def _patch_listing(root, files, folders, with_ancestors, relatives):
|
||||
"""The listing with each settled output in *relatives* added (blocking).
|
||||
|
||||
Returns ``(files, folders, with_ancestors, unresolved)``; *files* is the
|
||||
same list when nothing changed. ``unresolved`` is True when an output
|
||||
could not be placed exactly, so the root needs a walk.
|
||||
"""
|
||||
unresolved = False
|
||||
for relative in relatives:
|
||||
try:
|
||||
entry = _settled_entry(root, relative)
|
||||
except OSError:
|
||||
unresolved = True
|
||||
continue
|
||||
if entry is None:
|
||||
continue # a walk would not list it either
|
||||
patched = gallery_listing.insert_entry(files, entry, MAX_GALLERY_FILES)
|
||||
if patched is None:
|
||||
unresolved = True
|
||||
elif patched is not files:
|
||||
files = patched
|
||||
folders, with_ancestors = gallery_listing.add_folder(
|
||||
root, entry[0], folders, with_ancestors
|
||||
)
|
||||
return files, folders, with_ancestors, unresolved
|
||||
|
||||
|
||||
def _iter_media_files(output_path, max_files):
|
||||
"""Yield the ``Path`` of each media file under *output_path* (see _walk_media)."""
|
||||
for entry in _walk_media(output_path, max_files):
|
||||
@@ -1527,7 +1408,6 @@ class ImageRoutesMixin:
|
||||
even if a previous listing exists.
|
||||
"""
|
||||
record = self._root_listing(output_path)
|
||||
record.retired = False # a configured root: read, so listened to
|
||||
needed = record.walk_seq + 1 if fresh else 0
|
||||
discarded = 0
|
||||
while True:
|
||||
@@ -1553,172 +1433,28 @@ class ImageRoutesMixin:
|
||||
)
|
||||
|
||||
def _refresh_if_stale(self, output_path, record):
|
||||
"""Ask for a background walk when the listing is stale or dirty.
|
||||
"""Walk a stale listing again, REFRESH_DELAY after this read (on the loop).
|
||||
|
||||
A stale listing (older than REFRESH_AFTER) is walked REFRESH_DELAY
|
||||
after this read, so the walk does not share the GIL with the page's
|
||||
own requests; a dirty one as soon as its interval allows.
|
||||
At most one pending refresh per root, none while a walk runs.
|
||||
"""
|
||||
age = _time.monotonic() - (record.walked_at or 0.0)
|
||||
if age >= gallery_listing.REFRESH_AFTER:
|
||||
self._schedule_walk(output_path, gallery_listing.REFRESH_DELAY)
|
||||
elif record.dirty:
|
||||
self._schedule_walk(output_path)
|
||||
|
||||
def output_settled(self, path):
|
||||
"""The monitor settled an output at *path*: list it (on the loop).
|
||||
|
||||
*path* is resolved by the monitor's thread (``notify_output_settled``).
|
||||
A root with a stored listing and no walk running gets the file added
|
||||
in the executor (``_run_adds``), without a walk. Every other root
|
||||
holding it is marked dirty (``mark_root_dirty``): one with a walk
|
||||
running gets the trailing walk, as that walk may have passed the
|
||||
file's folder already.
|
||||
"""
|
||||
for key, record in self._settled_records(path):
|
||||
relative = None
|
||||
if (
|
||||
record.files is not None
|
||||
and record.flight is None
|
||||
and record.resolved_root is not None
|
||||
):
|
||||
relative = gallery_listing.relative_part(record.resolved_root, path)
|
||||
if relative is None:
|
||||
self._mark_dirty(key, record)
|
||||
continue
|
||||
record.pending_adds.append(relative)
|
||||
if record.add_task is None:
|
||||
task = asyncio.get_running_loop().create_task(
|
||||
self._run_adds(Path(key), record)
|
||||
)
|
||||
record.add_task = task
|
||||
self._listing_tasks.add(task)
|
||||
task.add_done_callback(self._listing_tasks.discard)
|
||||
|
||||
def mark_root_dirty(self, path):
|
||||
"""A settled output at *path*: walk every root holding it again (on the loop).
|
||||
|
||||
*path* is resolved by the monitor's thread (``notify_output_settled``),
|
||||
so this does no file I/O: it compares ``normcase`` strings with each
|
||||
walk's ``resolved_root``. Only roots with a listing, or a walk running,
|
||||
are marked; a root without either is walked by its next read anyway.
|
||||
"""
|
||||
for key, record in self._settled_records(path):
|
||||
self._mark_dirty(key, record)
|
||||
|
||||
def _settled_records(self, path):
|
||||
"""``(key, record)`` of every listed or walking root holding *path*."""
|
||||
if self._listings_closed:
|
||||
return []
|
||||
found = []
|
||||
for key, record in list(self._root_listings.items()):
|
||||
if record.retired or (record.files is None and record.flight is None):
|
||||
continue
|
||||
resolved = record.resolved_root or os.path.normcase(key)
|
||||
if gallery_listing.root_contains(resolved, path):
|
||||
found.append((key, record))
|
||||
return found
|
||||
|
||||
def _mark_dirty(self, key, record):
|
||||
if self._listings_closed or record.retired:
|
||||
if age < gallery_listing.REFRESH_AFTER:
|
||||
return
|
||||
if record.files is None and record.flight is None:
|
||||
return # walked by its next read anyway
|
||||
record.dirty = True
|
||||
self._schedule_walk(Path(key))
|
||||
|
||||
async def _run_adds(self, root, record):
|
||||
"""Add the root's pending settled outputs to its listing (on the loop).
|
||||
|
||||
Each batch is placed in the executor against the stored listing, and
|
||||
stored only when no walk started and no drop happened meanwhile: a
|
||||
walk started after the outputs settled lists them itself, and a
|
||||
dropped listing is walked by its next read. An output that cannot be
|
||||
placed exactly marks the root dirty instead.
|
||||
"""
|
||||
key = str(root)
|
||||
try:
|
||||
while record.pending_adds:
|
||||
relatives, record.pending_adds = record.pending_adds, []
|
||||
if record.files is None or record.flight is not None:
|
||||
continue # dropped, or a walk started after they settled
|
||||
base, generation, seq = record.files, record.generation, record.walk_seq
|
||||
files, folders, with_ancestors, unresolved = (
|
||||
await self._run_in_executor(
|
||||
_patch_listing,
|
||||
root,
|
||||
base,
|
||||
record.folders,
|
||||
record.folders_with_ancestors,
|
||||
relatives,
|
||||
)
|
||||
)
|
||||
if self._listings_closed or record.retired:
|
||||
return
|
||||
if record.generation != generation or record.walk_seq != seq:
|
||||
continue
|
||||
if files is not base:
|
||||
record.files = files
|
||||
record.folders = folders
|
||||
record.folders_with_ancestors = with_ancestors
|
||||
record.patched += 1
|
||||
self._schedule_workspace_warm()
|
||||
if unresolved:
|
||||
self._mark_dirty(key, record)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
self.logger.warning(f"Could not add a new output to the listing: {exc}")
|
||||
record.pending_adds = []
|
||||
self._mark_dirty(key, record)
|
||||
finally:
|
||||
record.add_task = None
|
||||
|
||||
def _schedule_walk(self, output_path, delay=0.0):
|
||||
"""Start or defer a background walk of *output_path* (on the loop).
|
||||
|
||||
- A walk is running: nothing; its end re-checks ``dirty``.
|
||||
- Otherwise the walk is due *delay* from now, and never sooner than
|
||||
``rewalk_interval`` after the previous walk started. Due now: start
|
||||
it. Later: one deferred walk per root, at least MIN_REWALK_DELAY
|
||||
away; when one is already pending, the earlier deadline wins.
|
||||
Every walk start cancels the pending one (``_start_walk``).
|
||||
"""
|
||||
if self._listings_closed:
|
||||
if record.flight is not None or record.rewalk_handle is not None:
|
||||
return
|
||||
record = self._root_listing(output_path)
|
||||
if record.flight is not None or record.retired:
|
||||
return
|
||||
now = _time.monotonic()
|
||||
due = now + delay
|
||||
if record.last_started_at is not None:
|
||||
interval = gallery_listing.rewalk_interval(record.last_duration)
|
||||
due = max(due, record.last_started_at + interval)
|
||||
if due <= now:
|
||||
self._start_walk(output_path, record)
|
||||
return
|
||||
wait = max(gallery_listing.MIN_REWALK_DELAY, due - now)
|
||||
loop = asyncio.get_running_loop()
|
||||
handle = record.rewalk_handle
|
||||
if handle is not None:
|
||||
if handle.when() - loop.time() <= wait:
|
||||
return # the pending walk comes first
|
||||
handle.cancel()
|
||||
record.rewalk_handle = loop.call_later(
|
||||
wait, self._delayed_walk, output_path, record
|
||||
record.rewalk_handle = asyncio.get_running_loop().call_later(
|
||||
gallery_listing.REFRESH_DELAY, self._delayed_walk, output_path, record
|
||||
)
|
||||
|
||||
def _delayed_walk(self, output_path, record):
|
||||
"""The deferred walk's timer: clear the handle first, then walk.
|
||||
|
||||
It was armed for a stale or dirty listing at the earliest moment the
|
||||
interval allowed. Only a walk start clears either reason, and every
|
||||
start cancels this timer, so when it fires the walk is still due.
|
||||
"""
|
||||
record.rewalk_handle = None
|
||||
if self._listings_closed or record.flight is not None or record.retired:
|
||||
return
|
||||
self._start_walk(output_path, record)
|
||||
self._schedule_walk(output_path)
|
||||
|
||||
def _schedule_walk(self, output_path):
|
||||
"""Start a background walk of *output_path* unless one is running."""
|
||||
record = self._root_listing(output_path)
|
||||
if record.flight is None:
|
||||
self._start_walk(output_path, record)
|
||||
|
||||
def _start_walk(self, output_path, record):
|
||||
"""Start the root's one walk; returns its flight future (on the loop)."""
|
||||
@@ -1726,11 +1462,7 @@ class ImageRoutesMixin:
|
||||
self._listing_loop = loop
|
||||
handle, record.rewalk_handle = record.rewalk_handle, None
|
||||
if handle is not None:
|
||||
handle.cancel() # this walk is the one it was waiting to start
|
||||
# This walk lists every output reported so far.
|
||||
record.dirty = False
|
||||
record.pending_adds = []
|
||||
record.last_started_at = _time.monotonic()
|
||||
handle.cancel() # this walk is the refresh it was waiting to start
|
||||
record.walk_seq += 1
|
||||
flight = loop.create_future()
|
||||
record.flight = flight
|
||||
@@ -1750,14 +1482,11 @@ class ImageRoutesMixin:
|
||||
error never reaches a future nobody awaits.
|
||||
"""
|
||||
outcome = None
|
||||
cancelled = False
|
||||
started = record.last_started_at
|
||||
try:
|
||||
files, folders, with_ancestors, resolved = await self._run_in_executor(
|
||||
files, folders, with_ancestors = await self._run_in_executor(
|
||||
walk_root, self._scan_gallery_files_sync, output_path, record.files
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
cancelled = True
|
||||
outcome = WalkOutcome.failed(RuntimeError(WALK_STOPPED))
|
||||
raise
|
||||
except Exception as exc:
|
||||
@@ -1767,15 +1496,11 @@ class ImageRoutesMixin:
|
||||
f"Could not list output folder {output_path}: {exc}"
|
||||
)
|
||||
record.failing = True
|
||||
# Retried by the 30 s rule, not by a trailing walk.
|
||||
record.dirty = False
|
||||
outcome = WalkOutcome.failed(exc)
|
||||
else:
|
||||
record.failing = False
|
||||
record.resolved_root = resolved
|
||||
if record.generation == generation:
|
||||
changed = files is not record.files
|
||||
record.patched = 0
|
||||
record.files = files
|
||||
record.folders = folders
|
||||
record.folders_with_ancestors = with_ancestors
|
||||
@@ -1787,14 +1512,10 @@ class ImageRoutesMixin:
|
||||
else:
|
||||
outcome = WalkOutcome.discarded()
|
||||
finally:
|
||||
record.last_duration = _time.monotonic() - started
|
||||
if record.flight is flight:
|
||||
record.flight = None
|
||||
if not flight.done():
|
||||
flight.set_result(outcome)
|
||||
if record.dirty and not cancelled:
|
||||
# An output settled during this walk: the trailing walk.
|
||||
self._schedule_walk(output_path)
|
||||
|
||||
async def get_output_images(self, request):
|
||||
"""Get all images from ComfyUI output folder(s)."""
|
||||
@@ -1923,10 +1644,6 @@ class ImageRoutesMixin:
|
||||
cached = self._workspace_listed_cache.get(key)
|
||||
if cached is not None and cached[0] is files:
|
||||
return cached[1]
|
||||
if cached is not None:
|
||||
delta = gallery_listing.listing_delta(cached[0], files)
|
||||
if delta is not None:
|
||||
return self._workspace_listed_step(key, cached, files, delta)
|
||||
previous = cached[2] if cached is not None else {}
|
||||
listed = []
|
||||
for path, mtime in files:
|
||||
@@ -1946,42 +1663,7 @@ class ImageRoutesMixin:
|
||||
self._workspace_card_cache[key] = {
|
||||
file: facts for file, facts in cards.items() if file in lookup
|
||||
}
|
||||
self._workspace_listed_cache[key] = (files, listed, lookup, None)
|
||||
return listed
|
||||
|
||||
def _workspace_listed_step(self, key, cached, files, delta) -> list:
|
||||
"""The _Listed values after a few files were added or replaced.
|
||||
|
||||
A listing patched by a settled output (``output_settled``) shares
|
||||
every other entry with the cached one, so only the added files are
|
||||
built; the order of the values does not matter (every sort is a total
|
||||
order). The step ``(previous values, added, removed)`` lets the sort
|
||||
orders follow it the same way (``_workspace_orders``).
|
||||
"""
|
||||
index, root = key[0], Path(key[1])
|
||||
_files, old_listed, old_lookup, _step = cached
|
||||
added, removed = delta
|
||||
if not added and not removed:
|
||||
self._workspace_listed_cache[key] = (files, old_listed, old_lookup, None)
|
||||
return old_listed
|
||||
lookup = dict(old_lookup)
|
||||
gone = [lookup.pop(item) for item in removed if item in lookup]
|
||||
fresh = [_listed(path, mtime, root, index) for path, mtime in added]
|
||||
for entry in fresh:
|
||||
lookup[(entry.path, entry.mtime)] = entry
|
||||
if gone:
|
||||
gone_ids = {id(entry) for entry in gone}
|
||||
listed = [entry for entry in old_listed if id(entry) not in gone_ids]
|
||||
cards = self._workspace_card_cache.get(key)
|
||||
if cards:
|
||||
self._workspace_card_cache[key] = {
|
||||
file: facts for file, facts in cards.items() if file in lookup
|
||||
}
|
||||
else:
|
||||
listed = list(old_listed)
|
||||
listed.extend(fresh)
|
||||
step = (old_listed, fresh, gone)
|
||||
self._workspace_listed_cache[key] = (files, listed, lookup, step)
|
||||
self._workspace_listed_cache[key] = (files, listed, lookup)
|
||||
return listed
|
||||
|
||||
def _workspace_orders(self, listings: list[tuple]) -> dict:
|
||||
@@ -1996,54 +1678,14 @@ class ImageRoutesMixin:
|
||||
old is new for old, new in zip(cached[0], parts, strict=True)
|
||||
):
|
||||
return cached[1]
|
||||
newest = None
|
||||
if cached is not None:
|
||||
newest = self._newest_step(listings, parts, cached)
|
||||
if newest is None:
|
||||
newest = sorted(
|
||||
(entry for part in parts for entry in part),
|
||||
key=_newest_key,
|
||||
)
|
||||
# The other orders, the wanted paths and the locators follow lazily.
|
||||
newest = sorted(
|
||||
(entry for part in parts for entry in part),
|
||||
key=lambda entry: _workspace_key(entry, "newest"),
|
||||
)
|
||||
orders = {"newest": newest}
|
||||
self._workspace_order_cache[key] = (parts, orders)
|
||||
return orders
|
||||
|
||||
def _newest_step(self, listings, parts, cached) -> list | None:
|
||||
"""The cached newest order with each part's last step applied, or None.
|
||||
|
||||
Only when every changed part is exactly the value list this call
|
||||
built (*parts*) and one step (``_workspace_listed_step``) past the
|
||||
part the cached order was built from: the added values are placed by
|
||||
bisection instead of sorting the whole listing again. Page reads and
|
||||
the warm build run in executor threads at once, so the per-file cache
|
||||
may already hold a newer step than *parts*; then this sorts in full.
|
||||
"""
|
||||
old_parts, orders = cached
|
||||
added, gone = [], []
|
||||
for (index, root, _files), part, old in zip(
|
||||
listings, parts, old_parts, strict=True
|
||||
):
|
||||
if part is old:
|
||||
continue
|
||||
listed = self._workspace_listed_cache.get((index, str(root)))
|
||||
if listed is None or listed[1] is not part:
|
||||
return None
|
||||
step = listed[3]
|
||||
if step is None or step[0] is not old:
|
||||
return None
|
||||
added.extend(step[1])
|
||||
gone.extend(step[2])
|
||||
newest = orders["newest"]
|
||||
if gone:
|
||||
gone_ids = {id(entry) for entry in gone}
|
||||
newest = [entry for entry in newest if id(entry) not in gone_ids]
|
||||
else:
|
||||
newest = list(newest)
|
||||
for entry in added:
|
||||
insort(newest, entry, key=_newest_key)
|
||||
return newest
|
||||
|
||||
def _workspace_wanted(self, listings: list[tuple]) -> frozenset:
|
||||
"""The selected listing's normalized paths (cached with its orders)."""
|
||||
orders = self._workspace_orders(listings)
|
||||
@@ -2272,8 +1914,7 @@ class ImageRoutesMixin:
|
||||
)
|
||||
|
||||
image_path = Path(image["image_path"])
|
||||
lora_dirs = await self._run_in_executor(self._lora_image_dirs)
|
||||
allowed_dirs = list(self._get_all_output_dirs()) + lora_dirs
|
||||
allowed_dirs = list(self._get_all_output_dirs()) + self._lora_image_dirs()
|
||||
|
||||
denied = _media_access_error(image_path, allowed_dirs)
|
||||
if denied is not None:
|
||||
@@ -2296,58 +1937,29 @@ 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.
|
||||
|
||||
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. The layout comes from the layout
|
||||
cache (``lora_utils.cached_lora_manager_layout``, #362), which reloads
|
||||
when LoraManager's settings files change; blocking, so the routes
|
||||
call it in the executor.
|
||||
"""
|
||||
"""Extra directories served when the LoRA integration is enabled."""
|
||||
try:
|
||||
from ..config import IntegrationConfig
|
||||
|
||||
if not IntegrationConfig.LORA_MANAGER_ENABLED:
|
||||
return []
|
||||
from ..lora_utils import (
|
||||
cached_lora_manager_layout,
|
||||
get_lora_image_cache_dir,
|
||||
)
|
||||
from ..lora_utils import find_lora_directories, get_lora_image_cache_dir
|
||||
|
||||
configured = IntegrationConfig.LORA_MANAGER_PATH or ""
|
||||
install, layout = cached_lora_manager_layout(configured)
|
||||
if configured.strip() and install is None:
|
||||
lora_root = IntegrationConfig.LORA_MANAGER_PATH or ""
|
||||
if lora_root.strip():
|
||||
# A configured path comes from config.json, which can be hand
|
||||
# edited: only a LoraManager install under custom_nodes may
|
||||
# widen the serving allow-list. An empty path keeps the
|
||||
# auto-detection, which itself only scans custom_nodes.
|
||||
return []
|
||||
from ..lora_utils import resolve_lora_manager_path
|
||||
|
||||
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)))
|
||||
lora_root = resolve_lora_manager_path(lora_root)
|
||||
if not lora_root:
|
||||
return []
|
||||
|
||||
dirs = [Path(d) for d in find_lora_directories(lora_root)]
|
||||
dirs.append(Path(get_lora_image_cache_dir()))
|
||||
return dirs
|
||||
except Exception:
|
||||
# LoRA integration is optional; never widen access on failure.
|
||||
@@ -2378,21 +1990,9 @@ class ImageRoutesMixin:
|
||||
return []
|
||||
|
||||
async def serve_output_image(self, request):
|
||||
"""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.
|
||||
"""
|
||||
"""Serve a media file from the ComfyUI output folder(s)."""
|
||||
try:
|
||||
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)
|
||||
roots = list(self._get_all_output_dirs())
|
||||
scope_error = self._gallery_scope_error(request.query, roots)
|
||||
if scope_error is not None:
|
||||
return scope_error
|
||||
@@ -2407,13 +2007,9 @@ class ImageRoutesMixin:
|
||||
):
|
||||
return _forbidden("Only media files can be served")
|
||||
|
||||
output_dirs = self._select_root(roots, root_param)
|
||||
output_dirs = self._select_root(roots, request.query.get("root"))
|
||||
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
|
||||
@@ -2445,46 +2041,6 @@ 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.
|
||||
|
||||
@@ -3339,9 +2895,7 @@ class ImageRoutesMixin:
|
||||
if error is not None:
|
||||
return error
|
||||
index, scope, relative = target
|
||||
# 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)
|
||||
roots = list(self._get_all_output_dirs())
|
||||
error = self._gallery_scope_error(
|
||||
{"root_scope": scope, "root": str(index)}, roots
|
||||
)
|
||||
@@ -3350,9 +2904,6 @@ 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(
|
||||
@@ -3364,16 +2915,12 @@ class ImageRoutesMixin:
|
||||
status=429,
|
||||
)
|
||||
try:
|
||||
return await self._reveal(relative, root, linked_only=index >= len(outputs))
|
||||
return await self._reveal(relative, root)
|
||||
finally:
|
||||
gate.leave()
|
||||
|
||||
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.
|
||||
"""
|
||||
async def _reveal(self, relative: str, root: Path):
|
||||
"""Resolve and open, inside the reveal gate."""
|
||||
try:
|
||||
path = await self._run_in_executor(self._reveal_path, relative, root)
|
||||
except (OSError, ValueError) as exc:
|
||||
@@ -3384,10 +2931,6 @@ 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
|
||||
@@ -3424,19 +2967,18 @@ class ImageRoutesMixin:
|
||||
cached = self._favorite_listing_cache.get(cache_key)
|
||||
if cached is not None and cached[0] is files:
|
||||
return cached[1]
|
||||
listing = None
|
||||
if cached is not None:
|
||||
delta = gallery_listing.listing_delta(cached[0], files)
|
||||
if delta is not None:
|
||||
listing = _favorite_listing_step(root, cached[1], *delta)
|
||||
if listing is None:
|
||||
# From the path strings, as _listed does: relative_to per file was
|
||||
# the largest cost of a cold first open (#348).
|
||||
prefix = os.path.join(str(root), "")
|
||||
listing = {}
|
||||
for path, _mtime in files:
|
||||
posix = _listing_posix(root, prefix, path)
|
||||
listing.setdefault(_favorite_relative_key(posix), posix)
|
||||
# From the path strings, as _listed does: relative_to per file was the
|
||||
# largest cost of a cold first open (#348).
|
||||
prefix = os.path.join(str(root), "")
|
||||
listing = {}
|
||||
for path, _mtime in files:
|
||||
text = str(path)
|
||||
if text.startswith(prefix):
|
||||
relative = text[len(prefix) :]
|
||||
posix = relative if os.sep == "/" else relative.replace(os.sep, "/")
|
||||
else:
|
||||
posix = path.relative_to(root).as_posix()
|
||||
listing.setdefault(_favorite_relative_key(posix), posix)
|
||||
self._favorite_listing_cache[cache_key] = (files, listing)
|
||||
return listing
|
||||
|
||||
@@ -3479,13 +3021,8 @@ class ImageRoutesMixin:
|
||||
return set()
|
||||
|
||||
async def inspect_image(self, request):
|
||||
"""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)
|
||||
"""Read prompt evidence for one explicitly rooted physical media file."""
|
||||
roots = list(self._get_all_output_dirs())
|
||||
scope_error = self._gallery_scope_error(request.query, roots)
|
||||
if scope_error is not None:
|
||||
return scope_error
|
||||
@@ -3510,8 +3047,6 @@ 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():
|
||||
|
||||
@@ -130,14 +130,6 @@ class JobRoutesMixin:
|
||||
async def start_image_tag_job_route(request):
|
||||
return await self.start_image_tag_job(request)
|
||||
|
||||
@routes.post("/prompt_manager/jobs/model-download")
|
||||
async def start_model_download_job_route(request):
|
||||
return await self.start_model_download_job(request)
|
||||
|
||||
@routes.post("/prompt_manager/jobs/autotag")
|
||||
async def start_autotag_job_route(request):
|
||||
return await self.start_autotag_job(request)
|
||||
|
||||
@routes.post("/prompt_manager/jobs/thumbnails")
|
||||
async def start_thumbnail_job_route(request):
|
||||
return await self.start_thumbnail_job(request)
|
||||
|
||||
@@ -87,8 +87,6 @@ SCAN_GROUPS = {
|
||||
# Record-changing: it writes image tag sets and new vocabulary names.
|
||||
IMAGE_TAG_JOB_KIND = "image-tags"
|
||||
IMAGE_TAG_GROUP_NAMES = ("tagged", "unchanged", "skipped", "errors")
|
||||
# Tag output images (#328) counts its prompts the same way.
|
||||
AUTOTAG_GROUP_NAMES = IMAGE_TAG_GROUP_NAMES
|
||||
|
||||
# Thumbnail generation (#249 System S5): previews written, already fresh, failed.
|
||||
THUMBNAIL_GROUP_NAMES = ("generated", "skipped", "errors")
|
||||
@@ -111,8 +109,6 @@ DUPLICATE_CLEANUP_JOB_KIND = "duplicate_cleanup"
|
||||
THUMBNAILS_JOB_KIND = "thumbnails"
|
||||
LORA_IMPORT_JOB_KIND = "lora_import"
|
||||
AUTOTAG_BATCH_JOB_KIND = "autotag"
|
||||
# A tagging model download (#328): read admission, not cancellable.
|
||||
MODEL_DOWNLOAD_JOB_KIND = "model_download"
|
||||
# Duplicate image files (#249 System S6, py/api/duplicate_files.py): the
|
||||
# find is a cancellable ``read`` job; the reviewed removal is ``records``.
|
||||
DUPLICATES_JOB_KIND = "duplicates"
|
||||
|
||||
+8
-36
@@ -234,22 +234,9 @@ class LibraryRoutesMixin:
|
||||
)
|
||||
|
||||
def _library_media(
|
||||
self,
|
||||
records: list[dict],
|
||||
roots: list[Path],
|
||||
*,
|
||||
linked: bool = True,
|
||||
outputs: int | None = None,
|
||||
self, records: list[dict], roots: list[Path], *, linked: bool = True
|
||||
) -> list[dict]:
|
||||
"""Workspace descriptors for a prompt's linked media under ``roots``.
|
||||
|
||||
Roots from index ``outputs`` on are linked-only (LoRA media, #355):
|
||||
their media keeps its rooted identity, but gets no preview file (the
|
||||
card loads the original; a video shows its label), no favorite star
|
||||
and is marked ``linked_only`` (no Tag action).
|
||||
"""
|
||||
page = []
|
||||
linked_only = set()
|
||||
for record in records:
|
||||
try:
|
||||
raw_path = Path(record["image_path"])
|
||||
@@ -263,25 +250,10 @@ class LibraryRoutesMixin:
|
||||
canonical_root = _canonical(root)
|
||||
if path.is_relative_to(canonical_root):
|
||||
page.append((path, 0, canonical_root, index))
|
||||
if outputs is not None and index >= outputs:
|
||||
relative = path.relative_to(canonical_root).as_posix()
|
||||
linked_only.add((index, relative))
|
||||
break
|
||||
except (OSError, ValueError, RuntimeError):
|
||||
continue
|
||||
images = self._format_workspace_page(page, linked=linked)
|
||||
for image in images:
|
||||
if (image["root_index"], image["path"]) in linked_only:
|
||||
image["thumbnail_url"] = None
|
||||
image["thumbnail_ready"] = False
|
||||
image["linked_only"] = True
|
||||
image.pop("favorite", None)
|
||||
return images
|
||||
|
||||
def _library_roots(self) -> tuple[list[Path], int]:
|
||||
"""``(roots, number of output roots)`` for a prompt's linked media."""
|
||||
outputs = list(self._get_all_output_dirs())
|
||||
return self._media_roots(outputs), len(outputs)
|
||||
return self._format_workspace_page(page, linked=linked)
|
||||
|
||||
@pair_response
|
||||
async def get_library_options(self, request: web.Request) -> web.Response:
|
||||
@@ -322,13 +294,13 @@ class LibraryRoutesMixin:
|
||||
return result, records
|
||||
|
||||
result, records = await self._run_in_executor(read)
|
||||
roots = list(self._get_all_output_dirs())
|
||||
|
||||
def previews() -> None:
|
||||
roots, outputs = self._library_roots()
|
||||
for parent in result["items"]:
|
||||
# Previews belong to their parent; a per-parent link read is N+1.
|
||||
parent["previews"] = self._library_media(
|
||||
records[parent["id"]], roots, linked=False, outputs=outputs
|
||||
records[parent["id"]], roots, linked=False
|
||||
)
|
||||
|
||||
await self._run_in_executor(previews)
|
||||
@@ -370,9 +342,9 @@ class LibraryRoutesMixin:
|
||||
result, records = await self._run_in_executor(read)
|
||||
if result is None:
|
||||
return _missing()
|
||||
roots, outputs = await self._run_in_executor(self._library_roots)
|
||||
roots = list(self._get_all_output_dirs())
|
||||
result["positive"]["previews"] = await self._run_in_executor(
|
||||
lambda: self._library_media(records, roots, outputs=outputs)
|
||||
self._library_media, records, roots
|
||||
)
|
||||
return web.json_response({"success": True, **result})
|
||||
|
||||
@@ -386,7 +358,7 @@ class LibraryRoutesMixin:
|
||||
archive = _archive_scope(request)
|
||||
limit = _page_limit(request)
|
||||
cursor = request.query.get("cursor")
|
||||
roots, outputs = await self._run_in_executor(self._library_roots)
|
||||
roots = list(self._get_all_output_dirs())
|
||||
canonical_roots = [str(_canonical(root)) for root in roots]
|
||||
|
||||
def read() -> dict | None:
|
||||
@@ -419,7 +391,7 @@ class LibraryRoutesMixin:
|
||||
if result is None:
|
||||
return _missing()
|
||||
result["images"] = await self._run_in_executor(
|
||||
lambda: self._library_media(result["images"], roots, outputs=outputs)
|
||||
self._library_media, result["images"], roots
|
||||
)
|
||||
return web.json_response({"success": True, **result})
|
||||
|
||||
|
||||
@@ -5,12 +5,9 @@ it lists the LoraManager metadata files and parses them, and counts the
|
||||
current imports, those a saved pair uses and the example images an import
|
||||
would download. It writes nothing and requests nothing. It answers a
|
||||
``Plan`` (the maintenance plan's shape, ``py/api/maintenance.py``) plus
|
||||
``{lora_files, with_metadata, skipped, with_local_media, current_imports,
|
||||
blocked_by_pairs, planned_downloads, network:{planned_downloads, api_key},
|
||||
layout}`` (#355: ``layout`` says which LoraManager settings file, storage mode
|
||||
and library the import read, and ``error`` when they cannot be used). A plan
|
||||
that saved pairs block, or whose layout is unusable, carries no token
|
||||
(nothing can be applied).
|
||||
``{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).
|
||||
|
||||
``POST /prompt_manager/lora/import/apply {token}`` starts the job (202
|
||||
``{job}``): kind ``lora_import``, titled "LoRA reference import", admission
|
||||
@@ -34,7 +31,6 @@ 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 (
|
||||
@@ -156,9 +152,8 @@ class LoraImportRoutesMixin:
|
||||
return None, ("lora-path", _LORA_PATH_ERROR)
|
||||
return allowed, None
|
||||
|
||||
async def stage_lora_import(self, discovery, *, on_record=None, cancelled=None):
|
||||
"""Stage ``discovery``'s sidecars off the loop: records, skips and
|
||||
download counts."""
|
||||
async def stage_lora_import(self, files, *, on_record=None, cancelled=None):
|
||||
"""Stage ``files`` 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)
|
||||
@@ -171,10 +166,9 @@ class LoraImportRoutesMixin:
|
||||
|
||||
return await self._run_in_executor(
|
||||
lambda: lora_import.stage(
|
||||
discovery,
|
||||
files,
|
||||
cache_dir,
|
||||
IntegrationConfig.CIVITAI_API_KEY,
|
||||
media_extensions=SERVABLE_EXTENSIONS,
|
||||
on_record=report,
|
||||
cancelled=cancelled or (lambda: False),
|
||||
logger=self.logger,
|
||||
@@ -182,8 +176,8 @@ class LoraImportRoutesMixin:
|
||||
)
|
||||
|
||||
async def lora_files_unchanged(self, lm_path, expected):
|
||||
"""Whether LoraManager's sidecars for ``lm_path`` still fingerprint as
|
||||
``expected`` (found again, so an added file counts too)."""
|
||||
"""Whether the metadata files under ``lm_path`` still fingerprint as
|
||||
``expected`` (listed again, so an added file counts too)."""
|
||||
|
||||
def current():
|
||||
return lora_import.fingerprint(lora_import.discover(lm_path))
|
||||
@@ -194,13 +188,8 @@ class LoraImportRoutesMixin:
|
||||
return False
|
||||
|
||||
async def replace_lora_imports(self, db, staged, *, expected_ids, allow_empty):
|
||||
"""Swap the collection for ``staged.records`` in one transaction.
|
||||
|
||||
Then the image routes' layout cache is dropped (#362): the import may
|
||||
have linked files under a folder LoraManager created after the cache
|
||||
stored its layout (its first example-image download creates one).
|
||||
"""
|
||||
result = await self._run_in_executor(
|
||||
"""Swap the collection for ``staged.records`` in one transaction."""
|
||||
return await self._run_in_executor(
|
||||
lambda: db.replace_category_prompts(
|
||||
CATEGORY,
|
||||
staged.records,
|
||||
@@ -208,8 +197,6 @@ class LoraImportRoutesMixin:
|
||||
allow_empty=allow_empty,
|
||||
)
|
||||
)
|
||||
lora_utils.clear_layout_cache()
|
||||
return result
|
||||
|
||||
# ── preview ───────────────────────────────────────────────────────
|
||||
|
||||
@@ -233,22 +220,8 @@ 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, SERVABLE_EXTENSIONS
|
||||
)
|
||||
found = await self._run_in_executor(lora_import.preview, lm_path, cache_dir)
|
||||
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)
|
||||
@@ -275,9 +248,7 @@ class LoraImportRoutesMixin:
|
||||
return _failure("The LoRA import could not be reviewed.", 500)
|
||||
token = None
|
||||
expires_at = None
|
||||
layout = found["layout"]
|
||||
# An unusable LoraManager layout found nothing it may replace with.
|
||||
if not blocked and layout["error"] is None:
|
||||
if not blocked:
|
||||
payload = {
|
||||
"path": lm_path,
|
||||
"fingerprint": found["fingerprint"],
|
||||
@@ -309,9 +280,6 @@ 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,
|
||||
@@ -374,8 +342,8 @@ class LoraImportRoutesMixin:
|
||||
store.discard(token)
|
||||
return _failure(INTEGRATION_CHANGED, 409, "plan-changed")
|
||||
try:
|
||||
discovery = await self._run_in_executor(lora_import.discover, lm_path)
|
||||
now = await self._run_in_executor(lora_import.fingerprint, discovery)
|
||||
files = await self._run_in_executor(lora_import.discover, lm_path)
|
||||
now = await self._run_in_executor(lora_import.fingerprint, files)
|
||||
backup = self._backup_promised(plan)
|
||||
_disk, short, _size = await self._run_in_executor(
|
||||
disk_plan, db.model.db_path, backup=backup, vacuum=False
|
||||
@@ -399,14 +367,9 @@ class LoraImportRoutesMixin:
|
||||
lora_import.JOB_TITLE,
|
||||
admission="records",
|
||||
cancellable=True,
|
||||
scope={
|
||||
"operation": LORA_OPERATION,
|
||||
"lora_files": len(discovery.sidecars),
|
||||
},
|
||||
scope={"operation": LORA_OPERATION, "lora_files": len(files)},
|
||||
groups=lora_import.GROUPS,
|
||||
runner=lambda job: self._run_lora_import(
|
||||
job, db, plan, discovery, backup
|
||||
),
|
||||
runner=lambda job: self._run_lora_import(job, db, plan, files, backup),
|
||||
)
|
||||
except JobBusy as busy:
|
||||
return busy_response(busy)
|
||||
@@ -415,8 +378,8 @@ class LoraImportRoutesMixin:
|
||||
|
||||
# ── the job ───────────────────────────────────────────────────────
|
||||
|
||||
async def _run_lora_import(self, job, db, plan, discovery, backup_on):
|
||||
total = len(discovery.sidecars)
|
||||
async def _run_lora_import(self, job, db, plan, files, backup_on):
|
||||
total = len(files)
|
||||
job.update(total=total, status_line="Preparing LoRA records")
|
||||
job.log_line(f"Found {_count(total, 'LoRA metadata file')}")
|
||||
|
||||
@@ -430,7 +393,7 @@ class LoraImportRoutesMixin:
|
||||
|
||||
try:
|
||||
staged = await self.stage_lora_import(
|
||||
discovery, on_record=on_record, cancelled=lambda: job.cancelled
|
||||
files, on_record=on_record, cancelled=lambda: job.cancelled
|
||||
)
|
||||
except lora_import.ImportCancelled:
|
||||
job.log_line(f"{CANCELLED} {CACHE_NOTE}")
|
||||
|
||||
@@ -50,15 +50,6 @@ _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."""
|
||||
|
||||
@@ -407,19 +398,13 @@ class LoraIntegrationMixin:
|
||||
return
|
||||
if isinstance(error, PlanChanged):
|
||||
message = _IMPORTS_CHANGED
|
||||
elif isinstance(error, (LoraFilesChanged, LoraSettingsUnusable)):
|
||||
elif isinstance(error, LoraFilesChanged):
|
||||
message = str(error)
|
||||
else:
|
||||
message = self._public_error(error)
|
||||
if isinstance(
|
||||
error,
|
||||
(
|
||||
PlanChanged,
|
||||
LoraFilesChanged,
|
||||
LoraSettingsUnusable,
|
||||
EmptyReplacement,
|
||||
PairDependencyError,
|
||||
),
|
||||
(PlanChanged, LoraFilesChanged, EmptyReplacement, PairDependencyError),
|
||||
):
|
||||
self.logger.warning(f"LoRA scan refused: {message}")
|
||||
else:
|
||||
@@ -447,11 +432,8 @@ class LoraIntegrationMixin:
|
||||
"progress": 2,
|
||||
}
|
||||
)
|
||||
discovery = await self._run_in_executor(lora_import.discover, lm_path)
|
||||
if discovery.layout.error is not None:
|
||||
raise LoraSettingsUnusable(discovery.layout.error)
|
||||
sidecars = discovery.sidecars
|
||||
total = len(sidecars)
|
||||
files = await self._run_in_executor(lora_import.discover, lm_path)
|
||||
total = len(files)
|
||||
await send_progress(
|
||||
{
|
||||
"type": "progress",
|
||||
@@ -461,7 +443,7 @@ class LoraIntegrationMixin:
|
||||
}
|
||||
)
|
||||
|
||||
reviewed = await self._run_in_executor(lora_import.fingerprint, discovery)
|
||||
reviewed = await self._run_in_executor(lora_import.fingerprint, files)
|
||||
events = asyncio.Queue()
|
||||
state = await self._run_in_executor(
|
||||
self.db.category_import_state, lora_import.LORA_CATEGORY
|
||||
@@ -485,7 +467,7 @@ class LoraIntegrationMixin:
|
||||
)
|
||||
if not existing:
|
||||
counts["imported"] += 1
|
||||
status = record["name"] if record else sidecars[index].metadata_path.name
|
||||
status = record["name"] if record else files[index].name
|
||||
if record and record["images"]:
|
||||
status += f" ({len(record['images'])} images)"
|
||||
await send_progress(
|
||||
@@ -508,7 +490,7 @@ class LoraIntegrationMixin:
|
||||
stop = threading.Event()
|
||||
staging = asyncio.ensure_future(
|
||||
self.stage_lora_import(
|
||||
discovery,
|
||||
files,
|
||||
on_record=on_record,
|
||||
cancelled=lambda: stop.is_set() or disconnected(),
|
||||
)
|
||||
|
||||
@@ -47,23 +47,11 @@ _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"
|
||||
|
||||
@@ -163,63 +151,6 @@ def _set_nested(config, sections, key, value):
|
||||
config[sections[0]] = merged
|
||||
|
||||
|
||||
def _validate_autotag_settings(data):
|
||||
"""The ``autotag_*`` settings as ``{AUTOTAG_ATTR: value}``.
|
||||
|
||||
The first bad value raises ValueError with a message that names the
|
||||
setting by its label in System › Auto-tagging.
|
||||
"""
|
||||
from ..config import (
|
||||
EXISTING_POLICIES,
|
||||
MAX_SUGGESTION_LIMIT,
|
||||
MODEL_TYPES,
|
||||
validate_saved_threshold,
|
||||
validate_template,
|
||||
)
|
||||
|
||||
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_saved_threshold(label, data[key])
|
||||
except ValueError as error:
|
||||
raise ValueError(f"{error}.") from None
|
||||
updates[key.upper()] = value
|
||||
if "autotag_template" in data:
|
||||
try:
|
||||
updates["AUTOTAG_TEMPLATE"] = validate_template(
|
||||
data["autotag_template"], "LLM tag template"
|
||||
)
|
||||
except ValueError as error:
|
||||
raise ValueError(f"{error}.") from None
|
||||
if "autotag_suggestion_limit" in data:
|
||||
value = data["autotag_suggestion_limit"]
|
||||
if type(value) is not int or not 1 <= value <= MAX_SUGGESTION_LIMIT:
|
||||
raise ValueError(
|
||||
"Suggestions per image must be a whole number from 1 to "
|
||||
f"{MAX_SUGGESTION_LIMIT}."
|
||||
)
|
||||
updates["AUTOTAG_SUGGESTION_LIMIT"] = value
|
||||
return updates
|
||||
|
||||
|
||||
async def _read_json(request):
|
||||
"""Parse a JSON object body through the package's capped reader."""
|
||||
from . import _read_json_body
|
||||
@@ -246,7 +177,7 @@ class SettingsRoutesMixin:
|
||||
async def get_settings(self, request):
|
||||
"""Get current settings."""
|
||||
try:
|
||||
from ..config import AutotagConfig, PromptManagerConfig, GalleryConfig
|
||||
from ..config import PromptManagerConfig, GalleryConfig
|
||||
|
||||
root_paths = [
|
||||
self._public_path(d) for d in GalleryConfig.MONITORING_DIRECTORIES
|
||||
@@ -272,14 +203,6 @@ 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 "",
|
||||
@@ -502,12 +425,7 @@ class SettingsRoutesMixin:
|
||||
"""
|
||||
from ..config import GalleryConfig, PromptManagerConfig
|
||||
|
||||
from ..config import AutotagConfig
|
||||
|
||||
for name, value in updates.items():
|
||||
if name in _AUTOTAG_SETTINGS:
|
||||
setattr(AutotagConfig, name[len(_AUTOTAG_PREFIX) :], value)
|
||||
continue
|
||||
owner = GalleryConfig if name in _GALLERY_SETTINGS else PromptManagerConfig
|
||||
setattr(owner, name, value)
|
||||
restart_fields = []
|
||||
@@ -646,7 +564,6 @@ 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):
|
||||
|
||||
+12
-31
@@ -488,16 +488,10 @@ class AutoThumbnailer:
|
||||
f"Automatic preview failed for {src.name}", exc_info=True
|
||||
)
|
||||
|
||||
def close(self, timeout=None):
|
||||
def close(self):
|
||||
"""Stop the worker at its next file boundary (queued files are
|
||||
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
|
||||
dropped), kill a running ffmpeg and wait up to CLOSE_WAIT_SECONDS;
|
||||
a worker still running after that is logged."""
|
||||
with self._changed:
|
||||
self._closed = True
|
||||
self._queue.clear()
|
||||
@@ -510,18 +504,14 @@ class AutoThumbnailer:
|
||||
process.kill()
|
||||
except OSError:
|
||||
pass
|
||||
if worker is None:
|
||||
return True
|
||||
if worker is threading.current_thread():
|
||||
return False
|
||||
worker.join(wait)
|
||||
if worker is None or worker is threading.current_thread():
|
||||
return
|
||||
worker.join(CLOSE_WAIT_SECONDS)
|
||||
if worker.is_alive():
|
||||
self._logger.warning(
|
||||
"The automatic preview worker is still running after "
|
||||
f"{wait} s; it stops at its next file"
|
||||
f"{CLOSE_WAIT_SECONDS} s; it stops at its next file"
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
# ── waits (tests and shutdown; never timers) ──────────────────────
|
||||
|
||||
@@ -619,19 +609,12 @@ class AutoThumbnailMixin:
|
||||
)
|
||||
return self._auto_thumbnailer
|
||||
|
||||
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):
|
||||
def close_auto_thumbnails(self, final=False):
|
||||
"""Stop the capture worker and the view pool, and forget them.
|
||||
|
||||
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.
|
||||
Bounded (see AutoThumbnailer.close). ``final`` (the app's cleanup)
|
||||
also keeps a late monitor event from starting a worker again; a
|
||||
test reset passes False and the next use starts fresh ones.
|
||||
"""
|
||||
with AutoThumbnailMixin._auto_thumbnailer_lock:
|
||||
if final:
|
||||
@@ -639,14 +622,12 @@ 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:
|
||||
closed = thumbnailer.close(timeout)
|
||||
thumbnailer.close()
|
||||
if pool is not None:
|
||||
pool.shutdown(wait=False, cancel_futures=True)
|
||||
if final:
|
||||
shutdown_encoders() # kill JPEG XL / ffmpeg encoders still running
|
||||
return closed
|
||||
|
||||
def _auto_thumbnail_failures(self):
|
||||
thumbnailer = self._auto_thumbnailer
|
||||
|
||||
+28
-30
@@ -25,23 +25,6 @@ except ImportError:
|
||||
from database.annotations import MAX_IMAGE_TAGS
|
||||
from database.vocabulary import validate_vocabulary_tag_name
|
||||
|
||||
# The tag template and thresholds live with the saved tagging defaults
|
||||
# (#328), so config.py never imports PIL.
|
||||
try:
|
||||
from .tagging_defaults import DEFAULT_TAG_TEMPLATE
|
||||
except ImportError:
|
||||
from py.tagging_defaults import DEFAULT_TAG_TEMPLATE
|
||||
|
||||
|
||||
def _autotag_config():
|
||||
"""The saved defaults; imported late, as config needs ComfyUI's server."""
|
||||
try:
|
||||
from .config import AutotagConfig
|
||||
except ImportError:
|
||||
from py.config import AutotagConfig
|
||||
return AutotagConfig
|
||||
|
||||
|
||||
# Try to import logging from utils, fallback to standard logging
|
||||
try:
|
||||
from ..utils.logging_config import get_logger
|
||||
@@ -62,12 +45,10 @@ except ImportError:
|
||||
return logger
|
||||
|
||||
|
||||
# Model configurations - same as standalone_tagger.py. ``label`` is the short
|
||||
# name System, Activity and a 409 model-busy show (#328).
|
||||
# Model configurations - same as standalone_tagger.py
|
||||
MODELS = {
|
||||
"gguf": {
|
||||
"name": "GGUF (Recommended)",
|
||||
"label": "GGUF JoyCaption",
|
||||
"description": "Quantized model, ~4.5GB, faster loading",
|
||||
"size": "~4.5GB",
|
||||
"repo": "mradermacher/llama-joycaption-beta-one-hf-llava-GGUF",
|
||||
@@ -78,7 +59,6 @@ MODELS = {
|
||||
},
|
||||
"hf": {
|
||||
"name": "HuggingFace",
|
||||
"label": "HuggingFace JoyCaption",
|
||||
"description": "Full precision model, ~16GB, higher quality",
|
||||
"size": "~16GB",
|
||||
"repo": "fancyfeast/llama-joycaption-beta-one-hf-llava",
|
||||
@@ -86,7 +66,6 @@ MODELS = {
|
||||
},
|
||||
"wd14-swinv2": {
|
||||
"name": "WD14 SwinV2",
|
||||
"label": "WD14 SwinV2",
|
||||
"description": "Fast classifier tagger, ~400MB, no prompt needed",
|
||||
"size": "~400MB",
|
||||
"repo": "SmilingWolf/wd-swinv2-tagger-v3",
|
||||
@@ -96,7 +75,6 @@ MODELS = {
|
||||
},
|
||||
"wd14-vit": {
|
||||
"name": "WD14 ViT",
|
||||
"label": "WD14 ViT",
|
||||
"description": "Fast classifier tagger, ~400MB, ViT architecture",
|
||||
"size": "~400MB",
|
||||
"repo": "SmilingWolf/wd-vit-tagger-v3",
|
||||
@@ -107,7 +85,17 @@ MODELS = {
|
||||
}
|
||||
|
||||
# Default prompts
|
||||
DEFAULT_PROMPT = DEFAULT_TAG_TEMPLATE
|
||||
DEFAULT_PROMPT = (
|
||||
"Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces).\n"
|
||||
"Include: subject_count (1girl/1boy/etc), gender, ethnicity (asian, caucasian, "
|
||||
"etc), age_approx (teen, young_woman), body_attributes, face_features, "
|
||||
"hairstyle, appearance, clothing, accessories, pose, expression, action, "
|
||||
"camera_angle, camera_shot, scene_composition (wide_shot, long_shot, "
|
||||
"dutch_angle), environment, background_elements, lighting, mood.\n"
|
||||
"Prioritize visual accuracy.\n"
|
||||
"No sentences, no commentary, no captions. Only tags. Keep length short but "
|
||||
"descriptive."
|
||||
)
|
||||
|
||||
# Upper bound of one suggestion request: the image tag set cap.
|
||||
MAX_SUGGESTION_LIMIT = MAX_IMAGE_TAGS
|
||||
@@ -191,6 +179,10 @@ class AutoTagService:
|
||||
self._current_model_type: Optional[str] = None
|
||||
self._custom_prompt: str = DEFAULT_PROMPT
|
||||
|
||||
# WD14 threshold defaults
|
||||
self._wd14_general_threshold: float = 0.35
|
||||
self._wd14_character_threshold: float = 0.85
|
||||
|
||||
@property
|
||||
def models_config(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Get the models configuration dictionary."""
|
||||
@@ -213,13 +205,19 @@ class AutoTagService:
|
||||
|
||||
@property
|
||||
def wd14_general_threshold(self) -> float:
|
||||
"""The saved General-tag confidence (``autotag.defaults``)."""
|
||||
return _autotag_config().GENERAL_THRESHOLD
|
||||
return self._wd14_general_threshold
|
||||
|
||||
@wd14_general_threshold.setter
|
||||
def wd14_general_threshold(self, value: float):
|
||||
self._wd14_general_threshold = max(0.0, min(1.0, value))
|
||||
|
||||
@property
|
||||
def wd14_character_threshold(self) -> float:
|
||||
"""The saved Character-tag confidence (``autotag.defaults``)."""
|
||||
return _autotag_config().CHARACTER_THRESHOLD
|
||||
return self._wd14_character_threshold
|
||||
|
||||
@wd14_character_threshold.setter
|
||||
def wd14_character_threshold(self, value: float):
|
||||
self._wd14_character_threshold = max(0.0, min(1.0, value))
|
||||
|
||||
def get_models_status(self) -> Dict[str, Dict[str, Any]]:
|
||||
"""Get availability status for all model types.
|
||||
@@ -862,12 +860,12 @@ class AutoTagService:
|
||||
gen_thresh = (
|
||||
general_threshold
|
||||
if general_threshold is not None
|
||||
else _autotag_config().GENERAL_THRESHOLD
|
||||
else self._wd14_general_threshold
|
||||
)
|
||||
char_thresh = (
|
||||
character_threshold
|
||||
if character_threshold is not None
|
||||
else _autotag_config().CHARACTER_THRESHOLD
|
||||
else self._wd14_character_threshold
|
||||
)
|
||||
return self._generate_wd14(
|
||||
session, tags_list, image, gen_thresh, char_thresh
|
||||
|
||||
+16
-98
@@ -563,92 +563,18 @@ class IntegrationConfig:
|
||||
DEFAULT_SUGGESTION_LIMIT = 10
|
||||
MAX_SUGGESTION_LIMIT = 100 # pinned to MAX_IMAGE_TAGS by tests/test_config.py
|
||||
|
||||
# The tagging engines, options, template and validators (server-free, #328).
|
||||
try:
|
||||
from .tagging_defaults import ( # noqa: F401 - re-exported
|
||||
DEFAULT_TAG_TEMPLATE,
|
||||
EXISTING_POLICIES,
|
||||
LLM_MODEL_TYPES,
|
||||
MAX_INSTRUCTION_BYTES,
|
||||
MAX_INSTRUCTION_CHARS,
|
||||
MODEL_TYPES,
|
||||
WD14_MODEL_TYPES,
|
||||
validate_saved_threshold,
|
||||
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_saved_threshold,
|
||||
validate_template,
|
||||
validate_threshold,
|
||||
)
|
||||
|
||||
|
||||
def _valid_model_type(value: Any) -> str:
|
||||
if not isinstance(value, str) or value not in MODEL_TYPES:
|
||||
raise ValueError("must be one of " + ", ".join(sorted(MODEL_TYPES)))
|
||||
return value
|
||||
|
||||
|
||||
def _valid_existing(value: Any) -> str:
|
||||
if not isinstance(value, str) or value not in EXISTING_POLICIES:
|
||||
raise ValueError("must be one of " + ", ".join(EXISTING_POLICIES))
|
||||
return value
|
||||
|
||||
|
||||
def _valid_suggestion_limit(value: Any) -> int:
|
||||
if type(value) is not int or not 1 <= value <= MAX_SUGGESTION_LIMIT:
|
||||
raise ValueError(f"must be an integer from 1 to {MAX_SUGGESTION_LIMIT}")
|
||||
return value
|
||||
|
||||
|
||||
class AutotagConfig:
|
||||
"""Persisted tagging defaults under ``autotag.defaults`` in config.json.
|
||||
|
||||
One tagging configuration (#328): the model, what a batch does with
|
||||
prompts already tagged, the two WD14 thresholds, the language-model
|
||||
template and how many suggestions Generate returns. Each key is checked
|
||||
alone on load; a hand-edited value that fails its check falls back to
|
||||
its default with one warning rather than failing the load.
|
||||
``suggestion_limit`` is how many suggestions Generate returns (the top N
|
||||
by score), an integer in ``1..MAX_SUGGESTION_LIMIT``. A hand-edited value
|
||||
outside that range, or of another type, falls back to the default with a
|
||||
warning rather than failing the load.
|
||||
"""
|
||||
|
||||
MODEL_TYPE = "wd14-swinv2"
|
||||
EXISTING = "skip"
|
||||
GENERAL_THRESHOLD = 0.35
|
||||
CHARACTER_THRESHOLD = 0.85
|
||||
TEMPLATE = DEFAULT_TAG_TEMPLATE
|
||||
SUGGESTION_LIMIT = DEFAULT_SUGGESTION_LIMIT
|
||||
|
||||
# config.json key -> (attribute, validator, default)
|
||||
_KEYS = {
|
||||
"model_type": ("MODEL_TYPE", _valid_model_type, "wd14-swinv2"),
|
||||
"existing": ("EXISTING", _valid_existing, "skip"),
|
||||
"general_threshold": (
|
||||
"GENERAL_THRESHOLD",
|
||||
lambda v: validate_saved_threshold("value", v),
|
||||
0.35,
|
||||
),
|
||||
"character_threshold": (
|
||||
"CHARACTER_THRESHOLD",
|
||||
lambda v: validate_saved_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
|
||||
@@ -661,14 +587,7 @@ class AutotagConfig:
|
||||
|
||||
@classmethod
|
||||
def get_config(cls) -> Dict[str, Any]:
|
||||
"""The saved defaults; the template only when it is not the built-in
|
||||
one, so a later built-in template still reaches who never set one."""
|
||||
defaults = {
|
||||
key: getattr(cls, attr) for key, (attr, _v, _d) in cls._KEYS.items()
|
||||
}
|
||||
if defaults["template"] == DEFAULT_TAG_TEMPLATE:
|
||||
del defaults["template"]
|
||||
return {"defaults": defaults}
|
||||
return {"defaults": {"suggestion_limit": cls.SUGGESTION_LIMIT}}
|
||||
|
||||
@classmethod
|
||||
def update_config(cls, new_config: Any) -> None:
|
||||
@@ -679,18 +598,17 @@ class AutotagConfig:
|
||||
if not isinstance(defaults, dict):
|
||||
config_logger.warning("Ignoring autotag.defaults that is not an object")
|
||||
return
|
||||
for key, (attr, validate, default) in cls._KEYS.items():
|
||||
if key not in defaults:
|
||||
continue
|
||||
try:
|
||||
setattr(cls, attr, validate(defaults[key]))
|
||||
except ValueError as error:
|
||||
config_logger.warning(
|
||||
f"autotag.defaults.{key} {error}; using {default!r}"
|
||||
if key != "template"
|
||||
else f"autotag.defaults.template {error}; using the built-in one"
|
||||
)
|
||||
setattr(cls, attr, default)
|
||||
if "suggestion_limit" not in defaults:
|
||||
return
|
||||
value = defaults["suggestion_limit"]
|
||||
if cls.valid_suggestion_limit(value):
|
||||
cls.SUGGESTION_LIMIT = value
|
||||
return
|
||||
config_logger.warning(
|
||||
"autotag.defaults.suggestion_limit must be an integer from 1 to "
|
||||
f"{MAX_SUGGESTION_LIMIT}; using {DEFAULT_SUGGESTION_LIMIT}"
|
||||
)
|
||||
cls.SUGGESTION_LIMIT = DEFAULT_SUGGESTION_LIMIT
|
||||
|
||||
|
||||
class PromptManagerConfig:
|
||||
|
||||
+55
-108
@@ -1,27 +1,18 @@
|
||||
"""The staged LoRA reference import (#221, #249 System S10, #355).
|
||||
"""The staged LoRA reference import (#221, #249 System S10).
|
||||
|
||||
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:
|
||||
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:
|
||||
|
||||
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 every scanned sidecar,
|
||||
kept or skipped (path, size, modification time, kept or not, #362) 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
|
||||
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
|
||||
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).
|
||||
@@ -33,15 +24,9 @@ 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 AbstractSet, Callable, Dict, List, Optional
|
||||
from typing import 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
|
||||
@@ -52,97 +37,69 @@ 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: Optional[str]) -> Discovery:
|
||||
"""The sidecars of every LoRA whose model exists inside a LoRA root.
|
||||
def discover(lora_manager_path: str) -> List[Path]:
|
||||
"""Every LoRA metadata file, sorted and without repeats.
|
||||
|
||||
``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.
|
||||
Raises whatever listing a folder raised: a discovery that cannot see
|
||||
every folder must not replace the collection with part of it.
|
||||
"""
|
||||
return discover_sidecars(lora_utils.lora_manager_layout(lora_manager_path))
|
||||
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
|
||||
|
||||
|
||||
def fingerprint(discovery: Discovery) -> str:
|
||||
"""What a review saw: the layout, then every scanned sidecar, kept or
|
||||
skipped, with its path, size, modification time and whether it was kept
|
||||
(#362). A LoraManager save that changes unrelated settings leaves it as
|
||||
it was."""
|
||||
layout = discovery.layout
|
||||
def fingerprint(files: List[Path]) -> str:
|
||||
"""What a review saw of ``files``: each path, size and modification time."""
|
||||
digest = hashlib.sha256()
|
||||
for value in (
|
||||
layout.settings_path,
|
||||
layout.source,
|
||||
layout.error,
|
||||
layout.mode,
|
||||
"\0".join(layout.lora_roots),
|
||||
layout.sidecar_root,
|
||||
layout.example_images_root,
|
||||
layout.library,
|
||||
):
|
||||
digest.update(f"{value}\n".encode("utf-8", "surrogatepass"))
|
||||
for name in discovery.unreadable:
|
||||
digest.update(f"unreadable\0{name}\n".encode("utf-8", "surrogatepass"))
|
||||
kept = {str(sidecar.metadata_path) for sidecar in discovery.sidecars}
|
||||
for path in discovery.scanned_files:
|
||||
flag = 1 if str(path) in kept else 0
|
||||
for path in files:
|
||||
try:
|
||||
stat = path.stat()
|
||||
entry = f"{path}\0{stat.st_size}\0{stat.st_mtime_ns}\0{flag}\n"
|
||||
entry = f"{path}\0{stat.st_size}\0{stat.st_mtime_ns}\n"
|
||||
except OSError:
|
||||
entry = f"{path}\0missing\0{flag}\n"
|
||||
entry = f"{path}\0missing\n"
|
||||
digest.update(entry.encode("utf-8", "surrogatepass"))
|
||||
return "sha256:" + digest.hexdigest()
|
||||
|
||||
|
||||
def preview(
|
||||
lora_manager_path: Optional[str],
|
||||
cache_dir: Path,
|
||||
media_extensions: AbstractSet[str],
|
||||
) -> Dict:
|
||||
def preview(lora_manager_path: str, cache_dir: Path) -> Dict:
|
||||
"""The review's counts (blocking; reads only, requests nothing).
|
||||
|
||||
``{lora_files, with_metadata, skipped, with_local_media,
|
||||
planned_downloads, fingerprint, names, layout}``: the metadata files
|
||||
seen, the kept LoRAs whose sidecar parses, the rest, the LoRAs with a
|
||||
local image, the CivitAI images an import would download for the LoRAs
|
||||
without one (not yet in the cache), the first LoRA names and what the
|
||||
layout says (``layout_summary``).
|
||||
``{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.
|
||||
"""
|
||||
discovery = discover(lora_manager_path)
|
||||
layout = discovery.layout
|
||||
files = discover(lora_manager_path)
|
||||
with_metadata = 0
|
||||
with_local_media = 0
|
||||
planned = 0
|
||||
names = []
|
||||
for sidecar in discovery.sidecars:
|
||||
metadata = lora_utils.read_lora_metadata(sidecar.metadata_path)
|
||||
for path in files:
|
||||
metadata = lora_utils.read_lora_metadata(path)
|
||||
if not metadata:
|
||||
continue
|
||||
with_metadata += 1
|
||||
if local_media(sidecar, metadata, layout, media_extensions):
|
||||
with_local_media += 1
|
||||
else:
|
||||
planned += lora_utils.planned_civitai_downloads(
|
||||
metadata, sidecar.metadata_path, cache_dir
|
||||
)
|
||||
planned += lora_utils.planned_civitai_downloads(metadata, path, cache_dir)
|
||||
if len(names) < 50:
|
||||
names.append(lora_utils.get_model_name_from_metadata(metadata))
|
||||
return {
|
||||
"lora_files": discovery.scanned,
|
||||
"lora_files": len(files),
|
||||
"with_metadata": with_metadata,
|
||||
"skipped": discovery.scanned - with_metadata,
|
||||
"with_local_media": with_local_media,
|
||||
"planned_downloads": planned,
|
||||
"fingerprint": fingerprint(discovery),
|
||||
"fingerprint": fingerprint(files),
|
||||
"names": names,
|
||||
"layout": layout_summary(layout, Path.home()),
|
||||
}
|
||||
|
||||
|
||||
@@ -174,45 +131,35 @@ class StagedImport:
|
||||
|
||||
|
||||
def stage(
|
||||
discovery: Discovery,
|
||||
files: List[Path],
|
||||
cache_dir: Path,
|
||||
api_key: str = "",
|
||||
*,
|
||||
media_extensions: AbstractSet[str],
|
||||
on_record: Optional[Callable[[int, Dict, "StagedImport"], None]] = None,
|
||||
cancelled: Callable[[], bool] = lambda: False,
|
||||
logger=None,
|
||||
) -> StagedImport:
|
||||
"""Read each sidecar and build its record with its images.
|
||||
"""Read each file and build its record, downloading its example images.
|
||||
|
||||
Blocking. Local images (``local_media``) come first; a LoRA without any
|
||||
downloads its CivitAI example images. A sidecar that cannot be read, or
|
||||
whose record fails to build, is skipped by its name, as are the
|
||||
centralized sidecars discovery could not parse. ``cancelled()`` is asked
|
||||
before each sidecar (raises ImportCancelled); ``on_record(index,
|
||||
record_or_None, staged)`` follows each one.
|
||||
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.
|
||||
"""
|
||||
sidecars = discovery.sidecars
|
||||
staged = StagedImport(total=len(sidecars))
|
||||
staged.skipped.extend(discovery.unreadable)
|
||||
for index, sidecar in enumerate(sidecars):
|
||||
staged = StagedImport(total=len(files))
|
||||
for index, path in enumerate(files):
|
||||
if cancelled():
|
||||
raise ImportCancelled()
|
||||
record = None
|
||||
path = sidecar.metadata_path
|
||||
metadata = lora_utils.read_lora_metadata(path)
|
||||
if not metadata:
|
||||
staged.skipped.append(path.name)
|
||||
else:
|
||||
try:
|
||||
local = local_media(
|
||||
sidecar, metadata, discovery.layout, media_extensions
|
||||
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
|
||||
)
|
||||
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)
|
||||
|
||||
@@ -1,217 +0,0 @@
|
||||
"""A per-process cache of LoraManager's layout for the image hot paths (#362).
|
||||
|
||||
Every Library list, prompt detail, inspect, reveal and serve request needs the
|
||||
LoRA allow-list (``ImagesMixin._lora_image_dirs``). Building it means finding
|
||||
the LoraManager install, reading its plugin settings file and its settings
|
||||
file, and resolving every folder. ``LayoutCache`` keeps the last result and
|
||||
reuses it while all of these hold:
|
||||
|
||||
- the key (the configured install path, stripped) and the home folder the
|
||||
folder guard compares against are the same;
|
||||
- the entry is younger than ``MAX_AGE_SECONDS`` (monotonic clock);
|
||||
- the plugin settings file and the settings file have the same signature
|
||||
(``st_mtime_ns``, ``st_size``, or ``MISSING``) as when it was stored;
|
||||
- the install and every stored media root still resolve to themselves (a
|
||||
folder swapped for a symlink reloads, so the folder guard runs again).
|
||||
|
||||
A result is stored only when no watched file is "racy" (modified within
|
||||
``RACY_SECONDS`` of the load, git's rule for racily clean index entries),
|
||||
every watched file could be checked, the settings file was readable, and no
|
||||
``clear()`` ran during the load. What each rule covers: signatures catch a
|
||||
settings change on the next request; the racy rule catches a write during the
|
||||
load, a half-written file and coarse timestamps; the maximum age bounds the
|
||||
inputs no signature sees (``extra_model_paths.yaml``, environment, temporary
|
||||
failures); ``clear()`` makes a finished import's new folders servable at once.
|
||||
|
||||
Hits and loads take no lock: a slow ``os.stat`` on a network share holds up
|
||||
only its own request, and two threads that miss together both load (both
|
||||
loads only read). A short lock guards the compare-and-store and the logging
|
||||
state; no file I/O runs under it. Standard library only, so it is tested
|
||||
without ComfyUI or LoraManager.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Callable, Optional, Tuple
|
||||
|
||||
from .lora_manager_layout import UNREADABLE, LoraManagerLayout, _canonical
|
||||
|
||||
# A watched file modified this recently is not trusted: LoraManager writes its
|
||||
# settings in place, and FAT/exFAT store modification times in 2 s steps.
|
||||
RACY_SECONDS = 3
|
||||
# Bounds what no signature sees (unwatched inputs, temporary failures).
|
||||
MAX_AGE_SECONDS = 30
|
||||
|
||||
# The signature of a watched file that does not exist (a valid state).
|
||||
MISSING = object()
|
||||
# ERROR_BAD_NETPATH and ERROR_BAD_NET_NAME, which Windows raises as
|
||||
# FileNotFoundError.
|
||||
_OFFLINE_SHARE_ERRORS = frozenset({53, 67})
|
||||
|
||||
logger = logging.getLogger("prompt_manager.lora_layout_cache")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Entry:
|
||||
key: str
|
||||
home: str
|
||||
install: Optional[str]
|
||||
layout: LoraManagerLayout
|
||||
signatures: Tuple[Tuple[str, object], ...]
|
||||
stored_at: float
|
||||
|
||||
|
||||
class _CannotValidate(Exception):
|
||||
"""A watched file could not be checked (not merely missing)."""
|
||||
|
||||
|
||||
class LayoutCache:
|
||||
"""At most one ``(install, layout)`` for the active LoraManager config."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
resolve: Callable[[str], Optional[str]],
|
||||
load: Callable[[Optional[str]], LoraManagerLayout],
|
||||
*,
|
||||
wall_ns: Callable[[], int] = time.time_ns,
|
||||
monotonic: Callable[[], float] = time.monotonic,
|
||||
stat=os.stat,
|
||||
pathmod=os.path,
|
||||
home: Callable[[], str] = lambda: str(Path.home()),
|
||||
):
|
||||
self._resolve = resolve
|
||||
self._load = load
|
||||
self._wall_ns = wall_ns
|
||||
self._monotonic = monotonic
|
||||
self._stat = stat
|
||||
self._pathmod = pathmod
|
||||
self._home = home
|
||||
self._entry: Optional[_Entry] = None
|
||||
self._generation = 0
|
||||
self._lock = threading.Lock()
|
||||
# What was last logged; survives clear() so a reload of the same
|
||||
# layout stays quiet.
|
||||
self._logged_layout: Optional[LoraManagerLayout] = None
|
||||
self._logged_refused: Tuple[str, ...] = ()
|
||||
|
||||
def get(self, configured: str) -> Tuple[Optional[str], LoraManagerLayout]:
|
||||
"""``(install, layout)`` for the configured path ("" auto-detects)."""
|
||||
key = (configured or "").strip()
|
||||
home = self._home()
|
||||
entry = self._entry
|
||||
if entry is not None and self._valid(entry, key, home):
|
||||
return entry.install, entry.layout
|
||||
with self._lock:
|
||||
if entry is not None and self._entry is entry:
|
||||
self._entry = None
|
||||
generation = self._generation
|
||||
install = self._resolve(key)
|
||||
layout = self._load(install)
|
||||
self._log_refused(layout)
|
||||
try:
|
||||
signatures = tuple(
|
||||
(path, self._signature(path)) for path in self._watched(install, layout)
|
||||
)
|
||||
except _CannotValidate:
|
||||
return install, layout
|
||||
if layout.error == UNREADABLE or self._racy(signatures, self._wall_ns()):
|
||||
return install, layout
|
||||
stored = _Entry(key, home, install, layout, signatures, self._monotonic())
|
||||
changed = False
|
||||
with self._lock:
|
||||
if self._generation == generation:
|
||||
self._entry = stored
|
||||
changed = layout != self._logged_layout
|
||||
if changed:
|
||||
self._logged_layout = layout
|
||||
if changed:
|
||||
logger.info(
|
||||
f"LoraManager layout changed: {layout.source or 'no'} settings, "
|
||||
f"{layout.mode} storage, {len(layout.media_roots)} media folders"
|
||||
)
|
||||
return install, layout
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Drop the entry; a load that started before this stores nothing."""
|
||||
with self._lock:
|
||||
self._entry = None
|
||||
self._generation += 1
|
||||
|
||||
def _log_refused(self, layout: LoraManagerLayout) -> None:
|
||||
"""One warning per default LoRA folder newly refused by the guard."""
|
||||
refused = tuple(layout.refused_defaults)
|
||||
with self._lock:
|
||||
if refused == self._logged_refused:
|
||||
return
|
||||
new = [folder for folder in refused if folder not in self._logged_refused]
|
||||
self._logged_refused = refused
|
||||
for folder in new:
|
||||
logger.warning(
|
||||
f"ComfyUI LoRA folder '{folder}' is neither scanned nor served: "
|
||||
"it is a filesystem root, the home folder or a folder that "
|
||||
"contains it"
|
||||
)
|
||||
|
||||
# ── validation ──────────────────────────────────────────────────
|
||||
|
||||
def _valid(self, entry: _Entry, key: str, home: str) -> bool:
|
||||
if entry.key != key or entry.home != home:
|
||||
return False
|
||||
if self._monotonic() - entry.stored_at >= MAX_AGE_SECONDS:
|
||||
return False
|
||||
try:
|
||||
for path, signature in entry.signatures:
|
||||
if self._signature(path) != signature:
|
||||
return False
|
||||
except _CannotValidate:
|
||||
return False
|
||||
folders = entry.layout.media_roots
|
||||
if entry.install is not None:
|
||||
folders = (entry.install,) + tuple(folders)
|
||||
for folder in folders:
|
||||
try:
|
||||
if _canonical(folder, self._pathmod) != folder:
|
||||
return False
|
||||
except (OSError, ValueError):
|
||||
return False
|
||||
return True
|
||||
|
||||
def _watched(self, install, layout) -> Tuple[str, ...]:
|
||||
"""The plugin settings file, then the settings file, each once."""
|
||||
if install is None:
|
||||
return ()
|
||||
paths = [self._pathmod.join(install, "settings.json")]
|
||||
if layout.settings_path is not None:
|
||||
paths.append(str(layout.settings_path))
|
||||
watched, seen = [], set()
|
||||
for path in paths:
|
||||
identity = self._pathmod.normcase(self._pathmod.abspath(path))
|
||||
if identity not in seen:
|
||||
seen.add(identity)
|
||||
watched.append(path)
|
||||
return tuple(watched)
|
||||
|
||||
def _signature(self, path: str):
|
||||
try:
|
||||
result = self._stat(path)
|
||||
except (FileNotFoundError, NotADirectoryError) as error:
|
||||
# Windows reports an offline share as "not found": that is a
|
||||
# temporary failure, not a missing file.
|
||||
if getattr(error, "winerror", None) in _OFFLINE_SHARE_ERRORS:
|
||||
raise _CannotValidate(path) from error
|
||||
return MISSING
|
||||
except (OSError, ValueError) as error:
|
||||
raise _CannotValidate(path) from error
|
||||
return (result.st_mtime_ns, result.st_size)
|
||||
|
||||
@staticmethod
|
||||
def _racy(signatures, now_ns: int) -> bool:
|
||||
threshold = now_ns - RACY_SECONDS * 10**9
|
||||
return any(
|
||||
signature is not MISSING and signature[0] > threshold
|
||||
for _path, signature in signatures
|
||||
)
|
||||
@@ -1,653 +0,0 @@
|
||||
"""Where LoraManager keeps its files (#355).
|
||||
|
||||
LoraManager (``custom_nodes/comfyui-lora-manager``) stores one
|
||||
``<name>.metadata.json`` sidecar per model, a preview image or video beside
|
||||
it, and each model's downloaded CivitAI example images. Where it stores them
|
||||
is a user setting, and its settings file itself can live in three places.
|
||||
This module repeats LoraManager's own rules, reading its files only (never
|
||||
importing it), so the LoRA import, the trigger-word cache and the image
|
||||
server share one definition:
|
||||
|
||||
- ``load_layout`` finds the settings file (environment folder, portable
|
||||
plugin file, or the user config folder, in LoraManager's order), keeps only
|
||||
the seven keys the import needs, and resolves the LoRA roots, the sidecar
|
||||
root (centralized storage), the example-images root and the active
|
||||
LoraManager library. Every folder from the settings, and each of
|
||||
ComfyUI's default LoRA folders, passes the folder guard: it must exist and
|
||||
must not be a filesystem root, the home folder or a folder that contains
|
||||
it. Default folders the guard refuses are listed in ``refused_defaults``.
|
||||
- ``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 logging
|
||||
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, the home folder or a folder "
|
||||
"that contains it"
|
||||
)
|
||||
|
||||
logger = logging.getLogger("prompt_manager.lora_manager_layout")
|
||||
|
||||
|
||||
@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, ...]
|
||||
# ComfyUI's default LoRA folders that exist but that the folder guard
|
||||
# refused (canonical): the layout cache logs them (#362).
|
||||
refused_defaults: 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, every metadata file seen, and those unread."""
|
||||
|
||||
layout: LoraManagerLayout
|
||||
sidecars: List[Sidecar]
|
||||
# Every ``*.metadata.json`` seen, kept or skipped: resolved, once each,
|
||||
# sorted by ``str(path)``. The import's fingerprint covers all of them.
|
||||
scanned_files: Tuple[Path, ...]
|
||||
unreadable: List[str]
|
||||
|
||||
@property
|
||||
def scanned(self) -> int:
|
||||
"""How many metadata files were seen."""
|
||||
return len(self.scanned_files)
|
||||
|
||||
|
||||
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 _refused(paths, home: str) -> Tuple[str, ...]:
|
||||
"""The folders in ``paths`` that exist but that the guard refuses."""
|
||||
refused = []
|
||||
for path in paths:
|
||||
absolute = _absolute(path)
|
||||
if not absolute or not os.path.isdir(absolute):
|
||||
continue
|
||||
if _guard_folder(absolute, home) is None:
|
||||
try:
|
||||
canonical = _canonical(absolute)
|
||||
except (OSError, ValueError):
|
||||
continue
|
||||
if canonical not in refused:
|
||||
refused.append(canonical)
|
||||
return tuple(refused)
|
||||
|
||||
|
||||
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())
|
||||
default_list = list(default_roots())
|
||||
defaults = _guarded(default_list, home)
|
||||
refused_defaults = _refused(default_list, 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,
|
||||
refused_defaults=refused_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, refused_defaults
|
||||
)
|
||||
|
||||
|
||||
def _layout_from_settings(
|
||||
settings, settings_path, source, defaults, home, refused_defaults=()
|
||||
):
|
||||
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]),
|
||||
refused_defaults=refused_defaults,
|
||||
)
|
||||
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, [], (), [])
|
||||
if layout.mode == MODE_CENTRALIZED:
|
||||
folders = [layout.sidecar_root]
|
||||
else:
|
||||
folders = list(layout.lora_roots)
|
||||
seen = set()
|
||||
scanned: List[Path] = []
|
||||
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.append(resolved)
|
||||
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))
|
||||
scanned.sort(key=str)
|
||||
return Discovery(layout, sidecars, tuple(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():
|
||||
if os.path.isabs(url.strip()):
|
||||
yield url.strip()
|
||||
else:
|
||||
# LoraManager writes absolute paths; it would resolve a relative
|
||||
# one against its working directory, not the sidecar folder.
|
||||
logger.debug(
|
||||
"Relative preview_url skipped in %s", sidecar.metadata_path.name
|
||||
)
|
||||
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
-132
@@ -13,7 +13,6 @@ import json
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
@@ -28,17 +27,6 @@ except ImportError:
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
from utils.logging_config import get_logger
|
||||
|
||||
from dataclasses import replace
|
||||
|
||||
from .lora_layout_cache import LayoutCache
|
||||
from .lora_manager_layout import (
|
||||
MODE_ALONGSIDE,
|
||||
METADATA_SUFFIX,
|
||||
LoraManagerLayout,
|
||||
discover_sidecars,
|
||||
load_layout,
|
||||
)
|
||||
|
||||
logger = get_logger("prompt_manager.lora_utils")
|
||||
|
||||
# ── LoraManager detection ────────────────────────────────────────────
|
||||
@@ -246,22 +234,23 @@ def resolve_lora_manager_path(path) -> Optional[str]:
|
||||
# ── Metadata reading ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def default_lora_roots() -> List[str]:
|
||||
"""The LoRA folders ComfyUI knows, without LoraManager's settings.
|
||||
def find_lora_directories(lora_manager_path: str) -> List[str]:
|
||||
"""Find directories that contain LoRA models (with .metadata.json files).
|
||||
|
||||
ComfyUI's ``models/loras``, the ``loras`` entries of
|
||||
``extra_model_paths.yaml`` and, inside ComfyUI, every folder
|
||||
``folder_paths`` registers for ``loras``. ``lora_manager_layout`` adds
|
||||
the folders LoraManager's own settings name.
|
||||
Searches: ComfyUI models/loras, extra_model_paths.yaml lora dirs,
|
||||
and the LoraManager extension dir itself.
|
||||
"""
|
||||
dirs = set()
|
||||
lm_path = Path(lora_manager_path)
|
||||
|
||||
root = find_comfyui_root()
|
||||
if root:
|
||||
# Default models/loras
|
||||
models_loras = root / "models" / "loras"
|
||||
if models_loras.is_dir():
|
||||
dirs.add(str(models_loras.resolve()))
|
||||
|
||||
# Extra model paths from ComfyUI config
|
||||
for extra_dir in _get_extra_lora_paths(root):
|
||||
if extra_dir.is_dir():
|
||||
dirs.add(str(extra_dir.resolve()))
|
||||
@@ -278,53 +267,13 @@ def default_lora_roots() -> List[str]:
|
||||
# folder_paths unavailable — not running inside ComfyUI runtime
|
||||
pass
|
||||
|
||||
# Check for any .metadata.json in the LoraManager dir tree
|
||||
for meta in lm_path.rglob("*.metadata.json"):
|
||||
dirs.add(str(meta.parent.resolve()))
|
||||
|
||||
return sorted(dirs)
|
||||
|
||||
|
||||
def lora_manager_install(lora_manager_path: Optional[str]) -> Optional[str]:
|
||||
"""The LoraManager install to read: the configured path when it is an
|
||||
allowed install under ``custom_nodes``, auto-detection when it is empty."""
|
||||
if isinstance(lora_manager_path, str) and lora_manager_path.strip():
|
||||
return resolve_lora_manager_path(lora_manager_path)
|
||||
detected = detect_lora_manager()
|
||||
return resolve_lora_manager_path(detected) if detected else None
|
||||
|
||||
|
||||
def layout_for_install(install: Optional[str]) -> LoraManagerLayout:
|
||||
"""The layout of an already resolved install (None: none was found)."""
|
||||
return load_layout(install, default_roots=default_lora_roots)
|
||||
|
||||
|
||||
def lora_manager_layout(lora_manager_path: Optional[str]) -> LoraManagerLayout:
|
||||
"""Where LoraManager keeps its files (``py/lora_manager_layout.py``).
|
||||
|
||||
A fresh read: the import review, Start, the legacy scan and the trigger
|
||||
word cache use it. The image hot paths use ``cached_lora_manager_layout``.
|
||||
"""
|
||||
return layout_for_install(lora_manager_install(lora_manager_path))
|
||||
|
||||
|
||||
# The image hot paths' layout cache (#362). The lambdas look the functions up
|
||||
# when called, so a test that patches them here is seen by the cache.
|
||||
_layout_cache = LayoutCache(
|
||||
lambda path: lora_manager_install(path),
|
||||
lambda install: layout_for_install(install),
|
||||
)
|
||||
|
||||
|
||||
def cached_lora_manager_layout(
|
||||
configured: Optional[str],
|
||||
) -> Tuple[Optional[str], LoraManagerLayout]:
|
||||
"""``(install, layout)`` for the configured path, from the layout cache
|
||||
while LoraManager's settings files are unchanged (``py/lora_layout_cache.py``)."""
|
||||
return _layout_cache.get(configured or "")
|
||||
|
||||
|
||||
def clear_layout_cache() -> None:
|
||||
"""Drop the cached layout (after an import linked files in new folders)."""
|
||||
_layout_cache.clear()
|
||||
|
||||
|
||||
def _get_extra_lora_paths(comfyui_root: Path) -> List[Path]:
|
||||
"""Parse extra_model_paths.yaml for additional LoRA directories."""
|
||||
results = []
|
||||
@@ -368,7 +317,7 @@ def read_lora_metadata(metadata_path: Path) -> Optional[Dict]:
|
||||
try:
|
||||
with open(metadata_path, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
except (ValueError, OSError) as e: # JSONDecodeError, UnicodeDecodeError
|
||||
except (json.JSONDecodeError, OSError) as e:
|
||||
logger.debug(f"Failed to read {metadata_path}: {e}")
|
||||
return None
|
||||
|
||||
@@ -430,6 +379,46 @@ 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
|
||||
@@ -848,16 +837,30 @@ 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.
|
||||
|
||||
@@ -868,89 +871,52 @@ class TriggerWordCache:
|
||||
self._cache: Dict[str, List[str]] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._loaded = False
|
||||
self._failed_at: Optional[float] = None
|
||||
|
||||
def load(self, lora_manager_path: str) -> int:
|
||||
"""Build the trigger word map from the sidecars LoraManager uses.
|
||||
|
||||
Never raises (a node calls this while it runs): a folder that cannot
|
||||
be listed, or any other failure, is logged and the previous map is
|
||||
kept. When LoraManager's settings cannot be used, the reason is
|
||||
logged and the default LoRA folders are read in alongside mode.
|
||||
"""Scan LoRA metadata files and build the trigger word mapping.
|
||||
|
||||
Returns:
|
||||
Number of LoRAs with trigger words in the map.
|
||||
Number of LoRAs with trigger words found.
|
||||
"""
|
||||
try:
|
||||
new_cache = self._build(lora_manager_path)
|
||||
except OSError as error:
|
||||
logger.warning(f"Trigger word cache not refreshed: {error}")
|
||||
return self._failed()
|
||||
except Exception:
|
||||
# Not a disk error: keep the traceback (#362).
|
||||
logger.exception("Trigger word cache not refreshed")
|
||||
return self._failed()
|
||||
new_cache: Dict[str, List[str]] = {}
|
||||
|
||||
lora_dirs = find_lora_directories(lora_manager_path)
|
||||
for lora_dir in lora_dirs:
|
||||
dir_path = Path(lora_dir)
|
||||
for meta_file in dir_path.rglob("*.metadata.json"):
|
||||
metadata = read_lora_metadata(meta_file)
|
||||
if not metadata:
|
||||
continue
|
||||
|
||||
words = get_trigger_words_from_metadata(metadata)
|
||||
if not words:
|
||||
continue
|
||||
|
||||
# Key by filename stem (what appears in <lora:NAME:weight>)
|
||||
file_name = metadata.get("file_name", "")
|
||||
if file_name:
|
||||
stem = Path(file_name).stem
|
||||
new_cache[stem.lower()] = words
|
||||
|
||||
# Also key by the metadata file stem
|
||||
meta_stem = meta_file.name.replace(".metadata.json", "")
|
||||
if meta_stem.lower() not in new_cache:
|
||||
new_cache[meta_stem.lower()] = words
|
||||
|
||||
with self._lock:
|
||||
self._cache = new_cache
|
||||
self._loaded = True
|
||||
self._failed_at = None
|
||||
|
||||
logger.info(
|
||||
f"Trigger word cache loaded: {len(new_cache)} LoRAs with trigger words"
|
||||
)
|
||||
return len(new_cache)
|
||||
|
||||
def _failed(self) -> int:
|
||||
"""Keep the previous map and note when this load failed."""
|
||||
with self._lock:
|
||||
self._failed_at = time.monotonic()
|
||||
return len(self._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:
|
||||
@@ -960,7 +926,6 @@ class TriggerWordCache:
|
||||
with self._lock:
|
||||
self._cache.clear()
|
||||
self._loaded = False
|
||||
self._failed_at = None
|
||||
|
||||
|
||||
# Module-level singleton
|
||||
|
||||
@@ -1,74 +0,0 @@
|
||||
"""The tagging configuration's constants and validators (#328).
|
||||
|
||||
Server-free on purpose: ``py.api`` must import without ComfyUI's ``server``
|
||||
module, which ``py/config.py`` needs; config re-exports all of this.
|
||||
"""
|
||||
|
||||
import math
|
||||
from typing import Any
|
||||
|
||||
# The tagging engines (``py/autotag.py`` MODELS) and the options they take.
|
||||
WD14_MODEL_TYPES = frozenset({"wd14-swinv2", "wd14-vit"})
|
||||
LLM_MODEL_TYPES = frozenset({"gguf", "hf"})
|
||||
MODEL_TYPES = WD14_MODEL_TYPES | LLM_MODEL_TYPES
|
||||
MAX_INSTRUCTION_CHARS = 4000
|
||||
MAX_INSTRUCTION_BYTES = 16 * 1024
|
||||
# What a batch does with a prompt that already has tags (#328).
|
||||
EXISTING_POLICIES = ("skip", "merge", "replace_after_review")
|
||||
|
||||
# The language-model tag template every page starts from until one is saved.
|
||||
DEFAULT_TAG_TEMPLATE = (
|
||||
"Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces).\n"
|
||||
"Include: subject_count (1girl/1boy/etc), gender, ethnicity (asian, caucasian, "
|
||||
"etc), age_approx (teen, young_woman), body_attributes, face_features, "
|
||||
"hairstyle, appearance, clothing, accessories, pose, expression, action, "
|
||||
"camera_angle, camera_shot, scene_composition (wide_shot, long_shot, "
|
||||
"dutch_angle), environment, background_elements, lighting, mood.\n"
|
||||
"Prioritize visual accuracy.\n"
|
||||
"No sentences, no commentary, no captions. Only tags. Keep length short but "
|
||||
"descriptive."
|
||||
)
|
||||
|
||||
|
||||
def validate_threshold(name: str, value: Any) -> float:
|
||||
"""A WD14 confidence threshold in [0, 1], or ValueError naming ``name``."""
|
||||
if type(value) not in (int, float) or not math.isfinite(value):
|
||||
raise ValueError(f"{name} must be a finite number")
|
||||
if not 0 <= value <= 1:
|
||||
raise ValueError(f"{name} must be between 0 and 1")
|
||||
return float(value)
|
||||
|
||||
|
||||
# Every page shows the thresholds in steps of 0.05 (the concept's System rows,
|
||||
# the Tag image dialogs and the legacy sliders), so a saved one must be on that
|
||||
# grid (#369). Per-request thresholds keep their own checks.
|
||||
THRESHOLD_STEP = 0.05
|
||||
_STEPS_PER_UNIT = round(1 / THRESHOLD_STEP)
|
||||
# A computed float can miss the grid by a hair: (0.1 + 0.2) * 20 is 6.000000000000001.
|
||||
_STEP_TOLERANCE = 1e-12 # in steps; refuses 0.35000000001
|
||||
|
||||
|
||||
def validate_saved_threshold(name: str, value: Any) -> float:
|
||||
"""A saved WD14 threshold: in [0, 1] and a multiple of ``THRESHOLD_STEP``,
|
||||
returned exactly on the grid; else ValueError naming ``name``."""
|
||||
value = validate_threshold(name, value)
|
||||
steps = value * _STEPS_PER_UNIT
|
||||
if abs(steps - round(steps)) > _STEP_TOLERANCE:
|
||||
raise ValueError(f"{name} must be a multiple of {THRESHOLD_STEP}")
|
||||
return round(steps) / _STEPS_PER_UNIT
|
||||
|
||||
|
||||
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
@@ -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.12"
|
||||
requires-python = ">=3.10"
|
||||
dependencies = ["watchdog>=2.1.0", "Pillow>=8.0.0"]
|
||||
|
||||
[build-system]
|
||||
|
||||
@@ -1,27 +0,0 @@
|
||||
"""Final stops for the background threads a test starts (#301, #302).
|
||||
|
||||
Every real ``ImageGenerationHandler``, ``ImageMonitor`` or ``PromptTracker``
|
||||
a test constructs ends with its final stop and an assertion, before the
|
||||
test's folder and databases go: ``close()`` for a handler,
|
||||
``stop_monitoring()`` for a monitor, ``shutdown()`` for a tracker.
|
||||
"""
|
||||
|
||||
from utils.prompt_tracker import PromptTracker
|
||||
|
||||
|
||||
def stop_tracker(testcase, tracker):
|
||||
"""Shut a real tracker's cleanup thread down; mocks are left alone."""
|
||||
if isinstance(tracker, PromptTracker):
|
||||
testcase.assertTrue(tracker.shutdown(), "prompt cleanup thread still running")
|
||||
|
||||
|
||||
def close_handler(testcase, handler):
|
||||
"""Close ``handler`` for good, then stop the tracker it was built with."""
|
||||
testcase.assertTrue(handler.close(), "image worker still running")
|
||||
stop_tracker(testcase, getattr(handler, "prompt_tracker", None))
|
||||
|
||||
|
||||
def stop_monitor(testcase, monitor):
|
||||
"""Stop ``monitor`` (its observer and worker), then its tracker."""
|
||||
testcase.assertTrue(monitor.stop_monitoring(), "image monitor still running")
|
||||
stop_tracker(testcase, getattr(monitor, "prompt_tracker", None))
|
||||
@@ -1,152 +0,0 @@
|
||||
// #328: the legacy admin and gallery auto-tag tools keep working beside System › Auto-tagging. The model download is
|
||||
// the `model_download` Activity job; the pages' stream attaches to it. No model is ever downloaded: the fixture's
|
||||
// fake download waits until the test releases it (`/__test__/tagging/download`).
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from './fixtures.mjs';
|
||||
|
||||
const TITLE = 'Tagging model download · WD14 ViT';
|
||||
|
||||
async function post(request, url, data) {
|
||||
await expect(await request.post(url, data === undefined ? {} : { data })).toBeOK();
|
||||
}
|
||||
|
||||
async function downloads(request) {
|
||||
const response = await request.get('/prompt_manager/jobs');
|
||||
await expect(response).toBeOK();
|
||||
return (await response.json()).jobs.filter(job => job.kind === 'model_download');
|
||||
}
|
||||
|
||||
test.beforeEach(async ({ page, request, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
// The 3.2.0 update notice would cover the admin page's toolbar.
|
||||
await page.addInitScript(() => { localStorage.setItem('pm_update_notice_dismissed', '3.2.0'); });
|
||||
await post(request, '/__test__/jobs/restart');
|
||||
await post(request, '/__test__/tagging/reset');
|
||||
await post(request, '/__test__/tagging/models', { unavailable: ['wd14-vit'] });
|
||||
await post(request, '/__test__/tagging/download', { mode: 'hold' });
|
||||
});
|
||||
|
||||
test.afterEach(async ({ page, request, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
// Releases a download a failed test left held, and puts the raising default back.
|
||||
await post(request, '/__test__/tagging/reset');
|
||||
await post(request, '/__test__/jobs/restart');
|
||||
});
|
||||
|
||||
async function startDownload(page, request) {
|
||||
await page.locator('#autoTagBtn').click();
|
||||
const status = page.locator('#wd14VitModelStatus');
|
||||
await status.getByRole('button', { name: 'Download', exact: true }).click();
|
||||
const modal = page.locator('#autoTagDownloadModal');
|
||||
await expect(modal).toBeVisible();
|
||||
await expect(page.locator('#downloadStatus')).toHaveText('Downloading…');
|
||||
await expect(page.locator('#downloadProgressBar')).toHaveClass(/animate-pulse/);
|
||||
await expect(modal).toContainText('The download continues in Activity.');
|
||||
await expect(page.locator('#cancelDownloadBtn')).toHaveText('Close');
|
||||
const running = (await downloads(request)).filter(job => job.state === 'running');
|
||||
expect(running.map(job => [job.title, job.cancellable])).toEqual([[TITLE, false]]);
|
||||
return { modal, status };
|
||||
}
|
||||
|
||||
test('admin download modal closes on complete and Activity lists the job without Cancel', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/admin');
|
||||
const { modal, status } = await startDownload(page, request);
|
||||
|
||||
await page.goto('/prompt_manager/workspace#/activity');
|
||||
const item = page.locator('#workspace-main section.surface > .activity-item')
|
||||
.filter({ has: page.getByRole('link', { name: new RegExp(`${TITLE}$`) }) });
|
||||
await expect(item.locator('.chip')).toHaveText('Running');
|
||||
await expect(item.locator('progress')).not.toHaveAttribute('value');
|
||||
await expect(item.getByRole('button', { name: 'Cancel job', exact: true })).toHaveCount(0);
|
||||
|
||||
// Back on the admin page a second download of the same model attaches to the running job.
|
||||
await page.goto('/prompt_manager/admin');
|
||||
await startDownload(page, request);
|
||||
await post(request, '/__test__/tagging/download', { mode: 'succeed' });
|
||||
await expect(modal).toBeHidden();
|
||||
await expect(status).toContainText('Downloaded');
|
||||
expect((await downloads(request)).map(job => [job.title, job.state])).toEqual([[TITLE, 'finished']]);
|
||||
});
|
||||
|
||||
test('gallery download modal closes on complete', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
const { modal, status } = await startDownload(page, request);
|
||||
await post(request, '/__test__/tagging/download', { mode: 'succeed' });
|
||||
await expect(modal).toBeHidden();
|
||||
await expect(page.locator('#autoTagModal')).toBeVisible();
|
||||
await expect(status).toContainText('Downloaded');
|
||||
expect((await downloads(request)).map(job => job.state)).toEqual(['finished']);
|
||||
});
|
||||
|
||||
test('a failed download closes the modal and says so', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
const { modal, status } = await startDownload(page, request);
|
||||
await post(request, '/__test__/tagging/download', { mode: 'fail' });
|
||||
await expect(modal).toBeHidden();
|
||||
await expect(page.getByText('Download failed: Download failed')).toBeVisible();
|
||||
await expect(status.getByRole('button', { name: 'Download', exact: true })).toBeVisible();
|
||||
expect((await downloads(request)).map(job => job.state)).toEqual(['failed']);
|
||||
});
|
||||
|
||||
// (wd14-vit is not downloaded in these tests: beforeEach marks it so the download tests can offer it.)
|
||||
// The gallery's batch reader keyed on `data.status` values the server never sends, so its modal never moved or
|
||||
// closed (#328 Problem 5). It now reads `type` frames as the admin page does.
|
||||
async function watchProgress(page) {
|
||||
await page.evaluate(() => {
|
||||
window.progressSeen = [];
|
||||
const node = document.getElementById('autoTagProgressPercent');
|
||||
new MutationObserver(() => window.progressSeen.push(node.textContent))
|
||||
.observe(node, { childList: true, characterData: true, subtree: true });
|
||||
});
|
||||
}
|
||||
|
||||
async function runGalleryBatch(page) {
|
||||
await page.locator('#autoTagBtn').click();
|
||||
await page.locator('input[name="autoTagModel"][value="wd14-swinv2"]').check();
|
||||
await watchProgress(page);
|
||||
const done = page.waitForResponse(response => response.url().includes('/prompt_manager/autotag/start'));
|
||||
await page.locator('#startAutoTagBtn').click();
|
||||
await done;
|
||||
await expect(page.locator('#autoTagProgressModal')).toBeHidden();
|
||||
await expect(page.getByText(/^Auto tagging complete! Applied tags to \d+ prompts\.$/)).toBeVisible();
|
||||
}
|
||||
|
||||
test('gallery batch modal counts progress and closes on complete', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await runGalleryBatch(page);
|
||||
const seen = await page.evaluate(() => window.progressSeen);
|
||||
// The fixture's prompts are already tagged, so the batch skips them; its progress frames still move the bar.
|
||||
expect(seen).toContain('95%');
|
||||
expect((await downloads(request)).length).toBe(0);
|
||||
});
|
||||
|
||||
test('admin review opens the first image', async ({ page }) => {
|
||||
await page.goto('/prompt_manager/admin');
|
||||
await page.locator('#autoTagBtn').click();
|
||||
await page.locator('input[name="autoTagModel"][value="wd14-swinv2"]').check();
|
||||
await page.locator('#startReviewBtn').click();
|
||||
// The first image's prompt is already tagged: confirm a new pass, then the review opens on that image.
|
||||
await page.locator('#retagConfirmBtn').click();
|
||||
await expect(page.locator('#autoTagReviewModal')).toBeVisible();
|
||||
await expect(page.locator('#reviewImageName')).not.toHaveText(/^(image_name\.png|Unknown|)$/);
|
||||
await expect(page.getByText('Failed to start review mode')).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('gallery unload works and is refused while the gate is held', async ({ page, request }) => {
|
||||
await page.goto('/prompt_manager/gallery');
|
||||
await runGalleryBatch(page); // keeps the model in memory
|
||||
await page.locator('#autoTagBtn').click();
|
||||
const unload = page.locator('#unloadModelBtn');
|
||||
await expect(unload).toBeVisible();
|
||||
|
||||
// A held download owns the model gate: Unload is refused with the server's error.
|
||||
const held = await request.post('/prompt_manager/jobs/model-download', { data: { model_type: 'gguf' } });
|
||||
expect(held.status()).toBe(202);
|
||||
await unload.click();
|
||||
await expect(page.getByText(/^Failed to unload model: A tagging model is busy/)).toBeVisible();
|
||||
|
||||
await post(request, '/__test__/tagging/download', { mode: 'succeed' });
|
||||
await expect.poll(async () => (await downloads(request)).map(job => job.state)).toEqual(['finished']);
|
||||
await unload.click();
|
||||
await expect(page.getByText('WD14-SWINV2 model unloaded successfully')).toBeVisible();
|
||||
});
|
||||
+1
-17
@@ -72,8 +72,6 @@ class TestModel:
|
||||
# Scored names the concept comparisons answer instead of the fixed list
|
||||
# (``/__test__/tagging/suggestions``); None keeps the fixed list.
|
||||
suggestions: list | None = None
|
||||
# The engine's model table (``py/autotag.py`` MODELS), set at startup.
|
||||
models_config: dict = {}
|
||||
|
||||
def get_models_status(self) -> dict:
|
||||
return {
|
||||
@@ -81,7 +79,7 @@ class TestModel:
|
||||
"downloaded": name not in self.unavailable,
|
||||
"loaded": self.loaded == name,
|
||||
}
|
||||
for name in self.models_config
|
||||
for name in ("gguf", "hf", "wd14-vit", "wd14-convnext", "wd14-swinv2")
|
||||
}
|
||||
|
||||
def is_model_loaded(self) -> bool:
|
||||
@@ -168,20 +166,7 @@ def register_reveal_stub(routes, api, root: Path):
|
||||
return types.SimpleNamespace(reset=reset)
|
||||
|
||||
|
||||
def isolate_lora_manager_settings(root: Path) -> Path:
|
||||
"""LoraManager's settings lookup reads an empty folder under ``root`` (#355).
|
||||
|
||||
The server's temporary root is removed when it exits, and with it this
|
||||
folder: no page test reads the developer's real LoraManager settings.
|
||||
"""
|
||||
folder = root / "lora-manager-settings"
|
||||
folder.mkdir(exist_ok=True)
|
||||
os.environ["LORA_MANAGER_SETTINGS_DIR"] = str(folder)
|
||||
return folder
|
||||
|
||||
|
||||
def create_app(root: Path) -> web.Application:
|
||||
isolate_lora_manager_settings(root)
|
||||
os.environ["PROMPT_MANAGER_CONFIG_PATH"] = str(root / "config.json")
|
||||
os.chdir(root)
|
||||
app = web.Application()
|
||||
@@ -209,7 +194,6 @@ def create_app(root: Path) -> web.Application:
|
||||
from py import autotag
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
TestModel.models_config = autotag.MODELS
|
||||
model = TestModel()
|
||||
autotag.get_autotag_service = lambda: model
|
||||
api = PromptManagerAPI()
|
||||
|
||||
@@ -270,8 +270,7 @@ 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: [['LoraManager', 'environment settings · alongside storage · library default'],
|
||||
['Scope', plan.scope.label], ['Found', '2 metadata files · 2 LoRAs · 0 with local images'],
|
||||
grid: { ...reference.grid, rows: [['Scope', plan.scope.label], ['Found', '2 metadata files · 2 readable'],
|
||||
['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']] },
|
||||
@@ -313,7 +312,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, preview: false });
|
||||
await fakeLoraManager(request, { loras: 2, civitai: true });
|
||||
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`) });
|
||||
@@ -325,15 +324,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, preview: false });
|
||||
await fakeLoraManager(request, { loras: 2, civitai: true });
|
||||
const concept = await conceptReview(page, 'dark');
|
||||
const { plan } = await productionReview(page, request, 'dark');
|
||||
const expected = mapImportConcept(concept, plan);
|
||||
const paths = async () => diffInventory(expected, await page.evaluate(measureReview)).map(item => item.path);
|
||||
expect(await paths()).toEqual([]);
|
||||
await page.evaluate(() => { document.querySelectorAll('dialog#modal[open] dl.meta-grid dt')[4].textContent = 'Downloads'; });
|
||||
expect(await paths()).toEqual(['grid.rows.4.0']);
|
||||
await page.evaluate(() => { document.querySelectorAll('dialog#modal[open] dl.meta-grid dt')[4].textContent = 'Network'; });
|
||||
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'; });
|
||||
const undoGrid = await mutate(page, { css: 'dialog#modal dl.meta-grid { column-gap: 3px; }' });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['grid.columnGap']));
|
||||
await undoGrid();
|
||||
|
||||
@@ -1,385 +0,0 @@
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from './fixtures.mjs';
|
||||
import { SIZES, THEMES, seedConcept, openConcept, openProduction, diffInventory, mutate } from './concept-inventory.mjs';
|
||||
|
||||
// #328 System S11 Task 8: System › Auto-tagging (app.js renderSystem 685–686) and the Tag output images review against
|
||||
// the immutable concept, on the concept-shaped fixture with the fake tagging model (every engine downloaded, none
|
||||
// loaded). Each concept measurement is mapped only by the recorded deviations (D132, D133, D134), each commented where
|
||||
// it applies; after that the difference must be empty, and a mutated label, option, order, class, select width or
|
||||
// token must be noticed. The concept has no batch review, so the Tag output images review is compared with the
|
||||
// concept's own record-changing review, `maintenancePlan` (953–958), whose shell, grid, note and actions it shares.
|
||||
test.use({ locale: 'en-US', timezoneId: 'UTC' });
|
||||
|
||||
test.beforeEach(async ({ page, request, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
// The tagging reset restores the fake model (nothing loaded, every engine on disk) and the saved tagging defaults;
|
||||
// the concept fixture then replaces its rows and files.
|
||||
await expect(await request.post('/__test__/tagging/reset', { data: {} })).toBeOK();
|
||||
await seedConcept(request);
|
||||
await expect(await request.post('/__test__/system/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/jobs/restart')).toBeOK();
|
||||
page.conceptErrors = [];
|
||||
page.on('pageerror', error => page.conceptErrors.push(error.message));
|
||||
});
|
||||
|
||||
test.afterEach(async ({ page, request, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
await request.post('/__test__/jobs/restart');
|
||||
await request.post('/__test__/system/reset');
|
||||
await request.post('/__test__/tagging/reset', { data: {} });
|
||||
expect(page.conceptErrors).toEqual([]);
|
||||
});
|
||||
|
||||
const PREVIEW_URL = '/prompt_manager/autotag/batch/preview';
|
||||
const MODEL_OPTIONS = ['WD14 SwinV2 · ~400 MB', 'WD14 ViT · ~400 MB', 'GGUF JoyCaption · ~4.5 GB', 'HuggingFace JoyCaption · ~16 GB'];
|
||||
const SUGGESTIONS = 'From 1 to 100: the most suggestions Generate shows for one image.';
|
||||
const BATCH = 'Add tags to each prompt from its linked output images with the model above. Review the scope first.';
|
||||
|
||||
/** Runs in the page: Auto-tagging's surface under `mainSelector`: rows, field, actions, operation, geometry and tokens. */
|
||||
function measureTagging(mainSelector) {
|
||||
const layout = document.querySelector(`${mainSelector} .settings-layout`);
|
||||
const css = (node, ...keys) => node ? Object.fromEntries(keys.map(key => [key, getComputedStyle(node)[key]])) : null;
|
||||
const name = node => [node.localName, ...[...node.classList].sort()].join('.');
|
||||
const text = node => node?.textContent.trim() ?? null;
|
||||
const control = row => {
|
||||
const select = row.querySelector(':scope > select');
|
||||
if (select) return { kind: 'select', id: select.id, options: [...select.options].map(option => [option.textContent, option.disabled]),
|
||||
selected: select.selectedOptions[0]?.textContent ?? null, disabled: select.disabled };
|
||||
const input = row.querySelector(':scope > input');
|
||||
return input ? { kind: `input.${input.type}`, id: input.id, value: input.value, disabled: input.disabled } : null;
|
||||
};
|
||||
const surface = layout.querySelector(':scope > section.surface');
|
||||
const firstRow = surface.querySelector('.setting-row');
|
||||
const field = surface.querySelector(':scope > .field');
|
||||
const actions = surface.querySelector(':scope > .row-actions');
|
||||
return {
|
||||
layout: [...layout.children].map(name),
|
||||
tree: [...surface.children].map(name),
|
||||
heading: text(surface.querySelector(':scope > h2')),
|
||||
intro: text(surface.querySelector(':scope > p')),
|
||||
// Per row: whether it may wrap, and whether its control sits under the text (a wrapped row) or beside it.
|
||||
rows: [...surface.querySelectorAll(':scope > .setting-row')].map(row => ({ tree: [...row.children].map(name),
|
||||
text: [...row.firstElementChild.children].map(node => [name(node), text(node), node.localName === 'label' ? node.htmlFor : null]),
|
||||
control: control(row), flexWrap: getComputedStyle(row).flexWrap,
|
||||
controlBelow: row.lastElementChild.getBoundingClientRect().top >= row.firstElementChild.getBoundingClientRect().bottom - 0.5 })),
|
||||
field: { tree: [...field.children].map(name), label: [text(field.querySelector('label')), field.querySelector('label').htmlFor],
|
||||
textarea: { id: field.querySelector('textarea').id, value: field.querySelector('textarea').value,
|
||||
...css(field.querySelector('textarea'), 'fontSize', 'minHeight', 'backgroundColor', 'color', 'borderTopColor') },
|
||||
...css(field, 'display', 'rowGap', 'fontSize', 'marginTop', 'marginBottom') },
|
||||
actions: { buttons: [...actions.children].map(item => [text(item), item.className]),
|
||||
...css(actions, 'display', 'flexWrap', 'columnGap', 'justifyContent'),
|
||||
primary: css(actions.querySelector('.ca-primary'), 'backgroundColor', 'color', 'fontSize', 'paddingTop', 'paddingLeft'),
|
||||
plain: css(actions.querySelector('button:not(.ca-primary)'), 'backgroundColor', 'color', 'fontSize', 'borderTopColor') },
|
||||
operations: [...surface.querySelectorAll(':scope > .operation')].map(row => ({ tree: [...row.children].map(name),
|
||||
title: text(row.querySelector('h3')), description: text(row.querySelector('p')),
|
||||
actions: [...row.querySelectorAll('.row-actions > button')].map(item => [text(item), item.className]) })),
|
||||
surface: css(surface, 'paddingTop', 'paddingLeft', 'backgroundColor', 'borderTopWidth', 'borderTopColor', 'borderTopLeftRadius'),
|
||||
row: css(firstRow, 'display', 'alignItems', 'justifyContent', 'columnGap', 'paddingTop', 'paddingBottom', 'borderTopWidth',
|
||||
'borderTopColor', 'marginTop'),
|
||||
label: css(firstRow.querySelector('label'), 'fontSize', 'fontWeight', 'color'),
|
||||
description: css(firstRow.querySelector('p'), 'fontSize', 'color', 'marginTop', 'maxWidth'),
|
||||
select: css(firstRow.querySelector('select'), 'fontSize', 'minWidth', 'backgroundColor', 'color', 'borderTopColor'),
|
||||
number: css(surface.querySelector('.setting-row input[type="number"]'), 'fontSize', 'minWidth', 'maxWidth', 'backgroundColor',
|
||||
'color', 'borderTopColor'),
|
||||
// Whether each select is as wide as the browser lays out a copy holding only its chosen option, i.e. it shows
|
||||
// that option whole (the browser's own select sizing, arrow included).
|
||||
selectsFit: [...surface.querySelectorAll('.setting-row > select')].map(select => {
|
||||
const copy = select.cloneNode(false);
|
||||
copy.removeAttribute('id');
|
||||
copy.append(new Option(select.selectedOptions[0]?.textContent ?? ''));
|
||||
// Important inline values: the phone cap is an !important rule (concept.css), which a plain inline style loses to.
|
||||
for (const [key, value] of [['position', 'absolute'], ['visibility', 'hidden'], ['width', 'max-content'], ['min-width', '0'],
|
||||
['max-width', 'none'], ['flex', 'none']]) copy.style.setProperty(key, value, 'important');
|
||||
select.parentElement.append(copy);
|
||||
const needed = copy.getBoundingClientRect().width;
|
||||
copy.remove();
|
||||
return select.getBoundingClientRect().width + 0.5 >= needed;
|
||||
}),
|
||||
pageScroll: document.documentElement.scrollWidth > document.documentElement.clientWidth,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* The concept's Auto-tagging with the recorded deviations applied, for the fixture (all engines on disk, none loaded).
|
||||
* `phone` is a viewport at the concept's 700px breakpoint or below.
|
||||
*/
|
||||
function mapTaggingConcept(reference, { template, phone }) {
|
||||
const [model, existing, general, character] = reference.rows;
|
||||
const [preview] = reference.actions.buttons;
|
||||
return {
|
||||
...reference,
|
||||
// D133 adds a row after the thresholds; D134 adds the operation after the actions.
|
||||
tree: [...reference.tree.slice(0, 6), 'div.setting-row', ...reference.tree.slice(6), 'div.operation'],
|
||||
rows: [
|
||||
// D132 (OQ5, OQ6): the four real engines with their download size, and the chosen one's live status.
|
||||
{ ...model, text: [model.text[0], ['p', 'Downloaded. Not loaded.', null]],
|
||||
control: { ...model.control, options: MODEL_OPTIONS.map(label => [label, false]), selected: MODEL_OPTIONS[0] },
|
||||
// D132: at phone width the row wraps, so the engine select sits whole under its text.
|
||||
...(phone ? { flexWrap: 'wrap', controlBelow: true } : {}) },
|
||||
existing, general, character,
|
||||
// D133: the persisted suggestion count, on the concept's number row.
|
||||
{ ...general, text: [['label', 'Suggestions per image', 'tag-suggestions'], ['p', SUGGESTIONS, null]],
|
||||
control: { ...general.control, id: 'tag-suggestions', value: '10' } },
|
||||
],
|
||||
// The saved template (the server's default until one is saved) replaces the concept's sample text.
|
||||
field: { ...reference.field, textarea: { ...reference.field.textarea, value: template } },
|
||||
// D132: Download (for an engine already on disk, "Download again") and Unload replace the two simulated jobs.
|
||||
actions: { ...reference.actions, buttons: [preview, ['Download again', ''], ['Unload', '']] },
|
||||
// D134: the prompt-tag batch, reviewed first.
|
||||
operations: [{ tree: ['div', 'div.row-actions'], title: 'Tag output images', description: BATCH, actions: [['Review', '']] }],
|
||||
// D132: the engine select shows its whole chosen option at every width (at 390 the concept's 120px cap cuts it;
|
||||
// `#tag-model` alone may grow there). Existing tags keeps the concept's own fit, cut or not.
|
||||
selectsFit: reference.selectsFit.map((fit, index) => (index === 0 ? true : fit)),
|
||||
};
|
||||
}
|
||||
|
||||
async function conceptTagging(page, theme) {
|
||||
await openConcept(page, 'system', theme);
|
||||
await page.locator('#main .section-nav').getByRole('button', { name: 'Auto-tagging', exact: true }).click();
|
||||
await expect(page.locator('#main #tag-template')).toBeVisible();
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureTagging, '#main');
|
||||
}
|
||||
|
||||
async function productionTagging(page, theme) {
|
||||
await openProduction(page, 'system/tagging', theme);
|
||||
// Compare once the settings and the model status answered.
|
||||
await expect(page.locator('#workspace-main #tag-model')).toBeEnabled();
|
||||
await expect(page.locator('#workspace-main .setting-row').first().locator('p')).toHaveText('Downloaded. Not loaded.');
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureTagging, '#workspace-main');
|
||||
}
|
||||
|
||||
async function savedTemplate(request, page) {
|
||||
const response = await request.get('/prompt_manager/settings');
|
||||
await expect(response).toBeOK();
|
||||
return { template: (await response.json()).settings.autotag_template, phone: page.viewportSize().width <= 700 };
|
||||
}
|
||||
|
||||
for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
test(`auto-tagging matches the concept ${theme} ${width}x${height}`, async ({ page, request }, testInfo) => {
|
||||
await page.setViewportSize({ width, height });
|
||||
const concept = await conceptTagging(page, theme);
|
||||
await page.screenshot({ path: testInfo.outputPath(`concept-tagging-${theme}-${width}.png`), fullPage: true });
|
||||
const production = await productionTagging(page, theme);
|
||||
await page.screenshot({ path: testInfo.outputPath(`production-tagging-${theme}-${width}.png`), fullPage: true });
|
||||
expect(diffInventory(mapTaggingConcept(concept, await savedTemplate(request, page)), production)).toEqual([]);
|
||||
expect(production.pageScroll).toBe(false);
|
||||
});
|
||||
}
|
||||
|
||||
// D132 below the concept's smallest size: on a 320px phone the engine select still shows its whole option.
|
||||
for (const theme of THEMES) test(`auto-tagging engine select fits at 320 wide ${theme}`, async ({ page }) => {
|
||||
await page.setViewportSize({ width: 320, height: 700 });
|
||||
const production = await productionTagging(page, theme);
|
||||
for (const model of ['hf', 'gguf', 'wd14-swinv2']) {
|
||||
await page.locator('#workspace-main #tag-model').selectOption(model);
|
||||
expect((await page.evaluate(measureTagging, '#workspace-main')).selectsFit[0], model).toBe(true);
|
||||
}
|
||||
expect(production.pageScroll).toBe(false);
|
||||
});
|
||||
|
||||
test('auto-tagging comparison detects mutations', async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
const concept = await conceptTagging(page, 'dark');
|
||||
const saved = await savedTemplate(request, page);
|
||||
const expected = mapTaggingConcept(concept, saved);
|
||||
const production = await productionTagging(page, 'dark');
|
||||
expect(diffInventory(expected, production)).toEqual([]);
|
||||
// Unmapped, the section differs only where D132, D133 and D134 are recorded (and the saved template's text).
|
||||
expect(diffInventory(concept, production).map(item => item.path.split('.').slice(0, 3).join('.'))
|
||||
.filter((path, index, all) => all.indexOf(path) === index).sort()).toEqual([
|
||||
'actions.buttons.1', 'actions.buttons.2', 'field.textarea.value', 'operations.0.actions', 'operations.0.description',
|
||||
'operations.0.title', 'operations.0.tree', 'operations.length', 'rows.0.control', 'rows.0.text', 'rows.4.control', 'rows.4.controlBelow',
|
||||
'rows.4.flexWrap', 'rows.4.text', 'rows.4.tree', 'rows.length', 'tree.6', 'tree.7', 'tree.8', 'tree.9', 'tree.length'].sort());
|
||||
const measure = () => page.evaluate(measureTagging, '#workspace-main');
|
||||
const paths = async () => diffInventory(expected, await measure()).map(item => item.path);
|
||||
// A removed row.
|
||||
const undoRow = await mutate(page, { remove: ['#workspace-main section.surface > .setting-row:nth-of-type(2)'] });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['rows.length', 'tree.length']));
|
||||
await undoRow();
|
||||
// A renamed label.
|
||||
await page.evaluate(() => { document.querySelector('#workspace-main label[for="tag-existing"]').textContent = 'Existing'; });
|
||||
expect(await paths()).toEqual(['rows.1.text.0.1']);
|
||||
await page.evaluate(() => { document.querySelector('#workspace-main label[for="tag-existing"]').textContent = 'Existing tags'; });
|
||||
// A reordered option.
|
||||
await page.evaluate(() => { const select = document.querySelector('#workspace-main #tag-existing'); select.append(select.options[0]); });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['rows.1.control.options.0.0']));
|
||||
await page.evaluate(() => { const select = document.querySelector('#workspace-main #tag-existing'); select.prepend(select.options[2]); });
|
||||
// ca-primary moved to another action.
|
||||
await page.evaluate(() => {
|
||||
const [first, second] = document.querySelectorAll('#workspace-main section.surface > .row-actions > button');
|
||||
first.classList.remove('ca-primary');
|
||||
second.classList.add('ca-primary');
|
||||
});
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['actions.buttons.0.1', 'actions.buttons.1.1']));
|
||||
await page.evaluate(() => {
|
||||
const [first, second] = document.querySelectorAll('#workspace-main section.surface > .row-actions > button');
|
||||
second.classList.remove('ca-primary');
|
||||
first.classList.add('ca-primary');
|
||||
});
|
||||
// A select squeezed until its chosen option is cut.
|
||||
const undoSqueeze = await mutate(page, { css: '#workspace-main #tag-model { max-width: 60px !important; min-width: 0 !important; }' });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['selectsFit.0']));
|
||||
await undoSqueeze();
|
||||
// A row that wraps where the concept's does not.
|
||||
const undoWrap = await mutate(page, { css: '#workspace-main .setting-row:has(> #tag-existing) { flex-wrap: wrap; }' });
|
||||
expect(await paths()).toEqual(['rows.1.flexWrap']);
|
||||
await undoWrap();
|
||||
// A swapped token.
|
||||
const undoToken = await mutate(page, { css: '#workspace-main .settings-layout { --line: rgb(255, 0, 0); }' });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['row.borderTopColor']));
|
||||
await undoToken();
|
||||
expect(diffInventory(expected, await measure())).toEqual([]);
|
||||
});
|
||||
|
||||
// D132 at phone width, unmapped: only the engine row wraps; Existing tags keeps the concept's cut select.
|
||||
test('auto-tagging at 390 wraps only the engine row', async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 390, height: 844 });
|
||||
const concept = await conceptTagging(page, 'dark');
|
||||
const production = await productionTagging(page, 'dark');
|
||||
expect(diffInventory(mapTaggingConcept(concept, await savedTemplate(request, page)), production)).toEqual([]);
|
||||
expect(production.rows.map(row => row.flexWrap)).toEqual(['wrap', 'nowrap', 'nowrap', 'nowrap', 'nowrap']);
|
||||
expect(production.selectsFit[1]).toBe(concept.selectsFit[1]);
|
||||
});
|
||||
|
||||
/** Runs in the page: the open review dialog: its title, body and the concept's tokens. */
|
||||
function measureReview() {
|
||||
const dialog = document.querySelector('dialog#modal[open]');
|
||||
const css = (node, ...keys) => node ? Object.fromEntries(keys.map(key => [key, getComputedStyle(node)[key]])) : null;
|
||||
const name = node => [node.localName, ...[...node.classList].sort()].join('.');
|
||||
const text = node => node?.textContent.trim() ?? null;
|
||||
const body = dialog.querySelector('#modal-body');
|
||||
const grid = body.querySelector(':scope > dl.meta-grid');
|
||||
const note = [...body.querySelectorAll(':scope > p.small')].at(-1);
|
||||
const actions = body.querySelector(':scope > .dialog-actions');
|
||||
return {
|
||||
title: text(dialog.querySelector('#modal-title')),
|
||||
tree: [...body.children].map(name),
|
||||
heading: { text: text(body.querySelector(':scope > h3')),
|
||||
...css(body.querySelector(':scope > h3'), 'fontSize', 'fontWeight', 'color', 'marginBottom') },
|
||||
grid: { rows: [...grid.children].map(row => [text(row.querySelector('dt')), text(row.querySelector('dd'))]),
|
||||
...css(grid, 'display', 'gridTemplateColumns', 'columnGap', 'rowGap'),
|
||||
dt: css(grid.querySelector('dt'), 'fontSize', 'color', 'textTransform'), dd: css(grid.querySelector('dd'), 'fontSize', 'color') },
|
||||
backup: body.querySelector(':scope > .notice') && { tree: [...body.querySelector(':scope > .notice').children].map(name),
|
||||
text: [...body.querySelectorAll(':scope > .notice > div:first-child > *')].map(node => [node.localName, text(node)]),
|
||||
buttons: [...body.querySelectorAll(':scope > .notice button')].map(text) },
|
||||
note: { text: text(note), ...css(note, 'fontSize', 'color', 'marginTop') },
|
||||
actions: { tree: [...actions.children].map(name), ...css(actions, 'display', 'justifyContent', 'columnGap', 'marginTop'),
|
||||
buttons: [...actions.children].map(item => ({ text: text(item), className: item.className,
|
||||
...css(item, 'backgroundColor', 'color', 'fontSize', 'paddingTop', 'paddingLeft', 'borderTopColor') })) },
|
||||
};
|
||||
}
|
||||
|
||||
const number = value => value.toLocaleString('en-US');
|
||||
|
||||
/**
|
||||
* D134: the server's batch preview in the concept's review: "Review tag output images" with "Tag output images" as its
|
||||
* `h3`, the model, policy, counts and settings in the concept's `dl.meta-grid` (Backup in the concept's wording), the
|
||||
* upper-bound sentence in the concept's `p.small`, the D55 backup box and Cancel / "Start tagging".
|
||||
*/
|
||||
function mapBatchConcept(reference, plan) {
|
||||
const [cancel, start] = reference.actions.buttons;
|
||||
return {
|
||||
...reference,
|
||||
title: 'Review tag output images',
|
||||
tree: ['h3', 'dl.meta-grid', 'p.small', 'div.notice', 'div.dialog-actions'],
|
||||
heading: { ...reference.heading, text: 'Tag output images' },
|
||||
grid: { ...reference.grid, rows: [['Model', 'WD14 SwinV2 · ~400 MB'], ['Existing tags', 'Merge with existing'],
|
||||
['Images linked', number(plan.images)], ['Prompts', number(plan.prompts)], ['Already tagged', number(plan.already_tagged)],
|
||||
['Prompts to tag', `up to ${number(plan.to_tag)}`], ['Settings', 'General 0.35 · Character 0.85'],
|
||||
['Backup', 'Create and verify before changes']] },
|
||||
backup: { tree: ['div', 'div.row-actions'], text: [['h3', 'Back up the database first (recommended)'],
|
||||
['p', 'Downloads a copy of the prompts database before this operation changes any records.']], buttons: ['Back up now'] },
|
||||
note: { ...reference.note, text: 'Nothing is written until you start. Missing files and files outside the output folders '
|
||||
+ 'are counted, so the prompts to tag are an upper bound; Activity reports what was tagged.' },
|
||||
actions: { ...reference.actions, buttons: [cancel, { ...start, text: 'Start tagging' }] },
|
||||
};
|
||||
}
|
||||
|
||||
async function conceptReview(page, theme) {
|
||||
await openConcept(page, 'system', theme);
|
||||
await page.locator('#main .section-nav').getByRole('button', { name: 'Maintenance', exact: true }).click();
|
||||
await page.locator('#main [data-plan="duplicate prompts"]').click();
|
||||
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Review maintenance plan');
|
||||
await page.mouse.move(0, 0);
|
||||
return page.evaluate(measureReview);
|
||||
}
|
||||
|
||||
async function productionReview(page, theme) {
|
||||
await productionTagging(page, theme);
|
||||
// Every concept-fixture prompt with an image is already tagged, so under the saved Skip there is nothing to start
|
||||
// (the Cancel-only review). Merge reviews a batch with work: the full review the concept's shell is compared with.
|
||||
await page.locator('#workspace-main #tag-existing').selectOption('merge');
|
||||
const answer = page.waitForResponse(response => response.url().endsWith(PREVIEW_URL));
|
||||
await page.locator('#workspace-main .operation').getByRole('button', { name: 'Review', exact: true }).click();
|
||||
const plan = await (await answer).json();
|
||||
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Review tag output images');
|
||||
await page.mouse.move(0, 0);
|
||||
return { plan, production: await page.evaluate(measureReview) };
|
||||
}
|
||||
|
||||
for (const theme of THEMES) for (const [width, height] of SIZES) {
|
||||
test(`tag output images review matches the concept ${theme} ${width}x${height}`, async ({ page }, testInfo) => {
|
||||
await page.setViewportSize({ width, height });
|
||||
const concept = await conceptReview(page, theme);
|
||||
expect(concept.tree).toEqual(['p.demo-note', 'h3', 'dl.meta-grid', 'p.small', 'div.dialog-actions']);
|
||||
await page.screenshot({ path: testInfo.outputPath(`concept-plan-review-${theme}-${width}.png`) });
|
||||
const { plan, production } = await productionReview(page, theme);
|
||||
// Under Merge every prompt with a linked image is tagged again: the batch has work to review.
|
||||
expect(plan.to_tag).toBeGreaterThan(0);
|
||||
await page.screenshot({ path: testInfo.outputPath(`production-tagging-review-${theme}-${width}.png`) });
|
||||
expect(diffInventory(mapBatchConcept(concept, plan), production)).toEqual([]);
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* D134 (CUJ 5): Replace after review never runs as a batch: the same review shell with the counts, the reason as its
|
||||
* `p.small`, no backup box and only Cancel (the shared close-only review Integrations also uses).
|
||||
*/
|
||||
function mapCancelOnlyConcept(reference, plan) {
|
||||
const [cancel] = reference.actions.buttons;
|
||||
const mapped = mapBatchConcept(reference, plan);
|
||||
return {
|
||||
...mapped,
|
||||
tree: ['h3', 'dl.meta-grid', 'p.small', 'div.dialog-actions'],
|
||||
grid: { ...mapped.grid, rows: mapped.grid.rows.map(([term, value]) => [term, term === 'Existing tags' ? 'Replace after review' : value]) },
|
||||
backup: null,
|
||||
note: { ...reference.note, text: 'Replacing tags needs a review of each image. Use Preview tagging review, or choose Skip or '
|
||||
+ 'Merge for a batch run.' },
|
||||
actions: { ...reference.actions, tree: [reference.actions.tree[0]], buttons: [cancel] },
|
||||
};
|
||||
}
|
||||
|
||||
test('the Cancel-only review matches the concept review shell', async ({ page }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
const concept = await conceptReview(page, 'dark');
|
||||
await productionTagging(page, 'dark');
|
||||
await page.locator('#workspace-main #tag-existing').selectOption('replace_after_review');
|
||||
const answer = page.waitForResponse(response => response.url().endsWith(PREVIEW_URL));
|
||||
await page.locator('#workspace-main .operation').getByRole('button', { name: 'Review', exact: true }).click();
|
||||
const plan = await (await answer).json();
|
||||
expect(plan.review_required).toBe(true);
|
||||
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Review tag output images');
|
||||
await page.mouse.move(0, 0);
|
||||
expect(diffInventory(mapCancelOnlyConcept(concept, plan), await page.evaluate(measureReview))).toEqual([]);
|
||||
});
|
||||
|
||||
test('tag output images review comparison detects mutations', async ({ page }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
const concept = await conceptReview(page, 'dark');
|
||||
const { plan } = await productionReview(page, 'dark');
|
||||
const expected = mapBatchConcept(concept, plan);
|
||||
const paths = async () => diffInventory(expected, await page.evaluate(measureReview)).map(item => item.path);
|
||||
expect(await paths()).toEqual([]);
|
||||
await page.evaluate(() => { document.querySelectorAll('dialog#modal[open] dl.meta-grid dt')[5].textContent = 'To tag'; });
|
||||
expect(await paths()).toEqual(['grid.rows.5.0']);
|
||||
await page.evaluate(() => { document.querySelectorAll('dialog#modal[open] dl.meta-grid dt')[5].textContent = 'Prompts to tag'; });
|
||||
const undoGrid = await mutate(page, { css: 'dialog#modal dl.meta-grid { column-gap: 3px; }' });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['grid.columnGap']));
|
||||
await undoGrid();
|
||||
const undoBox = await mutate(page, { remove: ['dialog#modal[open] #modal-body > .notice'] });
|
||||
expect(await paths()).toEqual(expect.arrayContaining(['backup', 'tree.length']));
|
||||
await undoBox();
|
||||
expect(await paths()).toEqual([]);
|
||||
});
|
||||
@@ -236,7 +236,7 @@ test('paired Tag dialog captures', async ({ page, request }, testInfo) => {
|
||||
await expect(await request.post('/__test__/tagging/models', { data: { unavailable: ['wd14-swinv2'] } })).toBeOK();
|
||||
await openTagDialog(page, 'production', 'dark');
|
||||
await generateWith(page, 'production').catch(() => {});
|
||||
await expect(status()).toHaveText('WD14 SwinV2 is not downloaded. Download it in System › Auto-tagging.');
|
||||
await expect(status()).toHaveText('WD14 SwinV2 is not downloaded. Download it from the legacy Auto-tag page.');
|
||||
await shot('production-unavailable');
|
||||
await expect(await request.post('/__test__/tagging/models', { data: { unavailable: [] } })).toBeOK();
|
||||
|
||||
|
||||
@@ -257,18 +257,6 @@ async function chooseRestore(page, file) {
|
||||
await page.locator('#restoreFileInput').setInputFiles(file);
|
||||
}
|
||||
|
||||
// Confirms the chosen restore and waits for the server's answer, not for a timer: a restore is disk-bound and can take
|
||||
// longer than an expect timeout on a slow runner (#359). The modal closes only after that answer.
|
||||
async function confirmRestoreAndWait(page) {
|
||||
const restored = page.waitForResponse(response => response.url().endsWith('/prompt_manager/restore'));
|
||||
await page.locator('#confirmRestore').click();
|
||||
const response = await restored;
|
||||
const body = await response.text();
|
||||
expect(response.status(), body).toBe(200);
|
||||
expect(JSON.parse(body)).toMatchObject({ success: true, prompt_count: 2, backup_created: expect.any(String) });
|
||||
await expect(page.locator('#restoreModal')).toBeHidden();
|
||||
}
|
||||
|
||||
function confirmNextDialog(page, accept, message) {
|
||||
page.once('dialog', async dialog => {
|
||||
expect(dialog.type()).toBe('confirm');
|
||||
@@ -299,7 +287,12 @@ test('downloaded backup restores deleted prompts, tags, and image links through
|
||||
expect(await promptSnapshot(request)).not.toEqual(before);
|
||||
await chooseRestore(page, backup);
|
||||
confirmNextDialog(page, true, backup.name);
|
||||
await confirmRestoreAndWait(page);
|
||||
const restored = page.waitForResponse(response => response.url().endsWith('/prompt_manager/restore'));
|
||||
await page.locator('#confirmRestore').click();
|
||||
const response = await restored;
|
||||
expect(response.ok()).toBe(true);
|
||||
expect(await response.json()).toMatchObject({ success: true, prompt_count: 2, backup_created: expect.any(String) });
|
||||
await expect(page.locator('#restoreModal')).toBeHidden();
|
||||
await expect(page.locator('.prompt-text')).toHaveCount(2);
|
||||
await page.reload();
|
||||
await expect(page.locator('.prompt-text').filter({ hasText: 'Cat portrait' })).toBeVisible();
|
||||
@@ -342,7 +335,8 @@ test('restore validation and cancellation preserve live data and allow a valid r
|
||||
expect(await promptSnapshot(request)).toEqual(before);
|
||||
await page.locator('#restoreFileInput').setInputFiles(backup);
|
||||
confirmNextDialog(page, true, backup.name);
|
||||
await confirmRestoreAndWait(page);
|
||||
await page.locator('#confirmRestore').click();
|
||||
await expect(page.locator('#restoreModal')).toBeHidden();
|
||||
expect(await promptSnapshot(request)).toEqual(before);
|
||||
});
|
||||
|
||||
@@ -362,7 +356,8 @@ test('backup and restore recover from interrupted network requests without losin
|
||||
await expect(page.locator('#restoreModal')).toBeVisible();
|
||||
expect(await promptSnapshot(request)).toEqual(before);
|
||||
confirmNextDialog(page, true, backup.name);
|
||||
await confirmRestoreAndWait(page);
|
||||
await page.locator('#confirmRestore').click();
|
||||
await expect(page.locator('#restoreModal')).toBeHidden();
|
||||
expect(await promptSnapshot(request)).toEqual(before);
|
||||
});
|
||||
|
||||
|
||||
@@ -1088,36 +1088,6 @@ test('a held legacy job and the health check show no tiles, no Results and no ma
|
||||
await expect(toast(page)).toHaveText('Health check finished.');
|
||||
});
|
||||
|
||||
// #328: a tagging model download is a job of its own kind with no byte progress and no Cancel; Activity shows it
|
||||
// through the generic held view (its title, status line and an indeterminate bar), then its end toast.
|
||||
test('a held model download shows its title, an indeterminate bar and no Cancel; its end toast follows', async ({ page, request }) => {
|
||||
await page.setViewportSize({ width: 1600, height: 1100 });
|
||||
const title = 'Tagging model download · WD14 ViT';
|
||||
await expect(await request.post('/__test__/tagging/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/tagging/download', { data: { mode: 'hold' } })).toBeOK();
|
||||
try {
|
||||
const response = await request.post('/prompt_manager/jobs/model-download', { data: { model_type: 'wd14-vit' } });
|
||||
expect(response.status()).toBe(202);
|
||||
const job = (await response.json()).job;
|
||||
expect([job.kind, job.title, job.cancellable]).toEqual(['model_download', title, false]);
|
||||
await page.goto(`${workspace}#/activity`);
|
||||
const item = items(page).filter({ has: page.getByRole('link', { name: new RegExp(`${title}$`) }) });
|
||||
await expect(item.locator('progress')).not.toHaveAttribute('value');
|
||||
await expect(item.getByRole('button', { name: 'Cancel job', exact: true })).toHaveCount(0);
|
||||
await item.getByRole('link').click();
|
||||
await expect(heading(page).locator('h1')).toHaveText(title);
|
||||
await expect(progressSection(page).locator('.panel-heading > span.small')).toHaveText('Downloading WD14 ViT (~400 MB)');
|
||||
await expect(progressSection(page).locator('progress')).not.toHaveAttribute('value');
|
||||
await expect(progressSection(page).locator('dl.stat-row')).toHaveCount(0);
|
||||
await expect(heading(page).getByRole('button', { name: /^Cancel/ })).toHaveCount(0);
|
||||
await expect(await request.post('/__test__/tagging/download', { data: { mode: 'succeed' } })).toBeOK();
|
||||
await expect(toast(page)).toHaveText(`${title} finished.`);
|
||||
await expect(heading(page).locator('p.muted')).toHaveText(/^Finished \d\d:\d\d$/);
|
||||
} finally {
|
||||
await request.post('/__test__/tagging/reset');
|
||||
}
|
||||
});
|
||||
|
||||
// S5 review I3: a thumbnail job with nothing to generate ends on the server's final line ("No previews to generate · N
|
||||
// already up to date"), never on a stale "Processing file 0 of 0" under the Completed chip.
|
||||
test('a finished thumbnail job with nothing to do says so', async ({ page, request }) => {
|
||||
|
||||
@@ -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, preview: false, previous: 2 });
|
||||
await lora(request, { loras: 3, civitai: true, previous: 2 });
|
||||
const previews = [];
|
||||
const applies = [];
|
||||
page.on('request', item => {
|
||||
@@ -339,17 +339,13 @@ 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 LoRAs · 0 with local images'],
|
||||
['Found', '3 metadata files · 3 readable'],
|
||||
['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');
|
||||
@@ -386,7 +382,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))[3]).toEqual(['Protected', '1 imported prompt used by saved pairs']);
|
||||
await expect.poll(async () => (await gridOf(page))[2]).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']);
|
||||
@@ -394,51 +390,6 @@ 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);
|
||||
@@ -457,7 +408,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, preview: false, previous: 1 });
|
||||
await lora(request, { loras: 2, civitai: true, previous: 1 });
|
||||
await openIntegrations(page);
|
||||
await enable(page);
|
||||
const before = await serverJson(request, '/__test__/integrations/state');
|
||||
|
||||
@@ -429,9 +429,7 @@ test('Rebuild all is explicit and the plan line says every preview is regenerate
|
||||
await expect(page).toHaveURL(/#\/activity\/[0-9a-f]{32}$/);
|
||||
expect(seen.starts).toEqual([{ mode: 'rebuild' }]);
|
||||
const id = page.url().split('/').pop();
|
||||
// Rebuild regenerates all 24 previews with real image work: on a busy CI
|
||||
// runner that took over 5 s (12 of 24 done at the default poll budget).
|
||||
await expect.poll(async () => (await serverJson(request, `/prompt_manager/jobs/${id}`)).job.state, { timeout: 30000 }).toBe('finished');
|
||||
await expect.poll(async () => (await serverJson(request, `/prompt_manager/jobs/${id}`)).job.state).toBe('finished');
|
||||
});
|
||||
|
||||
test('a plan that cannot be read says so; Start still starts', async ({ page }) => {
|
||||
|
||||
@@ -1,385 +0,0 @@
|
||||
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
|
||||
import { test, expect } from './fixtures.mjs';
|
||||
|
||||
// System › Auto-tagging (#328, System S11 Tasks 6–7) on the tagging fixture (workspace_tagging_fixture.py) with the
|
||||
// injected fake model: no model file exists and nothing is downloaded or run for real. The fake download waits on the
|
||||
// fixture (`/__test__/tagging/download {mode}`), and every reset releases it. Every test asserts that no legacy
|
||||
// download stream (`/autotag/download/`) was requested: the section downloads only through the Activity job.
|
||||
const workspace = '/prompt_manager/workspace';
|
||||
const MODELS_URL = '/prompt_manager/autotag/models';
|
||||
const PREVIEW_URL = '/prompt_manager/autotag/batch/preview';
|
||||
const SETTINGS_URL = '/prompt_manager/settings';
|
||||
const MODEL_LABELS = ['WD14 SwinV2 · ~400 MB', 'WD14 ViT · ~400 MB', 'GGUF JoyCaption · ~4.5 GB', 'HuggingFace JoyCaption · ~16 GB'];
|
||||
const FIRST = 'tagging/dunes_00001.png';
|
||||
const SECOND = 'tagging/dunes_00002.png';
|
||||
|
||||
test.beforeEach(async ({ page, request, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
await expect(await request.post('/__test__/tagging/reset', { data: {} })).toBeOK();
|
||||
await expect(await request.post('/__test__/system/reset')).toBeOK();
|
||||
await expect(await request.post('/__test__/jobs/restart')).toBeOK();
|
||||
page.systemErrors = [];
|
||||
page.on('pageerror', error => page.systemErrors.push(error.message));
|
||||
page.downloads = [];
|
||||
page.modelCalls = [];
|
||||
page.on('request', item => {
|
||||
const { pathname } = new URL(item.url());
|
||||
if (pathname.includes('/autotag/download/')) page.downloads.push(pathname);
|
||||
if (pathname.startsWith('/prompt_manager/autotag/') || pathname.startsWith('/prompt_manager/image-annotations/suggest')
|
||||
|| pathname.startsWith('/prompt_manager/jobs/model-download')) page.modelCalls.push(`${item.method()} ${pathname}`);
|
||||
});
|
||||
});
|
||||
|
||||
test.afterEach(async ({ page, request, browserName }) => {
|
||||
await saveCoverage(page, browserName);
|
||||
// The reset releases a held fake download, so no runner outlives the test.
|
||||
await request.post('/__test__/tagging/reset', { data: {} });
|
||||
await request.post('/__test__/jobs/restart');
|
||||
await request.post('/__test__/system/reset');
|
||||
expect(page.systemErrors).toEqual([]);
|
||||
expect(page.downloads).toEqual([]);
|
||||
});
|
||||
|
||||
const toast = page => page.locator('#workspace-toast');
|
||||
const layoutOf = page => page.locator('#workspace-main .settings-layout');
|
||||
const dialog = page => page.locator('dialog#modal[open]');
|
||||
const status = page => page.locator('#settings-status');
|
||||
const save = page => page.getByRole('button', { name: 'Save preferences', exact: true });
|
||||
const surface = page => layoutOf(page).locator(':scope > section.surface');
|
||||
const settingOf = (page, label) => layoutOf(page).locator('.setting-row').filter({ has: page.locator(`label:text-is("${label}")`) });
|
||||
const modelRow = page => settingOf(page, 'Tagging model');
|
||||
const actionOf = (page, name) => surface(page).locator(':scope > .row-actions').getByRole('button', { name, exact: true });
|
||||
const reviewButton = page => surface(page).locator('.operation').getByRole('button', { name: 'Review', exact: true });
|
||||
const card = (page, path) => page.locator(`article[data-image-key="0:${path}"]`);
|
||||
const tagDialog = page => page.getByRole('dialog', { name: 'Tag image', exact: true });
|
||||
const gridOf = page => dialog(page).locator('dl.meta-grid > div').evaluateAll(items => items.map(item =>
|
||||
[item.querySelector('dt').textContent, item.querySelector('dd').textContent]));
|
||||
|
||||
async function serverJson(request, url) {
|
||||
const response = await request.get(url);
|
||||
await expect(response).toBeOK();
|
||||
return response.json();
|
||||
}
|
||||
|
||||
async function fakeDownload(request, mode) {
|
||||
await expect(await request.post('/__test__/tagging/download', { data: { mode } })).toBeOK();
|
||||
}
|
||||
|
||||
async function unavailable(request, models) {
|
||||
await expect(await request.post('/__test__/tagging/models', { data: { unavailable: models } })).toBeOK();
|
||||
}
|
||||
|
||||
async function ended(request, id) {
|
||||
await expect.poll(async () => (await serverJson(request, `/prompt_manager/jobs/${id}`)).job.state).not.toBe('running');
|
||||
return (await serverJson(request, `/prompt_manager/jobs/${id}`)).job;
|
||||
}
|
||||
|
||||
/** System › Auto-tagging on a fresh page (the drafts live for the page session, Q10), settled on the model status. */
|
||||
async function openTagging(page) {
|
||||
if (page.url().endsWith('#/system/tagging')) await page.reload();
|
||||
else await page.goto(`${workspace}#/system/tagging`);
|
||||
await expect(surface(page)).toHaveCount(1);
|
||||
await expect(page.locator('#tag-model')).toBeEnabled();
|
||||
await expect(modelRow(page).locator('p').first()).not.toHaveText('Loading model status…');
|
||||
}
|
||||
|
||||
/** Select `path` in Images (the concept's `state.selectedImage`), then open System › Auto-tagging from the sidebar. */
|
||||
async function selectThenOpen(page, path) {
|
||||
await page.goto(`${workspace}#/images`);
|
||||
await card(page, path).getByRole('button', { name: `Inspect ${path.split('/').pop()}`, exact: true }).click();
|
||||
await expect(card(page, path)).toHaveClass(/\bselected\b/);
|
||||
await page.locator('aside.sidebar a[href="#/system"]').click();
|
||||
await page.locator('#workspace-main .section-nav').getByRole('button', { name: 'Auto-tagging', exact: true }).click();
|
||||
await expect(page.locator('#tag-model')).toBeEnabled();
|
||||
await expect(modelRow(page).locator('p').first()).not.toHaveText('Loading model status…');
|
||||
}
|
||||
|
||||
test('model options are the four real engines with size and the description shows live status', async ({ page, request }) => {
|
||||
await unavailable(request, ['wd14-vit', 'hf']);
|
||||
await openTagging(page);
|
||||
expect(await surface(page).locator(':scope > h2').textContent()).toBe('Auto-tagging');
|
||||
expect(await page.locator('#tag-model option').allTextContents()).toEqual(MODEL_LABELS);
|
||||
await expect(page.locator('#tag-model')).toHaveValue('wd14-swinv2');
|
||||
const line = modelRow(page).locator('p').first();
|
||||
await expect(line).toHaveText('Downloaded. Not loaded.');
|
||||
await expect(actionOf(page, 'Download again')).toBeVisible();
|
||||
// Choosing another engine drafts it and tells its own status; nothing is downloaded or loaded.
|
||||
await page.locator('#tag-model').selectOption('wd14-vit');
|
||||
await expect(line).toHaveText('Not downloaded. Download size ~400 MB.');
|
||||
await expect(actionOf(page, 'Download')).toBeVisible();
|
||||
await page.locator('#tag-model').selectOption('hf');
|
||||
await expect(line).toHaveText('Not downloaded. Download size ~16 GB.');
|
||||
await expect(status(page)).toHaveText('Unsaved changes.');
|
||||
// The other rows keep the concept's copy and the saved values (D133 adds Suggestions per image).
|
||||
expect(await surface(page).locator(':scope > .setting-row label').allTextContents()).toEqual(['Tagging model', 'Existing tags',
|
||||
'General-tag confidence', 'Character-tag confidence', 'Suggestions per image']);
|
||||
await expect(page.locator('#tag-existing')).toHaveValue('skip');
|
||||
await expect(page.locator('#threshold-general')).toHaveValue('0.35');
|
||||
await expect(page.locator('#threshold-character')).toHaveValue('0.85');
|
||||
await expect(page.locator('#tag-suggestions')).toHaveValue('10');
|
||||
const settings = await serverJson(request, SETTINGS_URL);
|
||||
await expect(page.locator('#tag-template')).toHaveValue(settings.settings.autotag_template);
|
||||
// Unload is offered only while a model is loaded (none is).
|
||||
await expect(actionOf(page, 'Unload')).toBeDisabled();
|
||||
});
|
||||
|
||||
test('entering the section makes no model call besides status', async ({ page }) => {
|
||||
await openTagging(page);
|
||||
await page.locator('#workspace-main .section-nav').getByRole('button', { name: 'Integrations', exact: true }).click();
|
||||
await page.locator('#workspace-main .section-nav').getByRole('button', { name: 'Auto-tagging', exact: true }).click();
|
||||
await expect(page.locator('#tag-model')).toBeEnabled();
|
||||
await expect(modelRow(page).locator('p').first()).toHaveText('Downloaded. Not loaded.');
|
||||
// One status read per mount (two mounts), and nothing else that could load, run or download a model.
|
||||
expect(page.modelCalls).toEqual([`GET ${MODELS_URL}`, `GET ${MODELS_URL}`]);
|
||||
const models = await (await page.request.get(MODELS_URL)).json();
|
||||
expect(models.model_loaded).toBe(false);
|
||||
});
|
||||
|
||||
test('defaults save through Save preferences and the Tag image dialog uses them', async ({ page, request }) => {
|
||||
await openTagging(page);
|
||||
await page.locator('#tag-existing').selectOption('merge');
|
||||
await page.locator('#threshold-general').fill('0.4');
|
||||
await page.locator('#threshold-general').press('Tab');
|
||||
await page.locator('#tag-suggestions').fill('15');
|
||||
await page.locator('#tag-suggestions').press('Tab');
|
||||
const posted = page.waitForRequest(item => item.method() === 'POST' && new URL(item.url()).pathname === SETTINGS_URL);
|
||||
await save(page).click();
|
||||
const body = (await posted).postDataJSON();
|
||||
expect(body).toMatchObject({ autotag_existing: 'merge', autotag_general_threshold: 0.4, autotag_suggestion_limit: 15 });
|
||||
expect(typeof body.expected_revision).toBe('string');
|
||||
await expect(status(page)).toHaveText('Saved to the server configuration.');
|
||||
const saved = (await serverJson(request, SETTINGS_URL)).settings;
|
||||
expect([saved.autotag_existing, saved.autotag_general_threshold, saved.autotag_character_threshold, saved.autotag_suggestion_limit])
|
||||
.toEqual(['merge', 0.4, 0.85, 15]);
|
||||
const models = await serverJson(request, MODELS_URL);
|
||||
expect([models.existing, models.wd14_general_threshold, models.suggestion_limit]).toEqual(['merge', 0.4, 15]);
|
||||
// The Tag image dialog in a fresh page starts from them.
|
||||
await page.goto(`${workspace}#/images`);
|
||||
await card(page, FIRST).getByRole('button', { name: `Tag or auto-tag ${FIRST.split('/').pop()}`, exact: true }).click();
|
||||
await expect(tagDialog(page)).toBeVisible();
|
||||
await expect(tagDialog(page).locator('#image-general-threshold')).toHaveValue('0.4');
|
||||
await expect(tagDialog(page).locator('#image-character-threshold')).toHaveValue('0.85');
|
||||
await expect(tagDialog(page).locator('#image-suggestion-limit')).toHaveValue('15');
|
||||
});
|
||||
|
||||
test('a refused value shows the server\'s message', async ({ page, request }) => {
|
||||
const before = (await serverJson(request, SETTINGS_URL)).settings;
|
||||
// An image is selected, so Preview tagging review is held back only by the invalid value below.
|
||||
await selectThenOpen(page, FIRST);
|
||||
await page.locator('#threshold-general').fill('1.5');
|
||||
await page.locator('#threshold-general').press('Tab');
|
||||
await save(page).click();
|
||||
await expect(status(page)).toHaveText('General-tag confidence must be between 0 and 1.');
|
||||
await expect(page.locator('#threshold-general')).toHaveAttribute('aria-invalid', 'true');
|
||||
// Nothing was written; the draft is kept for the next Save.
|
||||
expect((await serverJson(request, SETTINGS_URL)).settings.autotag_general_threshold).toBe(before.autotag_general_threshold);
|
||||
await expect(page.locator('#threshold-general')).toHaveValue('1.5');
|
||||
// While the value is invalid, Preview tagging review is aria-disabled and explains instead of opening (Playwright
|
||||
// waits on an aria-disabled control, so it is pressed from the keyboard, as a user can).
|
||||
await expect(actionOf(page, 'Preview tagging review')).toHaveAttribute('aria-disabled', 'true');
|
||||
await actionOf(page, 'Preview tagging review').focus();
|
||||
await page.keyboard.press('Enter');
|
||||
await expect(toast(page)).toHaveText('General-tag confidence must be between 0 and 1.');
|
||||
await expect(tagDialog(page)).toHaveCount(0);
|
||||
await page.locator('#threshold-general').fill('0.5');
|
||||
await page.locator('#threshold-general').press('Tab');
|
||||
await expect(page.locator('#threshold-general')).not.toHaveAttribute('aria-invalid', 'true');
|
||||
await expect(status(page)).toHaveText('Unsaved changes.');
|
||||
await expect(actionOf(page, 'Preview tagging review')).not.toHaveAttribute('aria-disabled', 'true');
|
||||
});
|
||||
|
||||
test('a value off the 0.05 step is refused and never reaches the Tag image dialog (#369)', async ({ page, request }) => {
|
||||
const before = (await serverJson(request, SETTINGS_URL)).settings;
|
||||
const STEP = 'General-tag confidence must be a multiple of 0.05.';
|
||||
await selectThenOpen(page, FIRST);
|
||||
await page.locator('#threshold-general').fill('0.37');
|
||||
await page.locator('#threshold-general').press('Tab');
|
||||
// Preview tagging review explains instead of opening a dialog whose Generate the browser would block.
|
||||
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(STEP);
|
||||
await expect(tagDialog(page)).toHaveCount(0);
|
||||
// The server refuses the save with the same sentence and writes nothing.
|
||||
await save(page).click();
|
||||
await expect(status(page)).toHaveText(STEP);
|
||||
expect((await serverJson(request, SETTINGS_URL)).settings.autotag_general_threshold).toBe(before.autotag_general_threshold);
|
||||
// A value on the step opens the dialog, and its Generate passes the browser's own check and sends.
|
||||
await page.locator('#threshold-general').fill('0.4');
|
||||
await page.locator('#threshold-general').press('Tab');
|
||||
await expect(actionOf(page, 'Preview tagging review')).not.toHaveAttribute('aria-disabled', 'true');
|
||||
await actionOf(page, 'Preview tagging review').click();
|
||||
await expect(tagDialog(page).locator('#image-general-threshold')).toHaveValue('0.4');
|
||||
expect(await tagDialog(page).locator('#image-general-threshold').evaluate(input => input.checkValidity())).toBe(true);
|
||||
const suggested = page.waitForResponse(item => item.url().endsWith('/image-annotations/suggest'));
|
||||
await tagDialog(page).getByRole('button', { name: 'Generate suggestions', exact: true }).click();
|
||||
const answer = await suggested;
|
||||
expect(answer.request().postDataJSON().options.general_threshold).toBe(0.4);
|
||||
});
|
||||
|
||||
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([]);
|
||||
});
|
||||
@@ -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);
|
||||
// Diagnostics still hands off (S12), so its body is one known panel.
|
||||
await nav.getByRole('button', { name: 'Diagnostics & logs', exact: true }).click();
|
||||
// Auto-tagging still hands off (S11), so its body is one known panel.
|
||||
await nav.getByRole('button', { name: 'Auto-tagging', exact: true }).click();
|
||||
await expect(await request.post('/__test__/release')).toBeOK();
|
||||
await expect.poll(held).toBe(0);
|
||||
// Both late answers reached a section that had left: Diagnostics keeps its own body.
|
||||
// Both late answers reached a section that had left: Auto-tagging 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 in System › Auto-tagging.';
|
||||
const NOT_DOWNLOADED = 'WD14 SwinV2 is not downloaded. Download it from the legacy Auto-tag page.';
|
||||
|
||||
test.beforeEach(async ({ page, request, browserName }) => {
|
||||
await startCoverage(page, browserName);
|
||||
|
||||
@@ -246,7 +246,7 @@ test('unavailable model shows inline status and never downloads', async ({ page,
|
||||
await expect(await request.post('/__test__/tagging/models', { data: { unavailable: ['wd14-swinv2'] } })).toBeOK();
|
||||
await openTag(page);
|
||||
expect((await generated(page)).status()).toBe(409);
|
||||
await expect(status(page)).toHaveText('WD14 SwinV2 is not downloaded. Download it in System › Auto-tagging.');
|
||||
await expect(status(page)).toHaveText('WD14 SwinV2 is not downloaded. Download it from the legacy Auto-tag page.');
|
||||
await expect(chips(page)).toHaveCount(0);
|
||||
await expect(dialog(page).locator('#suggestion-count')).toHaveText('Not generated');
|
||||
// D116: Retry leads the actions while the controller allows it, and generates again.
|
||||
|
||||
@@ -413,7 +413,7 @@ for (const width of [1600, 390]) test(`overview and handoff render with concept
|
||||
|
||||
await sidebarLink(page, 'System').click();
|
||||
// System's Overview is native (S2 Task 6); a section not delivered yet still hands off (D94).
|
||||
await main.locator('.section-nav').getByRole('button', { name: 'Diagnostics & logs', exact: true }).click();
|
||||
await main.locator('.section-nav').getByRole('button', { name: 'Auto-tagging', exact: true }).click();
|
||||
const handoff = main.locator('section.handoff');
|
||||
await expect(handoff).toHaveClass(/\bsurface\b/);
|
||||
expect(await handoff.evaluate(surfaceStyle)).toMatchObject(SURFACE);
|
||||
@@ -465,9 +465,8 @@ 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), 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();
|
||||
// workspace-system-maintenance.spec.mjs) and Backup & export (S9 Task 24, workspace-system-backups.spec.mjs);
|
||||
// every other section hands off.
|
||||
if (hash === 'overview') await expect(layout.locator(':scope > .summary-metrics')).toBeVisible();
|
||||
if (hash === 'sources') await expect(layout.locator('#source-list .operation')).toHaveCount(2);
|
||||
expect(await layout.evaluate(node => [...node.children].map(child => `${child.localName}.${child.className}`)))
|
||||
@@ -475,10 +474,9 @@ 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'],
|
||||
tagging: ['section.surface', 'div.save-row'] }[hash]
|
||||
backups: ['div.two-col', 'section.surface', 'div.save-row'] }[hash]
|
||||
?? ['section.surface handoff', 'div.save-row']);
|
||||
if (['overview', 'appearance', 'sources', 'integrations', 'tagging', 'maintenance', 'backups'].includes(hash)) continue;
|
||||
if (['overview', 'appearance', 'sources', 'integrations', '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,18 +7,13 @@ 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, 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/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/gate`` ``{hold}`` holds every CivitAI
|
||||
download until it is released (``{"hold": false}``), so a test sees the job
|
||||
while it downloads;
|
||||
@@ -35,7 +30,6 @@ test ever opens a network connection.
|
||||
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import threading
|
||||
from pathlib import Path
|
||||
@@ -97,11 +91,7 @@ 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
|
||||
@@ -128,17 +118,11 @@ 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)
|
||||
@@ -157,10 +141,8 @@ 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))
|
||||
if preview:
|
||||
(loras / f"{stem}.png").write_bytes(_png())
|
||||
(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(
|
||||
|
||||
@@ -97,25 +97,15 @@ 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
|
||||
from py.api.settings_routes import (
|
||||
_AUTOTAG_SETTINGS,
|
||||
_GALLERY_SETTINGS,
|
||||
_SETTINGS_FILE_KEYS,
|
||||
host_values = (
|
||||
PromptManagerConfig.RESULT_TIMEOUT,
|
||||
PromptManagerConfig.WEBUI_DISPLAY_MODE,
|
||||
PromptManagerConfig.WORKER_THREADS,
|
||||
GalleryConfig.MONITORING_ENABLED,
|
||||
PromptManagerConfig.THUMBNAIL_FORMAT,
|
||||
PromptManagerConfig.THUMBNAIL_QUALITY,
|
||||
PromptManagerConfig.THUMBNAIL_EDGE,
|
||||
)
|
||||
|
||||
# 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)
|
||||
@@ -145,8 +135,15 @@ def register_system_commands(app, routes, api, root: Path):
|
||||
unplugged.rmdir()
|
||||
GalleryConfig.MONITORING_DIRECTORIES = list(configured)
|
||||
api._health_report = None
|
||||
for (owner, name), value in host_values.items():
|
||||
setattr(owner, name, value)
|
||||
(
|
||||
PromptManagerConfig.RESULT_TIMEOUT,
|
||||
PromptManagerConfig.WEBUI_DISPLAY_MODE,
|
||||
PromptManagerConfig.WORKER_THREADS,
|
||||
GalleryConfig.MONITORING_ENABLED,
|
||||
PromptManagerConfig.THUMBNAIL_FORMAT,
|
||||
PromptManagerConfig.THUMBNAIL_QUALITY,
|
||||
PromptManagerConfig.THUMBNAIL_EDGE,
|
||||
) = host_values
|
||||
thumbnail_encode.reset_encoders()
|
||||
image_monitor._monitor_instance = monitor
|
||||
api.__dict__.pop("_library_resolver", None)
|
||||
|
||||
@@ -12,7 +12,6 @@ import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
import threading
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
@@ -182,28 +181,6 @@ def register_tagging_commands(
|
||||
routes: web.RouteTableDef, api: PromptManagerAPI, root: Path, model
|
||||
) -> None:
|
||||
"""Register fixed seed, snapshot and fake-model commands; no paths or SQL input."""
|
||||
# The fake download (``/__test__/tagging/download``): waits on ``event``,
|
||||
# then answers ``result``. Every reset sets the event and puts back the
|
||||
# model's raising default, so no test can leave a worker hanging.
|
||||
# ``generation`` grows on every reset: a download a reset released (or one
|
||||
# that timed out) fails without touching the next test's model state.
|
||||
download = {"event": threading.Event(), "result": True, "generation": 0}
|
||||
|
||||
def restore_download() -> None:
|
||||
download["generation"] += 1
|
||||
download["event"].set()
|
||||
model.__dict__.pop("download_model", None)
|
||||
|
||||
def fake_download(model_type: str, progress_callback=None) -> bool:
|
||||
event, generation = download["event"], download["generation"]
|
||||
released = event.wait(DOWNLOAD_HOLD_LIMIT_S)
|
||||
if not released or download["generation"] != generation:
|
||||
return False
|
||||
if download["result"]:
|
||||
model.unavailable = model.unavailable - {model_type}
|
||||
return download["result"]
|
||||
|
||||
_set_defaults(model, DEFAULTS) # the server starts on the standard ones
|
||||
|
||||
@routes.post("/__test__/tagging/reset")
|
||||
async def reset(request: web.Request) -> web.Response:
|
||||
@@ -218,7 +195,6 @@ def register_tagging_commands(
|
||||
model.unavailable = set()
|
||||
model.suggestions = None
|
||||
model.loaded = None
|
||||
restore_download()
|
||||
_set_defaults(model, DEFAULTS)
|
||||
# A test that ends with the dialog open leaves its capture behind; drop
|
||||
# every capture so the next test starts below the store's limits.
|
||||
@@ -247,33 +223,6 @@ def register_tagging_commands(
|
||||
model.unavailable = set(data["unavailable"])
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.post("/__test__/tagging/download")
|
||||
async def download_mode(request: web.Request) -> web.Response:
|
||||
"""Switch the model's download to a fake (never a real download).
|
||||
|
||||
``{"mode": "hold"}``: the next download waits until a later
|
||||
``succeed`` or ``fail`` (or a reset) releases it. ``succeed`` and
|
||||
``fail`` release a held download with that result and answer every
|
||||
later one at once. A successful fake marks the model downloaded.
|
||||
"""
|
||||
data = await request.json() if request.can_read_body else None
|
||||
if (
|
||||
request.query
|
||||
or not isinstance(data, dict)
|
||||
or set(data) != {"mode"}
|
||||
or data["mode"] not in ("hold", "succeed", "fail")
|
||||
):
|
||||
raise web.HTTPBadRequest()
|
||||
if data["mode"] == "hold":
|
||||
download["event"].set()
|
||||
download["event"] = threading.Event()
|
||||
download["result"] = True
|
||||
else:
|
||||
download["result"] = data["mode"] == "succeed"
|
||||
download["event"].set()
|
||||
model.download_model = fake_download
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.post("/__test__/tagging/suggestions")
|
||||
async def suggestions(request: web.Request) -> web.Response:
|
||||
"""Answer Generate with fixed scored names (the concept's own sample).
|
||||
@@ -324,28 +273,18 @@ def register_tagging_commands(
|
||||
return web.json_response({"captures": count})
|
||||
|
||||
|
||||
# A held fake download gives up after this long, whatever the test does.
|
||||
DOWNLOAD_HOLD_LIMIT_S = 60
|
||||
|
||||
# The standard tagging defaults: WD14 general and character thresholds, and
|
||||
# the suggestion count (``AutotagConfig``).
|
||||
DEFAULTS = (0.35, 0.85, 10)
|
||||
|
||||
|
||||
def _set_defaults(model, values) -> None:
|
||||
"""Set the saved tagging defaults (``AutotagConfig``, #328); a reset also
|
||||
puts back the saved model, batch policy and template (the fake model's
|
||||
short instructions, which the concept comparisons show)."""
|
||||
from py.config import AutotagConfig
|
||||
|
||||
general, character, limit = values
|
||||
AutotagConfig.GENERAL_THRESHOLD = general
|
||||
AutotagConfig.CHARACTER_THRESHOLD = character
|
||||
model.wd14_general_threshold = general
|
||||
model.wd14_character_threshold = character
|
||||
AutotagConfig.SUGGESTION_LIMIT = limit
|
||||
if values == DEFAULTS:
|
||||
AutotagConfig.MODEL_TYPE = "wd14-swinv2"
|
||||
AutotagConfig.EXISTING = "skip"
|
||||
AutotagConfig.TEMPLATE = model.default_prompt
|
||||
|
||||
|
||||
def _valid_defaults(data: object) -> bool:
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
"""Deterministic nodes and queue controls installed only in an owned test host."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
@@ -292,9 +291,7 @@ async def image_monitor(request: web.Request) -> web.Response:
|
||||
config.MONITORING_ENABLED = True
|
||||
monitor.start_monitoring([folder_paths.get_output_directory()])
|
||||
else:
|
||||
# Joins the image worker and the observer (one shared 5 s deadline),
|
||||
# so it runs on a worker thread, not on the event loop.
|
||||
await asyncio.to_thread(monitor.stop_monitoring)
|
||||
monitor.stop_monitoring()
|
||||
config.MONITORING_ENABLED = False
|
||||
return web.json_response({"monitoring": monitor.is_monitoring})
|
||||
|
||||
|
||||
@@ -388,16 +388,6 @@ 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);
|
||||
|
||||
@@ -225,26 +225,3 @@ 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);
|
||||
});
|
||||
|
||||
@@ -197,52 +197,6 @@ 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 });
|
||||
@@ -588,7 +542,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 in System › Auto-tagging.'],
|
||||
['model-unavailable', 'WD14 SwinV2 is not downloaded. Download it from the legacy Auto-tag page.'],
|
||||
['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 in System › Auto-tagging.');
|
||||
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({ gguf: { downloaded: false, name: 'GGUF (Recommended)' } }, 'gguf'),
|
||||
'GGUF (Recommended) is not downloaded. Download it in System › Auto-tagging.');
|
||||
'GGUF (Recommended) is not downloaded. Download it from the legacy Auto-tag page.');
|
||||
assert.throws(() => modelUnavailable({}, 'nope'), TypeError);
|
||||
});
|
||||
|
||||
|
||||
@@ -108,17 +108,13 @@ 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: 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 },
|
||||
const plan = { lora_files: 3, with_metadata: 2, current_imports: 5, blocked_by_pairs: 0, planned_downloads: 4,
|
||||
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', '2,221 metadata files · 615 LoRAs · 476 with local images'],
|
||||
['Found', '3 metadata files · 2 readable'],
|
||||
['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>)'],
|
||||
@@ -126,34 +122,11 @@ 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[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(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(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 () => {
|
||||
|
||||
@@ -341,22 +341,6 @@ 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);
|
||||
|
||||
@@ -1,600 +0,0 @@
|
||||
// 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');
|
||||
// #369: every page shows the thresholds in steps of 0.05, so a value off that grid cannot be saved or previewed.
|
||||
assert.deepEqual(draftProblem({ ...HOST, autotag_general_threshold: 0.37 }),
|
||||
{ key: 'autotag_general_threshold', message: 'General-tag confidence must be a multiple of 0.05.' });
|
||||
assert.deepEqual(draftProblem({ ...HOST, autotag_character_threshold: 0.33 }),
|
||||
{ key: 'autotag_character_threshold', message: 'Character-tag confidence must be a multiple of 0.05.' });
|
||||
assert.equal(draftProblem({ ...HOST, autotag_general_threshold: 0.15, autotag_character_threshold: 0.95 }), null);
|
||||
assert.equal(draftProblem({ ...HOST, autotag_general_threshold: 0, autotag_character_threshold: 1 }), null);
|
||||
// A computed value a hair off the grid passes and is handed on snapped to it, as the server stores it.
|
||||
assert.equal(draftProblem({ ...HOST, autotag_general_threshold: 0.1 + 0.2 }), null);
|
||||
assert.equal(taggingPreset({ ...HOST, autotag_general_threshold: 0.1 + 0.2 }).options.general_threshold, 0.3);
|
||||
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);
|
||||
byId('threshold-general').value = '0.37'; 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.at(-1), 'General-tag confidence must be a multiple of 0.05.');
|
||||
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 } });
|
||||
// #369: a threshold a hair off the grid is sent snapped to it.
|
||||
assert.equal(batchBody({ ...HOST, autotag_character_threshold: 0.1 + 0.2 }).options.character_threshold, 0.3);
|
||||
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(); }
|
||||
});
|
||||
@@ -841,11 +841,6 @@ 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']]);
|
||||
|
||||
+10
-29
@@ -33,9 +33,7 @@ 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
|
||||
|
||||
@@ -43,9 +41,6 @@ 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):
|
||||
@@ -120,16 +115,6 @@ 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
|
||||
@@ -143,20 +128,16 @@ def main(argv=None):
|
||||
return USAGE_EXIT
|
||||
|
||||
faulthandler.dump_traceback_later(per_test_limit, exit=True, file=stream)
|
||||
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)
|
||||
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)
|
||||
if result.testsRun == 0 and not result.skipped:
|
||||
return NO_TESTS_EXIT
|
||||
return 0 if result.wasSuccessful() else 1
|
||||
|
||||
@@ -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 in System › Auto-tagging."
|
||||
UNAVAILABLE = "Download it from the legacy Auto-tag page."
|
||||
|
||||
SCORED = [("sandstone", 0.98), ("desert", 0.96)] + [
|
||||
(f"tag_{i:03d}", round(0.9 - i * 0.005, 4)) for i in range(150)
|
||||
|
||||
@@ -1145,10 +1145,8 @@ 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
|
||||
|
||||
+3
-1066
File diff suppressed because it is too large
Load Diff
+18
-1352
File diff suppressed because it is too large
Load Diff
@@ -1,12 +1,9 @@
|
||||
"""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:
|
||||
@@ -981,186 +978,6 @@ 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
|
||||
|
||||
+18
-83
@@ -70,16 +70,8 @@ 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),
|
||||
"LORA_MANAGER_SETTINGS_DIR": str(self.lm_settings),
|
||||
},
|
||||
os.environ, {"PROMPT_MANAGER_CONFIG_PATH": str(self.config_file)}
|
||||
)
|
||||
self._env_patch.start()
|
||||
|
||||
@@ -202,8 +194,6 @@ 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
|
||||
@@ -211,8 +201,7 @@ 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
|
||||
]
|
||||
(directory / f"{stem}.safetensors").write_bytes(b"model")
|
||||
path = directory / f"{stem}.metadata.json"
|
||||
path = directory / f"{stem}.safetensors.metadata.json"
|
||||
path.write_text(json.dumps(meta))
|
||||
return path
|
||||
|
||||
@@ -311,7 +300,9 @@ class TestLoraEnableRoute(LoraAPITestCase):
|
||||
_write_metadata(loras, "neon", words=["glow"])
|
||||
self.detect_mock.return_value = str(lm)
|
||||
|
||||
with patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]):
|
||||
with patch.object(
|
||||
lora_utils, "find_lora_directories", return_value=[str(loras)]
|
||||
):
|
||||
resp = await self._post_json(
|
||||
"/prompt_manager/lora/enable",
|
||||
{
|
||||
@@ -857,7 +848,9 @@ class TestLoraScanRoute(LoraAPITestCase):
|
||||
|
||||
cache_dir = self.tmp_root / "img_cache"
|
||||
with (
|
||||
patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]),
|
||||
patch.object(
|
||||
lora_utils, "find_lora_directories", return_value=[str(loras)]
|
||||
),
|
||||
patch.object(
|
||||
lora_utils, "get_lora_image_cache_dir", return_value=cache_dir
|
||||
),
|
||||
@@ -876,18 +869,6 @@ class TestLoraScanRoute(LoraAPITestCase):
|
||||
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
|
||||
self.assertEqual(resp.status, 400)
|
||||
|
||||
async def test_a_finished_scan_clears_the_layout_cache(self):
|
||||
# #362: the legacy scan replaces through the same method.
|
||||
self._enable()
|
||||
loras = self.tmp_root / "loras"
|
||||
loras.mkdir()
|
||||
_write_metadata(loras, "plain")
|
||||
(loras / "plain.png").write_bytes(_png_bytes())
|
||||
with patch("py.lora_utils.clear_layout_cache") as clear:
|
||||
resp, events, _calls = await self._scan(loras)
|
||||
self.assertEqual(events[-1]["type"], "complete", events)
|
||||
clear.assert_called_once_with()
|
||||
|
||||
async def test_scan_imports_loras_with_images_and_links_existing(self):
|
||||
self._enable()
|
||||
loras = self.tmp_root / "loras"
|
||||
@@ -899,10 +880,9 @@ class TestLoraScanRoute(LoraAPITestCase):
|
||||
images=["https://civitai.com/a.png", "https://evil.example/b.png"],
|
||||
prompt="neon street",
|
||||
)
|
||||
# No local image: its CivitAI images are downloaded (#355).
|
||||
(loras / "neon.png").write_bytes(_png_bytes())
|
||||
_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,
|
||||
@@ -930,56 +910,8 @@ 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"])), 1)
|
||||
self.assertEqual(len(self.api.db.get_prompt_images(imported[0]["id"])), 2)
|
||||
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()
|
||||
@@ -1017,7 +949,9 @@ class TestLoraTriggerWordRoutes(LoraAPITestCase):
|
||||
loras.mkdir()
|
||||
_write_metadata(loras, "neon", words=["glow", "bright"])
|
||||
|
||||
with patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]):
|
||||
with patch.object(
|
||||
lora_utils, "find_lora_directories", return_value=[str(loras)]
|
||||
):
|
||||
resp = await self.client.request(
|
||||
"GET", "/prompt_manager/lora/trigger-words?name=NEON"
|
||||
)
|
||||
@@ -1051,16 +985,17 @@ class TestLoraTriggerWordRoutes(LoraAPITestCase):
|
||||
loras.mkdir()
|
||||
_write_metadata(loras, "neon", words=["glow"])
|
||||
|
||||
with patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]):
|
||||
with patch.object(
|
||||
lora_utils, "find_lora_directories", 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 the model name and the sidecar name, which are the same
|
||||
# for LoraManager's alongside layout.
|
||||
self.assertEqual(data["loras_with_trigger_words"], 1)
|
||||
# keyed by both the file_name stem and the metadata-file stem
|
||||
self.assertEqual(data["loras_with_trigger_words"], 2)
|
||||
|
||||
async def test_refresh_error_is_500(self):
|
||||
IntegrationConfig.LORA_MANAGER_ENABLED = True
|
||||
|
||||
@@ -188,7 +188,6 @@ class TestWorkspaceImages(ImageAPITestCase):
|
||||
newest = self.api._workspace_order_cache[key][1]["newest"]
|
||||
favorites = self.api._favorite_listing_cache[str(self.output_dir)]
|
||||
record.walked_at -= gallery_listing.REFRESH_AFTER + 1
|
||||
record.last_started_at -= gallery_listing.REFRESH_AFTER + 1
|
||||
with patch.object(
|
||||
self.api,
|
||||
"_scan_gallery_files_sync",
|
||||
|
||||
@@ -19,15 +19,8 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from PIL import Image
|
||||
|
||||
# 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
|
||||
from py import autotag
|
||||
from py.autotag import MODELS, AutoTagService
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
@@ -261,19 +254,12 @@ class TestWd14Thresholding(EngineTestCase):
|
||||
self.assertEqual(list(feeds.keys()), ["input"])
|
||||
self.assertEqual(tuple(feeds["input"].shape), (1, 8, 8, 3))
|
||||
|
||||
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"])
|
||||
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)
|
||||
|
||||
|
||||
class TestGgufNormalisation(EngineTestCase):
|
||||
@@ -383,13 +369,10 @@ class TestTagSuggestions(EngineTestCase):
|
||||
|
||||
def test_thresholds_default_to_the_service_values(self):
|
||||
service = self._wd14()
|
||||
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
|
||||
)
|
||||
service.wd14_general_threshold = 0.95
|
||||
service.wd14_character_threshold = 0.95
|
||||
|
||||
suggestions = service.generate_tag_suggestions(str(self.image_path), limit=10)
|
||||
|
||||
self.assertEqual(suggestions, [("sandstone", 0.98), ("desert", 0.96)])
|
||||
|
||||
|
||||
@@ -79,21 +79,5 @@ 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()
|
||||
|
||||
+2
-125
@@ -413,21 +413,9 @@ 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):
|
||||
for attr, value in self._defaults.items():
|
||||
setattr(self.config_class, attr, value)
|
||||
self.config_class.SUGGESTION_LIMIT = 10
|
||||
PromptManagerConfig.update_config(self._orig)
|
||||
|
||||
def _load(self, value):
|
||||
@@ -458,23 +446,8 @@ class TestAutotagConfig(unittest.TestCase):
|
||||
self.assertEqual(self.config_class.suggestion_limit(), 10)
|
||||
self.assertEqual(
|
||||
PromptManagerConfig.get_config()["autotag"],
|
||||
{
|
||||
"defaults": {
|
||||
"model_type": "wd14-swinv2",
|
||||
"existing": "skip",
|
||||
"general_threshold": 0.35,
|
||||
"character_threshold": 0.85,
|
||||
"suggestion_limit": 10,
|
||||
}
|
||||
},
|
||||
{"defaults": {"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(
|
||||
@@ -493,102 +466,6 @@ 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],
|
||||
# #369: a saved threshold off the 0.05 step no page can show.
|
||||
"general_threshold": [1.5, -0.1, float("nan"), "0.4", True, 0.37],
|
||||
"character_threshold": [2, float("inf"), None, 0.33],
|
||||
# 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_saved_threshold_is_a_multiple_of_the_step(self):
|
||||
"""#369: every page shows the thresholds in steps of 0.05, so a saved
|
||||
one must be on that grid; it is stored exactly on it."""
|
||||
from py.tagging_defaults import validate_saved_threshold as validate
|
||||
|
||||
for value in (0, 0.05, 0.15, 0.35, 0.85, 1):
|
||||
with self.subTest(value=value):
|
||||
self.assertEqual(validate("x", value), value)
|
||||
self.assertEqual(validate("x", 0.1 + 0.2), 0.3)
|
||||
for value, message in (
|
||||
(0.37, "x must be a multiple of 0.05"),
|
||||
(0.351, "x must be a multiple of 0.05"),
|
||||
(0.35000000001, "x must be a multiple of 0.05"),
|
||||
(1.5, "x must be between 0 and 1"),
|
||||
("0.4", "x must be a finite number"),
|
||||
):
|
||||
with self.subTest(value=value):
|
||||
with self.assertRaises(ValueError) as caught:
|
||||
validate("x", value)
|
||||
self.assertEqual(str(caught.exception), message)
|
||||
|
||||
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):
|
||||
|
||||
@@ -1398,50 +1398,6 @@ 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()
|
||||
|
||||
|
||||
@@ -28,8 +28,6 @@ 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:
|
||||
@@ -211,9 +209,6 @@ 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)
|
||||
@@ -344,20 +339,16 @@ class HistoryHooksTests(unittest.TestCase):
|
||||
with patch.object(self.db, "write_execution_start", blocked):
|
||||
registration = self.register()
|
||||
self.dequeue(0)
|
||||
# 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))
|
||||
self.assertTrue(entered.wait(1))
|
||||
began = time.monotonic()
|
||||
self.queue.task_done(0, {}, status())
|
||||
self.dequeue(1)
|
||||
self.assertLess(time.monotonic() - began, CEILING_SECONDS)
|
||||
self.assertLess(time.monotonic() - began, 0.2)
|
||||
began = time.monotonic()
|
||||
self.assertFalse(registration.dispose())
|
||||
self.assertGreaterEqual(time.monotonic() - began, 1.8)
|
||||
self.assertLess(time.monotonic() - began, CEILING_SECONDS)
|
||||
self.assertLess(time.monotonic() - began, 2.3)
|
||||
release.set()
|
||||
# Drain first: the 2 s bound is not under test here (#370).
|
||||
self.drain(registration)
|
||||
self.assertTrue(registration.dispose(2))
|
||||
|
||||
def test_restore_disables_capture_and_clears_handles_before_fresh_epoch(self):
|
||||
|
||||
@@ -642,8 +642,6 @@ 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())
|
||||
|
||||
@@ -245,8 +245,6 @@ 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()), {})
|
||||
@@ -399,10 +397,6 @@ class OutputHooksTests(fixtures.unittest.TestCase):
|
||||
self.queue.task_done(1, history("rejected.png"), fixtures.status())
|
||||
registration.writer.invalidate()
|
||||
release.set()
|
||||
# Let the writer work through the 1,024 queued commands first
|
||||
# (#370's pattern): on a busy Windows runner that took longer
|
||||
# than dispose's 2 s bound (merge queue, #382).
|
||||
self.drain(registration)
|
||||
self.assertTrue(registration.dispose(2))
|
||||
self.assertEqual(self.outputs(0)["state"], "not_captured")
|
||||
self.assertEqual(len(self.rows()), 1)
|
||||
|
||||
@@ -17,7 +17,6 @@ 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
|
||||
@@ -28,11 +27,6 @@ except ImportError:
|
||||
# Upper bound a blocked fake write waits for its release. Tests prove ordering
|
||||
# with events, so this only caps a failing run and never decides a pass.
|
||||
BLOCK_SECONDS = 10
|
||||
# A restore is a full verify/migrate/backup/replace cycle and is bound by the
|
||||
# disk: a busy Windows runner took over 10 s (merge queue, #382). Waits that
|
||||
# span a restore use this ceiling, which stays under the runner's 120 s hang
|
||||
# watchdog, so a real hang still fails with the thread dump.
|
||||
RESTORE_SECONDS = 90
|
||||
|
||||
|
||||
class OutputWriterTests(fixtures.unittest.TestCase):
|
||||
@@ -167,8 +161,6 @@ 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"
|
||||
@@ -322,8 +314,8 @@ class OutputWriterTests(fixtures.unittest.TestCase):
|
||||
|
||||
Returns the patch and an event set once that rotation has returned on
|
||||
the worker. A restore is a full verify/migrate/backup/replace cycle, so
|
||||
the handoff waits on this event with the ``RESTORE_SECONDS`` ceiling
|
||||
instead of the short worker-behavior deadline (TESTING.md,
|
||||
the handoff waits on this event with the generous ``BLOCK_SECONDS``
|
||||
ceiling instead of the short worker-behavior deadline (TESTING.md,
|
||||
Windows portability preflight).
|
||||
"""
|
||||
backup = os.path.join(self.directory.name, f"restore-{start.run_id}.db")
|
||||
@@ -366,7 +358,7 @@ class OutputWriterTests(fixtures.unittest.TestCase):
|
||||
)
|
||||
with rotation:
|
||||
self.finalization(writer, start)
|
||||
self.assertTrue(rotated.wait(RESTORE_SECONDS))
|
||||
self.assertTrue(rotated.wait(BLOCK_SECONDS))
|
||||
self.wait_for(lambda: writer.status().pending == 0)
|
||||
self.assertEqual(
|
||||
self.db.get_execution_outputs(start.run_id)["state"],
|
||||
@@ -408,7 +400,7 @@ class OutputWriterTests(fixtures.unittest.TestCase):
|
||||
self.assertEqual(writer.status().pending, 1)
|
||||
self.assertEqual(writer.context_budget.charged_bytes, charge)
|
||||
release.set()
|
||||
self.assertTrue(rotated.wait(RESTORE_SECONDS))
|
||||
self.assertTrue(rotated.wait(BLOCK_SECONDS))
|
||||
self.wait_for(lambda: writer.status().pending == 0)
|
||||
self.assertEqual(writer.context_budget.charged_bytes, 0)
|
||||
self.assertFalse(writer.status().accepting)
|
||||
@@ -572,8 +564,6 @@ class OutputWriterTests(fixtures.unittest.TestCase):
|
||||
self.assertIs(replacement.context_budget, writer.context_budget)
|
||||
self.assertIsNone(replacement.context_budget.reserve(8388608))
|
||||
release.set()
|
||||
# Drain first: the 2 s bound is not under test here (#370).
|
||||
drain_writer(writer)
|
||||
self.assertTrue(writer.close(2))
|
||||
finally:
|
||||
release.set()
|
||||
@@ -615,8 +605,6 @@ class OutputWriterTests(fixtures.unittest.TestCase):
|
||||
writer.invalidate()
|
||||
self.assertEqual(writer.context_budget.charged_bytes, 0)
|
||||
release.set()
|
||||
# Drain first: the 2 s bound is not under test here (#370).
|
||||
drain_writer(writer)
|
||||
self.assertTrue(writer.close(2))
|
||||
|
||||
def test_malformed_context_omission_preserves_independent_valid_output(self):
|
||||
@@ -707,8 +695,6 @@ class OutputWriterTests(fixtures.unittest.TestCase):
|
||||
release.set()
|
||||
invalidator.join(BLOCK_SECONDS)
|
||||
self.assertFalse(invalidator.is_alive())
|
||||
# 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_outputs(start.run_id)["state"], "complete"
|
||||
|
||||
@@ -1,13 +1,10 @@
|
||||
"""Pure helpers of the shared Images listing (py/api/gallery_listing.py, #348)."""
|
||||
|
||||
import ntpath
|
||||
import os
|
||||
import posixpath
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path, PurePath, PurePosixPath, PureWindowsPath
|
||||
from types import SimpleNamespace
|
||||
from pathlib import Path, PurePath, PureWindowsPath
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
@@ -19,13 +16,7 @@ from py.api.gallery_listing import ( # noqa: E402
|
||||
STORED,
|
||||
RootListing,
|
||||
WalkOutcome,
|
||||
add_folder,
|
||||
folder_sets,
|
||||
insert_entry,
|
||||
listing_delta,
|
||||
relative_part,
|
||||
rewalk_interval,
|
||||
root_contains,
|
||||
same_listing,
|
||||
walk_root,
|
||||
)
|
||||
@@ -132,36 +123,24 @@ class WalkRootTests(unittest.TestCase):
|
||||
def test_returns_the_listing_and_both_folder_sets(self):
|
||||
files = [(self.root / "sub" / "a.png", 2.0), (self.root / "b.png", 1.0)]
|
||||
scan, calls = self._scan(files)
|
||||
listed, folders, with_ancestors, resolved = walk_root(scan, self.root, None)
|
||||
listed, folders, with_ancestors = walk_root(scan, self.root, None)
|
||||
self.assertEqual(calls, [self.root])
|
||||
self.assertEqual(listed, files)
|
||||
self.assertEqual(folders, frozenset({"sub"}))
|
||||
self.assertEqual(with_ancestors, frozenset({"sub"}))
|
||||
self.assertEqual(resolved, os.path.normcase(os.path.realpath(self.root)))
|
||||
|
||||
def test_the_resolved_root_follows_a_symlinked_root(self):
|
||||
link = Path(self._tmp.name) / "linked"
|
||||
try:
|
||||
os.symlink(self.root, link, target_is_directory=True)
|
||||
except (OSError, NotImplementedError) as error:
|
||||
self.skipTest(f"symlinks unavailable: {error}")
|
||||
scan, _calls = self._scan([])
|
||||
resolved = walk_root(scan, link, None)[3]
|
||||
self.assertEqual(resolved, os.path.normcase(os.path.realpath(self.root)))
|
||||
self.assertNotEqual(resolved, os.path.normcase(str(link)))
|
||||
|
||||
def test_an_unchanged_walk_returns_the_previous_list_object(self):
|
||||
files = [(self.root / "sub" / "a.png", 2.0)]
|
||||
previous = list(files)
|
||||
scan, _calls = self._scan(files)
|
||||
listed, _folders, _ancestors, _resolved = walk_root(scan, self.root, previous)
|
||||
listed, _folders, _ancestors = walk_root(scan, self.root, previous)
|
||||
self.assertIs(listed, previous)
|
||||
|
||||
def test_a_changed_walk_returns_the_new_list(self):
|
||||
previous = [(self.root / "sub" / "a.png", 2.0)]
|
||||
files = [(self.root / "sub" / "a.png", 3.0)]
|
||||
scan, _calls = self._scan(files)
|
||||
listed, _folders, _ancestors, _resolved = walk_root(scan, self.root, previous)
|
||||
listed, _folders, _ancestors = walk_root(scan, self.root, previous)
|
||||
self.assertIsNot(listed, previous)
|
||||
self.assertEqual(listed, files)
|
||||
|
||||
@@ -201,278 +180,6 @@ class RecordTests(unittest.TestCase):
|
||||
self.assertEqual(gallery_listing.REFRESH_DELAY, 1.0)
|
||||
self.assertIsNone(RootListing().rewalk_handle)
|
||||
|
||||
def test_a_new_record_is_clean_and_never_walked(self):
|
||||
record = RootListing()
|
||||
self.assertFalse(record.dirty)
|
||||
self.assertIsNone(record.resolved_root)
|
||||
self.assertIsNone(record.last_started_at)
|
||||
self.assertIsNone(record.last_duration)
|
||||
|
||||
def test_drop_forgets_pending_adds(self):
|
||||
record = RootListing(files=[], pending_adds=["a.png"], patched=2)
|
||||
record.drop()
|
||||
self.assertEqual((record.pending_adds, record.patched), ([], 0))
|
||||
|
||||
def test_drop_keeps_the_resolved_root_and_the_walk_timing(self):
|
||||
# A drop clears the listing, not what is known about the root itself.
|
||||
record = RootListing(resolved_root="/o", last_started_at=1.0)
|
||||
record.last_duration = 0.5
|
||||
record.dirty = True
|
||||
record.drop()
|
||||
self.assertEqual(record.resolved_root, "/o")
|
||||
self.assertEqual((record.last_started_at, record.last_duration), (1.0, 0.5))
|
||||
self.assertTrue(record.dirty)
|
||||
|
||||
|
||||
class RewalkIntervalTests(unittest.TestCase):
|
||||
def test_constants(self):
|
||||
self.assertEqual(
|
||||
(
|
||||
gallery_listing.REWALK_FLOOR,
|
||||
gallery_listing.REWALK_FACTOR,
|
||||
gallery_listing.REWALK_CAP,
|
||||
gallery_listing.MIN_REWALK_DELAY,
|
||||
),
|
||||
(5, 5, 15, 0.5),
|
||||
)
|
||||
|
||||
def test_short_walks_wait_the_floor(self):
|
||||
self.assertEqual(rewalk_interval(0.2), 5)
|
||||
self.assertEqual(rewalk_interval(0), 5)
|
||||
self.assertEqual(rewalk_interval(None), 5) # never walked
|
||||
|
||||
def test_the_interval_scales_with_the_last_walk(self):
|
||||
self.assertEqual(rewalk_interval(1.0), 5)
|
||||
self.assertEqual(rewalk_interval(2.0), 10)
|
||||
self.assertAlmostEqual(rewalk_interval(2.5), 12.5)
|
||||
|
||||
def test_slow_walks_are_capped(self):
|
||||
self.assertEqual(rewalk_interval(3.0), 15)
|
||||
self.assertEqual(rewalk_interval(60.0), 15)
|
||||
|
||||
def test_the_constants_are_read_when_called(self):
|
||||
with (
|
||||
patch.object(gallery_listing, "REWALK_FLOOR", 0),
|
||||
patch.object(gallery_listing, "REWALK_CAP", 0),
|
||||
):
|
||||
self.assertEqual(rewalk_interval(2.0), 0)
|
||||
|
||||
|
||||
class RootContainsTests(unittest.TestCase):
|
||||
def test_posix_paths_under_the_root_match(self):
|
||||
root = posixpath.normcase("/srv/out")
|
||||
self.assertTrue(root_contains(root, "/srv/out/a.png", posixpath))
|
||||
self.assertTrue(root_contains(root, "/srv/out/sub/deep/a.png", posixpath))
|
||||
|
||||
def test_a_sibling_sharing_the_prefix_does_not_match(self):
|
||||
self.assertFalse(root_contains("/srv/out", "/srv/out2/a.png", posixpath))
|
||||
self.assertFalse(root_contains("/srv/out", "/srv/outside.png", posixpath))
|
||||
self.assertFalse(root_contains("/srv/out", "/srv/other/a.png", posixpath))
|
||||
|
||||
def test_posix_keeps_case(self):
|
||||
self.assertFalse(root_contains("/srv/out", "/srv/OUT/a.png", posixpath))
|
||||
|
||||
def test_the_root_itself_is_not_an_output(self):
|
||||
self.assertFalse(root_contains("/srv/out", "/srv/out", posixpath))
|
||||
|
||||
def test_windows_folds_case_and_separators(self):
|
||||
root = ntpath.normcase("C:\\ComfyUI\\Output")
|
||||
self.assertTrue(root_contains(root, "c:\\comfyui\\output\\a.png", ntpath))
|
||||
self.assertTrue(root_contains(root, "C:/ComfyUI/OUTPUT/sub/a.PNG", ntpath))
|
||||
self.assertFalse(root_contains(root, "C:\\ComfyUI\\Output2\\a.png", ntpath))
|
||||
self.assertFalse(root_contains(root, "D:\\ComfyUI\\Output\\a.png", ntpath))
|
||||
|
||||
def test_a_windows_drive_root_and_a_unc_share(self):
|
||||
drive = ntpath.normcase("E:\\")
|
||||
self.assertTrue(root_contains(drive, "e:\\x.png", ntpath))
|
||||
self.assertFalse(root_contains(drive, "F:\\x.png", ntpath))
|
||||
share = ntpath.normcase("\\\\Server\\Share\\Out")
|
||||
self.assertTrue(root_contains(share, "\\\\server\\share\\out\\a.png", ntpath))
|
||||
self.assertFalse(root_contains(share, "\\\\server\\share\\a.png", ntpath))
|
||||
|
||||
def test_a_windows_extended_length_prefix_is_ignored(self):
|
||||
root = ntpath.normcase("C:\\ComfyUI\\Output")
|
||||
long_file = "\\\\?\\C:\\ComfyUI\\Output\\deep\\a.png"
|
||||
self.assertTrue(root_contains(root, long_file, ntpath))
|
||||
prefixed_root = ntpath.normcase("\\\\?\\C:\\ComfyUI\\Output")
|
||||
self.assertTrue(
|
||||
root_contains(prefixed_root, "C:\\ComfyUI\\Output\\a.png", ntpath)
|
||||
)
|
||||
share = ntpath.normcase("\\\\Server\\Share\\Out")
|
||||
unc_file = "\\\\?\\UNC\\server\\share\\out\\a.png"
|
||||
self.assertTrue(root_contains(share, unc_file, ntpath))
|
||||
self.assertFalse(
|
||||
root_contains(root, "\\\\?\\C:\\ComfyUI\\Output2\\a.png", ntpath)
|
||||
)
|
||||
|
||||
def test_posix_keeps_a_literal_question_mark_folder(self):
|
||||
self.assertFalse(root_contains("/srv/out", "//?/srv/out/a.png", posixpath))
|
||||
|
||||
def test_the_default_is_this_platforms_rule(self):
|
||||
root = os.path.normcase(os.path.realpath(tempfile.gettempdir()))
|
||||
inside = os.path.join(root, "x.png")
|
||||
self.assertTrue(root_contains(root, inside))
|
||||
self.assertTrue(root_contains(root, inside.upper()) == (os.name == "nt"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
|
||||
class RelativePartTests(unittest.TestCase):
|
||||
def test_posix_keeps_the_case_and_the_nested_folders(self):
|
||||
self.assertEqual(
|
||||
relative_part("/srv/out", "/srv/out/Day 1/A.png", posixpath),
|
||||
"Day 1/A.png",
|
||||
)
|
||||
|
||||
def test_posix_outside_the_root_a_sibling_or_the_root_itself_is_none(self):
|
||||
for path in ("/srv/other/a.png", "/srv/out2/a.png", "/srv/out", "/srv/out/"):
|
||||
self.assertIsNone(relative_part("/srv/out", path, posixpath), path)
|
||||
|
||||
def test_a_root_spelled_with_a_trailing_separator(self):
|
||||
self.assertEqual(relative_part("/", "/a.png", posixpath), "a.png")
|
||||
|
||||
def test_windows_matches_folded_and_slices_the_original_case(self):
|
||||
root = ntpath.normcase("C:\\ComfyUI\\Output")
|
||||
self.assertEqual(
|
||||
relative_part(root, "c:/COMFYUI/output/Sub/ComfyUI_001.PNG", ntpath),
|
||||
"Sub\\ComfyUI_001.PNG",
|
||||
)
|
||||
|
||||
def test_windows_extended_length_and_unc_forms(self):
|
||||
root = ntpath.normcase("C:\\Out")
|
||||
self.assertEqual(
|
||||
relative_part(root, "\\\\?\\C:\\Out\\Deep\\a.png", ntpath),
|
||||
"Deep\\a.png",
|
||||
)
|
||||
share = ntpath.normcase("\\\\Server\\Share\\Out")
|
||||
self.assertEqual(
|
||||
relative_part(share, "\\\\?\\UNC\\Server\\Share\\Out\\B.png", ntpath),
|
||||
"B.png",
|
||||
)
|
||||
drive = ntpath.normcase("D:\\")
|
||||
self.assertEqual(relative_part(drive, "D:\\X\\a.png", ntpath), "X\\a.png")
|
||||
|
||||
def test_a_case_fold_that_changes_the_length_is_not_sliced(self):
|
||||
# str.lower folds the Turkish dotted I to two characters (Windows'
|
||||
# own normcase keeps the length; the guard covers any fold that not).
|
||||
folding = SimpleNamespace(
|
||||
sep="\\", normcase=lambda text: text.replace("/", "\\").lower()
|
||||
)
|
||||
path = "C:\\Out\\\u0130mage.png"
|
||||
self.assertNotEqual(len(folding.normcase(path)), len(path))
|
||||
self.assertIsNone(relative_part("c:\\out", path, folding))
|
||||
self.assertEqual(relative_part("c:\\out", "C:\\Out\\Ab.png", folding), "Ab.png")
|
||||
|
||||
def test_dot_parts_are_never_sliced(self):
|
||||
for path in ("/srv/out/../x.png", "/srv/out/./x.png", "/srv/out//x.png"):
|
||||
self.assertIsNone(relative_part("/srv/out", path, posixpath), path)
|
||||
|
||||
|
||||
class InsertEntryTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.root = PurePosixPath("/out")
|
||||
self.files = [
|
||||
(self.root / "c.png", 30.0),
|
||||
(self.root / "b.png", 20.0),
|
||||
(self.root / "a.png", 10.0),
|
||||
]
|
||||
|
||||
def test_newest_first_position_at_the_front_middle_and_end(self):
|
||||
for mtime, index in ((40.0, 0), (25.0, 1), (15.0, 2), (5.0, 3)):
|
||||
entry = (self.root / "n.png", mtime)
|
||||
patched = insert_entry(self.files, entry, 100)
|
||||
self.assertIs(patched[index], entry)
|
||||
self.assertEqual(len(patched), 4)
|
||||
self.assertEqual(len(self.files), 3) # never changed in place
|
||||
|
||||
def test_unchanged_entries_are_the_same_objects(self):
|
||||
patched = insert_entry(self.files, (self.root / "n.png", 40.0), 100)
|
||||
self.assertTrue(all(new is old for new, old in zip(patched[1:], self.files)))
|
||||
|
||||
def test_the_same_path_with_a_new_mtime_moves(self):
|
||||
patched = insert_entry(self.files, (PurePosixPath("/out/a.png"), 35.0), 100)
|
||||
self.assertEqual(
|
||||
[(str(path), mtime) for path, mtime in patched],
|
||||
[("/out/a.png", 35.0), ("/out/c.png", 30.0), ("/out/b.png", 20.0)],
|
||||
)
|
||||
|
||||
def test_an_entry_already_listed_returns_the_listing_itself(self):
|
||||
self.assertIs(
|
||||
insert_entry(self.files, (PurePosixPath("/out/b.png"), 20.0), 3), self.files
|
||||
)
|
||||
|
||||
def test_a_tie_in_mtime_is_not_placed(self):
|
||||
self.assertIsNone(insert_entry(self.files, (self.root / "n.png", 20.0), 100))
|
||||
|
||||
def test_a_tie_with_the_replaced_entry_itself_is_placed(self):
|
||||
files = [(self.root / "b.png", 20.0), (self.root / "a.png", 10.0)]
|
||||
patched = insert_entry(files, (self.root / "a.png", 20.5), 100)
|
||||
self.assertEqual([mtime for _path, mtime in patched], [20.5, 20.0])
|
||||
|
||||
def test_the_cap_refuses_a_new_path_but_replaces_a_listed_one(self):
|
||||
self.assertIsNone(insert_entry(self.files, (self.root / "n.png", 40.0), 3))
|
||||
patched = insert_entry(self.files, (self.root / "a.png", 40.0), 3)
|
||||
self.assertEqual(str(patched[0][0]), "/out/a.png")
|
||||
|
||||
def test_an_empty_listing(self):
|
||||
entry = (self.root / "n.png", 1.0)
|
||||
self.assertEqual(insert_entry([], entry, 100), [entry])
|
||||
|
||||
|
||||
class AddFolderTests(unittest.TestCase):
|
||||
def test_a_new_nested_folder_adds_it_and_its_ancestors(self):
|
||||
root = PurePath("/out")
|
||||
folders, with_ancestors = folder_sets(root, [(root / "x" / "a.png", 1.0)])
|
||||
new = add_folder(root, root / "y" / "z" / "b.png", folders, with_ancestors)
|
||||
self.assertEqual(
|
||||
new,
|
||||
folder_sets(
|
||||
root, [(root / "x" / "a.png", 1.0), (root / "y" / "z" / "b.png", 2.0)]
|
||||
),
|
||||
)
|
||||
|
||||
def test_a_known_folder_or_the_root_keeps_the_same_sets(self):
|
||||
root = PurePath("/out")
|
||||
sets = folder_sets(root, [(root / "x" / "a.png", 1.0)])
|
||||
for path in (root / "x" / "b.png", root / "c.png"):
|
||||
new = add_folder(root, path, *sets)
|
||||
self.assertIs(new[0], sets[0])
|
||||
self.assertIs(new[1], sets[1])
|
||||
|
||||
def test_windows_spelling_matches_folder_sets(self):
|
||||
root = PureWindowsPath("C:/Out")
|
||||
path = root / "A" / "B" / "c.png"
|
||||
self.assertEqual(
|
||||
add_folder(root, path, frozenset(), frozenset()),
|
||||
folder_sets(root, [(path, 1.0)]),
|
||||
)
|
||||
|
||||
|
||||
class ListingDeltaTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.old = [(PurePath(f"/out/{i}.png"), float(i)) for i in range(10)]
|
||||
|
||||
def test_an_added_entry(self):
|
||||
entry = (PurePath("/out/n.png"), 99.0)
|
||||
self.assertEqual(listing_delta(self.old, [entry, *self.old]), ([entry], []))
|
||||
|
||||
def test_a_replaced_entry(self):
|
||||
entry = (PurePath("/out/3.png"), 99.0)
|
||||
new = insert_entry(self.old, entry, 100)
|
||||
self.assertEqual(listing_delta(self.old, new), ([entry], [self.old[3]]))
|
||||
|
||||
def test_the_same_listing_has_no_change(self):
|
||||
self.assertEqual(listing_delta(self.old, list(self.old)), ([], []))
|
||||
|
||||
def test_a_walk_builds_new_entries_and_is_past_the_limit(self):
|
||||
walked = [(PurePath(str(path)), mtime) for path, mtime in self.old]
|
||||
self.assertEqual(walked, self.old)
|
||||
self.assertIsNone(listing_delta(self.old, walked, limit=5))
|
||||
self.assertIsNone(listing_delta(self.old, walked)) # shares nothing
|
||||
self.assertIsNone(listing_delta(self.old, self.old[:2], limit=5))
|
||||
|
||||
def test_an_emptied_listing_and_a_first_file(self):
|
||||
self.assertEqual(listing_delta(self.old[:1], []), ([], self.old[:1]))
|
||||
self.assertIsNone(listing_delta([], self.old[:1]))
|
||||
|
||||
@@ -12,7 +12,6 @@ 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,
|
||||
@@ -50,7 +49,6 @@ 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 ")
|
||||
|
||||
@@ -8,7 +8,6 @@ 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
|
||||
@@ -53,7 +52,6 @@ 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,
|
||||
|
||||
@@ -136,21 +136,13 @@ class TestMonitorSingletonLifecycle(unittest.TestCase):
|
||||
|
||||
|
||||
class FakeObserver:
|
||||
"""Records what the monitor does with the watchdog observer.
|
||||
"""Records what the monitor does with the watchdog observer."""
|
||||
|
||||
``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):
|
||||
def __init__(self):
|
||||
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))
|
||||
@@ -159,18 +151,12 @@ 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
|
||||
|
||||
|
||||
@@ -185,8 +171,6 @@ 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):
|
||||
@@ -241,8 +225,9 @@ class TestStartStop(unittest.TestCase):
|
||||
worker = handler._worker
|
||||
self.assertTrue(worker.is_alive())
|
||||
|
||||
self.assertTrue(self.monitor.stop_monitoring())
|
||||
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
|
||||
@@ -318,89 +303,6 @@ 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)
|
||||
fake.stopped = True # its thread has ended; stop() still raises
|
||||
with self.assertLogs("prompt_manager.image_monitor", "WARNING") as logs:
|
||||
self.assertFalse(self.monitor.stop_monitoring())
|
||||
text = "\n".join(logs.output)
|
||||
self.assertIn("Error while stopping image monitoring", text)
|
||||
self.assertNotIn("watchdog observer still running", text)
|
||||
fake.stop_error = None # tearDown's stop calls stop() again
|
||||
|
||||
def test_the_still_running_warning_names_the_time_the_join_got(self):
|
||||
fake = FakeObserver(stay_alive=True)
|
||||
handler = self._start_with(fake)
|
||||
now = [100.0]
|
||||
recorded_close = handler.close
|
||||
|
||||
def close_that_uses_the_deadline(timeout=im_mod.WORKER_JOIN_SECONDS):
|
||||
ok = recorded_close(timeout)
|
||||
now[0] += 5.0 # a busy worker took the whole deadline
|
||||
return ok
|
||||
|
||||
handler.close = close_that_uses_the_deadline
|
||||
with (
|
||||
patch.object(
|
||||
im_mod, "time", types.SimpleNamespace(monotonic=lambda: now[0])
|
||||
),
|
||||
self.assertLogs("prompt_manager.image_monitor", "WARNING") as logs,
|
||||
):
|
||||
self.assertFalse(self.monitor.stop_monitoring(timeout=5.0))
|
||||
self.assertEqual(fake.calls[-1], ("join", 0.0))
|
||||
self.assertIn(
|
||||
"watchdog observer still running after 0.0 s", "\n".join(logs.output)
|
||||
)
|
||||
fake.stay_alive = False # it ends; tearDown's stop waits for it
|
||||
|
||||
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):
|
||||
@@ -433,10 +335,6 @@ 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:
|
||||
@@ -501,127 +399,23 @@ class TestHandlerScheduling(unittest.TestCase):
|
||||
empty = self._image("empty.png", content=b"")
|
||||
self.assertFalse(self.handler.wait_until_settled(empty))
|
||||
|
||||
def test_shutdown_joins_the_worker_within_the_bound(self):
|
||||
def test_shutdown_stops_the_worker_and_clears_pending(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.assertTrue(self.handler.shutdown())
|
||||
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)
|
||||
self.assertTrue(fresh.shutdown())
|
||||
self.assertTrue(fresh.shutdown())
|
||||
self.assertTrue(fresh.close())
|
||||
fresh.shutdown()
|
||||
fresh.shutdown()
|
||||
self.handler.schedule(self._image("s3.png"))
|
||||
self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0))
|
||||
self.handler.shutdown()
|
||||
@@ -640,7 +434,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.assertTrue(self.handler.close())
|
||||
self.handler.shutdown()
|
||||
|
||||
def test_processing_errors_do_not_kill_the_worker(self):
|
||||
def explode(path, prompt_snapshot=None):
|
||||
@@ -655,150 +449,6 @@ 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)."""
|
||||
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.tmp.cleanup)
|
||||
self.handler = im_mod.ImageGenerationHandler(MagicMock(), MagicMock())
|
||||
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
|
||||
self.reported = []
|
||||
self.handler.output_listener = self.reported.append
|
||||
|
||||
def _image(self, name, content=b"x" * 16):
|
||||
path = os.path.join(self.tmp.name, name)
|
||||
with open(path, "wb") as f:
|
||||
f.write(content)
|
||||
return path
|
||||
|
||||
def test_a_settled_output_is_reported_once_with_its_resolved_path(self):
|
||||
path = self._image("a.png")
|
||||
self.handler._process_pending(path)
|
||||
self.assertEqual(self.reported, [os.path.realpath(path)])
|
||||
|
||||
def test_a_symlinked_spelling_is_reported_resolved(self):
|
||||
real_dir = os.path.join(self.tmp.name, "real")
|
||||
os.makedirs(real_dir)
|
||||
link_dir = os.path.join(self.tmp.name, "link")
|
||||
try:
|
||||
os.symlink(real_dir, link_dir, target_is_directory=True)
|
||||
except (OSError, NotImplementedError) as error:
|
||||
self.skipTest(f"symlinks unavailable: {error}")
|
||||
with open(os.path.join(real_dir, "b.png"), "wb") as f:
|
||||
f.write(b"x" * 16)
|
||||
self.handler._process_pending(os.path.join(link_dir, "b.png"))
|
||||
self.assertEqual(
|
||||
self.reported, [os.path.realpath(os.path.join(real_dir, "b.png"))]
|
||||
)
|
||||
|
||||
def test_it_is_reported_even_when_linking_raises(self):
|
||||
def explode(path, prompt_snapshot=None):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
self.handler.process_new_image = explode
|
||||
path = self._image("c.png")
|
||||
with self.assertRaises(RuntimeError):
|
||||
self.handler._process_pending(path)
|
||||
self.assertEqual(self.reported, [os.path.realpath(path)])
|
||||
self.assertEqual(self.handler.pending_count(), 0)
|
||||
|
||||
def test_an_output_that_never_settles_is_not_reported(self):
|
||||
self.handler.settle_attempts = 3
|
||||
self.handler._process_pending(os.path.join(self.tmp.name, "gone.png"))
|
||||
empty = self._image("empty.png", content=b"")
|
||||
self.handler._process_pending(empty)
|
||||
self.assertEqual(self.reported, [])
|
||||
|
||||
def test_a_video_is_reported_when_its_extension_is_monitored(self):
|
||||
self.handler.supported_extensions = (".png", ".mp4")
|
||||
path = self._image("clip.mp4")
|
||||
self.handler.on_created(FakeEvent(path))
|
||||
self.assertTrue(wait_for(lambda: self.reported != []))
|
||||
self.assertEqual(self.reported, [os.path.realpath(path)])
|
||||
|
||||
def test_a_failing_listener_is_logged_and_the_worker_carries_on(self):
|
||||
def broken(_path):
|
||||
raise RuntimeError("listing gone")
|
||||
|
||||
self.handler.output_listener = broken
|
||||
first = self._image("d.png")
|
||||
with self.assertLogs("prompt_manager.image_monitor", "WARNING") as logs:
|
||||
self.handler._process_pending(first) # no exception
|
||||
self.assertIn("Images listing", "\n".join(logs.output))
|
||||
self.handler.output_listener = self.reported.append
|
||||
second = self._image("e.png")
|
||||
self.handler.schedule(second)
|
||||
self.assertTrue(wait_for(lambda: self.reported != []))
|
||||
self.assertEqual(self.reported, [os.path.realpath(second)])
|
||||
|
||||
def test_no_listener_reports_nothing(self):
|
||||
self.handler.output_listener = None
|
||||
self.handler._process_pending(self._image("f.png")) # no exception
|
||||
self.assertEqual(self.reported, [])
|
||||
|
||||
def test_the_monitor_wires_running_and_later_handlers(self):
|
||||
monitor = im_mod.ImageMonitor(MagicMock(), MagicMock())
|
||||
listener = self.reported.append
|
||||
monitor.set_output_listener(listener)
|
||||
with (
|
||||
patch.object(im_mod, "Observer", FakeObserver),
|
||||
patch.object(im_mod.ImageMonitor, "_gallery_config", return_value=None),
|
||||
):
|
||||
monitor.start_monitoring([self.tmp.name])
|
||||
try:
|
||||
self.assertIs(monitor.handler.output_listener, listener)
|
||||
monitor.set_output_listener(None)
|
||||
self.assertIsNone(monitor.output_listener)
|
||||
self.assertIsNone(monitor.handler.output_listener)
|
||||
finally:
|
||||
self.assertTrue(monitor.stop_monitoring())
|
||||
|
||||
|
||||
class TestModuleIsolation(unittest.TestCase):
|
||||
"""This module must not leave mocks behind for later test modules."""
|
||||
|
||||
|
||||
@@ -8,7 +8,6 @@ import ctypes
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
@@ -695,40 +694,6 @@ 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()
|
||||
|
||||
@@ -310,7 +310,7 @@ class TestAutoTag(ImageTagJobTestCase):
|
||||
self.assertTrue(
|
||||
any(
|
||||
line.endswith(
|
||||
"is not downloaded. Download it in System › Auto-tagging."
|
||||
"is not downloaded. Download it from the legacy Auto-tag page."
|
||||
)
|
||||
for line in job.log
|
||||
),
|
||||
|
||||
+12
-439
@@ -4,8 +4,6 @@ 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
|
||||
@@ -18,7 +16,6 @@ import threading
|
||||
import types
|
||||
import unittest
|
||||
import unittest.mock as mock
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
@@ -43,20 +40,13 @@ class FakePromptQueue:
|
||||
return self.results.pop(0) if self.results else None
|
||||
|
||||
|
||||
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
|
||||
def fake_server_module():
|
||||
"""A ``server`` module whose PromptServer.instance accepts routes and hooks."""
|
||||
instance = types.SimpleNamespace(
|
||||
routes=web.RouteTableDef(),
|
||||
prompt_queue=queue,
|
||||
prompt_queue=FakePromptQueue(),
|
||||
handlers=[],
|
||||
app=app if app is not None else types.SimpleNamespace(middlewares=[]),
|
||||
app=types.SimpleNamespace(middlewares=[]),
|
||||
)
|
||||
instance.add_on_prompt_handler = instance.handlers.append
|
||||
module = types.ModuleType("server")
|
||||
@@ -102,7 +92,7 @@ def _import_package(name, db_path, output_dir, server=None):
|
||||
try:
|
||||
spec.loader.exec_module(package)
|
||||
except BaseException:
|
||||
cleanup_package(package, require_stopped=False)
|
||||
cleanup_package(package)
|
||||
raise
|
||||
finally:
|
||||
for key, module in saved.items():
|
||||
@@ -113,34 +103,20 @@ def _import_package(name, db_path, output_dir, server=None):
|
||||
return package, server.PromptServer.instance
|
||||
|
||||
|
||||
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"
|
||||
def cleanup_package(package):
|
||||
"""Stop fixture-owned workers before releasing its other resources."""
|
||||
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, require_stopped=True):
|
||||
def load_package(name, db_path, output_dir, server=None):
|
||||
"""Dispose package workers before closing its isolated logging handlers."""
|
||||
root_logger = logging.getLogger("prompt_manager")
|
||||
loggers = [root_logger] + [
|
||||
@@ -168,7 +144,7 @@ def load_package(name, db_path, output_dir, server=None, require_stopped=True):
|
||||
try:
|
||||
yield package, instance
|
||||
finally:
|
||||
cleanup_package(package, require_stopped=require_stopped)
|
||||
cleanup_package(package)
|
||||
finally:
|
||||
current_loggers = [root_logger] + [
|
||||
logger
|
||||
@@ -580,406 +556,3 @@ 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=[]
|
||||
)
|
||||
|
||||
|
||||
WAIT_SECONDS = 30 # ceiling for every Event wait (TESTING.md test rules)
|
||||
|
||||
|
||||
def start_busy_worker(package, folder, name="busy.png"):
|
||||
"""The image worker alive in its processing delay."""
|
||||
handler = package._global_image_monitor.handler
|
||||
if handler is None:
|
||||
raise AssertionError("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(folder, name)
|
||||
with open(path, "wb") as image:
|
||||
image.write(b"x" * 16)
|
||||
if not handler.schedule(path) or not handler._worker.is_alive():
|
||||
raise AssertionError("the image worker did not start")
|
||||
return handler._worker
|
||||
|
||||
|
||||
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"):
|
||||
return start_busy_worker(package, self.tmp.name, name)
|
||||
|
||||
def hold_a_preview(self, package, *, advance, kill_releases):
|
||||
"""A preview worker held inside a fake encode that tracks a fake
|
||||
ffmpeg, and a package clock that the monitor step moves by
|
||||
``advance`` seconds. Cleanups release the encode first."""
|
||||
# Not the watched folder. Real path for both the file and the root:
|
||||
# macOS /var vs /private/var, Windows short names.
|
||||
previews = os.path.realpath(os.path.join(self.tmp.name, "previews"))
|
||||
os.makedirs(previews)
|
||||
from PIL import Image
|
||||
|
||||
image = os.path.join(previews, "held.png")
|
||||
Image.new("RGB", (8, 8), "red").save(image)
|
||||
held = types.SimpleNamespace(
|
||||
entered=threading.Event(),
|
||||
release=threading.Event(),
|
||||
killed=threading.Event(),
|
||||
stop_seen=[],
|
||||
timeouts=[],
|
||||
now=[100.0],
|
||||
)
|
||||
|
||||
class FakeFfmpeg:
|
||||
def poll(self): # _track keeps only processes still running
|
||||
return None
|
||||
|
||||
def kill(self):
|
||||
held.killed.set()
|
||||
if kill_releases:
|
||||
held.release.set()
|
||||
|
||||
def held_encode(src, dst, settings, is_video, *, stop, guard, track):
|
||||
track(FakeFfmpeg())
|
||||
held.entered.set()
|
||||
held.release.wait(WAIT_SECONDS)
|
||||
held.stop_seen.append(stop())
|
||||
return {"action": "skipped"}
|
||||
|
||||
api = package.api
|
||||
monitor = package._global_image_monitor
|
||||
real_stop = monitor.stop_monitoring
|
||||
real_close = api.close_auto_thumbnails
|
||||
|
||||
def stop_then_advance(timeout):
|
||||
ok = real_stop(timeout)
|
||||
held.now[0] += advance
|
||||
return ok
|
||||
|
||||
def recorded_close(final=False, timeout=None):
|
||||
held.timeouts.append(timeout)
|
||||
return real_close(final=final, timeout=timeout)
|
||||
|
||||
for target, name, value in (
|
||||
(api, "_get_all_output_dirs", lambda: [Path(previews)]),
|
||||
(api, "_generate_one", held_encode),
|
||||
(api, "close_auto_thumbnails", recorded_close),
|
||||
(monitor, "stop_monitoring", stop_then_advance),
|
||||
(package, "_clock", lambda: held.now[0]),
|
||||
):
|
||||
patcher = mock.patch.object(target, name, value)
|
||||
patcher.start()
|
||||
self.addCleanup(patcher.stop)
|
||||
self.assertEqual(api.queue_capture_thumbnail(image), "queued")
|
||||
self.addCleanup(held.release.set)
|
||||
self.assertTrue(held.entered.wait(WAIT_SECONDS), "the encode never started")
|
||||
held.thumbnailer = api._auto_thumbnailer # the teardown drops the API's
|
||||
held.worker = held.thumbnailer._worker
|
||||
# Cleanups run newest first: release the encode, then join the worker
|
||||
# (even when an assertion failed first), then end the patches.
|
||||
self.addCleanup(held.worker.join, WAIT_SECONDS)
|
||||
self.addCleanup(held.release.set)
|
||||
self.assertTrue(held.worker.is_alive())
|
||||
return held
|
||||
|
||||
|
||||
class TestPackageTeardown(PackageTeardownCase):
|
||||
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_kills_ffmpeg_and_stops_the_preview_worker(self):
|
||||
package = self.load()
|
||||
held = self.hold_a_preview(package, advance=4.0, kill_releases=True)
|
||||
|
||||
result = package._stop_background_work()
|
||||
|
||||
self.assertTrue(held.killed.is_set())
|
||||
self.assertEqual(held.timeouts, [6.0]) # min(30, 10 - 4)
|
||||
self.assertTrue(result["thumbnails"], result)
|
||||
self.assertFalse(held.worker.is_alive())
|
||||
self.assertEqual(held.stop_seen, [True])
|
||||
|
||||
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 ListHandler(logging.Handler):
|
||||
def __init__(self):
|
||||
super().__init__(logging.INFO)
|
||||
self.lines = []
|
||||
|
||||
def emit(self, record):
|
||||
self.lines.append(f"{record.levelname}:{record.getMessage()}")
|
||||
|
||||
|
||||
class TestPackageTeardownOnce(PackageTestCase):
|
||||
"""The tests that only observe one plain teardown share one package.
|
||||
|
||||
The teardown runs once, so it runs here in ``setUpClass`` (per-test
|
||||
``assertLogs`` cannot wrap it); the tests read what it recorded.
|
||||
"""
|
||||
|
||||
package_name = "pm_teardown_once"
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
with contextlib.redirect_stdout(io.StringIO()):
|
||||
super().setUpClass()
|
||||
package = cls.package
|
||||
cls.monitor = package._global_image_monitor
|
||||
cls.observer = cls.monitor.observer
|
||||
cls.tracker_thread = package._global_prompt_tracker.cleanup_thread
|
||||
cls.worker = start_busy_worker(package, cls.tmp.name)
|
||||
cls.recorder_before = cls.monitor.review_recorder
|
||||
logger = logging.getLogger("prompt_manager.init")
|
||||
handler, level = ListHandler(), logger.level
|
||||
logger.addHandler(handler)
|
||||
logger.setLevel(logging.INFO) # otherwise set by the config
|
||||
try:
|
||||
with mock.patch("atexit.unregister", wraps=atexit.unregister) as unregister:
|
||||
cls.result = package._stop_background_work()
|
||||
finally:
|
||||
logger.removeHandler(handler)
|
||||
logger.setLevel(level)
|
||||
cls.log = "\n".join(handler.lines)
|
||||
cls.unregister_calls = unregister.call_args_list
|
||||
|
||||
def test_teardown_stops_every_part_once(self):
|
||||
expected = {"monitor": True, "thumbnails": True, "tracker": True}
|
||||
self.assertEqual(self.result, expected)
|
||||
self.assertEqual(list(self.result), ["monitor", "thumbnails", "tracker"])
|
||||
self.assertIn("Background work stopped in", self.log)
|
||||
self.assertIn("monitor=ok thumbnails=ok tracker=ok", self.log)
|
||||
self.assertFalse(self.worker.is_alive())
|
||||
self.assertFalse(self.observer.is_alive())
|
||||
self.assertFalse(self.tracker_thread.is_alive())
|
||||
copy = self.package._stop_background_work()
|
||||
copy["monitor"] = False # a copy: the stored result is unchanged
|
||||
with mock.patch.object(self.monitor, "stop_monitoring") as again:
|
||||
self.assertEqual(self.package._stop_background_work(), expected)
|
||||
again.assert_not_called()
|
||||
|
||||
def test_teardown_detaches_the_monitor_callbacks(self):
|
||||
self.assertIsNotNone(self.recorder_before)
|
||||
self.assertIsNone(self.monitor.review_recorder)
|
||||
self.assertIsNone(self.monitor.thumbnail_queue)
|
||||
self.assertIsNone(self.monitor.output_listener)
|
||||
|
||||
def test_teardown_removes_its_atexit_callback(self):
|
||||
self.assertEqual(
|
||||
self.unregister_calls, [mock.call(self.package._stop_background_work)]
|
||||
)
|
||||
|
||||
|
||||
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_an_interrupt_between_steps_still_records_the_teardown(self):
|
||||
package = self.load(require_stopped=False)
|
||||
real_step = package._teardown_step
|
||||
|
||||
def interrupted_before_the_previews(name, stop, logger):
|
||||
if name == "the preview worker":
|
||||
raise KeyboardInterrupt # outside the step's own try
|
||||
return real_step(name, stop, logger)
|
||||
|
||||
with (
|
||||
mock.patch.object(
|
||||
package, "_teardown_step", interrupted_before_the_previews
|
||||
),
|
||||
mock.patch("atexit.unregister", wraps=atexit.unregister) as unregister,
|
||||
self.assertRaises(KeyboardInterrupt),
|
||||
):
|
||||
package._stop_background_work()
|
||||
expected = {"monitor": True, "thumbnails": False, "tracker": False}
|
||||
self.assertEqual(package._background_work_result, expected)
|
||||
unregister.assert_called_once_with(package._stop_background_work)
|
||||
with mock.patch.object(
|
||||
package._global_image_monitor, "stop_monitoring"
|
||||
) as again:
|
||||
self.assertEqual(package._stop_background_work(), expected)
|
||||
again.assert_not_called()
|
||||
|
||||
def test_teardown_reports_a_preview_worker_still_held(self):
|
||||
package = self.load(require_stopped=False)
|
||||
held = self.hold_a_preview(package, advance=9.95, kill_releases=False)
|
||||
|
||||
with self.assertLogs("prompt_manager.init", "WARNING") as logs:
|
||||
result = package._stop_background_work()
|
||||
|
||||
self.assertTrue(held.killed.is_set())
|
||||
self.assertEqual(len(held.timeouts), 1)
|
||||
self.assertAlmostEqual(held.timeouts[0], 0.05, places=6)
|
||||
self.assertTrue(held.thumbnailer.closed)
|
||||
self.assertFalse(result["thumbnails"], result)
|
||||
self.assertRegex("\n".join(logs.output), r"not stopped in time .*thumbnails")
|
||||
held.release.set()
|
||||
held.worker.join(WAIT_SECONDS)
|
||||
self.assertFalse(held.worker.is_alive())
|
||||
self.assertEqual(held.stop_seen, [True])
|
||||
|
||||
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_}")
|
||||
|
||||
@@ -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", "Tag output images (legacy page)")
|
||||
self._run_job("autotag", "Auto-tag images (legacy page)")
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/autotag/start", json={}
|
||||
)
|
||||
|
||||
@@ -29,7 +29,6 @@ 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
|
||||
@@ -205,7 +204,6 @@ 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 = {}
|
||||
@@ -258,7 +256,6 @@ 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}, {})
|
||||
|
||||
+28
-407
@@ -252,7 +252,7 @@ class LoraImportAPITestCase(LoraAPITestCase):
|
||||
|
||||
self._patches = [
|
||||
patch.object(
|
||||
lora_utils, "default_lora_roots", return_value=[str(self.loras)]
|
||||
lora_utils, "find_lora_directories", return_value=[str(self.loras)]
|
||||
),
|
||||
patch.object(
|
||||
lora_utils, "get_lora_image_cache_dir", return_value=self.cache_dir
|
||||
@@ -276,8 +276,6 @@ 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",
|
||||
@@ -285,8 +283,8 @@ class LoraImportAPITestCase(LoraAPITestCase):
|
||||
images=["https://civitai.com/a.png"],
|
||||
prompt="neon street",
|
||||
)
|
||||
_write_metadata(self.loras, "plain", images=["https://civitai.com/p.png"])
|
||||
(self.loras / "plain.png").write_bytes(_png_bytes())
|
||||
(self.loras / "neon.png").write_bytes(_png_bytes())
|
||||
_write_metadata(self.loras, "plain")
|
||||
|
||||
async def _post(self, url, body=None):
|
||||
resp = await self._post_json(url, {} if body is None else body)
|
||||
@@ -323,34 +321,8 @@ class LoraImportRouteTests(LoraImportAPITestCase):
|
||||
|
||||
self.assertEqual(resp.status, 200, plan)
|
||||
self.assertEqual(
|
||||
{
|
||||
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,
|
||||
},
|
||||
{k: plan[k] for k in ("lora_files", "with_metadata", "current_imports")},
|
||||
{"lora_files": 3, "with_metadata": 2, "current_imports": 2},
|
||||
)
|
||||
self.assertEqual(plan["blocked_by_pairs"], 0)
|
||||
self.assertEqual(plan["planned_downloads"], 1)
|
||||
@@ -364,95 +336,7 @@ class LoraImportRouteTests(LoraImportAPITestCase):
|
||||
self.assertEqual(_size(db_path + "-shm"), shm)
|
||||
self.assertEqual(_snapshot(self.api.db), rows)
|
||||
self.assertNotIn(SECRET_KEY, 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))
|
||||
self.assertNotIn(str(self.tmp_root), json.dumps(plan))
|
||||
|
||||
async def test_preview_requires_the_integration(self):
|
||||
IntegrationConfig.LORA_MANAGER_ENABLED = False
|
||||
@@ -474,7 +358,7 @@ class LoraImportRouteTests(LoraImportAPITestCase):
|
||||
self._two_loras()
|
||||
resp, plan = await self._post(PREVIEW)
|
||||
rows = _snapshot(self.api.db)
|
||||
meta = self.loras / "plain.metadata.json"
|
||||
meta = self.loras / "plain.safetensors.metadata.json"
|
||||
meta.write_text(meta.read_text() + " ")
|
||||
|
||||
resp, body = await self._post(APPLY, {"token": plan["token"]})
|
||||
@@ -513,39 +397,19 @@ 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"])), 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.get_prompt_images(neon["id"])), 2)
|
||||
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"]})
|
||||
self.assertEqual((resp.status, body["code"]), (409, "plan-used"))
|
||||
|
||||
async def test_a_finished_import_clears_the_layout_cache(self):
|
||||
# #362: the import may link files under a folder LoraManager created
|
||||
# after the image routes cached its layout.
|
||||
self._two_loras()
|
||||
with patch("py.lora_utils.clear_layout_cache") as clear:
|
||||
_plan, job = await self._apply_reviewed()
|
||||
await job.task
|
||||
self.assertEqual(job.state, "finished", job.log)
|
||||
clear.assert_called_once_with()
|
||||
|
||||
async def test_discovery_failure_keeps_previous_imports(self):
|
||||
self._two_loras()
|
||||
before = _snapshot(self.api.db)
|
||||
|
||||
# Legacy scan: discovery raises once the stream is open.
|
||||
with patch.object(
|
||||
lora_import, "discover_sidecars", side_effect=OSError("share gone")
|
||||
lora_utils, "find_lora_directories", side_effect=OSError("share gone")
|
||||
):
|
||||
resp = await self.client.request("POST", SCAN)
|
||||
events = self._sse(await resp.text())
|
||||
@@ -701,7 +565,6 @@ 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())
|
||||
@@ -946,35 +809,6 @@ class LoraImportRefusalTests(LoraImportAPITestCase):
|
||||
self.assertEqual((job.state, job.scope["failure"]), ("failed", "plan-changed"))
|
||||
self.assertEqual(_snapshot(self.api.db), before)
|
||||
|
||||
async def test_a_refused_import_keeps_the_layout_cache(self):
|
||||
self._two_loras()
|
||||
resp, plan = await self._post(PREVIEW)
|
||||
self.api.db.save_prompt(
|
||||
text="late", category=CATEGORY, prompt_hash=generate_prompt_hash("late")
|
||||
)
|
||||
with patch("py.lora_utils.clear_layout_cache") as clear:
|
||||
resp, body = await self._post(APPLY, {"token": plan["token"]})
|
||||
job = self.api._jobs.get(body["job"]["id"])
|
||||
await job.task
|
||||
self.assertEqual(job.scope["failure"], "plan-changed")
|
||||
clear.assert_not_called()
|
||||
|
||||
async def test_start_refused_when_a_skipped_sidecar_changed(self):
|
||||
# #362 CUJ 4: the review counted the orphan as skipped.
|
||||
self._two_loras()
|
||||
orphan = self.loras / "orphan.metadata.json"
|
||||
orphan.write_text(json.dumps({"model_name": "orphan"}))
|
||||
resp, plan = await self._post(PREVIEW)
|
||||
self.assertEqual(resp.status, 200, plan)
|
||||
self.assertEqual((plan["lora_files"], plan["skipped"]), (3, 1))
|
||||
before = _snapshot(self.api.db)
|
||||
orphan.write_text(orphan.read_text() + " ")
|
||||
resp, body = await self._post(APPLY, {"token": plan["token"]})
|
||||
self.assertEqual(resp.status, 409, body)
|
||||
self.assertEqual(body["code"], "plan-changed")
|
||||
self.assertEqual(self.api._jobs.list(), [])
|
||||
self.assertEqual(_snapshot(self.api.db), before)
|
||||
|
||||
async def test_a_pair_saved_after_the_review_blocks_the_job(self):
|
||||
self._two_loras()
|
||||
resp, plan = await self._post(PREVIEW)
|
||||
@@ -1019,13 +853,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.metadata.json"
|
||||
meta = self.loras / "plain.safetensors.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.metadata.json"
|
||||
meta = self.loras / "plain.safetensors.metadata.json"
|
||||
await self._change_during_download(meta.unlink)
|
||||
|
||||
async def test_a_file_added_while_staging_fails_the_job(self):
|
||||
@@ -1035,7 +869,7 @@ class LoraImportRefusalTests(LoraImportAPITestCase):
|
||||
self._two_loras()
|
||||
before = _snapshot(self.api.db)
|
||||
self.gate = threading.Event()
|
||||
meta = self.loras / "plain.metadata.json"
|
||||
meta = self.loras / "plain.safetensors.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))
|
||||
@@ -1152,250 +986,37 @@ class LoraImportRestoreTests(LoraImportAPITestCase):
|
||||
|
||||
|
||||
class LoraImportStageTests(unittest.TestCase):
|
||||
"""The stages over a disposable LoraManager install (#355)."""
|
||||
|
||||
def setUp(self):
|
||||
self.tmp = Path(tempfile.mkdtemp()).resolve()
|
||||
self.tmp = Path(tempfile.mkdtemp())
|
||||
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.loras / "a"
|
||||
b = self.loras / "a" / "b"
|
||||
a = self.tmp / "a"
|
||||
b = self.tmp / "a" / "b"
|
||||
b.mkdir(parents=True)
|
||||
_write_metadata(b, "zeta")
|
||||
_write_metadata(a, "alpha")
|
||||
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, [])
|
||||
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)
|
||||
|
||||
def test_fingerprint_follows_size_and_mtime(self):
|
||||
meta = _write_metadata(self.loras, "one")
|
||||
first = lora_import.fingerprint(self.discover())
|
||||
self.assertEqual(lora_import.fingerprint(self.discover()), first)
|
||||
meta = _write_metadata(self.tmp, "one")
|
||||
first = lora_import.fingerprint([meta])
|
||||
self.assertEqual(lora_import.fingerprint([meta]), first)
|
||||
meta.write_text(meta.read_text() + " ")
|
||||
self.assertNotEqual(lora_import.fingerprint(self.discover()), first)
|
||||
changed = lora_import.fingerprint(self.discover())
|
||||
self.assertNotEqual(lora_import.fingerprint([meta]), first)
|
||||
meta.unlink()
|
||||
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)
|
||||
|
||||
# #362: a skipped sidecar is part of what the review counted.
|
||||
|
||||
def _assert_fingerprint_follows(self, change):
|
||||
first = lora_import.fingerprint(self.discover())
|
||||
self.assertEqual(lora_import.fingerprint(self.discover()), first)
|
||||
change()
|
||||
self.assertNotEqual(lora_import.fingerprint(self.discover()), first)
|
||||
|
||||
def _assert_add_edit_remove_change_it(self, path, text):
|
||||
self._assert_fingerprint_follows(lambda: path.write_text(text))
|
||||
self._assert_fingerprint_follows(
|
||||
lambda: path.write_text(path.read_text() + " ")
|
||||
)
|
||||
self._assert_fingerprint_follows(path.unlink)
|
||||
|
||||
def _centralized(self):
|
||||
sidecars = self.tmp / "sidecars"
|
||||
sidecars.mkdir()
|
||||
self.settings_file(
|
||||
{
|
||||
"sidecar_storage_mode": "centralized",
|
||||
"sidecar_storage_path": str(sidecars),
|
||||
}
|
||||
)
|
||||
model = self.loras / "kept.safetensors"
|
||||
model.write_bytes(b"m")
|
||||
(sidecars / "kept.metadata.json").write_text(
|
||||
json.dumps({"file_path": str(model)})
|
||||
)
|
||||
return sidecars
|
||||
|
||||
def test_fingerprint_follows_a_centralized_sidecar_whose_model_is_missing(self):
|
||||
sidecars = self._centralized()
|
||||
gone = self.loras / "gone.safetensors"
|
||||
self._assert_add_edit_remove_change_it(
|
||||
sidecars / "gone.metadata.json", json.dumps({"file_path": str(gone)})
|
||||
)
|
||||
self.assertEqual(len(self.discover().sidecars), 1)
|
||||
|
||||
def test_fingerprint_follows_a_checkpoint_sidecar_outside_the_lora_roots(self):
|
||||
sidecars = self._centralized()
|
||||
checkpoints = self.tmp / "checkpoints"
|
||||
checkpoints.mkdir()
|
||||
ckpt = checkpoints / "base.safetensors"
|
||||
ckpt.write_bytes(b"c")
|
||||
self._assert_add_edit_remove_change_it(
|
||||
sidecars / "base.metadata.json", json.dumps({"file_path": str(ckpt)})
|
||||
)
|
||||
|
||||
def test_fingerprint_follows_an_alongside_sidecar_without_its_model(self):
|
||||
_write_metadata(self.loras, "kept")
|
||||
self._assert_add_edit_remove_change_it(
|
||||
self.loras / "orphan.metadata.json", json.dumps({"model_name": "orphan"})
|
||||
)
|
||||
self.assertEqual(len(self.discover().sidecars), 1)
|
||||
|
||||
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)
|
||||
self.assertNotEqual(lora_import.fingerprint([meta]), first)
|
||||
|
||||
def test_stage_cancel_before_a_file_raises(self):
|
||||
_write_metadata(self.loras, "one")
|
||||
meta = _write_metadata(self.tmp, "one")
|
||||
with self.assertRaises(lora_import.ImportCancelled):
|
||||
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"})
|
||||
lora_import.stage([meta], self.tmp, cancelled=lambda: True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,590 +0,0 @@
|
||||
"""The LoraManager layout cache for the image hot paths (#362).
|
||||
|
||||
``py/lora_layout_cache.py``. Every test injects the install resolver, the
|
||||
layout loader (both counting calls), the wall clock, the monotonic clock and,
|
||||
where needed, ``stat`` and ``pathmod``, over a temporary folder. No test
|
||||
sleeps: modification times are set with ``os.utime`` and the clocks are
|
||||
plain attributes the test moves.
|
||||
"""
|
||||
|
||||
import ntpath
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from dataclasses import replace
|
||||
from pathlib import Path, PureWindowsPath
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from py import lora_layout_cache as cache_module
|
||||
from py.lora_layout_cache import (
|
||||
MAX_AGE_SECONDS,
|
||||
MISSING,
|
||||
RACY_SECONDS,
|
||||
LayoutCache,
|
||||
)
|
||||
from py.lora_manager_layout import UNREADABLE, LoraManagerLayout, _canonical
|
||||
|
||||
SECOND_NS = 10**9
|
||||
# NTFS keeps modification times in 100 ns steps and FAT in 2 s steps: every
|
||||
# time a test sets is a whole millisecond, and a test whose file system rounds
|
||||
# it anyway skips (``set_mtime``).
|
||||
TICK_NS = 10**6
|
||||
|
||||
|
||||
def make_layout(settings_path=None, media_roots=(), error=None, source="user config"):
|
||||
return LoraManagerLayout(
|
||||
settings_path=settings_path,
|
||||
source=source,
|
||||
error=error,
|
||||
mode="alongside",
|
||||
lora_roots=tuple(media_roots),
|
||||
sidecar_root=None,
|
||||
example_images_root=None,
|
||||
library="default",
|
||||
media_roots=tuple(media_roots),
|
||||
)
|
||||
|
||||
|
||||
class CacheTestCase(unittest.TestCase):
|
||||
"""An install, its plugin settings file, a settings file and a LoRA root."""
|
||||
|
||||
def setUp(self):
|
||||
self.tmp = Path(tempfile.mkdtemp()).resolve()
|
||||
self.addCleanup(shutil.rmtree, self.tmp)
|
||||
self.install_dir = self.tmp / "custom_nodes" / "ComfyUI-Lora-Manager"
|
||||
self.install_dir.mkdir(parents=True)
|
||||
self.plugin_settings = self.install_dir / "settings.json"
|
||||
self.settings = self.tmp / "user-config" / "settings.json"
|
||||
self.settings.parent.mkdir()
|
||||
self.root_dir = self.tmp / "loras"
|
||||
self.root_dir.mkdir()
|
||||
# Files are written ten minutes before the injected wall clock, so
|
||||
# nothing is racy unless a test sets a recent modification time.
|
||||
self.now_ns = time.time_ns() // TICK_NS * TICK_NS
|
||||
self.mono = 1000.0
|
||||
self.resolved = []
|
||||
self.loaded = []
|
||||
self.install = _canonical(str(self.install_dir))
|
||||
self.layout = make_layout(
|
||||
settings_path=self.settings,
|
||||
media_roots=(_canonical(str(self.root_dir)),),
|
||||
)
|
||||
self.write(self.settings, "{}")
|
||||
|
||||
def write(self, path, text, mtime_ns=None):
|
||||
Path(path).write_text(text, encoding="utf-8")
|
||||
stamp = mtime_ns if mtime_ns is not None else self.now_ns - 600 * SECOND_NS
|
||||
self.set_mtime(path, stamp)
|
||||
|
||||
def set_mtime(self, path, stamp):
|
||||
os.utime(path, ns=(stamp, stamp))
|
||||
if os.stat(path).st_mtime_ns != stamp:
|
||||
self.skipTest("this file system rounds modification times")
|
||||
|
||||
def resolve(self, key):
|
||||
self.resolved.append(key)
|
||||
return self.install
|
||||
|
||||
def load(self, install):
|
||||
self.loaded.append(install)
|
||||
return self.layout
|
||||
|
||||
def cache(self, **kwargs):
|
||||
kwargs.setdefault("wall_ns", lambda: self.now_ns)
|
||||
kwargs.setdefault("monotonic", lambda: self.mono)
|
||||
return LayoutCache(
|
||||
kwargs.pop("resolve", self.resolve), kwargs.pop("load", self.load), **kwargs
|
||||
)
|
||||
|
||||
|
||||
class ValidationTests(CacheTestCase):
|
||||
def test_hit_reuses_the_entry_without_resolving_or_loading(self):
|
||||
cache = self.cache()
|
||||
first = cache.get("")
|
||||
second = cache.get("")
|
||||
self.assertEqual(len(self.resolved), 1)
|
||||
self.assertEqual(len(self.loaded), 1)
|
||||
self.assertIs(first[1], second[1])
|
||||
self.assertIs(first[0], second[0])
|
||||
self.assertEqual(first, (self.install, self.layout))
|
||||
|
||||
def test_settings_size_change_reloads_once(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
stamp = self.settings.stat().st_mtime_ns
|
||||
self.write(self.settings, '{"a": 1}', mtime_ns=stamp)
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_settings_mtime_change_reloads_once(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
self.set_mtime(self.settings, self.settings.stat().st_mtime_ns + TICK_NS)
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_settings_deletion_reloads_once(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
self.settings.unlink()
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_settings_creation_reloads_once(self):
|
||||
self.settings.unlink()
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
self.write(self.settings, "{}")
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_plugin_settings_change_reloads_once(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
self.write(self.plugin_settings, '{"use_portable_settings": true}')
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_a_different_configured_path_reloads(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
cache.get("other")
|
||||
self.assertEqual(self.resolved, ["", "other"])
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_a_different_home_folder_reloads(self):
|
||||
# The folder guard compares against the home folder: a cached layout
|
||||
# must not outlive a change of it.
|
||||
home = {"value": str(self.tmp / "home-a")}
|
||||
cache = self.cache(home=lambda: home["value"])
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
home["value"] = str(self.tmp / "home-b")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_configured_path_is_stripped_for_the_key_and_the_resolver(self):
|
||||
cache = self.cache()
|
||||
cache.get(" x ")
|
||||
cache.get("x")
|
||||
self.assertEqual(self.resolved, ["x"])
|
||||
self.assertEqual(len(self.loaded), 1)
|
||||
|
||||
def test_a_recently_modified_file_is_not_stored(self):
|
||||
self.write(self.settings, "{}", mtime_ns=self.now_ns - SECOND_NS)
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
self.assertIsNone(cache._entry)
|
||||
self.now_ns += RACY_SECONDS * SECOND_NS
|
||||
cache.get("")
|
||||
self.assertIsNotNone(cache._entry)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 3)
|
||||
|
||||
def test_racy_boundary(self):
|
||||
# Exactly RACY_SECONDS old is trusted; anything newer is not.
|
||||
boundary = self.now_ns - RACY_SECONDS * SECOND_NS
|
||||
self.write(self.settings, "{}", mtime_ns=boundary + TICK_NS)
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
self.assertIsNone(cache._entry)
|
||||
self.set_mtime(self.settings, boundary)
|
||||
cache.get("")
|
||||
self.assertIsNotNone(cache._entry)
|
||||
|
||||
def test_a_modification_time_in_the_future_is_racy(self):
|
||||
self.write(self.plugin_settings, "{}", mtime_ns=self.now_ns + 60 * SECOND_NS)
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
self.assertIsNone(cache._entry)
|
||||
|
||||
def test_a_missing_file_is_not_racy(self):
|
||||
self.assertFalse(self.plugin_settings.exists())
|
||||
self.settings.unlink()
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
self.assertIsNotNone(cache._entry)
|
||||
signatures = dict(cache._entry.signatures)
|
||||
self.assertTrue(all(value is MISSING for value in signatures.values()))
|
||||
self.assertEqual(len(signatures), 2)
|
||||
|
||||
def test_entry_is_reused_below_the_maximum_age_and_reloaded_at_it(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
stored_at = self.mono
|
||||
self.mono = stored_at + MAX_AGE_SECONDS - 0.5
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 1)
|
||||
self.mono = stored_at + MAX_AGE_SECONDS
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_a_wall_clock_moved_backward_does_not_extend_the_age(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
self.now_ns -= 3600 * SECOND_NS
|
||||
self.mono += MAX_AGE_SECONDS
|
||||
# The old entry is too old; the new load sees files "from the future"
|
||||
# (racy) and is not stored, but it does load.
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_stat_failure_on_a_hit_reloads(self):
|
||||
failing = {"on": False}
|
||||
|
||||
def stat(path):
|
||||
if failing["on"]:
|
||||
raise PermissionError(13, "sharing violation", path)
|
||||
return os.stat(path)
|
||||
|
||||
cache = self.cache(stat=stat)
|
||||
cache.get("")
|
||||
failing["on"] = True
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
self.assertIsNone(cache._entry)
|
||||
|
||||
def test_stat_failure_on_a_miss_leaves_the_layout_unstored(self):
|
||||
def stat(path):
|
||||
raise PermissionError(13, "sharing violation", path)
|
||||
|
||||
cache = self.cache(stat=stat)
|
||||
self.assertEqual(cache.get(""), (self.install, self.layout))
|
||||
self.assertIsNone(cache._entry)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_an_unreadable_settings_layout_is_not_stored(self):
|
||||
self.layout = make_layout(settings_path=self.settings, error=UNREADABLE)
|
||||
cache = self.cache()
|
||||
self.assertEqual(cache.get("")[1].error, UNREADABLE)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
self.assertIsNone(cache._entry)
|
||||
|
||||
def test_other_layout_errors_are_stored(self):
|
||||
self.layout = make_layout(
|
||||
settings_path=self.settings, error="the settings file is not valid JSON"
|
||||
)
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 1)
|
||||
|
||||
def _symlink(self, link, target):
|
||||
try:
|
||||
os.symlink(target, link, target_is_directory=True)
|
||||
except (OSError, NotImplementedError) as error:
|
||||
self.skipTest(f"symlinks are not available here: {error}")
|
||||
|
||||
def test_a_media_root_swapped_for_a_symlink_reloads(self):
|
||||
elsewhere = self.tmp / "elsewhere"
|
||||
elsewhere.mkdir()
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
shutil.rmtree(self.root_dir)
|
||||
self._symlink(self.root_dir, elsewhere)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_an_install_swapped_for_a_symlink_reloads(self):
|
||||
elsewhere = self.tmp / "elsewhere-lm"
|
||||
elsewhere.mkdir()
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
moved = self.tmp / "moved-lm"
|
||||
self.install_dir.rename(moved)
|
||||
self._symlink(self.install_dir, elsewhere)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_a_root_whose_realpath_fails_reloads(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
with patch.object(
|
||||
cache_module, "_canonical", side_effect=OSError("realpath failed")
|
||||
):
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_a_stat_value_error_cannot_validate(self):
|
||||
def stat(path):
|
||||
raise ValueError("embedded null character in path")
|
||||
|
||||
cache = self.cache(stat=stat)
|
||||
cache.get("")
|
||||
self.assertIsNone(cache._entry)
|
||||
|
||||
def test_not_a_directory_is_missing(self):
|
||||
def stat(path):
|
||||
raise NotADirectoryError(20, "Not a directory", path)
|
||||
|
||||
cache = self.cache(stat=stat)
|
||||
cache.get("")
|
||||
self.assertTrue(
|
||||
all(value is MISSING for _path, value in cache._entry.signatures)
|
||||
)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 1)
|
||||
|
||||
def test_an_offline_windows_share_cannot_validate(self):
|
||||
def stat(path):
|
||||
error = FileNotFoundError(2, "The network path was not found", path)
|
||||
error.winerror = 53
|
||||
raise error
|
||||
|
||||
cache = self.cache(stat=stat)
|
||||
cache.get("")
|
||||
self.assertIsNone(cache._entry)
|
||||
|
||||
def test_a_failed_reload_leaves_no_stale_entry(self):
|
||||
failing = {"on": False}
|
||||
|
||||
def load(install):
|
||||
self.loaded.append(install)
|
||||
if failing["on"]:
|
||||
raise RuntimeError("boom")
|
||||
return self.layout
|
||||
|
||||
cache = self.cache(load=load)
|
||||
cache.get("")
|
||||
self.mono += MAX_AGE_SECONDS
|
||||
failing["on"] = True
|
||||
with self.assertRaises(RuntimeError):
|
||||
cache.get("")
|
||||
self.assertIsNone(cache._entry)
|
||||
|
||||
def test_a_resolve_failure_propagates_and_stores_nothing(self):
|
||||
def resolve(key):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
cache = self.cache(resolve=resolve)
|
||||
with self.assertRaises(RuntimeError):
|
||||
cache.get("")
|
||||
self.assertIsNone(cache._entry)
|
||||
|
||||
def test_without_a_settings_path_only_the_plugin_file_is_watched(self):
|
||||
self.layout = make_layout(media_roots=self.layout.media_roots)
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
self.assertEqual(
|
||||
[path for path, _sig in cache._entry.signatures],
|
||||
[os.path.join(self.install, "settings.json")],
|
||||
)
|
||||
|
||||
def test_portable_settings_are_watched_once(self):
|
||||
self.layout = make_layout(settings_path=self.install_dir / "settings.json")
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
self.assertEqual(len(cache._entry.signatures), 1)
|
||||
|
||||
def test_a_load_failure_propagates_and_stores_nothing(self):
|
||||
calls = []
|
||||
|
||||
def load(install):
|
||||
calls.append(install)
|
||||
if len(calls) == 1:
|
||||
raise RuntimeError("boom")
|
||||
return self.layout
|
||||
|
||||
cache = self.cache(load=load)
|
||||
with self.assertRaises(RuntimeError):
|
||||
cache.get("")
|
||||
self.assertIsNone(cache._entry)
|
||||
self.assertEqual(cache.get(""), (self.install, self.layout))
|
||||
self.assertEqual(len(calls), 2)
|
||||
|
||||
def test_without_an_install_nothing_is_watched_and_the_entry_ages_out(self):
|
||||
stats = []
|
||||
|
||||
def stat(path):
|
||||
stats.append(path)
|
||||
return os.stat(path)
|
||||
|
||||
self.layout = make_layout(media_roots=self.layout.media_roots, source=None)
|
||||
cache = self.cache(resolve=lambda key: self.resolved.append(key), stat=stat)
|
||||
self.assertEqual(cache.get(""), (None, self.layout))
|
||||
self.assertEqual(cache._entry.signatures, ())
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 1)
|
||||
self.assertEqual(stats, [])
|
||||
self.mono += MAX_AGE_SECONDS
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
def test_windows_paths_differing_only_in_case_are_watched_once(self):
|
||||
stats = []
|
||||
|
||||
def stat(path):
|
||||
stats.append(path)
|
||||
return SimpleNamespace(st_mtime_ns=self.now_ns - 600 * SECOND_NS, st_size=2)
|
||||
|
||||
install = _canonical("C:\\ComfyUI\\custom_nodes\\LM", ntpath)
|
||||
# The hit check needs ntpath's realpath to be idempotent here.
|
||||
self.assertEqual(install, _canonical(install, ntpath))
|
||||
self.layout = make_layout(
|
||||
settings_path=PureWindowsPath(
|
||||
"C:\\ComfyUI\\custom_nodes\\lm\\SETTINGS.JSON"
|
||||
)
|
||||
)
|
||||
cache = self.cache(resolve=lambda key: install, stat=stat, pathmod=ntpath)
|
||||
cache.get("")
|
||||
self.assertEqual(len(cache._entry.signatures), 1)
|
||||
self.assertEqual(len(stats), 1)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 1)
|
||||
self.assertEqual(len(stats), 2)
|
||||
|
||||
|
||||
class ConcurrencyTests(CacheTestCase):
|
||||
def test_a_load_that_started_before_clear_stores_nothing(self):
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
def load(install):
|
||||
self.loaded.append(install)
|
||||
if len(self.loaded) == 1:
|
||||
started.set()
|
||||
self.assertTrue(release.wait(10))
|
||||
return self.layout
|
||||
|
||||
cache = self.cache(load=load)
|
||||
results = []
|
||||
worker = threading.Thread(target=lambda: results.append(cache.get("")))
|
||||
worker.start()
|
||||
self.assertTrue(started.wait(10))
|
||||
cache.clear()
|
||||
release.set()
|
||||
worker.join(10)
|
||||
self.assertFalse(worker.is_alive())
|
||||
self.assertEqual(results, [(self.install, self.layout)])
|
||||
self.assertIsNone(cache._entry)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
self.assertIsNotNone(cache._entry)
|
||||
|
||||
def test_concurrent_misses_all_get_a_layout_and_one_entry_is_stored(self):
|
||||
cache = self.cache()
|
||||
count = 6
|
||||
barrier = threading.Barrier(count)
|
||||
results = []
|
||||
lock = threading.Lock()
|
||||
|
||||
def worker():
|
||||
barrier.wait(10)
|
||||
result = cache.get("")
|
||||
with lock:
|
||||
results.append(result)
|
||||
|
||||
threads = [threading.Thread(target=worker) for _ in range(count)]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join(10)
|
||||
self.assertEqual(results, [(self.install, self.layout)] * count)
|
||||
self.assertIsNotNone(cache._entry)
|
||||
loads = len(self.loaded)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), loads)
|
||||
|
||||
def test_a_valid_hit_takes_no_lock(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
results = []
|
||||
with cache._lock:
|
||||
worker = threading.Thread(target=lambda: results.append(cache.get("")))
|
||||
worker.start()
|
||||
worker.join(10)
|
||||
alive = worker.is_alive()
|
||||
worker.join(10)
|
||||
self.assertFalse(alive)
|
||||
self.assertEqual(results, [(self.install, self.layout)])
|
||||
|
||||
def test_clear_forces_a_reload(self):
|
||||
cache = self.cache()
|
||||
cache.get("")
|
||||
cache.clear()
|
||||
self.assertIsNone(cache._entry)
|
||||
cache.get("")
|
||||
self.assertEqual(len(self.loaded), 2)
|
||||
|
||||
|
||||
class LoggingTests(CacheTestCase):
|
||||
LOGGER = "prompt_manager.lora_layout_cache"
|
||||
|
||||
def get(self, cache):
|
||||
with self.assertLogs(self.LOGGER, "DEBUG") as logs:
|
||||
cache.get("")
|
||||
# assertLogs needs one record; this marker is not counted.
|
||||
cache_module.logger.debug("marker")
|
||||
return [r for r in logs.records if r.getMessage() != "marker"]
|
||||
|
||||
def test_a_new_layout_logs_one_info_line(self):
|
||||
cache = self.cache()
|
||||
records = self.get(cache)
|
||||
self.assertEqual([r.levelname for r in records], ["INFO"])
|
||||
self.assertEqual(
|
||||
records[0].getMessage(),
|
||||
"LoraManager layout changed: user config settings, alongside storage, "
|
||||
"1 media folders",
|
||||
)
|
||||
self.mono += MAX_AGE_SECONDS
|
||||
self.assertEqual(self.get(cache), []) # an equal layout, reloaded
|
||||
cache.clear()
|
||||
self.assertEqual(self.get(cache), [])
|
||||
self.layout = replace(self.layout, media_roots=(), lora_roots=())
|
||||
self.mono += MAX_AGE_SECONDS
|
||||
self.assertEqual(len(self.get(cache)), 1)
|
||||
|
||||
def test_no_settings_source_prints_no(self):
|
||||
self.layout = replace(self.layout, source=None)
|
||||
records = self.get(self.cache())
|
||||
self.assertIn("layout changed: no settings,", records[0].getMessage())
|
||||
|
||||
def test_an_unstored_load_logs_no_info_line(self):
|
||||
self.write(self.settings, "{}", mtime_ns=self.now_ns)
|
||||
self.assertEqual(
|
||||
[r for r in self.get(self.cache()) if r.levelname == "INFO"], []
|
||||
)
|
||||
|
||||
def test_a_newly_refused_default_folder_logs_one_warning(self):
|
||||
self.layout = replace(self.layout, refused_defaults=("/home",))
|
||||
cache = self.cache()
|
||||
records = self.get(cache)
|
||||
warnings = [r for r in records if r.levelname == "WARNING"]
|
||||
self.assertEqual(len(warnings), 1)
|
||||
self.assertIn("/home", warnings[0].getMessage())
|
||||
self.mono += MAX_AGE_SECONDS
|
||||
self.assertEqual(self.get(cache), [])
|
||||
cache.clear()
|
||||
self.assertEqual(self.get(cache), [])
|
||||
self.layout = replace(self.layout, refused_defaults=("/home", "/"))
|
||||
self.mono += MAX_AGE_SECONDS
|
||||
warnings = [r for r in self.get(cache) if r.levelname == "WARNING"]
|
||||
self.assertEqual([w.getMessage().count("'/'") for w in warnings], [1])
|
||||
|
||||
|
||||
class ModuleTests(unittest.TestCase):
|
||||
def test_constants(self):
|
||||
self.assertEqual(cache_module.RACY_SECONDS, 3)
|
||||
self.assertEqual(cache_module.MAX_AGE_SECONDS, 30)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,969 +0,0 @@
|
||||
"""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"
|
||||
BAD_SIDECAR_REASON = (
|
||||
"the sidecar folder is a filesystem root, the home folder or a folder "
|
||||
"that contains it"
|
||||
)
|
||||
|
||||
|
||||
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, BAD_SIDECAR_REASON)
|
||||
self.assertNotIn(os.path.normcase(os.path.realpath("/")), layout.media_roots)
|
||||
|
||||
def test_a_centralized_sidecar_root_that_contains_home(self):
|
||||
home = self.tmp / "outer" / "me"
|
||||
home.mkdir(parents=True)
|
||||
self.write_settings(
|
||||
self.user_config,
|
||||
{
|
||||
"sidecar_storage_mode": "centralized",
|
||||
"sidecar_storage_path": str(self.tmp / "outer"),
|
||||
},
|
||||
)
|
||||
with patch.object(layout_module.Path, "home", return_value=home):
|
||||
layout = self.load()
|
||||
self.assertEqual(layout.error, BAD_SIDECAR_REASON)
|
||||
self.assertEqual(
|
||||
layout.media_roots, (os.path.normcase(str(self.default_loras)),)
|
||||
)
|
||||
|
||||
|
||||
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_every_scanned_file_is_listed_resolved_and_sorted(self):
|
||||
# #362: the fingerprint covers every scanned sidecar, kept or skipped.
|
||||
ckpt = self.model(self.checkpoints, "base")
|
||||
alpha = self.model(self.loras, "alpha")
|
||||
self.sidecar(self.sidecars / "z", "base", {"file_path": str(ckpt)})
|
||||
self.sidecar(self.sidecars / "a", "alpha", {"file_path": str(alpha)})
|
||||
self.sidecar(self.sidecars / "m", "nopath", {"model_name": "x"})
|
||||
discovery = discover_sidecars(self.load())
|
||||
files = discovery.scanned_files
|
||||
self.assertIsInstance(files, tuple)
|
||||
self.assertEqual(len(files), 3)
|
||||
self.assertEqual(discovery.scanned, len(files))
|
||||
self.assertEqual([str(f) for f in files], sorted(str(f) for f in files))
|
||||
self.assertEqual(
|
||||
[f.name for f in files],
|
||||
[
|
||||
f"alpha{METADATA_SUFFIX}",
|
||||
f"nopath{METADATA_SUFFIX}",
|
||||
f"base{METADATA_SUFFIX}",
|
||||
],
|
||||
)
|
||||
for path in files:
|
||||
self.assertEqual(str(path), os.path.realpath(path))
|
||||
self.assertEqual(len(discovery.sidecars), 1)
|
||||
|
||||
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_a_relative_preview_url_is_logged_and_skipped(self):
|
||||
fallback = self.touch(self.default_loras / "alpha.webp")
|
||||
self.touch(self.default_loras / "relative.png")
|
||||
with self.assertLogs("prompt_manager.lora_manager_layout", "DEBUG") as logs:
|
||||
media = self.media({"preview_url": "relative.png"})
|
||||
self.assertEqual(media, [fallback])
|
||||
self.assertEqual(len(logs.records), 1)
|
||||
self.assertEqual(logs.records[0].levelname, "DEBUG")
|
||||
self.assertIn(self.meta.name, logs.output[0])
|
||||
|
||||
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 RefusedDefaultsTests(LayoutTestCase):
|
||||
"""#362: ComfyUI's default LoRA folders pass the same guard; the refused
|
||||
ones that exist are listed so the layout cache can warn about them."""
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.outer = self.tmp / "outer"
|
||||
self.home = self.outer / "me"
|
||||
self.home.mkdir(parents=True)
|
||||
self.fs_root = os.path.abspath(os.sep)
|
||||
self.defaults = [
|
||||
str(self.home),
|
||||
str(self.outer),
|
||||
self.fs_root,
|
||||
str(self.tmp / "missing"),
|
||||
str(self.default_loras),
|
||||
]
|
||||
self.expected = (
|
||||
layout_module._canonical(str(self.home)),
|
||||
layout_module._canonical(str(self.outer)),
|
||||
layout_module._canonical(self.fs_root),
|
||||
)
|
||||
home_patch = patch.object(layout_module.Path, "home", return_value=self.home)
|
||||
home_patch.start()
|
||||
self.addCleanup(home_patch.stop)
|
||||
|
||||
def load_defaults(self, path="__plugin__"):
|
||||
return self.load(path, default_roots=lambda: list(self.defaults))
|
||||
|
||||
def test_refused_existing_defaults_are_listed_on_a_usable_layout(self):
|
||||
self.write_settings(self.user_config, {})
|
||||
layout = self.load_defaults()
|
||||
self.assertIsNone(layout.error)
|
||||
self.assertEqual(layout.refused_defaults, self.expected)
|
||||
self.assertEqual(
|
||||
layout.lora_roots, (os.path.normcase(str(self.default_loras)),)
|
||||
)
|
||||
|
||||
def test_they_are_listed_on_an_unusable_centralized_layout(self):
|
||||
self.write_settings(
|
||||
self.user_config,
|
||||
{"sidecar_storage_mode": "centralized", "sidecar_storage_path": "/"},
|
||||
)
|
||||
layout = self.load_defaults()
|
||||
self.assertEqual(layout.error, BAD_SIDECAR_REASON)
|
||||
self.assertEqual(layout.refused_defaults, self.expected)
|
||||
|
||||
def test_they_are_listed_without_an_install(self):
|
||||
layout = self.load_defaults(None)
|
||||
self.assertEqual(layout.error, "no LoraManager install found")
|
||||
self.assertEqual(layout.refused_defaults, self.expected)
|
||||
|
||||
def test_nothing_is_listed_when_every_default_is_accepted(self):
|
||||
self.write_settings(self.user_config, {})
|
||||
self.assertEqual(self.load().refused_defaults, ())
|
||||
|
||||
|
||||
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()
|
||||
+52
-209
@@ -224,14 +224,6 @@ class TestReadLoraMetadata(unittest.TestCase):
|
||||
result = read_lora_metadata(Path("/nonexistent/file.metadata.json"))
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_returns_none_for_a_file_that_is_not_utf8(self):
|
||||
# #362: one malformed sidecar is skipped; it must not fail the whole
|
||||
# trigger-word load (now logged with a traceback) or the review.
|
||||
with tempfile.TemporaryDirectory() as folder:
|
||||
path = Path(folder) / "bad.metadata.json"
|
||||
path.write_bytes(b'{"model_name": "\xff\xfe"}')
|
||||
self.assertIsNone(read_lora_metadata(path))
|
||||
|
||||
|
||||
class TestGetLoraImageCacheDir(unittest.TestCase):
|
||||
"""Test get_lora_image_cache_dir — returns and creates cache path."""
|
||||
@@ -1271,12 +1263,11 @@ class TestTriggerWordCache(unittest.TestCase):
|
||||
"file_name": "my_lora.safetensors",
|
||||
"civitai": {"trainedWords": ["trigger1", "trigger2"]},
|
||||
}
|
||||
meta_path = Path(tmpdir) / "my_lora.metadata.json"
|
||||
meta_path = Path(tmpdir) / "my_lora.safetensors.metadata.json"
|
||||
meta_path.write_text(json.dumps(meta))
|
||||
(Path(tmpdir) / "my_lora.safetensors").write_bytes(b"m")
|
||||
|
||||
# No LoraManager install: the default LoRA folders, alongside.
|
||||
with patch("py.lora_utils.default_lora_roots", return_value=[tmpdir]):
|
||||
# Patch find_lora_directories to return our temp dir
|
||||
with patch("py.lora_utils.find_lora_directories", return_value=[tmpdir]):
|
||||
count = self.cache.load(tmpdir)
|
||||
|
||||
self.assertTrue(self.cache.is_loaded)
|
||||
@@ -1292,10 +1283,11 @@ class TestTriggerWordCache(unittest.TestCase):
|
||||
"file_name": "MyLoRA.safetensors",
|
||||
"civitai": {"trainedWords": ["word1"]},
|
||||
}
|
||||
(Path(tmpdir) / "MyLoRA.metadata.json").write_text(json.dumps(meta))
|
||||
(Path(tmpdir) / "MyLoRA.safetensors").write_bytes(b"m")
|
||||
(Path(tmpdir) / "MyLoRA.safetensors.metadata.json").write_text(
|
||||
json.dumps(meta)
|
||||
)
|
||||
|
||||
with patch("py.lora_utils.default_lora_roots", return_value=[tmpdir]):
|
||||
with patch("py.lora_utils.find_lora_directories", return_value=[tmpdir]):
|
||||
self.cache.load(tmpdir)
|
||||
|
||||
self.assertEqual(self.cache.get_trigger_words("mylora"), ["word1"])
|
||||
@@ -1520,7 +1512,7 @@ class TestDetectLoraManager(unittest.TestCase):
|
||||
self.assertEqual(lora_utils.detect_lora_manager(), str(lm.resolve()))
|
||||
|
||||
|
||||
class TestDefaultLoraRoots(unittest.TestCase):
|
||||
class TestFindLoraDirectories(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.tmp = Path(tempfile.mkdtemp())
|
||||
@@ -1528,7 +1520,7 @@ class TestDefaultLoraRoots(unittest.TestCase):
|
||||
self.root = _make_comfy_root(self.tmp)
|
||||
self.lm = _make_lora_manager(self.root / "custom_nodes")
|
||||
|
||||
def test_collects_default_extra_and_runtime_dirs_only(self):
|
||||
def test_collects_default_extra_runtime_and_metadata_dirs(self):
|
||||
extra_abs = self.tmp / "extra_abs"
|
||||
extra_abs.mkdir()
|
||||
(self.root / "rel_loras").mkdir()
|
||||
@@ -1550,14 +1542,14 @@ class TestDefaultLoraRoots(unittest.TestCase):
|
||||
{"folder_paths": _FakeFolderPaths(str(self.root), [str(runtime)])},
|
||||
),
|
||||
):
|
||||
dirs = lora_utils.default_lora_roots()
|
||||
dirs = lora_utils.find_lora_directories(str(self.lm))
|
||||
|
||||
# #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))
|
||||
@@ -1568,21 +1560,12 @@ class TestDefaultLoraRoots(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_nothing_is_found(self):
|
||||
def test_without_root_only_metadata_dirs_are_found(self):
|
||||
with (
|
||||
patch.object(lora_utils, "find_comfyui_root", return_value=None),
|
||||
patch.dict(sys.modules, {"folder_paths": None}),
|
||||
):
|
||||
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)
|
||||
self.assertEqual(lora_utils.find_lora_directories(str(self.lm)), [])
|
||||
|
||||
|
||||
class TestModelNameAndPreviews(unittest.TestCase):
|
||||
@@ -1598,6 +1581,44 @@ 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):
|
||||
|
||||
@@ -1682,189 +1703,11 @@ 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, "default_lora_roots", return_value=[str(tmp)]):
|
||||
with patch.object(lora_utils, "find_lora_directories", 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_an_os_error_is_a_warning_without_a_traceback(self):
|
||||
cache = TriggerWordCache()
|
||||
cache._cache = {"old": ["word"]}
|
||||
with (
|
||||
patch.object(TriggerWordCache, "_build", side_effect=OSError("denied")),
|
||||
self.assertLogs("prompt_manager.lora_utils", level="WARNING") as logs,
|
||||
):
|
||||
self.assertEqual(cache.load(str(self.lm)), 1)
|
||||
self.assertEqual([r.levelname for r in logs.records], ["WARNING"])
|
||||
self.assertIsNone(logs.records[0].exc_info)
|
||||
self.assertEqual(cache.get_trigger_words("old"), ["word"])
|
||||
self.assertFalse(cache.retry_due())
|
||||
|
||||
def test_an_unexpected_failure_logs_its_traceback(self):
|
||||
# #362: a programming error must not look like a disk error.
|
||||
cache = TriggerWordCache()
|
||||
cache._cache = {"old": ["word"]}
|
||||
with (
|
||||
patch.object(TriggerWordCache, "_build", side_effect=ValueError("bug")),
|
||||
self.assertLogs("prompt_manager.lora_utils", level="WARNING") as logs,
|
||||
):
|
||||
self.assertEqual(cache.load(str(self.lm)), 1)
|
||||
self.assertEqual(len(logs.records), 1)
|
||||
self.assertIsNotNone(logs.records[0].exc_info)
|
||||
self.assertIs(logs.records[0].exc_info[0], ValueError)
|
||||
self.assertEqual(cache.get_trigger_words("old"), ["word"])
|
||||
self.assertFalse(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()
|
||||
|
||||
@@ -315,43 +315,6 @@ 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
|
||||
|
||||
@@ -24,7 +24,6 @@ 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
|
||||
@@ -225,7 +224,6 @@ 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()
|
||||
|
||||
@@ -1,142 +0,0 @@
|
||||
"""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()
|
||||
+23
-5
@@ -26,9 +26,27 @@ def _read(name):
|
||||
|
||||
|
||||
def _pyproject():
|
||||
import tomllib
|
||||
|
||||
return tomllib.loads(_read("pyproject.toml"))
|
||||
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,
|
||||
}
|
||||
|
||||
|
||||
class PyprojectContract(unittest.TestCase):
|
||||
@@ -47,8 +65,8 @@ class PyprojectContract(unittest.TestCase):
|
||||
}
|
||||
self.assertEqual(deps, wanted)
|
||||
|
||||
def test_requires_python_floor_is_3_12(self):
|
||||
self.assertEqual(_pyproject()["project"]["requires-python"], ">=3.12")
|
||||
def test_requires_python_floor_is_3_10(self):
|
||||
self.assertEqual(_pyproject()["project"]["requires-python"], ">=3.10")
|
||||
|
||||
def test_build_system_is_declared(self):
|
||||
self.assertIsNotNone(_pyproject().get("build-system"))
|
||||
|
||||
@@ -23,7 +23,6 @@ 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
|
||||
@@ -279,7 +278,6 @@ 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"])
|
||||
|
||||
@@ -5,12 +5,10 @@ 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
|
||||
|
||||
@@ -314,93 +312,5 @@ 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()
|
||||
|
||||
@@ -27,7 +27,6 @@ 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
|
||||
@@ -791,7 +790,6 @@ 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):
|
||||
@@ -955,7 +953,7 @@ class TestMonitorRecording(ReviewTestCase):
|
||||
monitor.set_review_recorder(None)
|
||||
self.assertIsNone(monitor.handler.review_recorder)
|
||||
finally:
|
||||
stop_monitor(self, monitor)
|
||||
monitor.stop_monitoring()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
@@ -1094,7 +1092,6 @@ 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")
|
||||
|
||||
@@ -17,7 +17,6 @@ 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
|
||||
@@ -159,8 +158,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,
|
||||
@@ -171,7 +170,6 @@ class ReviewToolTests(ToolFixture):
|
||||
):
|
||||
verdict = tool.judge(db, handler, row, self.output)
|
||||
finally:
|
||||
close_handler(self, handler)
|
||||
db.close_all()
|
||||
self.assertEqual(verdict["verdict"], "conflict")
|
||||
|
||||
|
||||
@@ -28,7 +28,6 @@ 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
|
||||
@@ -438,7 +437,6 @@ 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
|
||||
|
||||
@@ -762,7 +760,6 @@ 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},
|
||||
|
||||
@@ -38,7 +38,6 @@ 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
|
||||
@@ -300,7 +299,6 @@ 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
|
||||
|
||||
|
||||
@@ -11,7 +11,6 @@ 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
|
||||
@@ -152,7 +151,6 @@ 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"]
|
||||
@@ -334,7 +332,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(close_handler, self, self.handler)
|
||||
self.addCleanup(self.handler.shutdown)
|
||||
self.linked = []
|
||||
self.handler.link_image_to_prompt = lambda path, prompt, meta: (
|
||||
self.linked.append((os.path.basename(path), prompt["id"]))
|
||||
@@ -652,7 +650,7 @@ class TestLinkedOutputRoot(RecordedRuns, unittest.TestCase):
|
||||
handler = ImageGenerationHandler(db, make_tracker(db))
|
||||
handler.processing_delay = 0
|
||||
handler.settle_interval = 0.01
|
||||
self.addCleanup(close_handler, self, handler)
|
||||
self.addCleanup(handler.shutdown)
|
||||
linked = []
|
||||
handler.link_image_to_prompt = lambda path, prompt, meta: linked.append(
|
||||
prompt["id"]
|
||||
|
||||
@@ -152,30 +152,6 @@ 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)
|
||||
|
||||
@@ -976,156 +976,3 @@ 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_threshold_near_the_grid_is_saved_on_it(self):
|
||||
"""#369: a computed value a hair off the 0.05 grid is stored on it, so
|
||||
every page's step check accepts what /settings and the file return."""
|
||||
revision = (await self._get())["revision"]
|
||||
status, data = await self._post(
|
||||
"/prompt_manager/settings",
|
||||
{"expected_revision": revision, "autotag_general_threshold": 0.1 + 0.2},
|
||||
)
|
||||
self.assertEqual(status, 200, data)
|
||||
self.assertEqual(
|
||||
repr(self._read_config()["autotag"]["defaults"]["general_threshold"]),
|
||||
"0.3",
|
||||
)
|
||||
self.assertEqual(
|
||||
repr((await self._get())["settings"]["autotag_general_threshold"]), "0.3"
|
||||
)
|
||||
|
||||
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_general_threshold",
|
||||
0.37,
|
||||
"General-tag confidence must be a multiple of 0.05.",
|
||||
),
|
||||
(
|
||||
"autotag_character_threshold",
|
||||
0.33,
|
||||
"Character-tag confidence must be a multiple of 0.05.",
|
||||
),
|
||||
(
|
||||
"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)
|
||||
|
||||
@@ -32,7 +32,6 @@ 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/"
|
||||
@@ -738,27 +737,11 @@ class TestShutdown(AutoThumbnailCase):
|
||||
patch.object(auto, "CLOSE_WAIT_SECONDS", 0),
|
||||
self.assertLogs("prompt_manager.api", "WARNING") as logs,
|
||||
):
|
||||
self.assertFalse(self.api.close_auto_thumbnails())
|
||||
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):
|
||||
@@ -997,7 +980,6 @@ 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
|
||||
@@ -1040,7 +1022,7 @@ class TestMonitorQueuesThumbnails(AutoThumbnailCase):
|
||||
monitor.set_thumbnail_queue(None)
|
||||
self.assertIsNone(monitor.handler.thumbnail_queue)
|
||||
finally:
|
||||
stop_monitor(self, monitor)
|
||||
monitor.stop_monitoring()
|
||||
|
||||
|
||||
class TestScanHandsOffThumbnails(AutoThumbnailCase):
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user