Compare commits

..
Author SHA1 Message Date
Vito Sansevero ca4e259169 fix(nodes): pin the Classic wildcard behaviour and match queued ids by type (#343)
- Classic PromptManager/PromptManagerText expose _pmTextSerializerReady, settled
  once their own text serializer is installed; the real-host specs wait on it
  and assert the Classic text is queued as written, as on main.
- The Primitive fill only touches an output entry whose class_type matches the
  graph node at that id.
- TESTING.md notes the host's harmless "Invalid widget" warning per Run.
2026-10-25 17:22:47 -07:00
Vito Sansevero 10da6fc3a0 fix(nodes): resolve a legacy Primitive's text like the host, once per run (#343)
- Inputs still bound to their widget (Classic PromptManager/PromptManagerText)
  are left as the host resolved them; the injection no longer overwrites them.
- For the editors' detached text inputs, the Primitive's value goes through its
  "Run widget replace on values" replacement and the hidden widget's own
  serializer, so {a|b} wildcards and %date% tokens resolve as on main.
- Each queued output is filled once (graphToPrompt wrapper, or the queue hook
  for an output that did not come through it): a wildcard is drawn once.
- Editors inside subgraphs are found under their <subgraph>:<inner> id.
- Real-host runPrompt installs its request route before the run starts.
- Shared parser case pins the widget-table fallback for a core KSampler.
2026-10-25 17:02:36 -07:00
Vito Sansevero fef1fcf35f Merge remote-tracking branch 'origin/main' into feat/343-linkable-text 2026-10-25 16:26:41 -07:00
Vito Sansevero 2fadd765f6 fix(nodes): linked text review round for #343
- A legacy frontend Primitive linked to an editor's text or the Pair's
  negative now runs with its value: the graphToPrompt wrapper and the queue
  hook fill it in (through virtual reroutes), in new and pre-#343 workflows.
- Moving a linked text and releasing it on the editor keeps it on `text`:
  a link being moved counts as free, as the frontend counts it.
- A link dropped on a prompt box (or its link note) fills that box's input,
  so the Pair's Negative prompt box takes `negative_text` (Classic through
  getInputOnPos, Nodes 2.0 through the box's slot key).
- Only exact STRING links are steered; wildcards keep the frontend choice.
- A linked Classic text box gives its row back and gets it back on
  disconnect; the PromptManagerText note uses the theme colour.
- The metadata parsers map positional widget values through the widget
  table when an editor's widget-bound inputs do not line up with it.
- Real-host cases for the legacy Primitive (new and saved workflows, a real
  Run), the moved link, the Negative box, paste/clone/undo and Split.
2026-10-25 15:08:28 -07:00
Vito Sansevero d5d7f5e33d test(nodes): PromptSearchList text takes a dragged link like a core widget 2026-10-25 13:44:40 -07:00
Vito Sansevero 0e9cbde453 test(coverage): floor for the linked-text module 2026-10-25 13:33:00 -07:00
Vito Sansevero 0a80cfd9e5 Merge remote-tracking branch 'origin/main' into feat/343-linkable-text 2026-10-25 13:24:24 -07:00
Vito Sansevero bd2556f32f feat(nodes): let a link drive the prompt text of every PromptManager node
The Focused, Split and Pair editors hide their widgets, and ComfyUI only
draws a widget's input socket while the widget is visible. `text` (and the
Pair's `negative_text`) therefore had no socket, and a STRING link dropped
on the editor fell through to the first free STRING input, `prepend_text`,
which is how #342's workflow was wired.

- Detach those inputs from their hidden widgets on create and after a
  workflow load, so they are ordinary sockets in Classic and Nodes 2.0. The
  schema, widget order and saved values are unchanged.
- A STRING link dropped on the editor body fills the main prompt, then the
  negative, before the frontend's own choice.
- Each editor says when its text comes from a link; the typed text is kept
  for disconnect. The Classic nodes now save the typed text under a link
  instead of null, and "use prompt" no longer overwrites it while linked.
- Real-host tests: mouse drags in both renderers, a workflow saved before
  this change, and executions proving the linked text is saved, counted,
  recorded in history and linked to the run's image.

Closes #343
2026-10-25 13:24:19 -07:00
77 changed files with 683 additions and 10990 deletions
+6 -3
View File
@@ -2,15 +2,18 @@
"required_status_checks": {
"strict": true,
"contexts": [
"py3.10 on ubuntu-latest",
"py3.11 on ubuntu-latest",
"py3.12 on ubuntu-latest",
"py3.13 on ubuntu-latest",
"py3.14 on ubuntu-latest",
"py3.10 on windows-latest",
"py3.11 on windows-latest",
"py3.12 on windows-latest",
"py3.13 on windows-latest",
"py3.14 on windows-latest",
"py3.10 on macos-latest",
"py3.11 on macos-latest",
"py3.12 on macos-latest",
"py3.13 on macos-latest",
"py3.14 on macos-latest",
"coverage gate (85%)",
"black and flake8",
"frontend modules",
+2 -5
View File
@@ -19,17 +19,14 @@ jobs:
security:
name: bandit (high severity)
runs-on: ubuntu-latest
# ceil(1.5 x slowest of three green main runs) + 3, never under 5 (test.yml): runs 40257217665,
# 40254137380, 40240706583 took 23s / 16s / 23s -> 5 min.
timeout-minutes: 5
steps:
- uses: actions/checkout@v7
- name: Set up Python 3.12
- name: Set up Python 3.10
uses: actions/setup-python@v7
with:
python-version: "3.12"
python-version: "3.10"
- name: Install bandit
run: pip install bandit[toml]
+4 -21
View File
@@ -19,28 +19,15 @@ concurrency:
env:
PYTHONFAULTHANDLER: "1"
# Every job has a time limit: ceil(1.5 x the slowest of three recent green runs on main) + 3
# minutes, never under 5. The 3 minutes leave room for the Python runner's watchdog (120 s, see
# tests/run_suite.py) to print its stack dump before GitHub cancels the job. A limit changes only
# with new measurements written here (TESTING.md "Time limits and hangs").
# Measured 2026-10-26 on runs 40257217681 (3be0ada), 40240706568 (fb9134f), 40228167328 (f9835aa):
# py matrix, slowest py3.10 on windows-latest: 12m43s / 15m54s / 14m53s -> 27 min
# (ubuntu 5m27s-6m27s, macos 6m14s-8m30s, windows 11m42s-15m54s)
# 3.10 and 3.11 left the matrix and 3.14 joined it in #356; re-measure after three green main runs.
# real ComfyUI compatibility: 13m32s / 8m55s / 11m46s -> 24 min
# coverage gate (85%): 7m26s / 7m22s / 7m54s -> 15 min
# perf budget 30s / 29s / 27s, black and flake8 39s / 37s / 39s,
# frontend modules 23s / 19s / 29s -> 5 min (the floor)
jobs:
test:
name: py${{ matrix.python-version }} on ${{ matrix.os }}
runs-on: ${{ matrix.os }}
timeout-minutes: 27
strategy:
fail-fast: false
matrix:
os: [ubuntu-latest, windows-latest, macos-latest]
python-version: ["3.12", "3.13", "3.14"]
python-version: ["3.10", "3.11", "3.12", "3.13"]
steps:
- uses: actions/checkout@v7
@@ -58,12 +45,12 @@ jobs:
pip install -r requirements-dev.txt
- name: Run tests
run: python -m tests.run_suite -v
run: python -m unittest discover -s tests -v
comfyui:
name: real ComfyUI compatibility
runs-on: ubuntu-latest
timeout-minutes: 24
timeout-minutes: 25
steps:
- uses: actions/checkout@v7
with:
@@ -140,7 +127,6 @@ jobs:
coverage:
name: coverage gate (85%)
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- uses: actions/checkout@v7
@@ -157,7 +143,7 @@ jobs:
- name: Measure coverage
run: |
python -m coverage run -m tests.run_suite
python -m coverage run -m unittest discover -s tests
python -m coverage combine
python -m coverage report --fail-under=85
@@ -168,7 +154,6 @@ jobs:
name: perf budget
if: github.event_name == 'push' || github.event_name == 'workflow_dispatch'
runs-on: ubuntu-latest
timeout-minutes: 5
steps:
- uses: actions/checkout@v7
@@ -191,7 +176,6 @@ jobs:
lint:
name: black and flake8
runs-on: ubuntu-latest
timeout-minutes: 5
steps:
- uses: actions/checkout@v7
- uses: actions/setup-python@v7
@@ -204,7 +188,6 @@ jobs:
js-test:
name: frontend modules
runs-on: ubuntu-latest
timeout-minutes: 5
steps:
- uses: actions/checkout@v7
+1 -1
View File
@@ -22,7 +22,7 @@ repos:
- id: tests
name: tests
entry: python -m tests.run_suite -q
entry: python -m unittest discover tests/ -v
language: system
pass_filenames: false
stages: [pre-push]
+6 -10
View File
@@ -1,7 +1,7 @@
# ComfyUI Prompt Manager
![License](https://img.shields.io/badge/license-MIT-blue.svg)
![Python](https://img.shields.io/badge/python-3.12%2B-blue.svg)
![Python](https://img.shields.io/badge/python-3.10%2B-blue.svg)
![ComfyUI](https://img.shields.io/badge/ComfyUI-compatible-green.svg)
![Tests](https://github.com/ComfyAssets/ComfyUI_PromptManager/workflows/Tests/badge.svg)
![Code Quality](https://github.com/ComfyAssets/ComfyUI_PromptManager/workflows/Code%20Quality/badge.svg)
@@ -483,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.
![LoRA Results](images/pm-lora-filtered.png)
After import, filter by the `lora-manager` category to browse your LoRA collection.
@@ -928,12 +924,12 @@ db.model.backup_database("backup_prompts.db")
```bash
cd ComfyUI_PromptManager
pip install -r requirements-dev.txt
python -m tests.run_suite # Python (what CI runs: unittest discovery plus a hang watchdog)
python -m unittest discover -s tests # Python (what CI runs)
node --test tests/js/*.test.js # frontend pure modules
python -m coverage run -m tests.run_suite && python -m coverage combine && python -m coverage report
python -m coverage run -m unittest discover -s tests && python -m coverage combine && python -m coverage report
```
CI runs the suite on Linux, Windows and macOS for Python 3.12 to 3.14 with an 85% coverage gate. `python -m unittest discover -s tests` still works; see `TESTING.md` ("Time limits and hangs") for the watchdog. Tests that delete a database file first check that nothing in the process still holds it open (`tests/open_handles.py`), so a connection leak fails on Linux the way it would on Windows.
CI runs the suite on Linux, Windows and macOS for Python 3.10 to 3.12 with an 85% coverage gate. Tests that delete a database file first check that nothing in the process still holds it open (`tests/open_handles.py`), so a connection leak fails on Linux the way it would on Windows.
### Code Style
+9 -323
View File
@@ -31,7 +31,7 @@ its reset endpoint exists only in the test harness.
Use a Python environment containing `requirements-dev.txt` and Node.js 24:
```sh
python -m tests.run_suite -q # what CI runs; see "Time limits and hangs"
python -m unittest discover -s tests -q
npm ci --prefix tools/vendor --ignore-scripts
npx --prefix tools/vendor playwright install --with-deps chromium firefox webkit
npm run --prefix tools/vendor test
@@ -39,9 +39,6 @@ npm run --prefix tools/vendor test:coverage
npm run --prefix tools/vendor test:workflows -- --project firefox --project webkit
```
`python -m unittest discover -s tests` still works and finds the same tests;
it just has no hang watchdog.
`PM_TEST_PYTHON` selects the Python executable; set it when Python is outside
PATH. Every Playwright worker starts its own `tests/browser/server.py --port 0`,
which binds a port the OS picks and prints it as `PM_TEST_SERVER_PORT=<n>`, with
@@ -86,18 +83,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`.
@@ -219,62 +204,6 @@ not certify every existing feature, arbitrary host versions,
GPU/model changes, cross-pack installation or visual design changes. Those still
need targeted validation.
## Time limits and hangs
CI and the pre-push hook run the Python suite as `python -m tests.run_suite`
(`-v`, `-q`, `-s START` and `-p PATTERN` as in `unittest discover`), and the
coverage job as `python -m coverage run -m tests.run_suite`. The runner
(`tests/run_suite.py`) discovers the same tests with the same `sys.path` and
warning filters as `python -m unittest discover -s tests`, exits 0, 1, or 5 when
no tests ran (Python 3.12+ behaviour, on every version), and adds a watchdog: it writes `start: <test id>` to stderr before every test, at every
verbosity, and arms `faulthandler.dump_traceback_later`. If a test, its teardown,
its event-loop shutdown or interpreter exit stops making progress, the stack of
every thread goes to stderr (`Timeout (0:02:00)!` then `Thread 0x...` blocks) and
the process exits non-zero.
- `PM_TEST_HANG_LIMIT` (default 120 s) bounds each test from its `start:` line to
the next one, including the class and module fixtures (`tearDownClass`,
`setUpModule`, ...) that run in between. `PM_TEST_SHUTDOWN_LIMIT` (default
120 s) bounds the end of the run: the summary, non-daemon threads, atexit
handlers and coverage's save. Both must be
positive integers; anything else stops the runner with exit 2. 120 s stays below
asyncio's 300 s wait for default-executor threads on Python 3.12 and 3.13 (3.10
and 3.11 wait without limit), so a hang in an `IsolatedAsyncioTestCase` loop
shutdown is dumped inside the test that hung.
- Reading a dump: the last `start:` line names the test. An `ok` after it means
the test body passed and the hang came after it: in its teardown or event-loop
shutdown (where `IsolatedAsyncioTestCase` joins executor threads), or in a
class or module fixture before the next test; the stack shows which. A dump
after the summary is a hang at interpreter exit. The dump shows thread stacks, not asyncio task
stacks.
- When the watchdog fires under coverage, that run's coverage data is lost; the
job fails anyway and the dump is what matters.
Every job in `test.yml`, `code-quality.yml` and `viewerjs-check.yml` has
`timeout-minutes` (`tests/vendor/ci-workflow.test.mjs` checks it). The rule is
ceil(1.5 x the slowest of three recent green runs on `main`) + 3 minutes, never
under 5; the three run ids, their times and the resulting limits are written
in a comment above `jobs:`. The 3 minutes leave room for one watchdog firing, so
the dump prints before GitHub cancels the job. A job that hits its limit, or a
watchdog dump, is a failure like any other: read the dump, find the blocking
call, fix it. Never raise a limit or rerun to get green; a limit changes only
with new measurements written next to it.
### After a push
Follow the run for the commit you pushed, never an older one:
```sh
gh run list --commit <pushed sha> --workflow test.yml # the run id for that commit
gh run watch <run id> --exit-status # returns when the run ends
gh run view <run id> --json jobs # shows a job that already failed
gh run view --job <job id> --log-failed # that job's failed log, at once
```
Start fixing as soon as a failed log shows the cause, but push the fix only after
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.
## Performance budgets
A spec performance budget gets two tests:
@@ -306,122 +235,6 @@ change the budget in the spec with a reason) in a follow-up PR. Never rerun it t
green. It is not a required check, and it cannot be one while it skips pull
requests (`.github/branch-protection.json`).
## Images listing cache (#348)
`/prompt_manager/images/output`, `/prompt_manager/gallery/subfolders` and
`POST /prompt_manager/gallery/rescan` read one shared listing per output root
(`RootListing` in `py/api/gallery_listing.py`, `_listing` in
`py/api/images.py`). The tests are `tests/test_gallery_listing.py` (pure
helpers) and `TestSharedListing` in `tests/test_api_images.py` (routes); they
hold a walk open with a gated scan instead of sleeping.
- **Single flight.** At most one walk per root runs at a time; every reader
that needs that root while it runs waits for it (shielded, so a client that
disconnects never cancels it). The walk, the comparison with the previous
listing and both folder sets run in the executor (`walk_root`). A walk that
finds the same `(path, mtime)` list keeps the previous list object, so the
favorites lookup and the sort orders are not rebuilt.
- **30 s background refresh.** A listing older than `REFRESH_AFTER` (30 s) is
served at once and one background walk replaces it. Files added or removed
by hand therefore show up one open later than before; Rescan shows them at
once.
- **Blocking invalidation after deletes.** Duplicate cleanup and the legacy
duplicate delete still call `invalidate_gallery_cache()`: the listing is
dropped and the next read waits for a walk, so a deleted file is never
listed. A walk running during the drop is discarded (the record's
`generation` outlives every drop); a cold reader gives up with a 500 after 3
discarded walks in a row. Thumbnail writes (jobs, conversion, the legacy
generate and clear routes) no longer drop the listing: thumbnail URLs are
computed per request.
- **Startup warm-up.** The root `__init__.py` registers
`api.warm_gallery_listings` on the host's `on_startup`, which starts one
background walk per root, and `api.cleanup_gallery_listings` on
`on_cleanup`, which cancels running walks. A warm-up failure never stops
startup; the first read retries.
- **Failures and Rescan.** `walk_root` opens the root with `os.scandir` first,
so a deleted or unreadable root fails the walk instead of listing nothing.
A failed walk keeps the previous listing, logs one warning per streak of
failures and is retried at most once per `REFRESH_AFTER`; a reader with no
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
`node tools/vendor/perf/images_load_probe.mjs http://127.0.0.1:8188` against
a running ComfyUI. It prints the time to the first decoded thumbnail, the
time to network idle and the 10 slowest requests, from a fresh browser
context each run.
## Windows portability preflight
Before changing paths, SQLite, streams, restore or cleanup:
@@ -989,8 +802,8 @@ Recorded shell deviations:
`file_size` equal to the file's size in this walk. Both paths go through
`normalize_image_path` (absolute, normalized, `os.path.normcase`); a row without
`file_path` matches by its `image_path`.
- Read again: files linked without a pair (the monitor's queue and run-key links,
fallback links made before #342, the legacy positive-only import), linked to more than one prompt or pair, not linked, a
- Read again: files linked without a pair (the monitor's fallback and FIFO links, the
legacy positive-only import), linked to more than one prompt or pair, not linked, a
file whose size changed (one replaced at the same path, e.g. a reused ComfyUI
counter name, is read and its conflict lands under Needs review), and a row with no
stored size (unknown is not known). The scan now stores each file's size when it
@@ -3422,40 +3235,14 @@ refused while a restore runs; the review is a request lease ("LoRA import review
accepts the optional `expected_revision`. LoraManager detection also searches every
`custom_nodes` folder ComfyUI registers through `folder_paths` (#144), main folder first.
**Where LoraManager keeps its files (#355).** `py/lora_manager_layout.py` repeats
LoraManager's rules without importing it (`tests/test_lora_manager_layout.py`). The
settings file is `<LORA_MANAGER_SETTINGS_DIR>/settings.json`, else the plugin's own
`settings.json` when portable mode is on (`LORA_MANAGER_PORTABLE` `1`/`0` decides, other
values are ignored; otherwise the JSON value `use_portable_settings: true`), else
`platformdirs.user_config_dir("ComfyUI-LoRA-Manager")`. It is read as UTF-8, at most 5 MB,
and only seven keys are kept (`SETTINGS_KEYS`); a missing file means LoraManager's
defaults. Discovery: in `alongside` mode every LoRA root is searched and a sidecar is kept
when `<base>.safetensors` exists beside it; in `centralized` mode only the sidecar root is
searched and a sidecar is kept when its `file_path` exists inside a LoRA root (unparseable
ones are skipped by name). `lora_files` counts the metadata files seen, `with_metadata`
the kept, readable LoRAs and `skipped` the rest. Media per LoRA, without repeats: the
preview (`preview_url`, else LoraManager's `PREVIEW_EXTENSIONS` order), then every file in
`<example-images root>[/<library>]/<sha256>/` in natural order; a leading video yields to the
first still. A file counts only with a servable extension and inside the guarded media
roots. CivitAI downloads run only for a LoRA without local media. The folder guard keeps a
settings folder only when it exists and is neither a filesystem root (`/`, `C:\`, a UNC
share root) nor the home folder or one of its ancestors (realpath + `normcase`). An
unusable layout (no install, unreadable / non-JSON / non-object / oversized settings, no
`platformdirs`, a missing or refused centralized sidecar root) gives a review without a
token ("LoraManager settings: not usable: <reason>. Nothing was changed."), a legacy scan
error with that reason, the default LoRA roots for serving, and trigger words from the
default roots in alongside mode. `tests/run_suite.py` and `tests/browser/server.py` point
`LORA_MANAGER_SETTINGS_DIR` at a disposable folder, so no test reads the developer's real
LoraManager settings; tests that call `load_layout` inject their own folder too.
`tests/browser/workspace_integrations_fixture.py` (`/__test__/integrations/*`) makes a fake
LoraManager install and LoRA folder under the server's temp root, holds or fails the
import, saves a blocking pair and answers CivitAI downloads with a stub opener, so no
test reaches the network.
```sh
python3 -m unittest tests.test_lora_manager_layout tests.test_lora_import tests.test_api_lora \
tests.test_lora_database tests.test_lora_utils -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 \
@@ -3565,14 +3352,14 @@ generation did not change meanwhile (`PromptModel.get_connection`).
Writes racing the switch: a background write that was waiting on the old file's lock
when the switch closed its connection fails, and run counting (`count_runs`,
`count_pair_runs`) and the image monitor's link (`ImageGenerationHandler.
link_image_to_prompt`: the resolved-pair ingest and the plain prompt link)
link_image_to_prompt`: the resolved-pair ingest and the fallback or recent-prompt link)
run it once more through `PromptDatabase.survive_library_switch` when the model's
generation changed during the attempt; the retry opens the new library. The monitor
links through `PromptDatabase.link_image_to_prompt_or_raise`, which raises database
errors (`link_image_to_prompt` keeps its log-and-return-0 contract for other callers);
each attempt is one transaction, so a failed attempt rolls back and a committed one
does not retry, and the image is linked once. `TestBackgroundWritesSurviveTheSwitch`
races a real UPDATE (run counting) and a real plain image link against the lock and
races a real UPDATE (run counting) and a real fallback image link against the lock and
asserts each lands in the new file only. Any other failure is not retried. Known limit: that
blocked writer, and the mover with it, waits out the writer's busy timeout (30 s)
before its connection closes: SQLite holds the connection's mutex during the busy
@@ -4171,106 +3958,6 @@ Viewer, including nested modal focus, mobile layouts and per-role sending.
real Classic and Nodes 2.0 widgets and verifies workflow serialization without
loading a workflow or queueing generation.
## Image attribution: certain or unlinked (#342)
The image monitor and the output scan link an output only to a prompt that
provably produced it; anything uncertain stays unlinked and the monitor logs it
at info ("No certain prompt for <file>; left unlinked"). There is no recency
fallback (`_get_fallback_prompt` is gone) and no unproven tracker entry.
- **Empty editor.** Prepend and append text are never stored as the prompt
(maintainer's product decision; #343 tracks letting a link drive `text`). A
PromptManager, Text, Focused, Split or Pair node whose `text` is empty saves,
counts and tracks nothing for the positive side, logs at info (no "Could not
save or track" warning), and marks its run key empty, so no earlier entry of
that key can be used and a batch item keeps its queue place
(`PromptTracker.mark_empty`). The queue hook counts nothing for it either. A
save or tracking failure of a positive node marks its run key the same way, so
an older entry of that key never stands in for the failed run.
- **Run keys** (`utils/run_keys.py`): sha256 of a prompt node plus its whole
upstream closure (`class_type` and `inputs`) in the API graph. Nodes register
tracker entries under it; the monitor computes it from the image's embedded
graph. The snapshot taken when the file appears names the key's latest run; a
queued entry is taken only by key and from that run (batch items in order, a
placeholder leaves its image unlinked, same-key entries of earlier runs ahead
of it are stale and dropped, entries of later runs that executed during the
monitor's delay wait for their own images). Without a queued entry the
snapshot's entry itself is used only when its run id is known
(`comfy_execution.utils.get_executing_context`) and that run registered one
entry for the key, which covers a cached re-run (sampler seed changed) without
applying a batch's last item to every image. When the key's latest entry is a
placeholder the image stays unlinked without consulting history (the same
rule as a popped placeholder). More than one positive PromptManager node in
the image graph, or any positive node without a key: never linked from the
tracker. Links are followed exactly as ComfyUI's `is_link` defines them (a
`str` node id and an `int` or `float` slot in a two-item list). Equal keys do
not always mean equal text: the key can't see `IS_CHANGED`, so an upstream
node that forces re-runs gives two runs one key; that is why the run named by
the snapshot is required as well. A file event delivered after the next run
already re-registered the key is the one remaining exposure (see the ledger).
- **Recorded history** (`database/run_attribution.py`): a file is attributed when
exactly one recorded run reported its output location (root key, subfolder,
name), its mtime lies within that run (dequeue to terminal or callback, +-2 s),
and every positive context names the same stored prompt (the live
`positive_prompt_id` hint). The monitor tries it last; the scan tries it when the
#170 browser-parser floor finds nothing, and also for files with no prompt
metadata at all (`--disable-metadata`, savers that embed none); it never
creates a prompt. Lookups are batched (`stored_run_prompts`): one
parameterized `filename IN (...)` query per chunk of at most 500 names
(`LOOKUP_CHUNK`), matched in memory, so a scan batch asks history once per
chunk, never once per file. The history schema is unchanged (no extra index),
so older builds still open the library.
- **After the run ends.** An image with no run keys (no embedded API graph) that
nothing identifies waits for its run: when the run's outputs are recorded
(`write_execution_outputs_for_finish` notifies `subscribe_outputs_sealed`
listeners after commit), each waiting image the run reported gets exactly one
history-only re-check on the monitor's worker. At most 1000 images wait, each
for at most an hour; shutdown forgets them. The image is registered before the
first history check, so a run sealed in between is not missed; a re-check that
can't be scheduled yet (the image's own pass still pending, or the pending set
full) is kept due and scheduled when a pending image finishes. A keyless image
outside ComfyUI's output folder never waits (history can't attribute it).
- **Spellings.** History keeps ComfyUI's own spelling of its output folder (the
root key). A file is matched through that spelling first, then with both sides
resolved (`realpath`), so an output folder reached through a symlink or a
Windows junction still attributes the scan's resolved file paths; waiting
files are keyed by their resolved path for the same reason. The scan makes
its batch's history lookup only when the first file in the batch needs it.
- **Blank positive.** A resolved pair whose stored positive is empty (the #342
image with its Pair execution record) is `unresolved` for the scan and for
`ingest_resolved_pair` (it used to raise in `_save_pair`).
- **Parity floor.** Metadata linking, the scan's resolver and parser floor are
unchanged for every non-blank result; only guesses are gone.
Tests: `tests/test_prompt_attribution.py` (nodes, tracker, monitor, run keys),
`tests/test_run_attribution.py` (history attribution, monitor and scan),
`tests/test_review_image_links_tool.py`, and the real-host
`tests/comfyui/image-attribution.spec.mjs` (both renderers): a typed Pair image
links, a Pair image whose positive is only linked prepend text from an unknown
node stays unlinked with nothing stored, an unknown node feeding `text` links
through its run key, and its cached re-run with a new sampler seed links again
(the host's `execution_cached` event names the PromptManager node, under the
fixture's Classic cache). A second case saves with metadata disabled and holds
the run open behind the fixture gate after SaveImage: the output is waiting and
unlinked while the run is open, and links to the stored prompt once the run's
outputs are recorded. The managed host config disables monitoring, so these
cases start and stop the plugin's monitor through the owned fixture route
`/pm_execution_fixture/monitor` (its GET lists the waiting files).
On origin/main the prepend image was linked to an unrelated prompt.
Repair: `tools/review_image_links.py <db> --output-dir <ComfyUI output> --ids
<ids>` (or `--prompt <id>`) judges links by these rules (`keep`, `relink`,
`unlink`, `conflict`, `missing`) on a backup copy and changes nothing.
`--apply` writes `relink`/`unlink` only for ids named with `--ids`, after saving
`<db>.pre-relink_<stamp>`, and only while each image is still linked to the
prompt it was judged on (otherwise it is counted as skipped). `relink` points the
link record at the other prompt and drops its pair links. `unlink` removes only
the image's link record (its `generated_images` row, the metadata snapshot stored
in it and its pair links): that is how the app represents an unlinked image, the
same as a file the scan could not attribute, so the file stays on disk and in the
gallery as unlinked, and its favorite and tags (kept by file) are untouched. Every
run prints this. Queue-linked batch images can't be proven after the fact and read
`unlink`, so review before applying. Stop ComfyUI first.
## Linked prompt text (#343)
A link drives `text` on every PromptManager node, and `negative_text` on the Pair,
@@ -4353,9 +4040,8 @@ PM_COMFY_ROOT=... PM_TEST_PYTHON=... npm run --prefix tools/vendor test:executio
## Durable execution ledger
`npm run --prefix tools/vendor test:executions` runs 52 real execution cases
(26 each in Classic and Nodes 2.0, retaining all 48 prior execution cases and adding
the two #342 image-attribution cases),
`npm run --prefix tools/vendor test:executions` runs 46 real execution cases
(23 each in Classic and Nodes 2.0, retaining all 42 prior execution cases),
alongside the existing 80 stock host cases (the original 78 plus the native
editor host case) and six pack compatibility cases. Run
these suites serially because they share report cleanup and owned host resources. Use the same pinned disposable `PM_COMFY_ROOT`,
-9
View File
@@ -100,11 +100,6 @@ try:
if host_app is not None and hasattr(host_app, "on_cleanup"):
host_app.on_cleanup.append(api.cleanup_history_reads)
host_app.on_cleanup.append(api.cleanup_annotation_captures)
host_app.on_cleanup.append(api.cleanup_gallery_listings)
if host_app is not None and hasattr(host_app, "on_startup"):
# Build each output root's Images listing once the server's loop runs,
# so the first open does not wait for the walk (#348).
host_app.on_startup.append(api.warm_gallery_listings)
if host_app is not None and hasattr(host_app, "on_shutdown"):
# Open Needs review change streams end before the server waits on them.
host_app.on_shutdown.append(api.close_review_events)
@@ -149,9 +144,6 @@ try:
# Every settled output's preview is queued for the API's background
# worker (#249): the workspace never asks to generate thumbnails.
_global_image_monitor.set_thumbnail_queue(_review_api.queue_capture_thumbnail)
# Every settled output walks its Images listing again (#348): the
# API hops from the monitor's thread onto the server's loop.
_global_image_monitor.set_output_listener(_review_api.notify_output_settled)
_review_host = getattr(config.server_instance, "app", None)
if _review_host is not None and hasattr(_review_host, "on_cleanup"):
@@ -160,7 +152,6 @@ try:
):
monitor.set_review_recorder(None)
monitor.set_thumbnail_queue(None)
monitor.set_output_listener(None)
api.close_auto_thumbnails(final=True)
_review_host.on_cleanup.insert(0, _stop_review_recording)
+1 -4
View File
@@ -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})
-6
View File
@@ -4,8 +4,6 @@ import sqlite3
from contextlib import ExitStack
from typing import Callable, ContextManager
from .run_attribution import notify_outputs_sealed
try:
from ..utils.execution_history_events import RunFinish, encode_command
from ..utils.execution_outputs import OutputBundle, validate_output_bundle
@@ -115,10 +113,6 @@ class ExecutionOutputOperations:
"VALUES(?,?,?,?,?,?,?,?)",
[(run_id, *reference) for reference in prepared],
)
# Committed: images waiting for this run may now be attributed (#342)
notify_outputs_sealed(
run_id, [dict(zip(_REFERENCE_FIELDS, row)) for row in prepared]
)
return True
def get_execution_outputs(self, run_id: str) -> dict:
-16
View File
@@ -23,26 +23,10 @@ except ImportError:
from utils.prompt_pair_resolver import RESOLVER_VERSION
def blank_positive(resolution: Dict[str, Any]) -> bool:
"""A resolved pair whose stored positive text is empty (#342).
A PromptManager node whose editor is empty encodes only its prepend and
append text, which is never stored as a prompt: such a result names no
prompt to link.
"""
if resolution.get("state") != "resolved":
return False
positive = resolution.get("positive") or {}
text = positive.get("storage_text", positive.get("text"))
return not (isinstance(text, str) and text.strip())
def ingest_resolved_pair(
db: Any, image_path: str, metadata: Dict[str, Any], resolution: Dict[str, Any]
) -> Dict[str, Any]:
state = resolution.get("state", "unresolved")
if state == "resolved" and blank_positive(resolution):
state = "unresolved"
if state != "resolved":
return {"status": state, "pairs_added": 0, "links_added": 0}
with db._pair_transaction() as conn:
-244
View File
@@ -1,244 +0,0 @@
"""Attribute an output file to a stored prompt through recorded history (#342).
Execution history records, per run, the files ComfyUI reported
(``execution_run_outputs``) and the PromptManager texts the run used
(``execution_run_contexts``, whose ``positive_prompt_id`` hint is set only
while the raw text matches exactly one stored prompt). A file is attributed
only when all of this holds:
- exactly one recorded run reported the file's output location (same output
root, subfolder and name);
- the file was last written while that run was executing, from its dequeue
to its terminal event (else the finish callback): a later run that history
did not record may have reused the name;
- every positive context of that run names the same stored prompt.
Anything else returns None: the caller leaves the file unlinked. Read-only.
Lookups are batched (``stored_run_prompts``): one parameterized query per
chunk of at most ``LOOKUP_CHUNK`` distinct file names, matched in memory. The
history schema is left exactly as it is (no extra index), so older builds can
still open the library.
"""
import logging
import os
import sqlite3
import threading
import weakref
from typing import Any, Callable, Dict, Iterable, List, Optional
try:
from ..utils.execution_outputs import snapshot_output_roots
from ..utils.library_location import import_folder_paths
except ImportError:
from utils.execution_outputs import snapshot_output_roots
from utils.library_location import import_folder_paths
# File times may trail or lead the recorded run times a little (coarse file
# system clocks, the write finishing just after the host reported it).
MTIME_SLACK_MS = 2000
# Distinct file names per lookup query: well under SQLite's bound-parameter
# limit (999 on old builds), so a scan batch costs one query per chunk.
LOOKUP_CHUNK = 500
_RUNS_FOR_NAMES = (
"SELECT DISTINCT o.filename, o.subfolder, o.root_key, o.path_style, o.run_id, "
"r.observed_dequeue_at, COALESCE(r.terminal_event_at, r.observed_callback_at) "
"FROM execution_run_outputs o JOIN execution_runs r ON r.run_id = o.run_id "
"WHERE o.folder_type = 'output' AND o.filename IN ({marks})"
)
def host_output_directory() -> Optional[str]:
"""ComfyUI's output directory, or None outside ComfyUI."""
folder_paths = import_folder_paths()
if folder_paths is None:
return None
try:
value = folder_paths.get_output_directory()
except Exception:
return None
return value if isinstance(value, str) and value else None
def _inside(image_path: str, output_root: str) -> Optional[str]:
"""``image_path`` relative to ``output_root``, or None when outside it."""
try:
relative = os.path.relpath(image_path, output_root)
except ValueError: # another drive on Windows
return None
if os.path.isabs(relative) or relative == os.pardir:
return None
if relative.startswith(os.pardir + os.sep):
return None
return relative
def _resolved(path: str) -> str:
try:
return os.path.realpath(path)
except (OSError, ValueError):
return os.path.abspath(path)
def _location(image_path: str, output_root: str, path_style: str):
"""(root key, subfolder, filename) as history records them, or None.
The root key is always the host's own spelling of its output folder,
because history recorded that one. The file may be spelled through the
same folder or through its resolved path (a symlinked or junction output
folder, as scan roots are): the lexical spelling is tried first, then
both sides resolved.
"""
root_key = snapshot_output_roots(path_style, output_root, output_root).output_key
if root_key is None:
return None
relative = _inside(os.path.abspath(image_path), os.path.abspath(output_root))
if relative is None:
relative = _inside(_resolved(image_path), _resolved(output_root))
if relative is None:
return None
subfolder, filename = os.path.split(relative)
if not filename:
return None
return root_key, subfolder, filename
def in_output_folder(image_path: str, output_root: Optional[str]) -> bool:
"""Whether history could ever attribute ``image_path`` (inside the host's
output folder, in either spelling)."""
if not image_path or not output_root:
return False
style = "windows" if os.name == "nt" else "posix"
return _location(image_path, output_root, style) is not None
def _written_during(image_path: str, started: Any, finished: Any) -> bool:
if not isinstance(started, int) or not isinstance(finished, int):
return False
try:
written = os.stat(image_path).st_mtime_ns // 1_000_000
except OSError:
return False
return started - MTIME_SLACK_MS <= written <= finished + MTIME_SLACK_MS
def stored_run_prompt(
db: Any,
image_path: str,
output_root: Optional[str],
path_style: Optional[str] = None,
) -> Optional[int]:
"""The stored prompt id history attributes to ``image_path``, else None."""
if not image_path:
return None
return stored_run_prompts(db, [image_path], output_root, path_style).get(image_path)
def _runs_by_location(db: Any, names: List[str]) -> Optional[Dict[tuple, dict]]:
"""``{(filename, subfolder, root_key, style): {run_id: (start, end)}}``.
One query per ``LOOKUP_CHUNK`` names; None without history tables.
"""
found: Dict[tuple, dict] = {}
conn = db.model.get_connection()
for start in range(0, len(names), LOOKUP_CHUNK):
chunk = names[start : start + LOOKUP_CHUNK]
query = _RUNS_FOR_NAMES.format(marks=",".join("?" for _ in chunk))
try:
rows = conn.execute(query, chunk).fetchall() # nosec B608
except sqlite3.Error:
return None # no history tables (a library without history)
for name, subfolder, root_key, style, run_id, began, ended in rows:
runs = found.setdefault((name, subfolder, root_key, style), {})
runs[run_id] = (began, ended)
return found
def _named_prompt(db: Any, run_id: str) -> Optional[int]:
contexts = db.get_execution_contexts(run_id).get("contexts") or []
named = {
context.get("positive_prompt_id")
for context in contexts
if context.get("positive") is not None
}
if len(named) != 1:
return None
(prompt_id,) = named
return prompt_id if isinstance(prompt_id, int) else None
def stored_run_prompts(
db: Any,
image_paths: Iterable[str],
output_root: Optional[str],
path_style: Optional[str] = None,
) -> Dict[str, Optional[int]]:
"""``{path: stored prompt id or None}`` for many files in few queries.
The output lookup is one query per ``LOOKUP_CHUNK`` distinct names; each
run that alone reported a file is then read once for its contexts.
"""
paths = list(dict.fromkeys(path for path in image_paths if path))
result: Dict[str, Optional[int]] = dict.fromkeys(paths)
if not paths or not output_root:
return result
style = path_style or ("windows" if os.name == "nt" else "posix")
located = {}
for path in paths:
location = _location(path, output_root, style)
if location is not None:
located[path] = location
if not located:
return result
names = sorted({filename for _key, _sub, filename in located.values()})
runs = _runs_by_location(db, names)
if not runs:
return result
prompts: Dict[str, Optional[int]] = {}
for path, (root_key, subfolder, filename) in located.items():
reported = runs.get((filename, subfolder, root_key, style), {})
if len(reported) != 1:
continue
((run_id, (began, ended)),) = reported.items()
if not _written_during(path, began, ended):
continue
if run_id not in prompts:
prompts[run_id] = _named_prompt(db, run_id)
result[path] = prompts[run_id]
return result
_LOGGER = logging.getLogger("prompt_manager.run_attribution")
_SEALED_LISTENERS: List[weakref.WeakMethod] = []
_SEALED_LOCK = threading.Lock()
def subscribe_outputs_sealed(method: Callable[[str, List[dict]], None]) -> None:
"""Call bound ``method(run_id, references)`` after a run's outputs are recorded.
Held weakly: a listener that is garbage collected is dropped.
"""
with _SEALED_LOCK:
_SEALED_LISTENERS.append(weakref.WeakMethod(method))
def notify_outputs_sealed(run_id: str, references: Iterable[dict]) -> None:
"""Tell listeners a run's output references were just committed.
A listener failure is logged and never reaches the history writer.
"""
references = list(references)
with _SEALED_LOCK:
alive = [ref for ref in _SEALED_LISTENERS if ref() is not None]
_SEALED_LISTENERS[:] = alive
listeners = [ref() for ref in alive]
for listener in listeners:
if listener is None:
continue
try:
listener(run_id, references)
except Exception:
_LOGGER.warning("Output seal listener failed", exc_info=True)
-3
View File
@@ -238,9 +238,6 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
except Exception as e:
# Log error but don't fail the encoding
self.logger.warning(f"Failed to save prompt to database: {e}")
self._mark_untracked(prompt, unique_id)
else:
self._skip_empty_prompt(prompt, unique_id)
# Perform standard CLIP text encoding using the combined text
self.logger.debug(
+3 -60
View File
@@ -24,7 +24,6 @@ try:
from .utils.prompt_tracker import get_prompt_tracker
from .utils.hashing import generate_prompt_hash
from .utils.prompt_graph import resolve_text, run_prompt_nodes
from .utils.run_keys import executing_run_id, prompt_node_key
from .utils.usage_tracking import PENDING_FIRST_USE
except ImportError:
import sys
@@ -36,7 +35,6 @@ except ImportError:
from utils.prompt_tracker import get_prompt_tracker
from utils.hashing import generate_prompt_hash
from utils.prompt_graph import resolve_text, run_prompt_nodes
from utils.run_keys import executing_run_id, prompt_node_key
from utils.usage_tracking import PENDING_FIRST_USE
@@ -170,67 +168,14 @@ class PromptManagerBase:
prompt_id = self._save_prompt_to_database(
text=text, category=category, tags=tags, count_run=batch_item
)
if role["is_run_prompt"] and not prompt_id:
# The save failed: an older entry of this key must not stand in (#342)
self._mark_untracked(prompt_graph, unique_id)
elif role["is_run_prompt"]:
if prompt_id and role["is_run_prompt"]:
self.prompt_tracker.set_current_prompt(
prompt_text=encoding_text,
additional_data={**additional_data, "prompt_id": prompt_id},
push_to_queue=batch_item or not role["known"],
**self._run_identity(prompt_graph, unique_id),
)
return prompt_id
@staticmethod
def _run_identity(prompt_graph: Any, unique_id: Any) -> Dict[str, Any]:
"""The tracker's proof that an image belongs to this execution (#342)."""
return {
"run_key": prompt_node_key(prompt_graph, unique_id),
"run_id": executing_run_id(),
}
def _skip_empty_prompt(self, prompt_graph: Any, unique_id: Any) -> None:
"""The ``text`` editor is empty: save and track nothing (#342).
Prepend and append text are deliberately never stored as the prompt, so
an empty editor has no prompt to save, count or link. For the run's
positive node this is logged at info and the node's run key is marked
empty (``_mark_untracked``).
"""
self._mark_untracked(
prompt_graph,
unique_id,
"Prompt text is empty (prepend and append text are not stored): "
"nothing saved or tracked, and this run's images stay unlinked",
)
def _mark_untracked(
self, prompt_graph: Any, unique_id: Any, message: Optional[str] = None
) -> None:
"""This execution of a positive node has no prompt to link (#342).
An empty editor, or a save or tracking failure. The node's run key is
marked with a placeholder, so no earlier prompt of that key can be
linked to this run's images, and a batch item keeps its queue place.
``message`` is logged at info for a positive node.
"""
try:
role = self._prompt_role(prompt_graph, unique_id, "")
if not role["is_run_prompt"]:
if message:
self.logger.debug("Empty text on a non-positive node; not saved")
return
if message:
self.logger.info(message)
self.prompt_tracker.mark_empty(
push_to_queue=role["known"],
**self._run_identity(prompt_graph, unique_id),
)
except Exception as error:
# Tracking must never block generation
self.logger.warning(f"Could not mark the untracked prompt: {error}")
def _record_use(self, prompt_id: int, times: int = 1) -> None:
"""Count runs; a tracking failure must never block generation."""
try:
@@ -268,10 +213,8 @@ class PromptManagerBase:
cache = get_trigger_cache()
# Lazy-load cache on first use; a failed load waits before the next try
# (#355: never a rescan of LoraManager's folders on every run).
retry_due = getattr(cache, "retry_due", lambda: True)
if not cache.is_loaded and IntegrationConfig.LORA_MANAGER_PATH and retry_due():
# Lazy-load cache on first use
if not cache.is_loaded and IntegrationConfig.LORA_MANAGER_PATH:
cache.load(IntegrationConfig.LORA_MANAGER_PATH)
modified, injected = inject_trigger_words(text, cache)
+77 -117
View File
@@ -117,29 +117,85 @@ class PromptManagerPair(PromptManagerFocused):
)
is_positive = "text" in used_fields
pair_used = node_id in intrinsic_pair_nodes(prompt) if known else True
if not text.strip():
# Nothing to store: prepend/append are never the stored prompt (#342)
self._skip_empty_prompt(prompt, unique_id)
else:
try:
pair_id, prompt_id = self._save_and_track_pair(
text,
try:
pair_id = self.db.save_prompt_pair(
text,
negative_text,
category=category.strip() or None,
tags=self._parse_tags(tags),
)
pair = self.db.get_prompt_pair(pair_id)
prompt_id = pair["positive_prompt_id"]
resolved_positive = (
resolve_text_input(prompt, node_id, "text") if known else None
)
resolved_negative = (
resolve_text_input(prompt, node_id, "negative_text") if known else None
)
positive_seen = resolved_positive is not None and generate_prompt_hash(
resolved_positive
) == generate_prompt_hash(text)
pair_seen = (
positive_seen
and resolved_negative is not None
and generate_prompt_hash(resolved_negative)
== generate_prompt_hash(negative_text)
)
positive_pending = getattr(self, "pending_first_use", PENDING_FIRST_USE)
pair_pending = getattr(
self, "pending_pair_first_use", PENDING_PAIR_FIRST_USE
)
counted = set()
for field, component_text, resolved_text, component_id in (
("text", text, resolved_positive, prompt_id),
(
"negative_text",
negative_text,
positive,
category=category,
tags=tags,
prepend_text=prepend_text,
append_text=append_text,
prompt=prompt,
unique_id=unique_id,
used_fields=used_fields,
is_positive=is_positive,
pair_used=pair_used,
resolved_negative,
pair["negative_prompt_id"],
),
):
component_hash = generate_prompt_hash(component_text)
if (
field not in used_fields
or not component_id
or component_hash in counted
):
continue
counted.add(component_hash)
seen_at_queue = (
resolved_text is not None
and generate_prompt_hash(resolved_text) == component_hash
)
except Exception as error:
self.logger.warning(f"Could not save or track paired prompt: {error}")
if is_positive:
self._mark_untracked(prompt, unique_id)
runs = positive_pending.consume(component_hash) + int(not seen_at_queue)
if runs:
self.db.record_prompt_use(component_id, times=runs)
if pair_used:
runs = pair_pending.consume(pair_run_key(text, negative_text)) + int(
not pair_seen
)
if runs:
self.db.record_prompt_pair_use(pair_id, times=runs)
if is_positive:
self.prompt_tracker.set_current_prompt(
prompt_text=positive,
additional_data={
"prompt_id": prompt_id,
"pair_id": pair_id if pair_used else None,
"positive_text": positive,
"negative_text": negative_text,
"positive_storage_text": text.strip(),
"negative_storage_text": negative_text.strip(),
"prepend_text": prepend_text.strip(),
"append_text": append_text.strip(),
"final_text": positive,
"category": category.strip() or None,
"tags": self._parse_tags(tags) or [],
},
push_to_queue=not positive_seen or (pair_used and not pair_seen),
)
except Exception as error:
self.logger.warning(f"Could not save or track paired prompt: {error}")
try:
self.comfyui_integration.register_prompt(
f"promptmanager_pair_{time.time_ns()}",
@@ -167,102 +223,6 @@ class PromptManagerPair(PromptManagerFocused):
)
return result
def _save_and_track_pair(
self,
text: str,
negative_text: str,
positive: str,
*,
category: str,
tags: str,
prepend_text: str,
append_text: str,
prompt: Optional[Dict[str, Any]],
unique_id: Optional[str],
used_fields: set,
is_positive: bool,
pair_used: bool,
) -> Tuple[Optional[int], Optional[int]]:
"""Save the recipe, count its runs and register the positive for linking."""
known = isinstance(prompt, dict) and unique_id is not None
node_id = str(unique_id)
pair_id = self.db.save_prompt_pair(
text,
negative_text,
category=category.strip() or None,
tags=self._parse_tags(tags),
)
pair = self.db.get_prompt_pair(pair_id)
prompt_id = pair["positive_prompt_id"]
resolved_positive = (
resolve_text_input(prompt, node_id, "text") if known else None
)
resolved_negative = (
resolve_text_input(prompt, node_id, "negative_text") if known else None
)
positive_seen = resolved_positive is not None and generate_prompt_hash(
resolved_positive
) == generate_prompt_hash(text)
pair_seen = (
positive_seen
and resolved_negative is not None
and generate_prompt_hash(resolved_negative)
== generate_prompt_hash(negative_text)
)
positive_pending = getattr(self, "pending_first_use", PENDING_FIRST_USE)
pair_pending = getattr(self, "pending_pair_first_use", PENDING_PAIR_FIRST_USE)
counted = set()
for field, component_text, resolved_text, component_id in (
("text", text, resolved_positive, prompt_id),
(
"negative_text",
negative_text,
resolved_negative,
pair["negative_prompt_id"],
),
):
component_hash = generate_prompt_hash(component_text)
if (
field not in used_fields
or not component_id
or component_hash in counted
):
continue
counted.add(component_hash)
seen_at_queue = (
resolved_text is not None
and generate_prompt_hash(resolved_text) == component_hash
)
runs = positive_pending.consume(component_hash) + int(not seen_at_queue)
if runs:
self.db.record_prompt_use(component_id, times=runs)
if pair_used:
runs = pair_pending.consume(pair_run_key(text, negative_text)) + int(
not pair_seen
)
if runs:
self.db.record_prompt_pair_use(pair_id, times=runs)
if is_positive:
self.prompt_tracker.set_current_prompt(
prompt_text=positive,
additional_data={
"prompt_id": prompt_id,
"pair_id": pair_id if pair_used else None,
"positive_text": positive,
"negative_text": negative_text,
"positive_storage_text": text.strip(),
"negative_storage_text": negative_text.strip(),
"prepend_text": prepend_text.strip(),
"append_text": append_text.strip(),
"final_text": positive,
"category": category.strip() or None,
"tags": self._parse_tags(tags) or [],
},
push_to_queue=not positive_seen or (pair_used and not pair_seen),
**self._run_identity(prompt, unique_id),
)
return pair_id, prompt_id
@classmethod
def IS_CHANGED(
cls,
-3
View File
@@ -216,9 +216,6 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC):
except Exception as e:
# Log error but don't fail the processing
self.logger.warning(f"Failed to save prompt to database: {e}")
self._mark_untracked(prompt, unique_id)
else:
self._skip_empty_prompt(prompt, unique_id)
# Register with ComfyUI integration for standard metadata compatibility
node_id = f"promptmanagertext_{int(time.time() * 1000)}"
+13 -148
View File
@@ -31,7 +31,6 @@ import math
import ntpath
import os
import re
import threading
from pathlib import Path
from aiohttp import web
@@ -48,7 +47,6 @@ from .review_events import ReviewEventsMixin
from .saved_views import SavedViewRoutesMixin
from .history import HistoryRoutesMixin
from .images import ImageRoutesMixin
from .gallery_listing import WALK_STOPPED, RootListing, WalkOutcome
from .admin import AdminRoutesMixin
from .settings_routes import SettingsRoutesMixin
from .job_routes import JobRoutesMixin
@@ -89,18 +87,6 @@ def _get_project_root():
return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# How long an executor thread waits for the loop to drop the listings.
_OFF_LOOP_DROP_WAIT = 30
def _running_loop():
"""The event loop running in this thread, or None."""
try:
return asyncio.get_running_loop()
except RuntimeError:
return None
def _public_path(path) -> str:
"""Render a server path for an API response without revealing the layout.
@@ -358,20 +344,8 @@ class PromptManagerAPI(
self._captures_closed = False
self._cached_output_dir = None # Lazy-cached by _find_comfyui_output_dir()
self._html_cache = {} # Cached HTML file contents keyed by path
# One shared listing per output root, keyed by str(root), and the
# background walks refreshing them (py/api/gallery_listing.py, #348).
self._root_listings: dict[str, RootListing] = {}
self._listing_tasks: set = set()
# The loop the walks run on, set by the startup warm-up or the first
# walk; the monitor's thread hops onto it (notify_output_settled).
self._listing_loop = None
# Set by cleanup_gallery_listings: no background walk starts after it.
self._listings_closed = False
# The workspace caches rebuilt behind a stored walk that changed the
# listing, once the host's startup warm-up has run (#348).
self._workspace_warm_enabled = False
self._workspace_warm_task = None
self._workspace_warm_again = False
self._gallery_cache = {} # dict: path_str -> (files, timestamp)
self._gallery_cache_ttl = 30 # Cache TTL in seconds
# Workspace listing values and sort orders derived from a cached scan
# (py/api/images.py _workspace_listed/_workspace_orders).
self._workspace_listed_cache = {}
@@ -491,122 +465,23 @@ class PromptManagerAPI(
return await loop.run_in_executor(None, call)
return await loop.run_in_executor(None, func, *args)
def _root_listing(self, root) -> RootListing:
"""The shared listing record of one output root (created empty)."""
return self._root_listings.setdefault(str(root), RootListing())
def invalidate_gallery_cache(self):
"""Invalidate the cached gallery file listing.
def invalidate_gallery_cache(self, retire=False):
"""Drop every root's listing, so the next read waits for a new walk.
Called after files are deleted through the app (duplicate cleanup and
the legacy duplicate delete), never for thumbnail writes. A walk
running now is not stored (its generation moved on). The workspace
keeps its per-file values and card facts, so the next request
rebuilds only new or changed files (see ImageRoutesMixin.
_workspace_listed and _card_facts); its sort orders are dropped, as
the walk replaces them anyway, and so is the favorites lookup, which
holds the old list.
``retire`` (the roots changed): no background walk of a root starts
again until a reader or the warm-up asks for it.
Called by the image monitor when new files are detected. The workspace
keeps its per-file values and card facts, so the next request rebuilds
only new or changed files (see ImageRoutesMixin._workspace_listed and
_card_facts); its sort orders
are dropped, as the rescan replaces them anyway, and so is the
favorites lookup, which holds the old scan list.
"""
loop = self._listing_loop
if loop is not None and loop.is_running() and _running_loop() is not loop:
# Called from an executor thread (a settings save, a test seed):
# listing records are touched only on the loop, and the caller
# expects the drop done when this returns.
dropped = threading.Event()
def drop():
try:
self._drop_listings(retire)
finally:
dropped.set()
loop.call_soon_threadsafe(drop)
if not dropped.wait(_OFF_LOOP_DROP_WAIT):
self.logger.warning("The Images listing drop is still queued")
return
self._drop_listings(retire)
def _drop_listings(self, retire=False):
for record in list(self._root_listings.values()):
record.drop()
if retire:
# The roots changed: a removed root must not come back through
# a pending timer, a trailing walk or a monitor mark.
record.retired = True
record.dirty = False
handle, record.rewalk_handle = record.rewalk_handle, None
if handle is not None:
handle.cancel()
self._gallery_cache = {}
self._workspace_order_cache = {}
self._favorite_listing_cache = {}
async def warm_gallery_listings(self, app=None):
"""Start one background walk per output root (aiohttp on_startup).
Registered by the root ``__init__.py``, so the first Images open after
ComfyUI starts finds the listing built. It never fails startup: a
root that cannot be walked is retried by the first read.
"""
# Each stored walk that changes a listing then rebuilds the workspace's
# per-file values, orders and favorites behind it (_warm_workspace).
self._workspace_warm_enabled = True
self._listing_loop = asyncio.get_running_loop()
self._listings_closed = False
try:
for root in self._get_all_output_dirs():
self._root_listing(root).retired = False
self._schedule_walk(root)
except Exception as e:
self.logger.warning(f"Could not warm the Images listing: {e}")
def notify_output_settled(self, path):
"""The image monitor settled an output at *path* (any thread).
Called on the monitor's worker thread with the resolved path. It only
hands the path to the event loop (``output_settled``) and never
touches a listing from this thread; before the loop is known, or once
the listings are cleaned up, it does nothing.
"""
loop = self._listing_loop
if loop is None or self._listings_closed or not loop.is_running():
return
try:
loop.call_soon_threadsafe(self.output_settled, path)
except RuntimeError:
pass # the loop closed meanwhile (shutdown)
async def cleanup_gallery_listings(self, app=None):
"""Stop the background listing walks (aiohttp on_cleanup)."""
self._workspace_warm_enabled = False
# No mark, deferred walk or trailing walk starts a walk from now on.
self._listings_closed = True
# Before any await, so no delayed refresh can start a walk meanwhile.
for record in self._root_listings.values():
handle, record.rewalk_handle = record.rewalk_handle, None
if handle is not None:
handle.cancel()
warm, self._workspace_warm_task = self._workspace_warm_task, None
if warm is not None:
warm.cancel()
await asyncio.gather(warm, return_exceptions=True)
tasks = list(self._listing_tasks)
for task in tasks:
task.cancel()
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
self._listing_tasks.clear()
# A walk cancelled before its first step never ran its own cleanup.
for record in self._root_listings.values():
flight, record.flight = record.flight, None
if flight is not None and not flight.done():
flight.set_result(WalkOutcome.failed(RuntimeError(WALK_STOPPED)))
def reset_gallery_roots_cache(self):
"""Drop every cached listing and derived workspace value (roots changed)."""
self.invalidate_gallery_cache(retire=True)
self.invalidate_gallery_cache()
self._workspace_listed_cache = {}
self._workspace_card_cache = {}
@@ -750,12 +625,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 +635,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
+4 -102
View File
@@ -35,12 +35,7 @@ except ImportError:
from utils.parallel import map_parallel
try:
from ...database.pair_ingestion import blank_positive, ingest_resolved_pair
from ...database.run_attribution import (
host_output_directory,
stored_run_prompt,
stored_run_prompts,
)
from ...database.pair_ingestion import ingest_resolved_pair
from ...utils.generation_params import (
PLACEHOLDERS,
GenerationParamsError,
@@ -50,12 +45,7 @@ try:
from ...utils.prompt_pair_resolver import resolve_prompt_pair
from ...utils.review_capture import reviewable
except ImportError:
from database.pair_ingestion import blank_positive, ingest_resolved_pair
from database.run_attribution import (
host_output_directory,
stored_run_prompt,
stored_run_prompts,
)
from database.pair_ingestion import ingest_resolved_pair
from utils.generation_params import (
PLACEHOLDERS,
GenerationParamsError,
@@ -558,24 +548,6 @@ def _current_image_monitor():
return getattr(_image_monitor_module(), "_monitor_instance", None)
class _BatchHistory:
"""A scan batch's recorded-history lookup, made on the first file needing it.
Most files link through their metadata and never ask; the first one that
does triggers one lookup for the whole batch (``_recorded_runs_for``).
"""
def __init__(self, api, media_files):
self._api = api
self._files = media_files
self._found = None
def get(self, path):
if self._found is None:
self._found = self._api._recorded_runs_for(self._files)
return self._found.get(path)
class AdminRoutesMixin:
"""Mixin providing admin, diagnostics, and maintenance API endpoints."""
@@ -1592,8 +1564,6 @@ class AdminRoutesMixin:
if outcomes is not None:
outcomes.append(_scan_outcome(media_file, roots, status, title))
batch_results = list(batch_results)
history = _BatchHistory(self, [media for media, _meta in batch_results])
for media_file, metadata in batch_results:
if self._scan_job and self._scan_job.cancelled:
break
@@ -1607,7 +1577,6 @@ class AdminRoutesMixin:
ambiguous=ambiguous,
untraced=untraced,
records=records,
history=history,
)
except Exception as e:
self.logger.error(f"Error processing {media_file.name}: {e}")
@@ -1666,7 +1635,6 @@ class AdminRoutesMixin:
ambiguous=None,
untraced=None,
records=None,
history=None,
):
"""Enrich exact pairs; retain legacy positive-only import for absent pairs.
@@ -1684,9 +1652,7 @@ class AdminRoutesMixin:
a dict.
"""
if not metadata:
# No prompt metadata (--disable-metadata, bare savers): only the
# recorded run can say which prompt made it (#342)
return self._link_recorded_run(media_file, {}, detailed, history)
return None
canonical = dict(metadata)
for key in ("prompt", "workflow"):
if isinstance(canonical.get(key), str):
@@ -1696,12 +1662,9 @@ class AdminRoutesMixin:
# Preserve malformed metadata for conservative resolver rejection.
pass
if not canonical.get("prompt") and not canonical.get("parameters"):
return self._link_recorded_run(media_file, canonical, detailed, history)
return None
_with_file_size(canonical, media_file)
resolution = resolve_prompt_pair(canonical)
if blank_positive(resolution):
# Only prepend/append text: no stored prompt can own it (#342)
resolution = {**resolution, "state": "unresolved"}
# Ambiguous without a choice: unresolved unless something links it.
unchoosable = resolution["state"] == "ambiguous" and not reviewable(resolution)
if unchoosable:
@@ -1730,11 +1693,6 @@ class AdminRoutesMixin:
# Graphs without a recognized sampler keep the legacy extractor below.
prompt_text = self._graph_prompt_text(metadata)
if not prompt_text:
linked = self._link_recorded_run(
media_file, canonical, detailed, history
)
if linked is not None:
return linked
return "unresolved" if unchoosable else resolution["state"]
elif resolution["state"] == "ambiguous":
return "unresolved" if unchoosable else "ambiguous"
@@ -1758,62 +1716,6 @@ class AdminRoutesMixin:
return status
return {"status": status, "title": self._scan_title(media_file, prompt_text)}
def _recorded_runs_for(self, media_files):
"""``{path: prompt id or None}`` from recorded history for a scan batch.
One output lookup per chunk of names (``stored_run_prompts``), so a
batch never queries history once per file. An empty dict on failure.
"""
try:
return stored_run_prompts(
self.db, [str(path) for path in media_files], host_output_directory()
)
except Exception:
self.logger.warning(
"History attribution failed for a scan batch", exc_info=True
)
return {}
def _link_recorded_run(self, media_file, canonical, detailed=False, history=None):
"""Link a file recorded history attributes to one stored prompt (#342).
Returns the scan status (with ``detailed``, a stored or linked file's
``{"status", "title"}``), or None when history names no single stored
prompt for the file (``database.run_attribution``). Never creates a
prompt. Works without any prompt metadata in the file. ``history`` is
the batch's precomputed lookup; without it the file is looked up alone.
"""
if history is not None:
prompt_id = history.get(str(media_file))
else:
try:
prompt_id = stored_run_prompt(
self.db, str(media_file), host_output_directory()
)
except Exception:
self.logger.warning(
f"History attribution failed for {media_file.name}",
exc_info=True,
)
return None
if prompt_id is None:
return None
try:
image_id = self.db.link_image_to_prompt(
prompt_id, str(media_file), metadata=canonical
)
except Exception as e:
self.logger.error(f"Failed to link {media_file.name} from history: {e}")
return "found"
status = "linked" if image_id else "unchanged"
if not detailed or status not in SCAN_TITLED_DETAILS:
return status
prompt = self.db.get_prompt_by_id(prompt_id) or {}
return {
"status": status,
"title": self._scan_title(media_file, prompt.get("text") or ""),
}
def _scan_title(self, media_file, text):
"""The linked prompt's Library title (its stored title, else ``text``'s
first line), as the Library and the inspector name it."""
-337
View File
@@ -1,337 +0,0 @@
"""One shared listing per output root, and the pure helpers behind it (#348).
The API keeps one ``RootListing`` per output root (``_root_listings`` in
``py/api/__init__.py``). The record is read and written only on the event
loop; the walk itself, the comparison with the previous listing and both
folder sets run in the executor through ``walk_root``. Routes live in
``py/api/images.py`` (``_listing``); this module has no aiohttp import so
tests can exercise it without a server.
A settled output the image monitor reports is added to its root's stored
listing without a walk (``output_settled``, ``insert_entry``), in the place
a walk would list it. When that cannot be exact, the root is marked dirty
(``mark_root_dirty``) instead; a dirty root gets one background walk, no
sooner than ``rewalk_interval`` after its previous walk started (PR 2).
"""
from __future__ import annotations
import asyncio
import os
from bisect import bisect_left
from dataclasses import dataclass, field
from pathlib import PurePath
from typing import Callable, NamedTuple
# A listing older than this (seconds) is served as it is, and one background
# walk replaces it. Read as ``gallery_listing.REFRESH_AFTER`` so tests can
# patch it.
REFRESH_AFTER = 30
# The stale read is served at once; its refresh walk starts this many seconds
# later, so the walk does not share the GIL with the page's own requests
# (#348 live test: reads during the walk took 2.2 s instead of 0.07 s).
REFRESH_DELAY = 1.0
# No background walk of a root starts sooner than rewalk_interval(last walk's
# duration) after its previous walk started: FACTOR times that duration, at
# least FLOOR and at most CAP seconds. Walks up to CAP / FACTOR (3 s) then
# hold the GIL at most ~20% of the time; the cap keeps a new output listed
# within ~20 s, at the cost of a higher share (back to back from 15 s) on
# slower walks. Read through this module so tests can patch them.
REWALK_FLOOR = 5
REWALK_FACTOR = 5
REWALK_CAP = 15
# A deferred walk is never armed sooner than this: no zero-delay loop.
MIN_REWALK_DELAY = 0.5
# A reader with no listing gives up (500) after this many walks in a row were
# discarded by invalidations made while they ran.
MAX_DISCARDED_WALKS = 3
WALK_STOPPED = "The listing walk was stopped"
STORED = "stored"
DISCARDED = "discarded"
FAILED = "failed"
class WalkOutcome(NamedTuple):
"""What one walk ended with; a walk's flight future resolves to this."""
kind: str
error: BaseException | None = None
@classmethod
def stored(cls) -> "WalkOutcome":
return cls(STORED)
@classmethod
def discarded(cls) -> "WalkOutcome":
return cls(DISCARDED)
@classmethod
def failed(cls, error: BaseException) -> "WalkOutcome":
return cls(FAILED, error)
@dataclass
class RootListing:
"""The shared listing of one output root.
A record is never deleted: ``drop()`` clears the listing and bumps
``generation``, so a walk that started before a drop can never be stored
after it.
"""
# (path, mtime) newest first, as _scan_gallery_files_sync returns it.
files: list | None = None
folders: frozenset | None = None
folders_with_ancestors: frozenset | None = None
# Monotonic time the last walk finished, stored or failed.
walked_at: float | None = None
# Number of the last walk started; a running walk is always that one.
walk_seq: int = 0
# Number of the walk whose files are stored (Rescan waits on it).
stored_seq: int = 0
generation: int = 0
# The running walk's future (resolves to a WalkOutcome), or None.
flight: asyncio.Future | None = None
# True while walks keep failing: one warning per streak.
failing: bool = False
# The one deferred background walk not started yet (loop.call_later), or
# None: a stale refresh or a dirty root's rewalk, whichever is due first.
rewalk_handle: asyncio.TimerHandle | None = None
# normcase(realpath(root)), computed by the walk in the executor: monitor
# paths are matched against it without file I/O on the loop.
resolved_root: str | None = None
# Monotonic start and duration of the last walk (rewalk_interval).
last_started_at: float | None = None
last_duration: float | None = None
# A settled output was reported under this root; cleared when a walk starts.
dirty: bool = False
# Dropped because the output roots changed: no background walk until a
# reader (or the warm-up) asks for this root again. The monitor keeps
# watching a removed root until ComfyUI restarts.
retired: bool = False
# Relative paths of settled outputs waiting to be added to ``files``
# without a walk, and the one task adding them (``output_settled``).
pending_adds: list = field(default_factory=list)
add_task: asyncio.Task | None = None
# Settled outputs added to the stored listing since the last walk stored.
patched: int = 0
def drop(self) -> None:
self.files = None
self.folders = None
self.folders_with_ancestors = None
self.walked_at = None
self.pending_adds = []
self.patched = 0
self.generation += 1
def folder_sets(root, files) -> tuple[frozenset, frozenset]:
"""``(folders, folders_with_ancestors)`` of the files listed under *root*.
Folders keep the spelling ``/gallery/subfolders`` has always answered:
``str(relative.parent)`` (native separators) for the folders holding
files, plus every ancestor joined by ``/``.
"""
# From the path strings: relative_to per file was half of a large walk.
sep, prefix = _root_prefix(root)
folders = set()
for path, _mtime in files:
rel_dir = _folder_of(root, path, sep, prefix)
if rel_dir:
folders.add(rel_dir)
with_ancestors = set(folders)
for folder in folders:
with_ancestors.update(_ancestors(folder))
return frozenset(folders), frozenset(with_ancestors)
def add_folder(root, path, folders, with_ancestors) -> tuple[frozenset, frozenset]:
"""Both folder sets once *path* is listed too (the same sets when it adds none)."""
sep, prefix = _root_prefix(root)
rel_dir = _folder_of(root, path, sep, prefix)
if not rel_dir or rel_dir in folders:
return folders, with_ancestors
return folders | {rel_dir}, with_ancestors | {rel_dir, *_ancestors(rel_dir)}
def _root_prefix(root) -> tuple[str, str]:
# The separator is the path flavour's own (tests use PureWindowsPath).
sep = str(type(root)("a", "b"))[1]
prefix = str(root)
if not prefix.endswith(sep):
prefix += sep
return sep, prefix
def _folder_of(root, path, sep, prefix) -> str:
"""The folder ``/gallery/subfolders`` names for *path*; "" for the root."""
text = str(path)
if text.startswith(prefix):
rel_dir = text[len(prefix) :].rpartition(sep)[0]
else:
rel_dir = str(path.relative_to(root).parent)
return "" if rel_dir == "." else rel_dir
def _ancestors(folder: str) -> list[str]:
parts = folder.replace("\\", "/").split("/")
return ["/".join(parts[:i]) for i in range(1, len(parts))]
def relative_part(resolved_root: str, path, pathmod=os.path) -> str | None:
"""The part of the resolved *path* below *resolved_root*, as a walk spells it.
Kept in *path*'s own case (``realpath`` gives a file's on-disk case on
Windows) with the native separator. None when *path* is not below the
root, or when the case folding changed the length so the slice would
not be exact.
"""
sep = pathmod.sep
text = str(path)
if sep == "\\":
text = _plain(text.replace("/", sep), pathmod)
folded = pathmod.normcase(text)
prefix = _plain(resolved_root, pathmod)
if not prefix.endswith(sep):
prefix += sep
if len(folded) != len(text) or not folded.startswith(prefix):
return None
relative = text[len(prefix) :]
if any(part in ("", ".", "..") for part in relative.split(sep)):
return None
return relative
def insert_entry(files: list, entry: tuple, max_files: int) -> list | None:
"""A new listing with *entry* ``(path, mtime)`` where a walk would list it.
*files* is newest first; an entry for the same path is replaced. Returns
*files* itself when it already holds *entry*, and None when the place is
not exact: another file has the same mtime (a walk orders those by its
own visit order), or the listing is at the walk's *max_files* cap.
Blocking (one pass over *files*): run it in the executor.
"""
path, mtime = entry
text = str(path)
kept = files
for index, (listed, listed_mtime) in enumerate(files):
if str(listed) == text:
if listed_mtime == mtime:
return files
kept = files[:index] + files[index + 1 :]
break
else:
if len(files) >= max_files:
return None
index = bisect_left(kept, -mtime, key=lambda item: -item[1])
if index < len(kept) and kept[index][1] == mtime:
return None
return [*kept[:index], entry, *kept[index:]]
# A listing that differs from the cached one by more entries than this is
# rebuilt in full by the workspace caches (listing_delta).
MAX_DELTA = 256
def listing_delta(old: list, new: list, limit: int = MAX_DELTA):
"""``(added, removed)`` entries from *old* to *new*, or None past *limit*.
A listing patched by ``insert_entry`` shares every unchanged entry object
with the one it came from, so entries are told apart by identity: a walk
builds new tuples, shares none and is rebuilt in full. *old* must still
be alive (the
caller's cache holds it), so no identity was reused.
"""
if abs(len(new) - len(old)) > limit:
return None
old_ids = set(map(id, old))
added = []
for entry in new:
if id(entry) not in old_ids:
added.append(entry)
if len(added) > limit:
return None
if new and len(added) == len(new):
return None # nothing shared: a walk
if len(old) + len(added) == len(new):
return added, []
new_ids = set(map(id, new))
return added, [entry for entry in old if id(entry) not in new_ids]
def rewalk_interval(last_duration) -> float:
"""Seconds from a walk's start to the next background walk of its root."""
scaled = REWALK_FACTOR * (last_duration or 0.0)
return min(REWALK_CAP, max(REWALK_FLOOR, scaled))
def root_contains(resolved_root: str, path, pathmod=os.path) -> bool:
"""True when the resolved *path* lies under *resolved_root* (no file I/O).
*resolved_root* is already ``normcase``d; *path* is folded the same way,
so Windows matches case-insensitively and with either separator. The
match ends at a separator: ``/out2/x.png`` is not under ``/out``.
"""
target = _plain(pathmod.normcase(str(path)), pathmod)
prefix = _plain(resolved_root, pathmod)
if not prefix.endswith(pathmod.sep):
prefix += pathmod.sep
return target.startswith(prefix)
def _plain(path: str, pathmod) -> str:
"""*path* (normcased or not) without a Windows extended-length prefix.
``realpath`` may keep ``\\\\?\\`` on a file past MAX_PATH while its shorter
root resolves without it.
"""
if pathmod.sep != "\\":
return path
if path[:8].lower() == "\\\\?\\unc\\":
return "\\\\" + path[8:]
if path.startswith("\\\\?\\"):
return path[4:]
return path
def resolve_root(root) -> str:
"""``normcase(realpath(root))``, the spelling monitor paths are matched in.
File I/O (realpath follows links): call it in the executor.
"""
return os.path.normcase(os.path.realpath(str(root)))
def same_listing(old, new) -> bool:
"""True when two listings hold the same ``(path, mtime)`` pairs in order."""
if old is None or new is None:
return False
return old is new or old == new
def walk_root(
scan_fn: Callable[[PurePath], list], root, previous_files
) -> tuple[list, frozenset, frozenset, str]:
"""The single executor job of a walk.
Returns ``(files, folders, with_ancestors, resolved_root)``.
Opens the root first, so a missing or unreadable root raises ``OSError``
instead of listing nothing. Returns *previous_files* itself when nothing
changed, which keeps the identity-keyed caches built on it warm.
"""
with os.scandir(root):
pass
files = scan_fn(root)
if same_listing(previous_files, files):
files = previous_files
folders, with_ancestors = folder_sets(root, files)
return files, folders, with_ancestors, resolve_root(root)
+68 -747
View File
File diff suppressed because it is too large Load Diff
+8 -36
View File
@@ -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})
+17 -47
View File
@@ -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))
@@ -226,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)
@@ -268,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"],
@@ -302,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,
@@ -367,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
@@ -392,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)
@@ -408,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')}")
@@ -423,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}")
+7 -25
View File
@@ -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(),
)
+1
View File
@@ -175,6 +175,7 @@ class ThumbnailConvertMixin:
job.log_line(f"Conversion finished · {NOTHING_OUTDATED.lower()}")
return "finished"
totals, left = await self._convert_in_batches(job, options, plan, output_dirs)
self.invalidate_gallery_cache()
if totals["stopped"]:
job.update(status_line=FORMAT_CHANGED.format(converted=totals["converted"]))
job.log_line(
+1
View File
@@ -261,6 +261,7 @@ class ThumbnailJobsMixin:
stats, start, left, attempted = await self._generate_in_batches(
job, options, plan, roots, output_dirs
)
self.invalidate_gallery_cache()
generated = _plural(stats["generated"], "preview")
if job.cancelled:
return self._end_cancelled(
+53 -103
View File
@@ -1,26 +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 the kept sidecars (path,
size and modification time) and the layout; ``preview`` adds the counts a
review shows. Nothing is written and nothing is requested.
2. **stage**: read each sidecar and build its record. Its LoraManager preview
and example images are linked where they are (``local_media``); only a
LoRA without any local image downloads its CivitAI example images into the
image cache first (best effort, counted). A cancelled or failed stage
leaves the database untouched; downloaded images stay in the cache.
3. **replace**: the sidecars are found and fingerprinted once more (a file
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).
@@ -32,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
@@ -51,44 +37,35 @@ 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 each kept sidecar's path, size and
modification time (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"))
for sidecar in discovery.sidecars:
path = sidecar.metadata_path
for path in files:
try:
stat = path.stat()
entry = f"{path}\0{stat.st_size}\0{stat.st_mtime_ns}\n"
@@ -98,48 +75,31 @@ def fingerprint(discovery: Discovery) -> str:
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()),
}
@@ -171,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)
-602
View File
@@ -1,602 +0,0 @@
"""Where LoraManager keeps its files (#355).
LoraManager (``custom_nodes/comfyui-lora-manager``) stores one
``<name>.metadata.json`` sidecar per model, a preview image or video beside
it, and each model's downloaded CivitAI example images. Where it stores them
is a user setting, and its settings file itself can live in three places.
This module repeats LoraManager's own rules, reading its files only (never
importing it), so the LoRA import, the trigger-word cache and the image
server share one definition:
- ``load_layout`` finds the settings file (environment folder, portable
plugin file, or the user config folder, in LoraManager's order), keeps only
the seven keys the import needs, and resolves the LoRA roots, the sidecar
root (centralized storage), the example-images root and the active
LoraManager library. Every folder from the settings passes the folder
guard (exists, not a filesystem root, not the home folder itself).
- ``discover_sidecars`` finds every sidecar whose model file exists inside a
LoRA root: beside the models in ``alongside`` mode, under the sidecar root
(matched to its model by ``file_path``) in ``centralized`` mode.
- ``local_media`` lists a LoRA's preview and example images that the image
server may serve.
Standard library plus ``platformdirs`` only (LoraManager depends on it), so
it is testable outside ComfyUI.
"""
import json
import os
import re
from dataclasses import dataclass, replace
from pathlib import Path
from typing import (
AbstractSet,
Callable,
Dict,
Iterator,
List,
Mapping,
Optional,
Sequence,
Tuple,
)
# py/utils/settings_paths.py: SETTINGS_DIR_ENV, _LM_PORTABLE_ENV, APP_NAME.
SETTINGS_DIR_ENV = "LORA_MANAGER_SETTINGS_DIR"
PORTABLE_ENV = "LORA_MANAGER_PORTABLE"
APP_NAME = "ComfyUI-LoRA-Manager"
# py/utils/constants.py PREVIEW_EXTENSIONS, in LoraManager's order.
PREVIEW_EXTENSIONS = (
".webp",
".preview.webp",
".preview.png",
".preview.jpeg",
".preview.jpg",
".preview.mp4",
".png",
".jpeg",
".jpg",
".mp4",
".gif",
".webm",
".avif",
".jxl",
)
# LoraManager's LoRA scanner accepts .safetensors only (py/services/lora_scanner.py).
MODEL_SUFFIX = ".safetensors"
# py/utils/sidecar_paths.py METADATA_SUFFIX.
METADATA_SUFFIX = ".metadata.json"
MAX_SETTINGS_BYTES = 5 * 1024 * 1024
# The only settings the import reads; every other key (the file can hold API
# keys) is dropped right after parsing.
SETTINGS_KEYS = frozenset(
{
"sidecar_storage_mode",
"sidecar_storage_path",
"example_images_path",
"folder_paths",
"extra_folder_paths",
"libraries",
"active_library",
}
)
# py/utils/sidecar_paths.py STORAGE_MODE_ALONGSIDE / STORAGE_MODE_CENTRALIZED.
MODE_ALONGSIDE = "alongside"
MODE_CENTRALIZED = "centralized"
# utils/video_metadata.py VIDEO_EXTENSIONS (not imported: that package needs
# the node's relative imports).
VIDEO_SUFFIXES = frozenset({".mp4", ".mov", ".m4v", ".webm", ".mkv", ".avi", ".wmv"})
SOURCE_ENVIRONMENT = "environment"
SOURCE_PORTABLE = "portable"
SOURCE_USER_CONFIG = "user config"
NO_INSTALL = "no LoraManager install found"
UNREADABLE = "the settings file could not be read"
NOT_JSON = "the settings file is not valid JSON"
NOT_OBJECT = "the settings file is not a JSON object"
TOO_LARGE = "the settings file is larger than 5 MB"
NO_PLATFORMDIRS = "platformdirs is not available"
NO_SIDECAR_ROOT = "the sidecar folder does not exist"
BAD_SIDECAR_ROOT = "the sidecar folder is a filesystem root or the home folder"
@dataclass(frozen=True)
class LoraManagerLayout:
"""What LoraManager's settings say about where its files are."""
settings_path: Optional[Path]
source: Optional[str]
error: Optional[str]
mode: str
lora_roots: Tuple[str, ...]
sidecar_root: Optional[str]
example_images_root: Optional[str]
library: str
media_roots: Tuple[str, ...]
@dataclass(frozen=True)
class Sidecar:
"""One LoRA's metadata file and the model file it describes."""
metadata_path: Path
model_path: Path
@dataclass(frozen=True)
class Discovery:
"""The sidecars kept, how many metadata files were seen, and those unread."""
layout: LoraManagerLayout
sidecars: List[Sidecar]
scanned: int
unreadable: List[str]
class _SettingsError(Exception):
"""The settings file cannot be used, for the reason it carries."""
# ── paths ────────────────────────────────────────────────────────────
def _canonical(path: str, pathmod=os.path) -> str:
# py/config.py _canonical_path (not imported: py.config imports ComfyUI's
# ``server`` at import time).
return pathmod.normcase(pathmod.realpath(path))
def _is_filesystem_root(path: str, pathmod=os.path) -> bool:
# py/config.py _is_filesystem_root: '/', 'C:\\' and UNC share roots.
return pathmod.dirname(path) == path
def _guard_folder(
path: str,
home: str,
*,
pathmod=os.path,
isdir: Optional[Callable[[str], bool]] = None,
) -> Optional[str]:
"""The canonical form of ``path`` when it may be read and served, else None.
Kept only when it is a folder, not a filesystem root, and neither the
home folder nor one of its ancestors (``/home``, ``C:\\Users``: a
security review ruled that a folder holding the home folder must never
be read or served).
"""
isdir = isdir or os.path.isdir
try:
canonical = _canonical(path, pathmod)
if not isdir(canonical) or _is_filesystem_root(canonical, pathmod):
return None
if _inside(home, (canonical,), pathmod):
return None
except (OSError, ValueError):
return None
return canonical
def _inside(path, roots: Sequence[str], pathmod=os.path) -> bool:
"""True when ``path`` (canonical) is one of ``roots`` or below one."""
try:
candidate = _canonical(str(path), pathmod)
except (OSError, ValueError):
return False
for root in roots:
prefix = root if root.endswith(pathmod.sep) else root + pathmod.sep
if candidate == root or candidate.startswith(prefix):
return True
return False
def _absolute(value) -> Optional[str]:
"""A settings path, stripped, ``~``-expanded and absolute; None when unset."""
if not isinstance(value, str) or not value.strip():
return None
return os.path.abspath(os.path.expanduser(value.strip()))
# ── the settings file ───────────────────────────────────────────────
def _platform_user_config_dir() -> str:
from platformdirs import user_config_dir
return user_config_dir(APP_NAME, appauthor=False)
def _read_bytes(path: Path) -> bytes:
with open(path, "rb") as handle:
return handle.read(MAX_SETTINGS_BYTES + 1)
def _read_settings(path: Path) -> Optional[Dict]:
"""The parsed object in ``path``; None when it does not exist."""
try:
raw = _read_bytes(path)
except FileNotFoundError:
return None
except OSError:
raise _SettingsError(UNREADABLE) from None
if len(raw) > MAX_SETTINGS_BYTES:
raise _SettingsError(TOO_LARGE)
try:
data = json.loads(raw.decode("utf-8"))
except (UnicodeDecodeError, ValueError):
raise _SettingsError(NOT_JSON) from None
if not isinstance(data, dict):
raise _SettingsError(NOT_OBJECT)
return data
def _settings_subset(data: Mapping) -> Dict:
"""Only the keys the import reads."""
return {key: data[key] for key in SETTINGS_KEYS if key in data}
def _portable_flag(environ: Mapping[str, str], plugin_settings: Path) -> bool:
"""LoraManager's portable mode (settings_paths._should_use_portable_settings)."""
forced = environ.get(PORTABLE_ENV)
if forced == "1":
return True
if forced == "0":
return False
try:
data = _read_settings(plugin_settings)
except _SettingsError:
return False
return bool(data) and data.get("use_portable_settings") is True
def _settings_location(
lora_manager_path: str,
environ: Mapping[str, str],
user_config_dir: Callable[[], str],
) -> Tuple[Path, str]:
"""``(settings file, source)`` in LoraManager's lookup order."""
override = environ.get(SETTINGS_DIR_ENV)
if override:
folder = os.path.abspath(os.path.expanduser(override))
return Path(folder) / "settings.json", SOURCE_ENVIRONMENT
plugin_settings = Path(lora_manager_path) / "settings.json"
if _portable_flag(environ, plugin_settings):
return plugin_settings, SOURCE_PORTABLE
try:
folder = user_config_dir()
except ImportError:
raise _SettingsError(NO_PLATFORMDIRS) from None
return Path(folder) / "settings.json", SOURCE_USER_CONFIG
# ── libraries and folders ───────────────────────────────────────────
def _sanitize_library_name(name: str) -> str:
# py/utils/example_images_paths.py sanitize_library_name.
safe = re.sub(r"[^A-Za-z0-9_.-]", "_", name or "")
return safe or "default"
def _active_library(settings: Mapping) -> Tuple[str, Optional[Dict], int]:
"""``(name, library dict or None, number of libraries)``."""
libraries = settings.get("libraries")
if not isinstance(libraries, dict) or not libraries:
return "default", None, 0
active = settings.get("active_library")
if not isinstance(active, str) or active not in libraries:
active = next(iter(libraries))
library = libraries[active]
return str(active), library if isinstance(library, dict) else None, len(libraries)
def _lora_list(settings: Mapping, library: Optional[Dict], key: str) -> List[str]:
"""The ``loras`` folders of ``key`` (the library's, else the top level)."""
source = library if library is not None and key in library else settings
section = source.get(key)
if not isinstance(section, dict):
return []
value = section.get("loras")
if isinstance(value, str):
value = [value]
if not isinstance(value, list):
return []
return [entry for entry in value if isinstance(entry, str) and entry.strip()]
def _guarded(paths, home: str) -> Tuple[str, ...]:
kept = []
for path in paths:
absolute = _absolute(path)
guarded = _guard_folder(absolute, home) if absolute else None
if guarded and guarded not in kept:
kept.append(guarded)
return tuple(kept)
def _example_images_root(
settings: Mapping, library: str, libraries: int, home: str
) -> Optional[str]:
"""py/utils/example_images_paths.py get_library_root, when the folder exists."""
configured = _absolute(settings.get("example_images_path"))
root = _guard_folder(configured, home) if configured else None
if root is None or libraries <= 1:
return root
# Guarded (canonical) like its root: a library name keeps its case, and
# the media check compares normcased paths.
return _guard_folder(os.path.join(root, _sanitize_library_name(library)), home)
def _media_roots(*groups) -> Tuple[str, ...]:
roots = []
for group in groups:
for root in group:
if root and root not in roots:
roots.append(root)
return tuple(roots)
def load_layout(
lora_manager_path: Optional[str],
*,
default_roots: Callable[[], List[str]],
environ: Mapping[str, str] = os.environ,
user_config_dir: Callable[[], str] = _platform_user_config_dir,
) -> LoraManagerLayout:
"""Read LoraManager's settings the way LoraManager does (read only).
``lora_manager_path`` is the resolved install (``None`` when none was
found). An unusable layout carries ``error``, alongside mode and the
default LoRA roots, so serving and trigger words keep working.
"""
home = str(Path.home())
defaults = _guarded(default_roots(), home)
fallback = LoraManagerLayout(
settings_path=None,
source=None,
error=NO_INSTALL,
mode=MODE_ALONGSIDE,
lora_roots=defaults,
sidecar_root=None,
example_images_root=None,
library="default",
media_roots=defaults,
)
if not lora_manager_path:
return fallback
try:
settings_path, source = _settings_location(
lora_manager_path, environ, user_config_dir
)
except _SettingsError as error:
return replace(fallback, error=str(error))
fallback = replace(fallback, settings_path=settings_path, source=source)
try:
raw = _read_settings(settings_path)
except _SettingsError as error:
return replace(fallback, error=str(error))
settings = _settings_subset(raw or {})
return _layout_from_settings(settings, settings_path, source, defaults, home)
def _layout_from_settings(settings, settings_path, source, defaults, home):
library, entry, libraries = _active_library(settings)
lora_roots = _media_roots(
defaults,
_guarded(_lora_list(settings, entry, "folder_paths"), home),
_guarded(_lora_list(settings, entry, "extra_folder_paths"), home),
)
mode = settings.get("sidecar_storage_mode")
mode = MODE_CENTRALIZED if mode == MODE_CENTRALIZED else MODE_ALONGSIDE
example_root = _example_images_root(settings, library, libraries, home)
layout = LoraManagerLayout(
settings_path=settings_path,
source=source,
error=None,
mode=mode,
lora_roots=lora_roots,
sidecar_root=None,
example_images_root=example_root,
library=library,
media_roots=_media_roots(lora_roots, [example_root]),
)
if mode != MODE_CENTRALIZED:
return layout
configured = _absolute(settings.get("sidecar_storage_path"))
configured = configured or os.path.join(str(settings_path.parent), "sidecars")
sidecar_root = _guard_folder(configured, home)
if sidecar_root is None:
reason = NO_SIDECAR_ROOT if not os.path.isdir(configured) else BAD_SIDECAR_ROOT
return replace(
layout,
error=reason,
lora_roots=defaults,
example_images_root=None,
media_roots=defaults,
)
return replace(
layout,
sidecar_root=sidecar_root,
media_roots=_media_roots(lora_roots, [sidecar_root], [example_root]),
)
# ── sidecars ────────────────────────────────────────────────────────
def _metadata_files(folder: str) -> Iterator[Path]:
"""Every ``*.metadata.json`` under ``folder``; an unlistable folder raises.
Symlinked folders are followed, each real folder once (no loops).
"""
def fail(error):
raise error
visited = set()
for current, dirs, files in os.walk(folder, onerror=fail, followlinks=True):
real = _canonical(current)
if real in visited:
dirs[:] = []
continue
visited.add(real)
dirs.sort()
for name in sorted(files):
if name.endswith(METADATA_SUFFIX):
yield Path(current) / name
def _read_sidecar(path: Path) -> Optional[Dict]:
try:
with open(path, "r", encoding="utf-8") as handle:
data = json.load(handle)
except (OSError, ValueError):
return None
return data if isinstance(data, dict) else None
def _model_for(path: Path, layout: LoraManagerLayout, unreadable: List[str]):
"""The model file a sidecar describes, when it exists inside a LoRA root."""
if layout.mode == MODE_ALONGSIDE:
base = path.name[: -len(METADATA_SUFFIX)]
model = path.parent / f"{base}{MODEL_SUFFIX}"
return model if model.is_file() else None
metadata = _read_sidecar(path)
if metadata is None:
unreadable.append(path.name)
return None
file_path = metadata.get("file_path")
if not isinstance(file_path, str) or not file_path.strip():
return None
model = Path(file_path)
if not model.is_file() or not _inside(model, layout.lora_roots):
return None
return model
def discover_sidecars(layout: LoraManagerLayout) -> Discovery:
"""Every sidecar whose model exists inside a LoRA root, sorted, once each.
Raises ``OSError`` when a folder cannot be listed: a discovery that
cannot see every folder must not replace the collection with part of it.
"""
if layout.error is not None:
return Discovery(layout, [], 0, [])
if layout.mode == MODE_CENTRALIZED:
folders = [layout.sidecar_root]
else:
folders = list(layout.lora_roots)
seen = set()
scanned = 0
unreadable: List[str] = []
sidecars = []
for folder in folders:
for path in _metadata_files(folder):
resolved = Path(os.path.realpath(path))
key = os.path.normcase(str(resolved))
if key in seen:
continue
seen.add(key)
scanned += 1
model = _model_for(resolved, layout, unreadable)
if model is not None:
sidecars.append(Sidecar(resolved, Path(os.path.realpath(model))))
sidecars.sort(key=lambda sidecar: str(sidecar.metadata_path))
return Discovery(layout, sidecars, scanned, unreadable)
# ── images ──────────────────────────────────────────────────────────
def _natural_key(name: str):
parts = re.split(r"(\d+)", name)
return [int(part) if part.isdigit() else part.lower() for part in parts], name
def _example_files(folder: str) -> List[str]:
try:
with os.scandir(folder) as entries:
names = [entry.name for entry in entries if entry.is_file()]
except OSError:
return []
return [os.path.join(folder, name) for name in sorted(names, key=_natural_key)]
def _preview_candidates(sidecar: Sidecar, metadata: Dict) -> Iterator[str]:
"""The paths a preview is looked for at, in LoraManager's order."""
url = metadata.get("preview_url")
if isinstance(url, str) and url.strip() and os.path.isabs(url.strip()):
yield url.strip()
base = sidecar.metadata_path.name[: -len(METADATA_SUFFIX)]
for suffix in PREVIEW_EXTENSIONS:
yield str(sidecar.metadata_path.parent / f"{base}{suffix}")
def local_media(
sidecar: Sidecar,
metadata: Dict,
layout: LoraManagerLayout,
media_extensions: AbstractSet[str],
) -> List[str]:
"""The LoRA's preview, then its example images: absolute, resolved paths.
A file counts only when its extension is in ``media_extensions`` (what the
image server serves) and it lies inside ``layout.media_roots``.
"""
def usable(path: str) -> bool:
return (
os.path.splitext(path)[1].lower() in media_extensions
and os.path.isfile(path)
and _inside(path, layout.media_roots)
)
found: List[str] = []
for path in _preview_candidates(sidecar, metadata):
if usable(path):
found.append(path)
break
sha256 = metadata.get("sha256")
if layout.example_images_root and isinstance(sha256, str) and sha256.strip():
folder = os.path.join(layout.example_images_root, sha256.strip().lower())
found.extend(path for path in _example_files(folder) if usable(path))
media: List[str] = []
for path in found:
resolved = os.path.realpath(path)
if resolved not in media:
media.append(resolved)
return _cover_first(media)
def _cover_first(media: List[str]) -> List[str]:
"""A still image in front of a leading video (no thumbnails for these yet)."""
if not media or os.path.splitext(media[0])[1].lower() not in VIDEO_SUFFIXES:
return media
for index, path in enumerate(media):
if os.path.splitext(path)[1].lower() not in VIDEO_SUFFIXES:
return [path] + media[:index] + media[index + 1 :]
return media
# ── the review ──────────────────────────────────────────────────────
def layout_summary(layout: LoraManagerLayout, home: Path) -> Dict:
"""What the review shows about the layout (no settings values)."""
settings_file = None
if layout.settings_path is not None:
settings_file = str(layout.settings_path)
try:
relative = Path(settings_file).relative_to(home)
settings_file = os.path.join("~", str(relative))
except ValueError:
pass
return {
"source": layout.source,
"settings_file": settings_file,
"error": layout.error,
"mode": layout.mode,
"library": layout.library,
"example_images": layout.example_images_root is not None,
}
+96 -97
View File
@@ -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,16 +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_manager_layout import (
MODE_ALONGSIDE,
METADATA_SUFFIX,
LoraManagerLayout,
discover_sidecars,
load_layout,
)
logger = get_logger("prompt_manager.lora_utils")
# ── LoraManager detection ────────────────────────────────────────────
@@ -245,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()))
@@ -277,28 +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``)."""
return layout_for_install(lora_manager_install(lora_manager_path))
def _get_extra_lora_paths(comfyui_root: Path) -> List[Path]:
"""Parse extra_model_paths.yaml for additional LoRA directories."""
results = []
@@ -404,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
@@ -822,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.
@@ -842,81 +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 Exception as error:
logger.warning(f"Trigger word cache not refreshed: {error}")
with self._lock:
self._failed_at = time.monotonic()
return len(self._cache)
new_cache: Dict[str, List[str]] = {}
lora_dirs = find_lora_directories(lora_manager_path)
for lora_dir in lora_dirs:
dir_path = Path(lora_dir)
for meta_file in dir_path.rglob("*.metadata.json"):
metadata = read_lora_metadata(meta_file)
if not metadata:
continue
words = get_trigger_words_from_metadata(metadata)
if not words:
continue
# Key by filename stem (what appears in <lora:NAME:weight>)
file_name = metadata.get("file_name", "")
if file_name:
stem = Path(file_name).stem
new_cache[stem.lower()] = words
# Also key by the metadata file stem
meta_stem = meta_file.name.replace(".metadata.json", "")
if meta_stem.lower() not in new_cache:
new_cache[meta_stem.lower()] = words
with self._lock:
self._cache = new_cache
self._loaded = True
self._failed_at = None
logger.info(
f"Trigger word cache loaded: {len(new_cache)} LoRAs with trigger words"
)
return len(new_cache)
@staticmethod
def _build(lora_manager_path: str) -> Dict[str, List[str]]:
layout = lora_manager_layout(lora_manager_path)
if layout.error is not None:
logger.warning(
f"LoraManager settings not usable ({layout.error}); trigger words "
"come from the default LoRA folders"
)
layout = replace(
layout,
error=None,
mode=MODE_ALONGSIDE,
sidecar_root=None,
example_images_root=None,
)
new_cache: Dict[str, List[str]] = {}
for sidecar in discover_sidecars(layout).sidecars:
metadata = read_lora_metadata(sidecar.metadata_path)
if not metadata:
continue
words = get_trigger_words_from_metadata(metadata)
if not words:
continue
# Key by the model file name (what appears in <lora:NAME:weight>)
# and by the sidecar's name.
new_cache[sidecar.model_path.stem.lower()] = words
sidecar_stem = sidecar.metadata_path.name[: -len(METADATA_SUFFIX)]
new_cache.setdefault(sidecar_stem.lower(), words)
return new_cache
def get_trigger_words(self, lora_name: str) -> List[str]:
"""Look up trigger words for a LoRA by name (case-insensitive)."""
with self._lock:
return self._cache.get(lora_name.lower(), [])
def retry_due(self) -> bool:
"""Whether a node may try loading again: never failed, or the last
failure is TRIGGER_RETRY_SECONDS old (no rescan on every prompt run)."""
with self._lock:
failed_at = self._failed_at
return failed_at is None or time.monotonic() - failed_at >= (
TRIGGER_RETRY_SECONDS
)
@property
def is_loaded(self) -> bool:
with self._lock:
@@ -926,7 +926,6 @@ class TriggerWordCache:
with self._lock:
self._cache.clear()
self._loaded = False
self._failed_at = None
# Module-level singleton
+1 -1
View File
@@ -3,7 +3,7 @@ name = "promptmanager"
description = "A powerful ComfyUI custom node that extends the standard text encoder with persistent prompt storage, advanced search capabilities, and an automatic image gallery system using SQLite."
version = "3.3.1"
license = {file = "LICENSE"}
requires-python = ">=3.12"
requires-python = ">=3.10"
dependencies = ["watchdog>=2.1.0", "Pillow>=8.0.0"]
[build-system]
-13
View File
@@ -166,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()
@@ -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();
+10 -15
View File
@@ -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);
});
@@ -485,6 +485,5 @@ test('a safety backup file (no .db extension) can be picked and reviewed (M3)',
const answer = await chooseBackup(page, bytes, 'prompts.db.backup_20261025_093000_000000');
expect((await answer).status()).toBe(200);
await expect(dialog(page).locator('#modal-title')).toHaveText('Restore preview');
// The copy's size depends on what earlier tests on this server wrote: KB or MB.
expect((await gridOf(page))[0][1]).toMatch(/^prompts\.db\.backup_20261025_093000_000000 · [\d,.]+ (KB|MB) · .+$/);
expect((await gridOf(page))[0][1]).toMatch(/^prompts\.db\.backup_20261025_093000_000000 · \d+ KB · .+$/);
});
@@ -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');
@@ -385,8 +381,7 @@ test('saved pairs block the import: the review explains it and offers only Close
await openIntegrations(page);
await enable(page);
await layoutOf(page).getByRole('button', { name: 'Review import', exact: true }).click();
// gridOf reads once; poll until the review has rendered its rows.
await expect.poll(async () => (await gridOf(page))[3]).toEqual(['Protected', '1 imported prompt used by saved pairs']);
expect((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 +389,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 +407,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');
@@ -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(
-48
View File
@@ -1,6 +1,5 @@
"""Deterministic nodes and queue controls installed only in an owned test host."""
import os
import threading
import time
@@ -256,50 +255,3 @@ async def cache_mode(request: web.Request) -> web.Response:
raise web.HTTPBadRequest()
_CACHE_MODE = mode
return web.json_response({"configured": True})
def _plugin_module(suffix: str):
"""A loaded PromptManager module (ComfyUI names the package by its path)."""
import sys
for name, module in list(sys.modules.items()):
if name.endswith("promptmanager." + suffix):
return module
raise web.HTTPServiceUnavailable(text=f"PromptManager {suffix} not loaded")
@PromptServer.instance.routes.post("/pm_execution_fixture/monitor")
async def image_monitor(request: web.Request) -> web.Response:
"""Start or stop PromptManager's image monitor on the host output folder.
The managed host's config disables monitoring; only the attribution cases
(#342) turn it on, and they turn it off again.
"""
enabled = (await request.json())["enabled"]
if type(enabled) is not bool:
raise web.HTTPBadRequest()
import folder_paths
monitor_module = _plugin_module("utils.image_monitor")
config = _plugin_module("py.config").GalleryConfig
monitor = monitor_module._monitor_instance
if monitor is None:
database = _plugin_module("database.operations").PromptDatabase()
tracker = _plugin_module("utils.prompt_tracker").get_prompt_tracker(database)
monitor = monitor_module.get_image_monitor(database, tracker)
if enabled:
config.MONITORING_ENABLED = True
monitor.start_monitoring([folder_paths.get_output_directory()])
else:
monitor.stop_monitoring()
config.MONITORING_ENABLED = False
return web.json_response({"monitoring": monitor.is_monitoring})
@PromptServer.instance.routes.get("/pm_execution_fixture/monitor")
async def image_monitor_state(request: web.Request) -> web.Response:
"""File names of keyless images waiting for their run's recorded outputs."""
monitor = _plugin_module("utils.image_monitor")._monitor_instance
handler = getattr(monitor, "handler", None)
waiting = handler.awaiting_history() if handler is not None else []
return web.json_response({"awaiting": [os.path.basename(p) for p in waiting]})
-142
View File
@@ -1,142 +0,0 @@
// #342 on the real pinned host: the image monitor links an output only to a prompt
// that provably produced it. An unknown node (PMHistoryCached) stands in for the
// maintainer's LLM; KSampler is the fixture's deterministic stand-in.
import { randomUUID } from 'node:crypto';
import { execFile } from 'node:child_process';
import { promisify } from 'node:util';
import { fileURLToPath } from 'node:url';
import { expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs';
import { test } from './execution-fixture.mjs';
const run = promisify(execFile);
const reader = fileURLToPath(new URL('./read_execution_history.py', import.meta.url));
async function readLinks(databasePath, filenames, texts = []) {
const { stdout } = await run(process.env.PM_TEST_PYTHON || 'python',
[reader, databasePath, '--links', ...filenames, '--texts', ...texts], { timeout: 5000, maxBuffer: 1024 * 1024 });
return JSON.parse(stdout);
}
// kind: the positive node; text: a typed editor or 'llm' (linked from the unknown
// node); prepend: 'llm' links the unknown node into prepend_text.
function graph({ kind = 'PromptManagerPair', text = '', prepend = '', llm = '', seed = 1, prefix }) {
const linked = value => value === 'llm' ? ['11', 0] : value;
const nodes = {
'10': { class_type: 'PMHistoryInputs', inputs: {} },
'20': { class_type: kind, inputs: { text: linked(text), clip: ['10', 0], append_text: '',
...(prepend ? { prepend_text: linked(prepend) } : {}),
...(kind === 'PromptManagerPair' ? { negative_text: '' } : {}) } },
'30': { class_type: 'CLIPTextEncode', inputs: { text: '', clip: ['10', 0] } },
'40': { class_type: 'KSampler', inputs: { model: ['10', 1], latent_image: ['10', 2],
positive: ['20', 0], negative: kind === 'PromptManagerPair' ? ['20', 1] : ['30', 0],
seed, steps: 1, cfg: 1, sampler_name: 'euler', scheduler: 'normal', denoise: 1 } },
'60': { class_type: 'PMHistoryImages', inputs: { value: 64, batch: 1, latent: ['40', 0] } },
'70': { class_type: 'SaveImage', inputs: { images: ['60', 0], filename_prefix: prefix } },
};
if (llm) nodes['11'] = { class_type: 'PMHistoryCached', inputs: { value: llm } };
return nodes;
}
async function saved(request, prompt, history = {}) {
const id = randomUUID();
const response = await request.post('/prompt', { data: { prompt, prompt_id: id } });
expect(response.ok()).toBe(true);
let item;
await expect.poll(async () => {
item = (await (await request.get(`/history/${id}`)).json())[id];
return item?.status?.messages?.some(([kind]) => kind === 'execution_success') || false;
}, { timeout: 15000 }).toBe(true);
const [image] = item.outputs['70'].images;
expect(image.type).toBe('output');
history.item = item;
return image.filename;
}
async function linkedTo(execution, filename, expected) {
await expect.poll(async () => (await readLinks(execution.databasePath, [filename])).links[filename],
{ timeout: 20000, message: `${filename} is linked by the image monitor` }).toBe(expected);
}
async function monitoring(execution, enabled) {
const response = await execution.request.post('/pm_execution_fixture/monitor', { data: { enabled } });
expect(response.ok()).toBe(true);
expect((await response.json()).monitoring).toBe(enabled);
}
async function cacheMode(execution, mode) {
expect((await execution.request.post('/pm_execution_fixture/cache', { data: { mode } })).ok()).toBe(true);
}
test.beforeEach(async ({ execution }) => {
expect((await execution.request.post('/pm_execution_fixture/metadata', { data: { disabled: false } })).ok()).toBe(true);
await cacheMode(execution, 'CLASSIC'); // keeps the re-run's PromptManager node cached
await monitoring(execution, true);
});
test.afterEach(async ({ execution }) => {
await monitoring(execution, false);
await cacheMode(execution, 'RAM_PRESSURE');
});
test('an output whose positive is only linked prepend text stays unlinked, never given an earlier prompt', async ({ execution }, testInfo) => {
const tag = `${testInfo.project.name} ${randomUUID()}`;
const earlier = `a woman holding a coffee cup ${tag}`;
const prepended = `llm prepend landscape ${tag}`, llmText = `llm text landscape ${tag}`;
// A typed prompt runs first: before #342 the next unidentifiable image took it.
const first = await saved(execution.request, graph({ text: earlier, prefix: 'misslink/earlier' }));
await linkedTo(execution, first, earlier);
// The maintainer's workflow: empty editor, the generated text arrives in prepend.
const orphan = await saved(execution.request,
graph({ prepend: 'llm', llm: prepended, prefix: 'misslink/prepend' }));
// An LLM feeding the editor itself is stored and links through its run key; the
// monitor works in creation order, so the prepend image was processed first.
const llmGraph = graph({ kind: 'PromptManager', text: 'llm', llm: llmText, prefix: 'misslink/text' });
const proven = await saved(execution.request, llmGraph);
await linkedTo(execution, proven, llmText);
const state = await readLinks(execution.databasePath, [orphan], [prepended]);
expect(state).toEqual({ links: { [orphan]: null }, stored: { [prepended]: false } });
// Only the sampler seed changes: the PromptManager node is served from the cache
// without running (the host says so), so no queue entry exists for this run and
// only the snapshot's run match can link it.
const rerunHistory = {};
const rerun = await saved(execution.request, { ...llmGraph, '40': { ...llmGraph['40'],
inputs: { ...llmGraph['40'].inputs, seed: 2 } } }, rerunHistory);
const cached = rerunHistory.item.status.messages.find(([kind]) => kind === 'execution_cached');
expect(cached?.[1].nodes).toEqual(expect.arrayContaining(['11', '20']));
await linkedTo(execution, rerun, llmText);
});
test('an output saved without metadata links once its recorded run names a stored prompt', async ({ execution }, testInfo) => {
const llmText = `bare output landscape ${testInfo.project.name} ${randomUUID()}`;
const request = execution.request;
expect((await request.post('/pm_execution_fixture/metadata', { data: { disabled: true } })).ok()).toBe(true);
try {
// No embedded graph: no run key can prove the tracker entry, so only the
// run's recorded history can, and only once the run's outputs are sealed.
// The gate after SaveImage holds the run open while the monitor looks.
const prompt = graph({ kind: 'PromptManager', text: 'llm', llm: llmText, prefix: 'misslink/bare' });
prompt['80'] = { class_type: 'PMHistoryImageOutput', inputs: { images: ['70', 0], mode: 'gate' } };
expect((await request.post('/pm_execution_fixture/gate/hold')).ok()).toBe(true);
const id = randomUUID();
try {
expect((await request.post('/prompt', { data: { prompt, prompt_id: id } })).ok()).toBe(true);
let waiting = [];
await expect.poll(async () => {
waiting = (await (await request.get('/pm_execution_fixture/monitor')).json()).awaiting;
return waiting.filter(name => name.startsWith('bare_')).length;
}, { timeout: 20000, message: 'the monitor processed the bare output while its run was open' }).toBe(1);
const [bare] = waiting.filter(name => name.startsWith('bare_'));
expect((await readLinks(execution.databasePath, [bare])).links[bare]).toBeNull();
expect((await request.post('/pm_execution_fixture/gate/release')).ok()).toBe(true);
await linkedTo(execution, bare, llmText);
expect((await (await request.get('/pm_execution_fixture/monitor')).json()).awaiting).not.toContain(bare);
} finally {
expect((await request.post('/pm_execution_fixture/gate/release')).ok()).toBe(true);
}
} finally {
expect((await request.post('/pm_execution_fixture/metadata', { data: { disabled: false } })).ok()).toBe(true);
}
});
+5 -40
View File
@@ -1,7 +1,6 @@
"""Bounded read-only SQLite inspection for the owned execution fixture."""
import argparse
from contextlib import closing
import json
from pathlib import Path
import sqlite3
@@ -90,48 +89,14 @@ def read_runs(database: str, job_id: str, session_id: str | None = None) -> dict
}
def read_links(database: str, filenames: list, texts: list) -> dict:
"""Each output file's linked prompt text (None: unlinked), and which texts
are stored as prompts (#342)."""
uri = Path(database).resolve().as_uri() + "?mode=ro"
with closing(sqlite3.connect(uri, uri=True, timeout=0.25)) as connection:
connection.execute("PRAGMA query_only=ON")
links = {}
for name in filenames[:64]:
rows = connection.execute(
"SELECT p.text FROM generated_images g JOIN prompts p "
"ON p.id = g.prompt_id WHERE g.filename = ? LIMIT 2",
(name,),
).fetchall()
values = [row[0] for row in rows]
links[name] = values[0] if len(values) == 1 else None
if len(values) > 1:
links[name] = "<several links>"
stored = {
text: connection.execute(
"SELECT 1 FROM prompts WHERE text = ? LIMIT 1", (text,)
).fetchone()
is not None
for text in texts[:64]
}
return {"links": links, "stored": stored}
if __name__ == "__main__":
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("database")
parser.add_argument("job_id", nargs="?")
parser.add_argument("job_id")
parser.add_argument("--session-id")
parser.add_argument("--links", nargs="*", default=None)
parser.add_argument("--texts", nargs="*", default=[])
arguments = parser.parse_args()
if arguments.links is not None:
print(
json.dumps(read_links(arguments.database, arguments.links, arguments.texts))
)
else:
print(
json.dumps(
read_runs(arguments.database, arguments.job_id, arguments.session_id)
)
print(
json.dumps(
read_runs(arguments.database, arguments.job_id, arguments.session_id)
)
)
+4 -31
View File
@@ -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 () => {
-5
View File
@@ -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']]);
-166
View File
@@ -1,166 +0,0 @@
"""Run the Python suite with a hang watchdog: ``python -m tests.run_suite``.
Why: on PR #349 a macOS job stopped inside the test suite and sat silent until
GitHub cancelled it after six hours, with no stack to read. This runner does
the same discovery as ``python -m unittest discover -s tests`` (same
``sys.path`` when started with ``-m`` from the repo root, same warning
filters, and the exit codes of Python 3.12+ on every version: 0, 1, and 5
when no tests ran), and adds a watchdog built on
``faulthandler.dump_traceback_later``: when the timer is not re-armed in time,
every thread's stack goes to stderr and the process exits non-zero. It works
the same on Linux, macOS and Windows.
The timer is armed before discovery, re-armed at the start of every test with
PM_TEST_HANG_LIMIT seconds (default 120), and re-armed when the run ends with
PM_TEST_SHUTDOWN_LIMIT seconds (default 120) to cover interpreter exit
(non-daemon threads, atexit handlers, coverage's save). It is never
cancelled. Class and module fixtures (setUpClass, tearDownModule, ...) run
between two ``start:`` lines, so they count toward the test before them.
120 s stays below asyncio's 300 s wait for default-executor threads on Python
3.12 and 3.13 (3.10 and 3.11 wait without limit), so a hang in an
``IsolatedAsyncioTestCase`` loop shutdown is dumped inside the test that hung.
Reading a dump: the last ``start: <test id>`` line names the test. If its
"ok" line comes after it, the hang is after the test body: its teardown or
event-loop shutdown, or a class/module fixture before the next test (the
stack shows which). A dump after the summary is a hang at interpreter exit
(a non-daemon thread or an atexit handler).
A limit is a measurement, not a dial: change it only with new measurements
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
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):
"""An environment limit is not a positive integer."""
def read_limit(name, default):
"""Return the positive integer in env var ``name``, or ``default``."""
raw = os.environ.get(name)
if raw is None:
return default
text = raw.strip()
if not (text.isascii() and text.isdigit()) or int(text) <= 0:
raise LimitError(f"run_suite: {name} must be a positive integer, got {raw!r}")
return int(text)
def make_result_class(stream, per_test_limit, shutdown_limit):
"""Build a TextTestResult that names each test and re-arms the watchdog."""
def arm(seconds):
faulthandler.dump_traceback_later(seconds, exit=True, file=stream)
class WatchdogResult(unittest.TextTestResult):
def startTest(self, test):
# At default verbosity the progress dots share a line; start ours.
prefix = "\n" if self.dots else ""
stream.write(f"{prefix}start: {test.id()}\n")
stream.flush()
super().startTest(test)
arm(per_test_limit)
def stopTestRun(self):
super().stopTestRun()
arm(shutdown_limit)
return WatchdogResult
def parse_args(argv):
parser = argparse.ArgumentParser(
prog="python -m tests.run_suite",
description="Discover and run the tests with a hang watchdog.",
)
# As in unittest, -v and -q share one setting and the last one wins.
parser.set_defaults(verbosity=1)
parser.add_argument(
"-v",
"--verbose",
dest="verbosity",
action="store_const",
const=2,
help="verbose output",
)
parser.add_argument(
"-q",
"--quiet",
dest="verbosity",
action="store_const",
const=0,
help="quiet output",
)
parser.add_argument(
"-s",
"--start-directory",
default=str(Path(__file__).resolve().parent),
help="directory to start discovery (default: this tests folder)",
)
parser.add_argument(
"-p", "--pattern", default="test*.py", help="pattern to match test files"
)
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
# needs a stream with a file descriptor.
stream = sys.stderr
try:
per_test_limit = read_limit("PM_TEST_HANG_LIMIT", DEFAULT_PER_TEST_LIMIT)
shutdown_limit = read_limit("PM_TEST_SHUTDOWN_LIMIT", DEFAULT_SHUTDOWN_LIMIT)
except LimitError as error:
stream.write(f"{error}\n")
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)
if result.testsRun == 0 and not result.skipped:
return NO_TESTS_EXIT
return 0 if result.wasSuccessful() else 1
if __name__ == "__main__":
sys.exit(main())
-1
View File
@@ -2466,7 +2466,6 @@ class TestActivityJobs(AdminAPITestCase):
ambiguous=None,
untraced=None,
records=None,
history=None,
):
key = Path(media_file).relative_to(self.output_dir).as_posix()
outcome = names[key]
+2 -16
View File
@@ -687,23 +687,9 @@ class TestInstanceHelpers(FolderPathsFixture):
self.assertEqual(self._run(self.api._run_in_executor(add, 1, 2)), 3)
def test_invalidate_gallery_cache(self):
record = self.api._root_listing(self.output_dir)
record.files = [(self.output_dir / "a.png", 1.0)]
record.folders = record.folders_with_ancestors = frozenset()
record.walked_at = 1.0
self.api._favorite_listing_cache = {"x": 1}
self.api._workspace_order_cache = {"x": 1}
self.api._gallery_cache = {"x": 1}
self.api.invalidate_gallery_cache()
# The record stays: its generation outlives every drop (#348).
self.assertIs(self.api._root_listings[str(self.output_dir)], record)
self.assertIsNone(record.files)
self.assertIsNone(record.folders)
self.assertIsNone(record.walked_at)
self.assertEqual(record.generation, 1)
self.assertEqual(self.api._favorite_listing_cache, {})
self.assertEqual(self.api._workspace_order_cache, {})
self.api.reset_gallery_roots_cache()
self.assertEqual(record.generation, 2)
self.assertEqual(self.api._gallery_cache, {})
def test_enrich_images_adds_urls_and_thumbnails(self):
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
+11 -2020
View File
File diff suppressed because it is too large Load Diff
-183
View File
@@ -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 -71
View File
@@ -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
),
@@ -887,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,
@@ -918,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()
@@ -1005,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"
)
@@ -1039,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
-39
View File
@@ -171,45 +171,6 @@ class TestWorkspaceImages(ImageAPITestCase):
self.assertIsNone(second["next_cursor"])
self.assertEqual(second["total"], 4)
async def test_an_unchanged_refresh_keeps_the_listing_and_derived_caches(self):
from py.api import gallery_listing
self.media("a.png", 300)
self.media("b.png", 200)
_, first = await self.get()
response = await self.client.put(
"/prompt_manager/images/favorite",
json={"root_index": 0, "path": "a.png"},
)
self.assertEqual(response.status, 200)
record = self.api._root_listings[str(self.output_dir)]
files = record.files
key = ((0, str(self.output_dir)),)
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",
wraps=self.api._scan_gallery_files_sync,
) as scan:
with patch.object(gallery_listing, "REFRESH_DELAY", 0):
await self.get() # served at once; one refresh walk follows
while record.rewalk_handle is not None:
await asyncio.sleep(0.01)
while self.api._listing_tasks:
await asyncio.gather(*list(self.api._listing_tasks))
# A finished walk leaves the set in a done callback, which needs
# a loop turn; gather over finished tasks does not yield one.
await asyncio.sleep(0)
_, again = await self.get()
self.assertEqual(scan.call_count, 1)
self.assertIs(record.files, files)
self.assertIs(self.api._workspace_order_cache[key][1]["newest"], newest)
self.assertIs(self.api._favorite_listing_cache[str(self.output_dir)], favorites)
self.assertEqual(again["total"], first["total"])
async def test_filters_apply_before_limit_and_sorts_are_complete(self):
self.media("folder/Zebra.png", 200)
self.media("folder/apple.png", 300)
-16
View File
@@ -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()
+5 -12
View File
@@ -24,13 +24,6 @@ from utils.execution_history_events import (
)
from utils.execution_history_writer import HistoryWriter
# These calls must return while another thread holds the writer in a step that
# waits up to 3 s (release.wait(3)). 1 s still tells "returned at once" from
# "waited for the release" without failing on a loaded shared runner (#242:
# generous ceilings, never a short wall clock). A 0.1 s ceiling failed at
# 0.19 s on macOS py3.13 (PR #351).
NON_BLOCKING_CEILING = 1.0
class ManualClock:
def __init__(self):
@@ -597,10 +590,10 @@ class HistoryWriterTests(unittest.TestCase):
before = time.monotonic()
self.assertTrue(writer.status().accepting)
self.assertTrue(writer.submit(self.finish_command(start)))
self.assertLess(time.monotonic() - before, NON_BLOCKING_CEILING)
self.assertLess(time.monotonic() - before, 0.1)
before = time.monotonic()
self.assertFalse(writer.close(0))
self.assertLess(time.monotonic() - before, NON_BLOCKING_CEILING)
self.assertLess(time.monotonic() - before, 0.1)
self.assertGreater(writer.context_budget.charged_bytes, 0)
self.assertEqual(writer.status().pending, 1)
finally:
@@ -670,11 +663,11 @@ class HistoryWriterTests(unittest.TestCase):
before = time.monotonic()
self.assertFalse(writer.submit(self.start_command()))
self.assertIsNone(writer.status().epoch)
self.assertLess(time.monotonic() - before, NON_BLOCKING_CEILING)
self.assertLess(time.monotonic() - before, 0.1)
self.assertIsNone(self.db.get_execution_run(self.start_command().run_id))
before = time.monotonic()
self.assertFalse(writer.close(0.03))
self.assertLess(time.monotonic() - before, NON_BLOCKING_CEILING)
self.assertLess(time.monotonic() - before, 0.2)
release.set()
self.wait_for(lambda: not writer.status().alive)
self.assertFalse(writer.status().accepting)
@@ -965,7 +958,7 @@ class HistoryWriterTests(unittest.TestCase):
self.assertTrue(writer.submit(queued))
before = time.monotonic()
self.assertFalse(writer.close(0.03))
self.assertLess(time.monotonic() - before, NON_BLOCKING_CEILING)
self.assertLess(time.monotonic() - before, 0.2)
release.set()
self.wait_for(lambda: not writer.status().alive)
self.assertIsNone(self.db.get_execution_run(queued.run_id))
-478
View File
@@ -1,478 +0,0 @@
"""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 unittest.mock import patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from py.api import gallery_listing # noqa: E402
from py.api.gallery_listing import ( # noqa: E402
DISCARDED,
FAILED,
STORED,
RootListing,
WalkOutcome,
add_folder,
folder_sets,
insert_entry,
listing_delta,
relative_part,
rewalk_interval,
root_contains,
same_listing,
walk_root,
)
def _legacy_subfolders(root, files, include_ancestors):
"""``_root_subfolders`` as it was before #348, the reference for folder_sets."""
subfolders = set()
for f, _mtime in files:
rel_dir = str(f.relative_to(root).parent)
if rel_dir and rel_dir != ".":
subfolders.add(rel_dir)
if include_ancestors:
ancestors = set()
for folder in subfolders:
parts = folder.replace("\\", "/").split("/")
for i in range(1, len(parts)):
ancestors.add("/".join(parts[:i]))
subfolders.update(ancestors)
return subfolders
class FolderSetsTests(unittest.TestCase):
def test_nested_folders_match_the_previous_subfolder_listing(self):
root = Path("/out")
files = [
(root / "top.png", 1.0),
(root / "a" / "one.png", 2.0),
(root / "a" / "b" / "c" / "two.png", 3.0),
(root / "d" / "three.png", 4.0),
]
folders, with_ancestors = folder_sets(root, files)
self.assertEqual(folders, _legacy_subfolders(root, files, False))
self.assertEqual(with_ancestors, _legacy_subfolders(root, files, True))
self.assertIn("a/b", with_ancestors)
self.assertNotIn("a/b", folders)
self.assertIsInstance(folders, frozenset)
self.assertIsInstance(with_ancestors, frozenset)
def test_windows_separators_keep_the_previous_spelling(self):
root = PureWindowsPath("C:/out")
files = [
(root / "x" / "y" / "z.png", 1.0),
(root / "w.png", 2.0),
]
folders, with_ancestors = folder_sets(root, files)
self.assertEqual(folders, {"x\\y"})
self.assertEqual(with_ancestors, {"x\\y", "x"})
self.assertEqual(folders, _legacy_subfolders(root, files, False))
self.assertEqual(with_ancestors, _legacy_subfolders(root, files, True))
def test_an_empty_listing_has_no_folders(self):
self.assertEqual(folder_sets(Path("/out"), []), (frozenset(), frozenset()))
def test_folders_come_from_the_path_strings_not_relative_to(self):
# relative_to per file was half of a 20,000-file walk (#348 live test).
root = Path("/out")
files = [(root / "a" / "b" / "one.png", 1.0), (root / "two.png", 2.0)]
with patch.object(PurePath, "relative_to", side_effect=AssertionError):
folders, with_ancestors = folder_sets(root, files)
self.assertEqual(folders, {os.path.join("a", "b")})
self.assertEqual(with_ancestors, {os.path.join("a", "b"), "a"})
def test_a_path_not_spelled_like_the_root_falls_back_to_relative_to(self):
# Windows folds case: the string prefix misses, relative_to matches.
root = PureWindowsPath("C:/Out")
files = [(PureWindowsPath("c:/out/x/y.png"), 1.0)]
folders, with_ancestors = folder_sets(root, files)
self.assertEqual(folders, {"x"})
self.assertEqual(folders, _legacy_subfolders(root, files, False))
self.assertEqual(with_ancestors, _legacy_subfolders(root, files, True))
class SameListingTests(unittest.TestCase):
def test_equal_path_and_mtime_lists_are_the_same(self):
old = [(Path("/o/a.png"), 1.0), (Path("/o/b.png"), 2.0)]
self.assertTrue(same_listing(old, list(old)))
self.assertTrue(same_listing(old, old))
def test_a_changed_mtime_path_or_length_differs(self):
old = [(Path("/o/a.png"), 1.0), (Path("/o/b.png"), 2.0)]
self.assertFalse(same_listing(old, [(Path("/o/a.png"), 1.5), old[1]]))
self.assertFalse(same_listing(old, [(Path("/o/c.png"), 1.0), old[1]]))
self.assertFalse(same_listing(old, old[:1]))
self.assertFalse(same_listing(None, old))
class WalkRootTests(unittest.TestCase):
def setUp(self):
self._tmp = tempfile.TemporaryDirectory()
self.addCleanup(self._tmp.cleanup)
self.root = Path(self._tmp.name) / "output"
(self.root / "sub").mkdir(parents=True)
def _scan(self, files):
calls = []
def scan(root):
calls.append(root)
return list(files)
return scan, calls
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)
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)
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)
self.assertIsNot(listed, previous)
self.assertEqual(listed, files)
def test_a_deleted_root_raises_instead_of_listing_nothing(self):
scan, calls = self._scan([])
missing = self.root / "gone"
with self.assertRaises(OSError):
walk_root(scan, missing, None)
self.assertEqual(calls, [])
class RecordTests(unittest.TestCase):
def test_drop_clears_the_listing_and_bumps_the_generation(self):
record = RootListing()
record.files = [(Path("/o/a.png"), 1.0)]
record.folders = frozenset({"x"})
record.folders_with_ancestors = frozenset({"x"})
record.walked_at = 5.0
record.drop()
record.drop()
self.assertIsNone(record.files)
self.assertIsNone(record.folders)
self.assertIsNone(record.folders_with_ancestors)
self.assertIsNone(record.walked_at)
self.assertEqual(record.generation, 2)
def test_outcomes(self):
error = OSError("gone")
self.assertEqual(WalkOutcome.stored().kind, STORED)
self.assertEqual(WalkOutcome.discarded().kind, DISCARDED)
failed = WalkOutcome.failed(error)
self.assertEqual(failed.kind, FAILED)
self.assertIs(failed.error, error)
def test_refresh_constants(self):
self.assertEqual(gallery_listing.REFRESH_AFTER, 30)
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]))
+13 -17
View File
@@ -12,7 +12,6 @@ from database.operations import PromptDatabase
from utils.hashing import generate_prompt_hash
from utils.image_monitor import ImageGenerationHandler
from utils.prompt_tracker import PromptTracker
from utils.run_keys import prompt_node_key
LOADER = {
"class_type": "CheckpointLoaderSimple",
@@ -104,11 +103,8 @@ class TestMetadataLinking(LinkingTestCase):
batch_graph = graph()
batch_graph["134"]["inputs"]["text"] = ["7", 0]
batch_graph["7"] = {"class_type": "PromptSearchList", "inputs": {}}
key = prompt_node_key(batch_graph, "134")
for text, prompt_id in (("item one", first), ("item two", second)):
self.tracker.set_current_prompt(
text, {"prompt_id": prompt_id}, run_key=key, run_id="run-1"
)
self.tracker.set_current_prompt("item one", {"prompt_id": first})
self.tracker.set_current_prompt("item two", {"prompt_id": second})
self.assertEqual(self._process({"prompt": batch_graph}), first)
self.assertEqual(self._process({"prompt": batch_graph}), second)
@@ -126,10 +122,9 @@ class TestMetadataLinking(LinkingTestCase):
self.assertEqual(
self.handler._find_prompt_from_metadata(metadata), ("unknown", None)
)
self._save("most recent prompt") # never picked by recency (#342)
self.assertIsNone(self._process(metadata))
self.assertEqual(self.linked, [])
self.assertIsNotNone(negative)
latest = self._save("most recent prompt") # recency fallback picks this
self.assertEqual(self._process(metadata), latest)
self.assertNotEqual(latest, negative)
def test_workflow_fallback_still_used_without_an_api_graph(self):
positive = self._save("from workflow only")
@@ -205,14 +200,12 @@ class TestForeignWorkflows(LinkingTestCase):
self.handler._find_prompt_from_metadata(metadata), ("foreign", None)
)
def test_missing_metadata_is_unknown_and_left_unlinked(self):
def test_missing_metadata_is_unknown_and_uses_the_queue(self):
first = self._queue_batch_item("item one")
self.assertEqual(
self.handler._find_prompt_from_metadata(None), ("unknown", None)
)
# Nothing proves which run made it (#342): unlinked, queue untouched
self.assertIsNone(self._process(None))
self.assertEqual(self.tracker.pop_next_prompt()["id"], first)
self.assertEqual(self._process(None), first)
def test_tri_state_for_linked_unknown_and_foreign(self):
positive = self._save("a cat")
@@ -231,11 +224,14 @@ class TestForeignWorkflows(LinkingTestCase):
("unknown", None),
)
def test_unknown_image_is_never_given_the_latest_prompt(self):
def test_unknown_image_without_any_context_falls_back_to_the_latest_prompt(self):
self._save("older")
self._save("latest")
latest = self._save("latest")
self.assertEqual(self._process(None), latest)
def test_recent_prompt_fallback_handles_an_empty_database(self):
self.assertIsNone(self.handler._get_fallback_prompt())
self.assertIsNone(self._process(None))
self.assertIsNone(self._process({"prompt": graph("not saved")}))
def test_metadata_lookup_errors_are_unknown(self):
self._save("a cat")
+2 -112
View File
@@ -326,7 +326,7 @@ class TestHandlerScheduling(unittest.TestCase):
self.tmp = tempfile.TemporaryDirectory()
self.addCleanup(self.tmp.cleanup)
self.tracker = MagicMock()
self.tracker.keyed_snapshot.return_value = {"key": ({"id": 7}, "run", 1)}
self.tracker.get_current_prompt.return_value = {"id": 7}
self.handler = im_mod.ImageGenerationHandler(MagicMock(), self.tracker)
self.handler.processing_delay = 0
self.handler.settle_interval = 0.01
@@ -354,14 +354,7 @@ class TestHandlerScheduling(unittest.TestCase):
self.handler.on_created(FakeEvent(os.path.join(self.tmp.name, "c.txt")))
self.handler.on_created(FakeEvent(self.tmp.name, is_directory=True))
self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0))
self.assertEqual(self.processed, [(path, {"key": ({"id": 7}, "run", 1)})])
def test_a_failed_snapshot_still_schedules_the_image(self):
self.tracker.keyed_snapshot.side_effect = RuntimeError("tracker gone")
path = self._image("d.png")
self.handler.on_created(FakeEvent(path))
self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0))
self.assertEqual(self.processed, [(path, None)])
self.assertEqual(self.processed, [(path, {"id": 7})])
def test_images_are_processed_in_creation_order(self):
paths = [self._image(f"{i}.png") for i in range(5)]
@@ -449,109 +442,6 @@ class TestHandlerScheduling(unittest.TestCase):
self.assertEqual(self.processed, [first, second])
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(self.handler.shutdown)
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:
monitor.stop_monitoring()
class TestModuleIsolation(unittest.TestCase):
"""This module must not leave mocks behind for later test modules."""
-35
View File
@@ -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()
+28 -312
View File
@@ -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,17 +397,7 @@ 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"]})
@@ -535,7 +409,7 @@ class LoraImportRouteTests(LoraImportAPITestCase):
# 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())
@@ -691,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())
@@ -980,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):
@@ -996,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))
@@ -1113,194 +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)
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__":
-858
View File
@@ -1,858 +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"
class LayoutTestCase(unittest.TestCase):
"""A disposable LoraManager install and the folders its settings name."""
def setUp(self):
self.tmp = Path(tempfile.mkdtemp()).resolve()
self.addCleanup(shutil.rmtree, self.tmp, True)
self.plugin = self.tmp / "custom_nodes" / "ComfyUI-Lora-Manager"
(self.plugin / "py").mkdir(parents=True)
(self.plugin / "__init__.py").write_text("")
self.user_config = self.tmp / "user-config" / "ComfyUI-LoRA-Manager"
self.env_dir = self.tmp / "env-settings"
self.default_loras = self.tmp / "models" / "loras"
self.default_loras.mkdir(parents=True)
self.environ = {}
self.config_calls = []
def user_config_dir(self):
self.config_calls.append(True)
return str(self.user_config)
def write_settings(self, folder, data):
folder.mkdir(parents=True, exist_ok=True)
path = folder / "settings.json"
path.write_text(json.dumps(data), encoding="utf-8")
return path
def load(self, path="__plugin__", **kwargs):
lora_manager = str(self.plugin) if path == "__plugin__" else path
kwargs.setdefault("default_roots", lambda: [str(self.default_loras)])
kwargs.setdefault("environ", self.environ)
kwargs.setdefault("user_config_dir", self.user_config_dir)
return load_layout(lora_manager, **kwargs)
def model(self, folder, stem):
folder.mkdir(parents=True, exist_ok=True)
path = folder / f"{stem}.safetensors"
path.write_bytes(b"model")
return path
def sidecar(self, folder, stem, data=None):
folder.mkdir(parents=True, exist_ok=True)
path = folder / f"{stem}{METADATA_SUFFIX}"
path.write_text(json.dumps(data if data is not None else {}), "utf-8")
return path
class SettingsLookupTests(LayoutTestCase):
def test_environment_folder_wins_over_portable_and_user_config(self):
env_file = self.write_settings(self.env_dir, {"sidecar_storage_mode": "x"})
self.write_settings(self.plugin, {"use_portable_settings": True})
self.write_settings(self.user_config, {})
self.environ["LORA_MANAGER_SETTINGS_DIR"] = str(self.env_dir)
layout = self.load()
self.assertEqual(layout.source, "environment")
self.assertEqual(layout.settings_path, env_file)
self.assertIsNone(layout.error)
self.assertEqual(self.config_calls, [])
def test_an_empty_environment_folder_is_ignored(self):
self.environ["LORA_MANAGER_SETTINGS_DIR"] = ""
self.assertEqual(self.load().source, "user config")
def test_portable_flag_in_the_plugin_settings(self):
plugin_file = self.write_settings(self.plugin, {"use_portable_settings": True})
layout = self.load()
self.assertEqual(layout.source, "portable")
self.assertEqual(layout.settings_path, plugin_file)
self.assertEqual(self.config_calls, [])
def test_portable_environment_one_turns_it_on(self):
self.environ["LORA_MANAGER_PORTABLE"] = "1"
layout = self.load()
self.assertEqual(layout.source, "portable")
self.assertEqual(layout.settings_path, self.plugin / "settings.json")
def test_portable_environment_zero_turns_it_off_despite_the_file(self):
self.write_settings(self.plugin, {"use_portable_settings": True})
self.environ["LORA_MANAGER_PORTABLE"] = "0"
self.assertEqual(self.load().source, "user config")
def test_other_portable_environment_values_are_ignored(self):
for value in ("true", "yes"):
with self.subTest(value=value):
self.environ["LORA_MANAGER_PORTABLE"] = value
self.assertEqual(self.load().source, "user config")
self.write_settings(self.plugin, {"use_portable_settings": True})
self.assertEqual(self.load().source, "portable")
(self.plugin / "settings.json").unlink()
def test_a_string_portable_flag_is_ignored(self):
self.write_settings(self.plugin, {"use_portable_settings": "true"})
self.assertEqual(self.load().source, "user config")
def test_an_unreadable_plugin_file_means_not_portable(self):
(self.plugin / "settings.json").write_text("{not json")
self.assertEqual(self.load().source, "user config")
def test_user_config_folder_is_the_injected_one(self):
user_file = self.write_settings(self.user_config, {})
layout = self.load()
self.assertEqual(layout.source, "user config")
self.assertEqual(layout.settings_path, user_file)
self.assertEqual(self.config_calls, [True])
def test_the_real_user_config_folder_is_never_opened(self):
opened = []
real_open = open
def spy(path, *args, **kwargs):
opened.append(os.path.realpath(str(path)))
return real_open(path, *args, **kwargs)
with patch("builtins.open", spy):
layout = self.load()
self.assertEqual(layout.settings_path, self.user_config / "settings.json")
for path in opened:
self.assertTrue(path.startswith(str(self.tmp)), path)
def test_a_missing_settings_file_means_defaults(self):
layout = self.load()
self.assertIsNone(layout.error)
self.assertEqual(layout.mode, "alongside")
self.assertIsNone(layout.example_images_root)
self.assertIsNone(layout.sidecar_root)
self.assertEqual(layout.library, "default")
self.assertEqual(
layout.lora_roots, (os.path.normcase(str(self.default_loras)),)
)
def test_settings_dir_is_expanded_and_made_absolute(self):
self.write_settings(self.env_dir, {})
# Relative to the process folder (relpath fails across Windows drives).
# The cleanup returns before the temp folder is removed (LIFO).
self.addCleanup(os.chdir, os.getcwd())
os.chdir(self.tmp)
self.environ["LORA_MANAGER_SETTINGS_DIR"] = self.env_dir.name
self.assertEqual(self.load().settings_path, self.env_dir / "settings.json")
class UnusableLayoutTests(LayoutTestCase):
def test_no_install(self):
layout = self.load(path=None)
self.assertEqual(layout.error, "no LoraManager install found")
self.assertIsNone(layout.settings_path)
self.assertIsNone(layout.source)
# The default LoRA roots stay known for serving and trigger words.
self.assertEqual(
layout.lora_roots, (os.path.normcase(str(self.default_loras)),)
)
self.assertEqual(layout.media_roots, layout.lora_roots)
self.assertEqual(discover_sidecars(layout).sidecars, [])
def test_unreadable_settings(self):
self.write_settings(self.user_config, {})
with patch.object(layout_module, "_read_bytes", side_effect=PermissionError):
layout = self.load()
self.assertEqual(layout.error, "the settings file could not be read")
def test_settings_that_are_not_json(self):
self.user_config.mkdir(parents=True)
(self.user_config / "settings.json").write_text("{nope")
self.assertEqual(self.load().error, "the settings file is not valid JSON")
def test_settings_that_are_not_utf8(self):
self.user_config.mkdir(parents=True)
(self.user_config / "settings.json").write_bytes(b'{"a": "\xff"}')
self.assertEqual(self.load().error, "the settings file is not valid JSON")
def test_settings_that_are_not_an_object(self):
self.write_settings(self.user_config, ["a"])
self.assertEqual(self.load().error, "the settings file is not a JSON object")
def test_oversized_settings(self):
self.user_config.mkdir(parents=True)
with open(self.user_config / "settings.json", "wb") as handle:
handle.write(b" " * (layout_module.MAX_SETTINGS_BYTES + 1))
self.assertEqual(self.load().error, "the settings file is larger than 5 MB")
def test_platformdirs_import_failure(self):
def missing():
raise ImportError("no platformdirs")
layout = self.load(user_config_dir=missing)
self.assertEqual(layout.error, "platformdirs is not available")
self.assertEqual(layout.mode, "alongside")
def test_default_user_config_dir_needs_platformdirs(self):
with patch.dict(sys.modules, {"platformdirs": None}):
with self.assertRaises(ImportError):
layout_module._platform_user_config_dir()
def test_default_user_config_dir_uses_the_app_name(self):
fake = type(sys)("platformdirs")
calls = []
fake.user_config_dir = lambda *a, **k: calls.append((a, k)) or "/x"
with patch.dict(sys.modules, {"platformdirs": fake}):
self.assertEqual(layout_module._platform_user_config_dir(), "/x")
self.assertEqual(calls, [(("ComfyUI-LoRA-Manager",), {"appauthor": False})])
def test_missing_centralized_sidecar_root(self):
self.write_settings(
self.user_config,
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": str(self.tmp / "gone"),
},
)
layout = self.load()
self.assertEqual(layout.error, "the sidecar folder does not exist")
self.assertEqual(layout.mode, "centralized")
self.assertEqual(
layout.lora_roots, (os.path.normcase(str(self.default_loras)),)
)
self.assertEqual(layout.media_roots, layout.lora_roots)
self.assertEqual(discover_sidecars(layout).scanned, 0)
def test_a_centralized_sidecar_root_that_is_a_filesystem_root(self):
self.write_settings(
self.user_config,
{"sidecar_storage_mode": "centralized", "sidecar_storage_path": "/"},
)
layout = self.load()
self.assertEqual(
layout.error, "the sidecar folder is a filesystem root or the home folder"
)
self.assertNotIn(os.path.normcase(os.path.realpath("/")), layout.media_roots)
class SecretTests(LayoutTestCase):
def test_other_settings_keys_never_leave_the_module(self):
# Raw text keeps credential-named keys out of test variables.
self.user_config.mkdir(parents=True)
(self.user_config / "settings.json").write_text(
'{"civitai_api_key": "%s", "nested": {"token": "%s"}}' % (MARKER, MARKER),
encoding="utf-8",
)
layout = self.load()
self.assertNotIn(MARKER, repr(layout))
self.assertNotIn(MARKER, json.dumps(layout_summary(layout, Path.home())))
def test_only_the_seven_keys_are_kept(self):
self.assertEqual(
layout_module.SETTINGS_KEYS,
frozenset(
{
"sidecar_storage_mode",
"sidecar_storage_path",
"example_images_path",
"folder_paths",
"extra_folder_paths",
"libraries",
"active_library",
}
),
)
kept = layout_module._settings_subset(
json.loads('{"civitai_api_key": "%s", "libraries": {}}' % MARKER)
)
self.assertEqual(kept, {"libraries": {}})
class ActiveLibraryTests(LayoutTestCase):
def setUp(self):
super().setUp()
self.top = self.tmp / "top-loras"
self.lib_a = self.tmp / "a-loras"
self.lib_b = self.tmp / "b-loras"
self.extra = self.tmp / "b-extra"
for folder in (self.top, self.lib_a, self.lib_b, self.extra):
folder.mkdir()
self.examples = self.tmp / "examples"
self.examples.mkdir()
def settings(self, **extra):
data = {
"folder_paths": {"loras": [str(self.top)]},
"libraries": {
"first lib": {"folder_paths": {"loras": [str(self.lib_a)]}},
"second": {
"folder_paths": {"loras": [str(self.lib_b)]},
"extra_folder_paths": {"loras": [str(self.extra)]},
},
},
"example_images_path": str(self.examples),
}
data.update(extra)
self.write_settings(self.user_config, data)
def test_the_active_library_lists_are_used(self):
self.settings(active_library="second")
(self.examples / "second").mkdir()
layout = self.load()
self.assertEqual(layout.library, "second")
self.assertIn(os.path.normcase(str(self.lib_b)), layout.lora_roots)
self.assertIn(os.path.normcase(str(self.extra)), layout.lora_roots)
self.assertNotIn(os.path.normcase(str(self.top)), layout.lora_roots)
self.assertNotIn(os.path.normcase(str(self.lib_a)), layout.lora_roots)
self.assertIn(os.path.normcase(str(self.default_loras)), layout.lora_roots)
self.assertEqual(
layout.example_images_root, os.path.normcase(str(self.examples / "second"))
)
def test_an_unknown_active_library_falls_back_to_the_first(self):
self.settings(active_library="missing")
(self.examples / "first_lib").mkdir()
layout = self.load()
self.assertEqual(layout.library, "first lib")
self.assertIn(os.path.normcase(str(self.lib_a)), layout.lora_roots)
# The sanitized library name is the example-images subfolder.
self.assertEqual(
layout.example_images_root,
os.path.normcase(str(self.examples / "first_lib")),
)
def test_no_libraries_means_default_and_the_top_level_lists(self):
self.write_settings(
self.user_config,
{
"folder_paths": {"loras": str(self.top)},
"extra_folder_paths": {"loras": [str(self.extra), 7]},
"example_images_path": str(self.examples),
},
)
layout = self.load()
self.assertEqual(layout.library, "default")
self.assertIn(os.path.normcase(str(self.top)), layout.lora_roots)
self.assertIn(os.path.normcase(str(self.extra)), layout.lora_roots)
self.assertEqual(
layout.example_images_root, os.path.normcase(str(self.examples))
)
def test_one_library_has_no_library_subfolder(self):
self.write_settings(
self.user_config,
{
"libraries": {
"comfyui": {"folder_paths": {"loras": [str(self.lib_a)]}}
},
"active_library": "comfyui",
"example_images_path": str(self.examples),
},
)
layout = self.load()
self.assertEqual(layout.library, "comfyui")
self.assertEqual(
layout.example_images_root, os.path.normcase(str(self.examples))
)
def test_a_mixed_case_library_folder_keeps_its_example_images(self):
# Review finding: the library subfolder must be canonical like its
# root, or every example image fails the media check on Windows
# (normcase lowercases there; native proof is the Windows CI job).
self.write_settings(
self.user_config,
{
"libraries": {"Studio A": {}, "other": {}},
"active_library": "Studio A",
"example_images_path": str(self.examples),
},
)
image = self.examples / "Studio_A" / "abc" / "1.png"
image.parent.mkdir(parents=True)
image.write_bytes(b"x")
model = self.default_loras / "m.safetensors"
model.write_bytes(b"m")
(self.default_loras / "m.metadata.json").write_text("{}")
layout = self.load()
self.assertEqual(
layout.example_images_root,
os.path.normcase(os.path.realpath(self.examples / "Studio_A")),
)
entry = discover_sidecars(layout).sidecars[0]
self.assertEqual(
local_media(entry, {"sha256": "ABC"}, layout, MEDIA),
[os.path.realpath(image)],
)
def test_a_missing_library_subfolder_means_no_example_images(self):
self.settings(active_library="second")
self.assertIsNone(self.load().example_images_root)
def test_an_unset_or_missing_example_folder_means_none(self):
self.write_settings(self.user_config, {"example_images_path": " "})
self.assertIsNone(self.load().example_images_root)
self.write_settings(
self.user_config, {"example_images_path": str(self.tmp / "gone")}
)
self.assertIsNone(self.load().example_images_root)
def test_sanitize_library_name(self):
sanitize = layout_module._sanitize_library_name
self.assertEqual(sanitize("My Lib/2"), "My_Lib_2")
self.assertEqual(sanitize("ok-name_1.0"), "ok-name_1.0")
self.assertEqual(sanitize(""), "default")
def test_odd_library_shapes_are_ignored(self):
self.write_settings(
self.user_config,
{"libraries": {"x": "not a dict"}, "folder_paths": ["not", "a dict"]},
)
layout = self.load()
self.assertEqual(layout.library, "x")
self.assertEqual(
layout.lora_roots, (os.path.normcase(str(self.default_loras)),)
)
class CentralizedDiscoveryTests(LayoutTestCase):
def setUp(self):
super().setUp()
self.loras = self.tmp / "loras"
self.checkpoints = self.tmp / "checkpoints"
self.checkpoints.mkdir()
self.sidecars = self.tmp / "sidecars"
self.sidecars.mkdir()
self.write_settings(
self.user_config,
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": str(self.sidecars),
"folder_paths": {"loras": [str(self.loras)]},
},
)
self.loras.mkdir()
def test_component_folders_and_the_legacy_nested_form(self):
a = self.model(self.loras, "alpha")
b = self.model(self.loras / "sub", "beta")
self.sidecar(self.sidecars / "loras-1a2b3c4d", "alpha", {"file_path": str(a)})
self.sidecar(
self.sidecars / "comfyui" / "loras-1a2b3c4d" / "sub",
"beta",
{"file_path": b.as_posix()},
)
discovery = discover_sidecars(self.load())
self.assertEqual(discovery.scanned, 2)
self.assertEqual(
sorted(s.model_path.name for s in discovery.sidecars),
["alpha.safetensors", "beta.safetensors"],
)
self.assertEqual(discovery.unreadable, [])
self.assertEqual(
[str(s.metadata_path) for s in discovery.sidecars],
sorted(str(s.metadata_path) for s in discovery.sidecars),
)
def test_missing_models_and_checkpoints_are_skipped(self):
ckpt = self.model(self.checkpoints, "base")
self.sidecar(self.sidecars / "c", "base", {"file_path": str(ckpt)})
self.sidecar(
self.sidecars / "l",
"gone",
{"file_path": str(self.loras / "gone.safetensors")},
)
self.sidecar(self.sidecars / "l", "nopath", {"model_name": "x"})
discovery = discover_sidecars(self.load())
self.assertEqual(discovery.scanned, 3)
self.assertEqual(discovery.sidecars, [])
def test_leftover_sidecars_beside_models_are_not_scanned(self):
model = self.model(self.loras, "alpha")
self.sidecar(self.loras, "alpha", {"file_path": str(model)})
discovery = discover_sidecars(self.load())
self.assertEqual((discovery.scanned, discovery.sidecars), (0, []))
def test_unparseable_sidecars_are_recorded(self):
(self.sidecars / f"broken{METADATA_SUFFIX}").write_text("{{{")
(self.sidecars / f"list{METADATA_SUFFIX}").write_text("[1]")
discovery = discover_sidecars(self.load())
self.assertEqual(discovery.scanned, 2)
self.assertEqual(
sorted(discovery.unreadable),
[f"broken{METADATA_SUFFIX}", f"list{METADATA_SUFFIX}"],
)
def test_default_sidecar_root_is_beside_the_settings(self):
self.write_settings(
self.user_config,
{
"sidecar_storage_mode": "centralized",
"folder_paths": {"loras": [str(self.loras)]},
},
)
(self.user_config / "sidecars").mkdir()
layout = self.load()
self.assertEqual(
layout.sidecar_root, os.path.normcase(str(self.user_config / "sidecars"))
)
self.assertIn(layout.sidecar_root, layout.media_roots)
def test_an_unknown_mode_is_alongside(self):
self.write_settings(self.user_config, {"sidecar_storage_mode": "elsewhere"})
self.assertEqual(self.load().mode, "alongside")
def test_an_unlistable_folder_raises(self):
layout = self.load()
def refuse(error):
raise error
def failing_walk(top, onerror=None, followlinks=False):
onerror(PermissionError(13, "denied", top))
return iter(())
with patch.object(layout_module.os, "walk", failing_walk):
with self.assertRaises(OSError):
discover_sidecars(layout)
class AlongsideDiscoveryTests(LayoutTestCase):
def test_a_sidecar_with_its_model_is_kept(self):
model = self.model(self.default_loras, "alpha")
meta = self.sidecar(self.default_loras, "alpha")
self.sidecar(self.default_loras, "orphan")
discovery = discover_sidecars(self.load())
self.assertEqual(discovery.scanned, 2)
self.assertEqual(
discovery.sidecars,
[Sidecar(Path(os.path.realpath(meta)), Path(os.path.realpath(model)))],
)
def test_a_folder_reached_through_two_nested_roots_is_listed_once(self):
nested = self.default_loras / "nested"
self.model(nested, "alpha")
self.sidecar(nested, "alpha")
layout = self.load(default_roots=lambda: [str(self.default_loras), str(nested)])
discovery = discover_sidecars(layout)
self.assertEqual((discovery.scanned, len(discovery.sidecars)), (1, 1))
@unittest.skipUnless(hasattr(os, "symlink"), "symlinks")
def test_a_symlink_loop_is_walked_once(self):
self.model(self.default_loras, "alpha")
self.sidecar(self.default_loras, "alpha")
try:
os.symlink(self.default_loras, self.default_loras / "loop")
except (OSError, NotImplementedError):
self.skipTest("symlinks need privileges here")
discovery = discover_sidecars(self.load())
self.assertEqual((discovery.scanned, len(discovery.sidecars)), (1, 1))
class LocalMediaTests(LayoutTestCase):
def setUp(self):
super().setUp()
self.examples = self.tmp / "examples"
self.examples.mkdir()
self.write_settings(
self.user_config, {"example_images_path": str(self.examples)}
)
self.model(self.default_loras, "alpha")
self.meta = self.sidecar(self.default_loras, "alpha")
self.layout = self.load()
self.entry = discover_sidecars(self.layout).sidecars[0]
def media(self, metadata=None):
return local_media(self.entry, metadata or {}, self.layout, MEDIA)
def touch(self, path):
path.parent.mkdir(parents=True, exist_ok=True)
path.write_bytes(b"x")
return os.path.realpath(path)
def test_preview_url_comes_first(self):
preview = self.touch(self.default_loras / "custom.jpeg")
self.touch(self.default_loras / "alpha.webp")
self.assertEqual(self.media({"preview_url": preview}), [preview])
def test_preview_url_outside_media_roots_or_not_servable_is_ignored(self):
outside = self.touch(self.tmp / "elsewhere" / "p.png")
avif = self.touch(self.default_loras / "p.avif")
fallback = self.touch(self.default_loras / "alpha.png")
for url in (
outside,
avif,
"relative.png",
7,
str(self.default_loras / "x.png"),
):
with self.subTest(url=url):
self.assertEqual(self.media({"preview_url": url}), [fallback])
def test_suffix_order_follows_lora_manager(self):
self.assertEqual(
PREVIEW_EXTENSIONS,
(
".webp",
".preview.webp",
".preview.png",
".preview.jpeg",
".preview.jpg",
".preview.mp4",
".png",
".jpeg",
".jpg",
".mp4",
".gif",
".webm",
".avif",
".jxl",
),
)
mp4 = self.touch(self.default_loras / "alpha.mp4")
self.assertEqual(self.media(), [mp4])
webp = self.touch(self.default_loras / "alpha.preview.webp")
self.assertEqual(self.media(), [webp])
def test_an_avif_only_preview_gives_nothing(self):
self.touch(self.default_loras / "alpha.avif")
self.assertEqual(self.media(), [])
def test_example_images_in_numeric_order(self):
folder = self.examples / "abcdef"
names = ["image_10.png", "image_2.png", "image_1.jpg", "notes.txt"]
for name in names:
self.touch(folder / name)
(folder / "sub").mkdir()
preview = self.touch(self.default_loras / "alpha.png")
found = self.media({"sha256": "ABCDEF"})
self.assertEqual(
[os.path.basename(p) for p in found],
["alpha.png", "image_1.jpg", "image_2.png", "image_10.png"],
)
self.assertEqual(found[0], preview)
def test_missing_sha256_or_folder_gives_no_example_images(self):
self.touch(self.examples / "abc" / "1.png")
self.assertEqual(self.media(), [])
self.assertEqual(self.media({"sha256": "zzz"}), [])
self.assertEqual(self.media({"sha256": 5}), [])
def test_an_unlistable_example_folder_gives_what_was_found(self):
self.touch(self.examples / "abc" / "1.png")
with patch.object(layout_module.os, "scandir", side_effect=PermissionError):
self.assertEqual(self.media({"sha256": "abc"}), [])
def test_cover_moves_the_first_still_in_front_of_a_video(self):
video = self.touch(self.default_loras / "alpha.mp4")
still = self.touch(self.examples / "abc" / "2.png")
other = self.touch(self.examples / "abc" / "1.webm")
self.assertEqual(self.media({"sha256": "abc"}), [still, video, other])
def test_videos_only_keep_their_order(self):
video = self.touch(self.default_loras / "alpha.mp4")
self.assertEqual(self.media(), [video])
def test_repeats_are_dropped(self):
preview = self.touch(self.default_loras / "alpha.png")
self.assertEqual(self.media({"preview_url": preview}), [preview])
def test_video_suffixes_match_the_video_module(self):
from utils.video_metadata import VIDEO_EXTENSIONS
self.assertEqual(layout_module.VIDEO_SUFFIXES, frozenset(VIDEO_EXTENSIONS))
class FolderGuardTests(unittest.TestCase):
"""Portable path-flavor controls (TESTING.md Windows portability preflight)."""
def setUp(self):
self.tmp = Path(tempfile.mkdtemp()).resolve()
self.addCleanup(shutil.rmtree, self.tmp, True)
def guard_nt(self, path, home="C:\\Users\\Vito"):
return layout_module._guard_folder(
path, home, pathmod=ntpath, isdir=lambda _p: True
)
def test_filesystem_roots_and_home_are_rejected(self):
self.assertIsNone(layout_module._guard_folder("/", str(self.tmp)))
self.assertIsNone(self.guard_nt("C:\\"))
self.assertIsNone(self.guard_nt("c:/"))
self.assertIsNone(self.guard_nt("\\\\server\\share\\"))
self.assertIsNone(self.guard_nt("c:\\USERS\\vito"))
self.assertIsNone(self.guard_nt("C:\\Users\\Vito\\"))
self.assertEqual(
self.guard_nt("C:\\Users\\Vito\\Loras"), "c:\\users\\vito\\loras"
)
self.assertIsNone(layout_module._guard_folder(str(self.tmp), str(self.tmp)))
def test_ancestors_of_home_are_rejected(self):
# Security review (#355): a settings folder that contains the home
# folder would serve every user's files.
home = self.tmp / "home" / "me"
inside = home / "loras"
inside.mkdir(parents=True)
guard = layout_module._guard_folder
self.assertIsNone(guard(str(self.tmp / "home"), str(home)))
self.assertIsNone(guard(str(self.tmp), str(home)))
self.assertIsNone(guard(str(home), str(home)))
self.assertEqual(guard(str(inside), str(home)), os.path.normcase(str(inside)))
# A sibling sharing the name prefix is not an ancestor.
sibling = self.tmp / "home" / "me2"
sibling.mkdir()
self.assertEqual(guard(str(sibling), str(home)), os.path.normcase(str(sibling)))
def test_windows_ancestors_of_home_are_rejected(self):
self.assertIsNone(self.guard_nt("C:\\Users", home="C:\\Users\\me"))
self.assertIsNone(self.guard_nt("c:/users/", home="C:\\Users\\me"))
self.assertEqual(
self.guard_nt("C:\\Users\\me\\Loras", home="C:\\Users\\me"),
"c:\\users\\me\\loras",
)
self.assertEqual(
self.guard_nt("C:\\Users\\meta", home="C:\\Users\\me"),
"c:\\users\\meta",
)
def test_settings_folders_that_contain_home_are_not_served(self):
home = self.tmp / "home" / "me"
home.mkdir(parents=True)
settings = self.tmp / "settings"
settings.mkdir()
(settings / "settings.json").write_text(
json.dumps(
{
"folder_paths": {"loras": [str(self.tmp / "home")]},
"example_images_path": str(self.tmp),
}
)
)
with patch.object(layout_module.Path, "home", return_value=home):
layout = load_layout(
str(self.tmp),
default_roots=lambda: [],
environ={"LORA_MANAGER_SETTINGS_DIR": str(settings)},
user_config_dir=lambda: str(self.tmp / "unused"),
)
self.assertIsNone(layout.error)
self.assertEqual(layout.lora_roots, ())
self.assertIsNone(layout.example_images_root)
self.assertEqual(layout.media_roots, ())
def test_missing_folders_add_nothing(self):
self.assertIsNone(layout_module._guard_folder(str(self.tmp / "x"), "/home"))
self.assertIsNone(
layout_module._guard_folder("\\\\server\\share\\loras", str(self.tmp))
)
self.assertEqual(
layout_module._guard_folder(str(self.tmp), "/nonexistent-home"),
os.path.normcase(str(self.tmp)),
)
def test_windows_model_path_matches_a_lower_case_root(self):
self.assertTrue(
layout_module._inside(
"C:/Models/Loras/A.safetensors", ("c:\\models\\loras",), ntpath
)
)
self.assertFalse(
layout_module._inside(
"C:/Models/LorasX/A.safetensors", ("c:\\models\\loras",), ntpath
)
)
self.assertFalse(
layout_module._inside(
"\\\\other\\share\\A.safetensors", ("c:\\models\\loras",), ntpath
)
)
def test_mixed_case_native_containment(self):
inside = str(self.tmp / "Loras" / "a.png")
root = os.path.normcase(str(self.tmp / "Loras"))
self.assertTrue(layout_module._inside(inside, (root,)))
self.assertFalse(layout_module._inside(str(self.tmp / "other.png"), (root,)))
class SummaryTests(LayoutTestCase):
def test_summary_names_what_the_import_read(self):
examples = self.tmp / "ex"
examples.mkdir()
sidecars = self.plugin / "sidecars"
sidecars.mkdir()
self.write_settings(
self.plugin,
{
"use_portable_settings": True,
"sidecar_storage_mode": "centralized",
"example_images_path": str(examples),
"libraries": {"comfyui": {}},
},
)
layout = self.load()
summary = layout_summary(layout, self.tmp)
self.assertEqual(
summary,
{
"source": "portable",
"settings_file": os.path.join(
"~", "custom_nodes", "ComfyUI-Lora-Manager", "settings.json"
),
"error": None,
"mode": "centralized",
"library": "comfyui",
"example_images": True,
},
)
def test_summary_outside_home_and_without_settings(self):
summary = layout_summary(self.load(path=None), self.tmp / "home")
self.assertEqual(summary["settings_file"], None)
self.assertEqual(summary["error"], "no LoraManager install found")
self.assertFalse(summary["example_images"])
layout = self.load()
self.assertEqual(
layout_summary(layout, self.tmp / "home")["settings_file"],
str(self.user_config / "settings.json"),
)
if __name__ == "__main__":
unittest.main()
+52 -173
View File
@@ -1263,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)
@@ -1284,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"])
@@ -1512,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())
@@ -1520,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()
@@ -1542,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))
@@ -1560,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):
@@ -1590,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):
@@ -1674,161 +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_any_other_failure_never_raises(self):
cache = TriggerWordCache()
with (
patch.object(lora_utils, "load_layout", side_effect=RuntimeError("bad")),
self.assertLogs("prompt_manager.lora_utils", level="WARNING"),
):
self.assertEqual(cache.load(str(self.lm)), 0)
self.assertFalse(cache.is_loaded)
def test_an_unusable_layout_falls_back_to_the_default_roots(self):
(self.settings / "settings.json").write_text("{not json")
model = self.loras / "neon.safetensors"
model.write_bytes(b"m")
self.sidecar(self.loras, "neon", model, ["glow"])
cache = TriggerWordCache()
with (
patch.object(
lora_utils, "default_lora_roots", return_value=[str(self.loras)]
),
self.assertLogs("prompt_manager.lora_utils", level="WARNING") as logs,
):
self.assertEqual(cache.load(str(self.lm)), 1)
self.assertEqual(cache.get_trigger_words("neon"), ["glow"])
self.assertIn("not valid JSON", "\n".join(logs.output))
def test_layout_wraps_the_install_and_the_default_roots(self):
with patch.object(
lora_utils, "default_lora_roots", return_value=[str(self.loras)]
):
layout = lora_utils.lora_manager_layout(str(self.lm))
self.assertIsNone(layout.error)
self.assertEqual(layout.source, "environment")
self.assertIn(os.path.normcase(str(self.loras)), layout.lora_roots)
self.assertEqual(
lora_utils.lora_manager_layout(str(self.tmp)).error,
"no LoraManager install found",
)
if __name__ == "__main__":
unittest.main()
-37
View File
@@ -315,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
+23 -5
View File
@@ -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"))
-662
View File
@@ -1,662 +0,0 @@
"""#342: an image is linked only to a prompt that provably produced it.
The maintainer's run: an LLM writes the prompt into a PromptManagerPair's
``prepend_text`` while its ``text`` editor is empty. Prepend and append text are
never stored as the prompt (product decision), so nothing is saved, and the
image must stay unlinked instead of taking the newest prompt in the database or
a stale tracker entry.
"""
import copy
import logging
import os
import sys
import tempfile
import threading
import unittest
from unittest import mock
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
try:
from tests.open_handles import assert_closed
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed
from database.operations import PromptDatabase
from prompt_manager import PromptManager
from prompt_manager_pair import PromptManagerPair
from prompt_manager_text import PromptManagerText
from utils.hashing import generate_prompt_hash
from utils.image_monitor import ImageGenerationHandler
from utils.prompt_tracker import PromptTracker
from utils.run_keys import image_run_keys, prompt_node_key
from utils.usage_tracking import (
PendingFirstUse,
queued_pair_runs,
queued_prompt_runs,
)
LLM_TEXT = "Photorealistic aerial landscape of misty mountains at dawn"
EVIDENCE_KEY = "comfyassets_prompt_pairs"
def issue_graph(text="", seed=1):
"""The workflow from #342 (subgraph ids kept), PromptManagerPair positive."""
return {
"18:35": {
"class_type": "PromptManagerPair",
"inputs": {
"text": text,
"negative_text": "",
"prepend_text": ["18:13", 0],
"append_text": "",
"category": "",
"tags": "",
"clip": ["4", 1],
},
},
"18:13": {"class_type": "PreviewAny", "inputs": {"source": ["18:16", 0]}},
"18:16": {
"class_type": "ComfySwitchNode",
"inputs": {
"switch": True,
"on_false": ["18:17", 0],
"on_true": ["18:11", 0],
},
},
"18:11": {
"class_type": "TextGenerate",
"inputs": {"prompt": ["18:17", 0], "seed": 7},
},
"18:17": {
"class_type": "PrimitiveStringMultiline",
"inputs": {"value": "Describe an aerial mountain landscape"},
},
"18:5": {
"class_type": "CLIPTextEncode",
"inputs": {"text": "blurry", "clip": ["4", 1]},
},
"4": {
"class_type": "CheckpointLoaderSimple",
"inputs": {"ckpt_name": "moss.safetensors"},
},
"3": {
"class_type": "KSampler",
"inputs": {
"model": ["4", 0],
"positive": ["18:35", 0],
"negative": ["18:5", 0],
"latent_image": ["6", 0],
"seed": seed,
},
},
"6": {"class_type": "EmptyLatentImage", "inputs": {"batch_size": 1}},
}
def llm_text_graph(seed=1):
"""An LLM feeding a PromptManager's ``text``: stored, but not in the metadata."""
graph = issue_graph(seed=seed)
graph["18:35"] = {
"class_type": "PromptManager",
"inputs": {"text": ["18:13", 0], "clip": ["4", 1]},
}
return graph
def pair_evidence(text="", prepend=LLM_TEXT):
"""The execution record PromptManagerPair embeds in the saved image."""
return {
"version": 1,
"nodes": {
"18:35": {
"records": [
{
"text": text,
"negative_text": "",
"positive_text": " ".join(p for p in (prepend, text) if p),
"negative_effective_text": "",
"prepend_text": prepend,
"append_text": "",
}
]
}
},
}
def make_tracker(db):
"""A real tracker without its cleanup thread or config import."""
with mock.patch.object(PromptTracker, "__init__", lambda self, db: None):
tracker = PromptTracker.__new__(PromptTracker)
tracker.logger = logging.getLogger("test.prompt_attribution.tracker")
tracker.db_manager = db
tracker._local = threading.local()
tracker.active_prompts = {}
tracker.lock = threading.Lock()
tracker._prompt_queue = []
tracker._queue_lock = threading.Lock()
tracker.cleanup_interval = 300
tracker.prompt_timeout = 600
return tracker
class DatabaseCase(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
self.addCleanup(self.tmp.cleanup)
self.addCleanup(assert_closed, self, self.tmp.name)
self.db = PromptDatabase(os.path.join(self.tmp.name, "prompts.db"))
self.addCleanup(self.db.close_all)
self.tracker = make_tracker(self.db)
def save(self, text):
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
class NodeCase(DatabaseCase):
def node(self, cls):
node = cls.__new__(cls)
node.db = self.db
node.logger = logging.getLogger("test.prompt_attribution.node")
node.prompt_tracker = self.tracker
node.comfyui_integration = mock.Mock()
node._inject_lora_trigger_words = lambda text: text
node.pending_first_use = PendingFirstUse()
node.pending_pair_first_use = PendingFirstUse()
return node
def clip(self):
clip = mock.Mock()
clip.tokenize.side_effect = lambda text: text
clip.encode_from_tokens_scheduled.side_effect = lambda t: ("cond", t)
return clip
def count_rows(self, table):
conn = self.db.model.get_connection()
return conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]
class TestEmptyEditorSavesNothing(NodeCase):
"""Prepend/append are never the stored prompt; an empty editor saves nothing."""
def test_pair_with_only_prepend_text_saves_nothing_and_logs_info(self):
graph = issue_graph()
with self.assertLogs("test.prompt_attribution.node", "INFO") as logs:
result = self.node(PromptManagerPair).encode_prompt(
self.clip(),
"",
negative_text="",
prepend_text=LLM_TEXT,
prompt=graph,
unique_id="18:35",
)
self.assertEqual(result[2], LLM_TEXT, "the composed text is still encoded")
self.assertEqual(self.count_rows("prompts"), 0)
self.assertEqual(self.count_rows("prompt_pairs"), 0)
self.assertFalse(
any(record.levelno >= logging.WARNING for record in logs.records),
[record.getMessage() for record in logs.records],
)
self.assertTrue(any("Prompt text is empty" in line for line in logs.output))
def test_prompt_manager_and_text_nodes_with_only_prepend_save_nothing(self):
graph = issue_graph()
graph["18:35"] = {
"class_type": "PromptManager",
"inputs": {"text": "", "prepend_text": ["18:13", 0], "clip": ["4", 1]},
}
self.node(PromptManager).encode_prompt(
self.clip(), "", prepend_text=LLM_TEXT, prompt=graph, unique_id="18:35"
)
graph["18:35"]["class_type"] = "PromptManagerText"
self.node(PromptManagerText).process_text(
"", prepend_text=LLM_TEXT, prompt=graph, unique_id="18:35"
)
self.assertEqual(self.count_rows("prompts"), 0)
def test_empty_editor_marks_its_run_key_so_no_older_entry_is_used(self):
graph = issue_graph()
key = prompt_node_key(graph, "18:35")
older = self.save("a woman holding a coffee cup")
self.tracker.set_current_prompt(
"a woman holding a coffee cup",
{"prompt_id": older},
push_to_queue=True,
run_key=key,
run_id="run-1",
)
self.node(PromptManagerPair).encode_prompt(
self.clip(), "", prepend_text=LLM_TEXT, prompt=graph, unique_id="18:35"
)
latest, _run, _count = self.tracker.keyed_snapshot()[key]
self.assertTrue(latest.get("placeholder"))
self.assertIsNone(self.tracker.pop_for_key(key, "another run"))
popped = self.tracker.pop_for_key(key)
self.assertEqual(popped["id"], older, "an earlier queued item stays first")
self.assertTrue(self.tracker.pop_for_key(key).get("placeholder"))
def test_empty_negative_node_is_not_marked(self):
graph = issue_graph(text="a typed positive")
graph["18:5"] = {
"class_type": "PromptManager",
"inputs": {"text": "", "clip": ["4", 1]},
}
self.node(PromptManager).encode_prompt(
self.clip(), "", prompt=graph, unique_id="18:5"
)
self.assertEqual(self.tracker.keyed_snapshot(), {})
def test_queue_hook_counts_nothing_for_an_empty_editor(self):
graph = issue_graph()
self.assertEqual(queued_prompt_runs(graph), [])
self.assertEqual(queued_pair_runs(graph), [])
def test_typed_text_with_prepend_still_stores_only_the_editor_text(self):
graph = issue_graph(text="mountains")
self.node(PromptManagerPair).encode_prompt(
self.clip(),
"mountains",
prepend_text=LLM_TEXT,
prompt=graph,
unique_id="18:35",
)
self.assertIsNotNone(
self.db.get_prompt_by_hash(generate_prompt_hash("mountains"))
)
self.assertIsNone(self.db.get_prompt_by_hash(generate_prompt_hash(LLM_TEXT)))
context, _run, _count = self.tracker.keyed_snapshot()[
prompt_node_key(graph, "18:35")
]
self.assertEqual(context["text"], f"{LLM_TEXT} mountains")
class MonitorCase(DatabaseCase):
def setUp(self):
super().setUp()
fd, self.image = tempfile.mkstemp(suffix=".png", dir=self.tmp.name)
os.close(fd)
self.handler = ImageGenerationHandler(self.db, self.tracker)
self.linked = []
self.handler.link_image_to_prompt = lambda path, prompt, meta: (
self.linked.append(prompt["id"])
)
def process(self, metadata, snapshot=None):
before = len(self.linked)
with mock.patch.object(
self.handler.metadata_extractor, "extract_metadata", return_value=metadata
):
self.handler.process_new_image(self.image, prompt_snapshot=snapshot)
return self.linked[-1] if len(self.linked) > before else None
class TestIssueReproduction(MonitorCase):
def test_llm_prepend_image_is_left_unlinked_not_given_the_newest_prompt(self):
"""Failing before #342: linked to the newest prompt in the database."""
self.save("a woman holding a coffee cup")
metadata = {"prompt": issue_graph(), EVIDENCE_KEY: pair_evidence()}
with self.assertLogs("prompt_manager.image_monitor", "INFO") as logs:
self.assertIsNone(self.process(metadata))
self.assertTrue(any("left unlinked" in line for line in logs.output))
def test_a_stale_entry_of_another_workflow_is_never_used(self):
"""Failing before #342: the live tracker's entry was taken."""
stale = self.save("a woman holding a coffee cup")
other = issue_graph(text="a woman holding a coffee cup")
self.tracker.set_current_prompt(
"a woman holding a coffee cup",
{"prompt_id": stale},
push_to_queue=True,
run_key=prompt_node_key(other, "18:35"),
run_id="run-0",
)
metadata = {"prompt": issue_graph(), EVIDENCE_KEY: pair_evidence()}
self.assertIsNone(self.process(metadata, self.tracker.keyed_snapshot()))
self.assertIsNone(self.process(metadata))
self.assertEqual(len(self.tracker._prompt_queue), 1, "not consumed")
def test_an_entry_without_a_run_key_is_never_used(self):
stale = self.save("previous run")
self.tracker.set_current_prompt("previous run", {"prompt_id": stale})
self.assertIsNone(self.process({"prompt": llm_text_graph()}))
def test_an_image_without_metadata_is_left_unlinked(self):
"""Failing before #342: popped the queue, else took the newest prompt."""
self.save("latest")
self.assertIsNone(self.process(None))
def test_the_recency_fallback_is_gone(self):
self.assertFalse(hasattr(self.handler, "_get_fallback_prompt"))
class TestProvenRuns(MonitorCase):
"""Entries the image's own graph proves still link (LLM into ``text``)."""
def track(self, graph, text, run_id="run-1", push=True):
prompt_id = (
self.save(text)
if not self.db.get_prompt_by_hash(generate_prompt_hash(text))
else self.db.get_prompt_by_hash(generate_prompt_hash(text))["id"]
)
self.tracker.set_current_prompt(
text,
{"prompt_id": prompt_id},
push_to_queue=push,
run_key=prompt_node_key(graph, "18:35"),
run_id=run_id,
)
return prompt_id
def test_queued_entry_of_the_same_run_key_links(self):
graph = llm_text_graph()
prompt_id = self.track(graph, LLM_TEXT)
self.assertEqual(self.process({"prompt": graph}), prompt_id)
def test_cached_node_rerun_links_through_the_snapshot(self):
graph = llm_text_graph(seed=1)
prompt_id = self.track(graph, LLM_TEXT)
self.assertEqual(self.process({"prompt": graph}), prompt_id)
# Only the sampler seed changed: the node was cached and did not run,
# but its upstream (the run key) is unchanged.
rerun = llm_text_graph(seed=2)
self.assertEqual(
prompt_node_key(rerun, "18:35"), prompt_node_key(graph, "18:35")
)
self.assertEqual(
self.process({"prompt": rerun}, self.tracker.keyed_snapshot()), prompt_id
)
def test_snapshot_needs_a_known_run(self):
graph = llm_text_graph()
self.track(graph, LLM_TEXT, run_id=None, push=False)
self.assertIsNone(self.process({"prompt": graph}))
def test_batch_items_link_in_order_and_the_latest_alone_proves_nothing(self):
graph = llm_text_graph()
first = self.track(graph, "item one")
second = self.track(graph, "item two")
self.assertEqual(self.process({"prompt": graph}), first)
self.assertEqual(self.process({"prompt": graph}), second)
# Queue drained: two entries in one run, so the latest is not proof
self.assertIsNone(self.process({"prompt": graph}))
def test_an_empty_batch_item_keeps_the_queue_aligned(self):
graph = llm_text_graph()
key = prompt_node_key(graph, "18:35")
first = self.track(graph, "item one")
self.tracker.mark_empty(key, "run-1", push_to_queue=True)
third = self.track(graph, "item three")
self.assertEqual(self.process({"prompt": graph}), first)
self.assertIsNone(self.process({"prompt": graph}))
self.assertEqual(self.process({"prompt": graph}), third)
def test_a_newer_run_drops_stale_entries_of_the_same_key(self):
graph = llm_text_graph()
# run-1 saved no image, so its entry was never taken
self.track(graph, "earlier run text", run_id="run-1")
newer = self.track(graph, "newer run text", run_id="run-2")
self.assertEqual(self.process({"prompt": graph}), newer)
self.assertEqual(self.tracker._prompt_queue, [])
def test_a_later_run_queued_before_processing_keeps_its_own_entry(self):
# The image appears (snapshot), the next queued run executes the same
# node again before the monitor's delay is over, then the image is
# processed: it takes its own run's entry, the later run keeps its own.
graph = llm_text_graph()
own = self.track(graph, "first run text", run_id="run-1")
snapshot = self.tracker.keyed_snapshot()
later = self.track(graph, "second run text", run_id="run-2")
self.assertEqual(self.process({"prompt": graph}, snapshot), own)
self.assertEqual(self.process({"prompt": graph}), later)
def test_no_entry_when_the_image_appeared_means_no_queue_match(self):
graph = llm_text_graph()
empty_snapshot = self.tracker.keyed_snapshot()
self.track(graph, "queued after the image appeared")
self.assertIsNone(self.process({"prompt": graph}, empty_snapshot))
self.assertEqual(len(self.tracker._prompt_queue), 1)
def test_hosts_without_run_ids_still_link_queued_entries_in_order(self):
graph = llm_text_graph()
first = self.track(graph, "item one", run_id=None)
second = self.track(graph, "item two", run_id=None)
self.assertEqual(self.process({"prompt": graph}), first)
self.assertEqual(self.process({"prompt": graph}), second)
self.assertIsNone(self.process({"prompt": graph}))
def test_other_keys_stay_queued_for_their_own_images(self):
mine, other = llm_text_graph(), copy.deepcopy(llm_text_graph())
other["18:17"]["inputs"]["value"] = "Describe a beach"
theirs = self.track(other, "a beach")
prompt_id = self.track(mine, LLM_TEXT)
self.assertEqual(self.process({"prompt": mine}), prompt_id)
self.assertEqual(self.process({"prompt": other}), theirs)
def test_two_positive_prompt_nodes_are_not_guessed(self):
graph = llm_text_graph()
self.track(graph, LLM_TEXT)
graph["30"] = {
"class_type": "PromptManager",
"inputs": {"text": ["18:13", 0], "clip": ["4", 1]},
}
graph["31"] = dict(graph["3"])
graph["31"]["inputs"] = {**graph["3"]["inputs"], "positive": ["30", 0]}
self.assertEqual(len(image_run_keys({"prompt": graph})), 2)
self.assertIsNone(self.process({"prompt": graph}))
def test_metadata_still_wins_and_leaves_the_queue_alone(self):
typed = self.save("a typed prompt")
graph = issue_graph(text="a typed prompt")
self.tracker.set_current_prompt(
"x",
{"prompt_id": typed},
run_key=prompt_node_key(graph, "18:35"),
run_id="r",
)
self.assertEqual(self.process({"prompt": graph}), typed)
self.assertEqual(len(self.tracker._prompt_queue), 1)
class TestRunKeys(unittest.TestCase):
def test_key_covers_the_upstream_and_ignores_downstream_and_titles(self):
graph = issue_graph()
key = prompt_node_key(graph, "18:35")
changed = copy.deepcopy(graph)
changed["3"]["inputs"]["seed"] = 99
changed["18:35"]["_meta"] = {"title": "renamed"}
self.assertEqual(prompt_node_key(changed, "18:35"), key)
changed["18:11"]["inputs"]["seed"] = 8
self.assertNotEqual(prompt_node_key(changed, "18:35"), key)
def test_unknowable_keys_are_none(self):
graph = issue_graph()
self.assertIsNone(prompt_node_key(None, "1"))
self.assertIsNone(prompt_node_key(graph, None))
self.assertIsNone(prompt_node_key(graph, "missing"))
broken = copy.deepcopy(graph)
broken["18:13"]["inputs"]["source"] = ["gone", 0]
self.assertIsNone(prompt_node_key(broken, "18:35"))
broken = copy.deepcopy(graph)
broken["18:17"]["inputs"]["value"] = object()
self.assertIsNone(prompt_node_key(broken, "18:35"))
def test_image_keys_read_string_graphs_and_reject_garbage(self):
import json
graph = issue_graph()
self.assertEqual(
image_run_keys({"prompt": json.dumps(graph)}),
{"18:35": prompt_node_key(graph, "18:35")},
)
self.assertEqual(image_run_keys({"prompt": "{not json"}), {})
self.assertEqual(image_run_keys(None), {})
self.assertEqual(image_run_keys({"prompt": []}), {})
def test_cyclic_and_huge_closures_have_no_key(self):
from utils.prompt_graph import MAX_VISITS
cyclic = {
"1": {"class_type": "A", "inputs": {"x": ["2", 0]}},
"2": {"class_type": "B", "inputs": {"x": ["1", 0]}},
}
self.assertIsNotNone(prompt_node_key(cyclic, "1"))
chain = {
str(i): {"class_type": "N", "inputs": {"x": [str(i + 1), 0]}}
for i in range(MAX_VISITS + 1)
}
chain[str(MAX_VISITS + 1)] = {"class_type": "N", "inputs": {}}
self.assertIsNone(prompt_node_key(chain, "0"))
class TestReviewFindings(MonitorCase, NodeCase):
"""Final-review findings M1-M4 (#344)."""
def test_m1_a_positive_node_without_a_key_makes_the_image_unkeyable(self):
graph = llm_text_graph()
graph["30"] = {
"class_type": "PromptManager",
"inputs": {"text": ["missing", 0], "clip": ["4", 1]},
}
graph["31"] = {
"class_type": "KSampler",
"inputs": {**graph["3"]["inputs"], "positive": ["30", 0]},
}
self.assertIsNone(prompt_node_key(graph, "30"))
self.assertEqual(image_run_keys({"prompt": graph}), {})
prompt_id = self.save(LLM_TEXT)
self.tracker.set_current_prompt(
LLM_TEXT,
{"prompt_id": prompt_id},
run_key=prompt_node_key(graph, "18:35"),
run_id="run-1",
)
self.assertIsNone(self.process({"prompt": graph}))
def test_m2_links_are_exactly_comfyui_is_link(self):
from utils.run_keys import _links
inputs = {
"string_slot_int": ["7", 0],
"string_slot_float": ["7", 0.0],
"string_slot_bool": ["7", True], # bool is an int to ComfyUI too
"int_source": [7, 0],
"tuple": ("7", 0),
"three": ["7", 0, 1],
"slot_text": ["7", "0"],
}
self.assertEqual(sorted(_links(inputs)), ["7", "7", "7"])
widget = issue_graph()
widget["18:17"]["inputs"]["size"] = [512, 768] # a value, not a link
self.assertIsNotNone(prompt_node_key(widget, "18:35"))
floated = issue_graph()
floated["18:35"]["inputs"]["prepend_text"] = ["18:13", 0.0]
changed = copy.deepcopy(floated)
changed["18:17"]["inputs"]["value"] = "another instruction"
self.assertNotEqual(
prompt_node_key(floated, "18:35"), prompt_node_key(changed, "18:35")
)
def _failing_save(self, node):
def explode(*args, **kwargs):
raise RuntimeError("database is locked")
node.db = mock.Mock(wraps=self.db)
for name in ("save_prompt", "save_prompt_pair", "get_prompt_by_hash"):
setattr(node.db, name, mock.Mock(side_effect=explode))
return node
def test_m3_a_failed_save_takes_over_its_run_key(self):
graph = llm_text_graph()
key = prompt_node_key(graph, "18:35")
older = self.save("previous run text")
self.tracker.set_current_prompt(
"previous run text", {"prompt_id": older}, run_key=key, run_id="run-1"
)
node = self._failing_save(self.node(PromptManager))
with mock.patch("prompt_manager_base.executing_run_id", return_value="run-2"):
node.encode_prompt(
self.clip(), "new run text", prompt=graph, unique_id="18:35"
)
latest, run_id, _count = self.tracker.keyed_snapshot()[key]
self.assertTrue(latest.get("placeholder"))
self.assertEqual(run_id, "run-2")
self.assertIsNone(self.process({"prompt": graph}))
def test_m3_a_failed_pair_save_takes_over_its_run_key(self):
graph = issue_graph(text="typed")
key = prompt_node_key(graph, "18:35")
node = self._failing_save(self.node(PromptManagerPair))
with self.assertLogs("test.prompt_attribution.node", "WARNING"):
node.encode_prompt(self.clip(), "typed", prompt=graph, unique_id="18:35")
self.assertTrue(self.tracker.keyed_snapshot()[key][0].get("placeholder"))
def test_m3_a_failed_text_node_save_takes_over_its_run_key(self):
graph = issue_graph()
graph["18:35"] = {
"class_type": "PromptManagerText",
"inputs": {"text": "typed"},
}
graph["40"] = {
"class_type": "CLIPTextEncode",
"inputs": {"text": ["18:35", 0], "clip": ["4", 1]},
}
graph["3"]["inputs"]["positive"] = ["40", 0]
node = self._failing_save(self.node(PromptManagerText))
with mock.patch.object(
node, "_track_prompt_execution", side_effect=RuntimeError("boom")
):
node.process_text("typed", prompt=graph, unique_id="18:35")
key = prompt_node_key(graph, "18:35")
self.assertTrue(self.tracker.keyed_snapshot()[key][0].get("placeholder"))
def test_m4_a_placeholder_never_falls_through_to_history(self):
graph = llm_text_graph()
key = prompt_node_key(graph, "18:35")
self.tracker.mark_empty(key, "run-1", push_to_queue=False)
with mock.patch.object(
self.handler, "_history_prompt", return_value=99
) as history:
self.assertIsNone(self.process({"prompt": graph})) # snapshot form
self.tracker.mark_empty(key, "run-2", push_to_queue=True)
self.assertIsNone(self.process({"prompt": graph})) # popped form
history.assert_not_called()
class TestRereviewFindings(MonitorCase, NodeCase):
"""Re-review findings R4 and R5 (#344)."""
def test_r4_a_logging_failure_after_tracking_never_doubles_a_queue_slot(self):
graph = llm_text_graph()
key = prompt_node_key(graph, "18:35")
node = self.node(PromptManager)
logger = mock.Mock(wraps=self.tracker.logger)
logger.debug.side_effect = RuntimeError("log handler broke")
self.tracker.logger = logger
node.encode_prompt(self.clip(), LLM_TEXT, prompt=graph, unique_id="18:35")
queued = [ctx for ctx in self.tracker._prompt_queue if ctx["run_key"] == key]
self.assertEqual(len(queued), 1)
def test_r5_a_keyless_image_outside_the_output_folder_does_not_wait(self):
outside = tempfile.mkdtemp()
self.addCleanup(lambda: __import__("shutil").rmtree(outside, True))
stray = os.path.join(outside, "copied.png")
with open(stray, "wb") as handle:
handle.write(b"png")
with mock.patch(
"utils.image_monitor.host_output_directory", return_value=self.tmp.name
):
with mock.patch.object(
self.handler.metadata_extractor, "extract_metadata", return_value=None
):
self.handler.process_new_image(stray)
self.handler.process_new_image(self.image) # inside: waits
self.assertEqual(
self.handler.awaiting_history(),
[os.path.normcase(os.path.realpath(self.image))],
)
if __name__ == "__main__":
unittest.main()
+1 -8
View File
@@ -291,15 +291,11 @@ class TestCancelledPreview(RestoreTestCase):
self.upload(["restored"], "up.db")
entered, release = threading.Event(), threading.Event()
real = self.api._live_fingerprint
done = threading.Event()
def slow(db):
entered.set()
release.wait(WAIT)
try:
return real(db)
finally:
done.set()
return real(db)
with patch.object(self.api, "_live_fingerprint", side_effect=slow):
task = asyncio.ensure_future(
@@ -314,9 +310,6 @@ class TestCancelledPreview(RestoreTestCase):
release.set()
with self.assertRaises(asyncio.CancelledError):
await task
# The cancel cannot stop the worker thread: it still opens the live
# library. Let it finish, or teardown can find the file open.
self.assertTrue(await loop.run_in_executor(None, done.wait, WAIT))
self.assertEqual(self.stages(), [])
self.assertEqual(len(self.api._restore_plans), 0)
-255
View File
@@ -1,255 +0,0 @@
"""tools/review_image_links.py: a dry run by default, changes only named images."""
import io
import json
import os
import shutil
import tempfile
import unittest
import unittest.mock
from PIL import Image
from PIL.PngImagePlugin import PngInfo
try:
from tests.open_handles import assert_closed
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed
from database.operations import PromptDatabase
from tests.test_api_images import ImageAPITestCase
from tests.test_prompt_attribution import issue_graph, pair_evidence
from tools import review_image_links as tool
from utils.hashing import generate_prompt_hash
def png(path, prompt=None, evidence=None):
info = PngInfo()
if prompt is not None:
info.add_text("prompt", json.dumps(prompt))
if evidence is not None:
info.add_text("comfyassets_prompt_pairs", json.dumps(evidence))
Image.new("RGB", (4, 4)).save(path, pnginfo=info)
class ToolFixture(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.mkdtemp()
self.addCleanup(shutil.rmtree, self.tmp, True)
self.addCleanup(assert_closed, self, self.tmp)
self.output = os.path.join(self.tmp, "output")
os.makedirs(self.output)
self.db_path = os.path.join(self.tmp, "prompts.db")
db = PromptDatabase(self.db_path)
try:
save = lambda text: db.save_prompt( # noqa: E731
text=text, prompt_hash=generate_prompt_hash(text)
)
self.coffee = save("a woman holding a coffee cup")
self.typed = save("snowy peaks")
# The six #342 images: only prepend text, linked by the old fallback
self.guessed = []
for n in range(9, 12):
path = os.path.join(self.output, f"moss_{n:05}_.png")
png(path, issue_graph(), pair_evidence())
self.guessed.append(
db.link_image_to_prompt(self.coffee, path, metadata={})
)
# A typed image linked to the wrong prompt, and a correct one
wrong = os.path.join(self.output, "typed_wrong.png")
png(wrong, issue_graph(text="snowy peaks"))
self.wrong = db.link_image_to_prompt(self.coffee, wrong, metadata={})
right = os.path.join(self.output, "typed_right.png")
png(right, issue_graph(text="snowy peaks"))
self.right = db.link_image_to_prompt(self.typed, right, metadata={})
gone = os.path.join(self.output, "gone.png")
self.gone = db.link_image_to_prompt(self.coffee, gone, metadata={})
finally:
db.close_all()
def run_tool(self, *args):
out = io.StringIO()
code = tool.main([self.db_path, "--output-dir", self.output, *args], out=out)
return code, out.getvalue()
def links(self):
db = PromptDatabase(self.db_path)
try:
rows = db.model.get_connection().execute(
"SELECT id, prompt_id FROM generated_images"
)
return {row[0]: row[1] for row in rows}
finally:
db.close_all()
class ReviewToolTests(ToolFixture):
def test_dry_run_judges_and_changes_nothing(self):
before = self.links()
ids = f"{self.guessed[0]}-{self.gone}"
code, out = self.run_tool("--ids", ids, "--json")
self.assertEqual(code, 0)
verdicts = {v["image_id"]: v for v in json.loads(out.split("--apply, for")[0])}
for image_id in self.guessed:
self.assertEqual(verdicts[image_id]["verdict"], "unlink")
self.assertEqual(verdicts[image_id]["sources"], {})
self.assertEqual(verdicts[self.wrong]["verdict"], "relink")
self.assertEqual(verdicts[self.wrong]["target"], self.typed)
self.assertEqual(verdicts[self.right]["verdict"], "keep")
self.assertEqual(verdicts[self.gone]["verdict"], "missing")
self.assertIn("Dry run: nothing was changed.", out)
self.assertEqual(self.links(), before)
self.assertEqual(
sorted(os.listdir(self.tmp)), ["output", "prompts.db"], "no side files"
)
def test_prompt_selection_and_text_report(self):
code, out = self.run_tool("--prompt", str(self.coffee))
self.assertEqual(code, 0)
self.assertIn("unlink: 3", out)
self.assertIn("relink: 1", out)
self.assertIn("missing: 1", out)
self.assertIn(f"-> {self.typed}", out)
def test_apply_backs_up_then_changes_only_named_images(self):
named = f"{self.guessed[0]}-{self.guessed[1]},{self.wrong}"
code, out = self.run_tool(
"--ids", named, "--prompt", str(self.coffee), "--apply"
)
self.assertEqual(code, 0)
self.assertIn("Applied: 1 relinked, 2 unlinked, 0 skipped", out)
links = self.links()
self.assertNotIn(self.guessed[0], links)
self.assertNotIn(self.guessed[1], links)
self.assertEqual(links[self.guessed[2]], self.coffee, "not named: kept")
self.assertEqual(links[self.wrong], self.typed)
self.assertEqual(links[self.gone], self.coffee, "missing files are kept")
backups = [n for n in os.listdir(self.tmp) if ".pre-relink_" in n]
self.assertEqual(len(backups), 1)
restored = PromptDatabase(os.path.join(self.tmp, backups[0]))
try:
count = (
restored.model.get_connection()
.execute("SELECT COUNT(*) FROM generated_images")
.fetchone()[0]
)
finally:
restored.close_all()
self.assertEqual(count, 6)
def test_arguments_are_validated(self):
for args in (
[],
["--ids", "5-2"],
["--ids", "x"],
["--prompt", "1", "--apply"],
):
with self.subTest(args=args), self.assertRaises(SystemExit):
with unittest.mock.patch("sys.stderr", io.StringIO()):
self.run_tool(*args)
with (
self.assertRaises(SystemExit),
unittest.mock.patch("sys.stderr", io.StringIO()),
):
tool.main(
[os.path.join(self.tmp, "none.db"), "--output-dir", "x", "--ids", "1"]
)
self.assertEqual(tool.parse_ids(" 3, 1-2 ,"), [1, 2, 3])
def test_conflicting_sources_are_left_to_the_maintainer(self):
db = PromptDatabase(self.db_path)
try:
handler = tool.ImageGenerationHandler(db, None)
row = {
"id": 1,
"prompt_id": self.coffee,
"image_path": os.path.join(self.output, "typed_right.png"),
}
with unittest.mock.patch.object(
tool, "certain_prompts", return_value={"scan": 1, "history": 2}
):
verdict = tool.judge(db, handler, row, self.output)
finally:
db.close_all()
self.assertEqual(verdict["verdict"], "conflict")
class ApplyGuardTests(ToolFixture):
"""M6: unlink clears only the link, and only the link it was judged on."""
def test_unlink_skips_a_row_linked_elsewhere_since_the_review(self):
db = PromptDatabase(self.db_path)
try:
verdicts = tool.review(db, [self.guessed[0]], None, self.output)
self.assertEqual(verdicts[0]["verdict"], "unlink")
db.model.get_connection().execute(
"UPDATE generated_images SET prompt_id=? WHERE id=?",
(self.typed, self.guessed[0]),
)
db.model.get_connection().commit()
result = tool.apply(db, verdicts, [self.guessed[0]])
finally:
db.close_all()
self.assertEqual(result["unlinked"], 0)
self.assertEqual(self.links()[self.guessed[0]], self.typed)
def test_unlink_keeps_the_file_and_its_favorite(self):
path = os.path.join(self.output, "moss_00009_.png")
db = PromptDatabase(self.db_path)
try:
db.set_image_favorite("a" * 64, "moss_00009_.png")
finally:
db.close_all()
self.run_tool("--ids", str(self.guessed[0]), "--apply")
self.assertTrue(os.path.isfile(path))
db = PromptDatabase(self.db_path)
try:
self.assertEqual(db.get_root_favorites("a" * 64), {"moss_00009_.png"})
summary = db.get_image_link_summaries([path]).get(path) or {}
finally:
db.close_all()
self.assertNotEqual(summary.get("status"), "linked")
def test_dry_run_states_what_apply_does(self):
_code, out = self.run_tool("--ids", str(self.guessed[0]))
self.assertIn("removes only the image's link record", out)
self.assertIn("file stays on disk and in the gallery as unlinked", out)
class GalleryAfterUnlinkTests(ImageAPITestCase):
"""M6: after --apply the file is still listed, as an unlinked image."""
async def test_unlinked_file_stays_in_the_workspace_gallery(self):
db = self.api.db
prompt = db.save_prompt(
text="a woman holding a coffee cup",
prompt_hash=generate_prompt_hash("a woman holding a coffee cup"),
)
path = self.output_dir / "moss_00009_.png"
png(str(path), issue_graph(), pair_evidence())
image_id = db.link_image_to_prompt(prompt, str(path), metadata={})
out = io.StringIO()
tool.main(
[
db.model.db_path,
"--output-dir",
str(self.output_dir),
"--ids",
str(image_id),
"--apply",
],
out=out,
)
self.assertIn("Applied: 0 relinked, 1 unlinked", out.getvalue())
response = await self.client.get(
"/prompt_manager/images/output", params={"mode": "workspace"}
)
self.assertEqual(response.status, 200)
images = {i["filename"]: i for i in (await response.json())["images"]}
self.assertIn("moss_00009_.png", images)
link = images["moss_00009_.png"].get("linked") or {}
self.assertNotEqual(link.get("status"), "linked")
if __name__ == "__main__":
unittest.main()
-674
View File
@@ -1,674 +0,0 @@
"""#342: recorded history attributes a file only to one certain, stored prompt."""
import json
import os
import sqlite3
import threading
import time
import unittest
from pathlib import Path
from unittest import mock
from database.run_attribution import stored_run_prompt
from tests import test_execution_history_database as history_tests
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
from utils.execution_contexts import ContextBundle, RecipeContext
from utils.execution_outputs import (
OutputBundle,
OutputReference,
snapshot_output_roots,
)
from utils.hashing import generate_prompt_hash
from utils.image_monitor import ImageGenerationHandler
STYLE = "windows" if os.name == "nt" else "posix"
LLM_TEXT = "Photorealistic aerial landscape of misty mountains at dawn"
def now_ms():
return time.time_ns() // 1_000_000
class RecordedRuns:
"""Write runs into execution history the way the capture pipeline does."""
start = history_tests.HistoryDatabaseTests.start
finish = history_tests.HistoryDatabaseTests.finish
def record_run(self, db, output_root, files, positives, item=0, window=None):
started, finished = window or (now_ms() - 5000, now_ms() + 5000)
start = self.start(db, item, started)
db.write_execution_start(start)
finish = self.finish(start, time=finished)
db.write_execution_finish(finish)
db.write_execution_contexts(
start.run_id,
start.storage_epoch,
ContextBundle(
tuple(
RecipeContext(str(4 + i), snapshot(text), snapshot(""), "inferred")
for i, text in enumerate(positives)
),
"complete",
"captured",
),
)
root = snapshot_output_roots(STYLE, output_root, output_root).output_key
references = tuple(
OutputReference("9", i, "output", sub, name, STYLE, root)
for i, (sub, name) in enumerate(files)
)
db.write_execution_outputs_for_finish(
finish, OutputBundle(references, "complete", "reported")
)
return start.run_id
class TestStoredRunPrompt(RecordedRuns, unittest.TestCase):
setUp = history_tests.HistoryDatabaseTests.setUp
close_databases = history_tests.HistoryDatabaseTests.close_databases
database = history_tests.HistoryDatabaseTests.database
def setUp(self): # noqa: F811 - the shared temp dir and database list first
history_tests.HistoryDatabaseTests.setUp(self)
self.db = self.database()
self.output = str(Path(self.temp.name) / "output")
os.makedirs(os.path.join(self.output, "sub"))
self.image = os.path.join(self.output, "sub", "img_00001_.png")
Path(self.image).write_bytes(b"png")
def save(self, text):
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
def attribute(self, path=None):
return stored_run_prompt(self.db, path or self.image, self.output)
def test_one_run_naming_one_stored_prompt_attributes_the_file(self):
prompt_id = self.save("an llm prompt")
self.record_run(
self.db, self.output, [("sub", "img_00001_.png")], ["an llm prompt"]
)
self.assertEqual(self.attribute(), prompt_id)
def test_text_that_is_not_stored_attributes_nothing(self):
"""The #342 run: raw editor text was empty, nothing was stored."""
self.save("a woman holding a coffee cup")
self.record_run(self.db, self.output, [("sub", "img_00001_.png")], [""])
self.assertIsNone(self.attribute())
def test_two_runs_reporting_the_same_file_attribute_nothing(self):
self.save("first")
self.save("second")
self.record_run(self.db, self.output, [("sub", "img_00001_.png")], ["first"])
self.record_run(
self.db, self.output, [("sub", "img_00001_.png")], ["second"], item=1
)
self.assertIsNone(self.attribute())
def test_a_file_written_outside_the_run_attributes_nothing(self):
self.save("an llm prompt")
self.record_run(
self.db,
self.output,
[("sub", "img_00001_.png")],
["an llm prompt"],
window=(1_000_000, 2_000_000),
)
self.assertIsNone(self.attribute())
def test_contexts_naming_different_prompts_attribute_nothing(self):
self.save("left")
self.save("right")
self.record_run(
self.db, self.output, [("sub", "img_00001_.png")], ["left", "right"]
)
self.assertIsNone(self.attribute())
def test_files_outside_the_output_root_or_unrecorded_attribute_nothing(self):
self.save("an llm prompt")
self.record_run(
self.db, self.output, [("sub", "img_00001_.png")], ["an llm prompt"]
)
outside = os.path.join(self.temp.name, "img_00001_.png")
Path(outside).write_bytes(b"png")
self.assertIsNone(self.attribute(outside))
other = os.path.join(self.output, "sub", "other.png")
Path(other).write_bytes(b"png")
self.assertIsNone(self.attribute(other))
self.assertIsNone(stored_run_prompt(self.db, self.image, None))
self.assertIsNone(stored_run_prompt(self.db, self.image, "relative/root"))
def test_a_library_without_history_tables_attributes_nothing(self):
broken = mock.Mock()
broken.model.get_connection.return_value.execute.side_effect = (
sqlite3.OperationalError("no such table")
)
self.assertIsNone(stored_run_prompt(broken, self.image, self.output))
def test_the_monitor_links_from_history_when_nothing_else_is_certain(self):
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))
linked = []
handler.link_image_to_prompt = lambda path, prompt, meta: linked.append(
prompt["id"]
)
with (
mock.patch(
"utils.image_monitor.host_output_directory", return_value=self.output
),
mock.patch.object(
handler.metadata_extractor,
"extract_metadata",
return_value={"prompt": llm_text_graph()},
),
):
handler.process_new_image(self.image)
self.assertEqual(linked, [prompt_id])
class TestScanUsesHistory(RecordedRuns, AdminAPITestCase):
"""The output scan applies the same rules as the monitor."""
def linked_prompt(self, path):
row = (
self.api.db.model.get_connection()
.execute(
"SELECT prompt_id FROM generated_images WHERE image_path=?",
(str(path),),
)
.fetchone()
)
return row[0] if row else None
def write(self, name):
path = self.output_dir / name
path.write_bytes(b"png")
return path
def test_an_llm_text_file_links_to_the_prompt_its_run_stored(self):
db = self.api.db
prompt_id = db.save_prompt(
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
)
path = self.write("llm_00001_.png")
self.record_run(db, str(self.output_dir), [("", path.name)], [LLM_TEXT])
meta = {"prompt": json.dumps(llm_text_graph())}
self.assertEqual(self.api._ingest_scanned_file(path, meta), "linked")
self.assertEqual(self.linked_prompt(path), prompt_id)
def test_a_metadata_less_file_links_through_its_recorded_run(self):
"""I1: --disable-metadata outputs carry no prompt; history still proves one."""
db = self.api.db
prompt_id = db.save_prompt(
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
)
path = self.write("bare_00001_.png")
self.record_run(db, str(self.output_dir), [("", path.name)], [LLM_TEXT])
self.assertEqual(self.api._ingest_scanned_file(path, {}), "linked")
self.assertEqual(self.linked_prompt(path), prompt_id)
self.assertEqual(self.api._ingest_scanned_file(path, None), "unchanged")
detailed = self.api._ingest_scanned_file(path, {}, detailed=True)
self.assertEqual(detailed["status"], "unchanged")
self.assertTrue(detailed["title"])
def test_a_metadata_less_file_without_a_recorded_run_is_still_skipped(self):
path = self.write("bare_00002_.png")
self.assertIsNone(self.api._ingest_scanned_file(path, {}))
self.assertIsNone(self.api._ingest_scanned_file(path, {"workflow": {}}))
self.assertIsNone(self.linked_prompt(path))
def test_the_issue_file_stays_unlinked_even_with_its_pair_record(self):
"""Resolved with an empty stored positive: before #342 the scan raised."""
db = self.api.db
db.save_prompt(
text="a woman holding a coffee cup",
prompt_hash=generate_prompt_hash("a woman holding a coffee cup"),
)
path = self.write("moss_00009_.png")
self.record_run(db, str(self.output_dir), [("", path.name)], [""])
evidence = {
"version": 1,
"nodes": {
"18:35": {
"records": [
{
"text": "",
"negative_text": "",
"positive_text": LLM_TEXT,
"negative_effective_text": "",
"prepend_text": LLM_TEXT,
"append_text": "",
}
]
}
},
}
meta = {
"prompt": json.dumps(issue_graph()),
"comfyassets_prompt_pairs": json.dumps(evidence),
}
self.assertEqual(self.api._ingest_scanned_file(path, meta), "unresolved")
self.assertIsNone(self.linked_prompt(path))
self.assertEqual(db.list_prompt_pairs()["pairs"], [])
class TestHostHelpers(unittest.TestCase):
"""The host lookups degrade to "no proof", never to an error."""
def test_host_output_directory(self):
import sys
import types
from database import run_attribution
cases = (
(types.SimpleNamespace(get_output_directory=lambda: "/out"), "/out"),
(types.SimpleNamespace(get_output_directory=lambda: ""), None),
(types.SimpleNamespace(get_output_directory=lambda: 1 / 0), None),
)
original = sys.modules.get("folder_paths")
try:
for module, expected in cases:
sys.modules["folder_paths"] = module
self.assertEqual(run_attribution.host_output_directory(), expected)
sys.modules["folder_paths"] = None # import fails: outside ComfyUI
self.assertIsNone(run_attribution.host_output_directory())
finally:
if original is None:
sys.modules.pop("folder_paths", None)
else:
sys.modules["folder_paths"] = original
def test_run_window_needs_integer_times_and_a_readable_file(self):
from database.run_attribution import _written_during
self.assertFalse(_written_during(__file__, None, 5))
self.assertFalse(_written_during(os.path.join("no", "such.png"), 0, 5))
def test_executing_run_id(self):
import sys
import types
from utils.run_keys import executing_run_id
def host(context):
def get_executing_context():
if isinstance(context, Exception):
raise context
return context
return types.SimpleNamespace(get_executing_context=get_executing_context)
original = sys.modules.get("comfy_execution.utils")
try:
for context, expected in (
(types.SimpleNamespace(prompt_id="run-1"), "run-1"),
(types.SimpleNamespace(prompt_id=""), None),
(None, None),
(RuntimeError("no context"), None),
):
sys.modules["comfy_execution.utils"] = host(context)
self.assertEqual(executing_run_id(), expected)
sys.modules["comfy_execution.utils"] = None
self.assertIsNone(executing_run_id())
finally:
if original is None:
sys.modules.pop("comfy_execution.utils", None)
else:
sys.modules["comfy_execution.utils"] = original
class TestMonitorHistoryRecheck(RecordedRuns, unittest.TestCase):
"""I1: a keyless image gets one history-only re-check when its run is sealed."""
def setUp(self):
history_tests.HistoryDatabaseTests.setUp(self)
self.db = history_tests.HistoryDatabaseTests.database(self)
self.output = str(Path(self.temp.name) / "output")
os.makedirs(self.output)
self.handler = ImageGenerationHandler(self.db, make_tracker(self.db))
self.handler.processing_delay = 0
self.handler.settle_interval = 0.01
self.addCleanup(self.handler.shutdown)
self.linked = []
self.handler.link_image_to_prompt = lambda path, prompt, meta: (
self.linked.append((os.path.basename(path), prompt["id"]))
)
patches = (
mock.patch(
"utils.image_monitor.host_output_directory", return_value=self.output
),
mock.patch.object(
self.handler.metadata_extractor, "extract_metadata", return_value=None
),
)
for patch in patches:
patch.start()
self.addCleanup(patch.stop)
close_databases = history_tests.HistoryDatabaseTests.close_databases
def image(self, name):
path = os.path.join(self.output, name)
Path(path).write_bytes(b"png")
return path
def save(self, text):
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
def wait_for(self, predicate, timeout=5.0):
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
if predicate():
return True
time.sleep(0.01)
return predicate()
def test_the_run_ending_after_processing_links_the_image_once(self):
prompt_id = self.save(LLM_TEXT)
path = self.image("late_00001_.png")
self.handler.process_new_image(path) # the run is still executing
self.assertEqual(self.linked, [])
self.assertIn(
os.path.normcase(os.path.realpath(path)), self.handler.awaiting_history()
)
self.record_run(self.db, self.output, [("", "late_00001_.png")], [LLM_TEXT])
self.assertTrue(self.wait_for(lambda: self.linked))
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
self.assertEqual(self.linked, [("late_00001_.png", prompt_id)])
self.assertEqual(self.handler.awaiting_history(), [])
# Another run reporting the same name later is not a second re-check
self.record_run(
self.db, self.output, [("", "late_00001_.png")], [LLM_TEXT], item=1
)
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
self.assertEqual(len(self.linked), 1)
def test_a_seal_during_the_first_pass_is_not_lost(self):
"""R1: the run seals while the image's own pass is still pending."""
prompt_id = self.save(LLM_TEXT)
path = self.image("race_00001_.png")
real = self.handler._history_prompt
sealed = []
def history(image_path):
result = real(image_path) # before the run's outputs commit: None
if not sealed:
sealed.append(True)
# The writer commits and notifies now; the path is still pending
self.record_run(
self.db, self.output, [("", "race_00001_.png")], [LLM_TEXT]
)
return result
self.handler._history_prompt = history
self.assertTrue(self.handler.schedule(path))
self.assertTrue(self.wait_for(lambda: self.linked))
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
self.assertEqual(self.linked, [("race_00001_.png", prompt_id)])
def test_a_seal_while_the_pending_set_is_full_is_rechecked_later(self):
from utils import image_monitor
prompt_id = self.save(LLM_TEXT)
path = self.image("full_00001_.png")
self.handler.process_new_image(path) # waiting, unlinked
blocker = self.image("blocker.png")
release = threading.Event()
original = self.handler.process_new_image
def slow(image_path, prompt_snapshot=None):
if image_path == blocker:
release.wait(5)
return original(image_path, prompt_snapshot=prompt_snapshot)
self.handler.process_new_image = slow
with mock.patch.object(image_monitor, "MAX_PENDING_IMAGES", 1):
self.assertTrue(self.handler.schedule(blocker))
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 1))
self.record_run(self.db, self.output, [("", "full_00001_.png")], [LLM_TEXT])
self.assertEqual(self.linked, [], "not scheduled while the set is full")
release.set()
self.assertTrue(self.wait_for(lambda: self.linked))
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
self.assertEqual(self.linked, [("full_00001_.png", prompt_id)])
def test_a_run_naming_nothing_stored_leaves_the_image_unlinked(self):
path = self.image("empty_00001_.png")
self.handler.process_new_image(path)
self.record_run(self.db, self.output, [("", "empty_00001_.png")], [""])
self.assertTrue(self.wait_for(lambda: not self.handler.awaiting_history()))
self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0))
self.assertEqual(self.linked, [])
def test_history_found_at_once_leaves_nothing_awaiting(self):
prompt_id = self.save(LLM_TEXT)
self.record_run(self.db, self.output, [("", "done_00001_.png")], [LLM_TEXT])
self.handler.process_new_image(self.image("done_00001_.png"))
self.assertEqual(self.linked, [("done_00001_.png", prompt_id)])
self.assertEqual(self.handler.awaiting_history(), [])
def test_keyed_images_do_not_wait_for_history(self):
with mock.patch.object(
self.handler.metadata_extractor,
"extract_metadata",
return_value={"prompt": llm_text_graph()},
):
self.handler.process_new_image(self.image("keyed_00001_.png"))
self.assertEqual(self.handler.awaiting_history(), [])
def test_the_waiting_set_is_bounded_and_expires(self):
from utils import image_monitor
with mock.patch.object(image_monitor, "MAX_AWAITING_HISTORY", 2):
for n in range(3):
self.handler.process_new_image(self.image(f"b_{n}.png"))
self.assertEqual(
[os.path.basename(p) for p in self.handler.awaiting_history()],
["b_1.png", "b_2.png"],
)
later = time.monotonic() + image_monitor.AWAIT_HISTORY_SECONDS + 1
with mock.patch.object(image_monitor.time, "monotonic", return_value=later):
self.assertEqual(self.handler.awaiting_history(), [])
def test_shutdown_forgets_waiting_images_and_a_dead_handler_is_dropped(self):
import gc
from database import run_attribution
self.handler.process_new_image(self.image("gone_00001_.png"))
self.handler.shutdown()
self.assertEqual(self.handler.awaiting_history(), [])
other = ImageGenerationHandler(self.db, make_tracker(self.db))
with mock.patch.object(
other.metadata_extractor, "extract_metadata", return_value=None
):
other.process_new_image(self.image("gone_00002_.png"))
del other
gc.collect()
run_attribution.notify_outputs_sealed("run", []) # no dead callback raises
class TestBoundedLookup(RecordedRuns, unittest.TestCase):
"""M5: history lookups are batched, and the history schema is unchanged."""
setUp = history_tests.HistoryDatabaseTests.setUp
close_databases = history_tests.HistoryDatabaseTests.close_databases
database = history_tests.HistoryDatabaseTests.database
def output_lookups(self, db, action):
statements = []
conn = db.model.get_connection()
conn.set_trace_callback(statements.append)
try:
result = action()
finally:
conn.set_trace_callback(None)
return result, [s for s in statements if "FROM execution_run_outputs" in s]
def test_one_query_per_chunk_not_per_file_and_same_answers(self):
from database import run_attribution
db = self.database()
output = str(Path(self.temp.name) / "output")
os.makedirs(output)
prompt_id = db.save_prompt(
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
)
names = [f"img_{n:05}_.png" for n in range(25)]
paths = [os.path.join(output, name) for name in names]
for path in paths:
Path(path).write_bytes(b"png")
self.record_run(db, output, [("", n) for n in names[:10]], [LLM_TEXT])
self.record_run(db, output, [("", n) for n in names[5:10]], [""], item=1)
with mock.patch.object(run_attribution, "LOOKUP_CHUNK", 10):
result, lookups = self.output_lookups(
db, lambda: run_attribution.stored_run_prompts(db, paths, output)
)
self.assertEqual(len(lookups), 3) # 25 names in chunks of 10
expected = {path: None for path in paths}
expected.update({path: prompt_id for path in paths[:5]})
self.assertEqual(result, expected)
for path in paths[:12]:
self.assertEqual(
run_attribution.stored_run_prompt(db, path, output), expected[path]
)
def test_the_history_schema_gains_no_index(self):
from database.execution_history_schema import history_schema_error
db = self.database()
conn = db.model.get_connection()
indexes = [
row[1]
for row in conn.execute("PRAGMA index_list(execution_run_outputs)")
if row[3] != "pk"
]
self.assertEqual(indexes, [])
self.assertIsNone(history_schema_error(conn, allow_missing=False))
class TestScanBatchLookup(RecordedRuns, AdminAPITestCase):
"""M5: a scan batch asks history once per chunk of names, not per file."""
def test_a_batch_its_metadata_links_never_asks_history(self):
"""R3: the batch lookup is made only when a file needs it."""
files = []
for n in range(3):
path = self.output_dir / f"typed_{n:05}_.png"
path.write_bytes(b"png")
files.append(path)
meta = {"prompt": json.dumps(issue_graph(text="typed mountains"))}
statements = []
conn = self.api.db.model.get_connection()
conn.set_trace_callback(statements.append)
try:
counts = self.api._ingest_scan_batch_sync([(f, meta) for f in files])
finally:
conn.set_trace_callback(None)
self.assertEqual(counts["added"] + counts["linked"], 3, counts)
lookups = [s for s in statements if "FROM execution_run_outputs" in s]
self.assertEqual(lookups, [])
def test_metadata_less_batch_links_with_one_history_lookup(self):
db = self.api.db
prompt_id = db.save_prompt(
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
)
files = []
for n in range(6):
path = self.output_dir / f"bare_{n:05}_.png"
path.write_bytes(b"png")
files.append(path)
self.record_run(
db, str(self.output_dir), [("", f.name) for f in files[:4]], [LLM_TEXT]
)
statements = []
conn = db.model.get_connection()
conn.set_trace_callback(statements.append)
try:
counts = self.api._ingest_scan_batch_sync([(f, {}) for f in files])
finally:
conn.set_trace_callback(None)
lookups = [s for s in statements if "FROM execution_run_outputs" in s]
self.assertEqual(len(lookups), 1)
self.assertEqual(counts["linked"], 4)
linked = {
row[0]
for row in conn.execute(
"SELECT filename FROM generated_images WHERE prompt_id=?",
(prompt_id,),
)
}
self.assertEqual(linked, {f.name for f in files[:4]})
class TestLinkedOutputRoot(RecordedRuns, unittest.TestCase):
"""R2: ComfyUI's output folder reached through a symlink or junction."""
setUp = history_tests.HistoryDatabaseTests.setUp
close_databases = history_tests.HistoryDatabaseTests.close_databases
database = history_tests.HistoryDatabaseTests.database
def linked_root(self):
self.real = str(Path(self.temp.name) / "big_drive" / "output")
os.makedirs(os.path.join(self.real, "sub"))
link = str(Path(self.temp.name) / "output_link")
try:
os.symlink(self.real, link, target_is_directory=True)
except (OSError, NotImplementedError) as error:
self.skipTest(f"symlinks unavailable: {error}")
return link
def test_either_spelling_of_the_file_is_attributed(self):
db = self.database()
link = self.linked_root()
prompt_id = db.save_prompt(
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
)
real_file = os.path.join(self.real, "sub", "img_00001_.png")
Path(real_file).write_bytes(b"png")
# History keeps the host's spelling of its output folder (the link)
self.record_run(db, link, [("sub", "img_00001_.png")], [LLM_TEXT])
link_file = os.path.join(link, "sub", "img_00001_.png")
self.assertEqual(stored_run_prompt(db, real_file, link), prompt_id)
self.assertEqual(stored_run_prompt(db, link_file, link), prompt_id)
# The real folder named as the root is another root key: not attributed
self.assertIsNone(stored_run_prompt(db, real_file, self.real))
def test_the_monitor_rechecks_a_file_seen_through_the_real_path(self):
db = self.database()
link = self.linked_root()
prompt_id = db.save_prompt(
text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT)
)
real_file = os.path.join(self.real, "late_00001_.png")
Path(real_file).write_bytes(b"png")
handler = ImageGenerationHandler(db, make_tracker(db))
handler.processing_delay = 0
handler.settle_interval = 0.01
self.addCleanup(handler.shutdown)
linked = []
handler.link_image_to_prompt = lambda path, prompt, meta: linked.append(
prompt["id"]
)
with (
mock.patch("utils.image_monitor.host_output_directory", return_value=link),
mock.patch.object(
handler.metadata_extractor, "extract_metadata", return_value=None
),
):
handler.process_new_image(real_file)
self.assertEqual(linked, [])
self.record_run(db, link, [("", "late_00001_.png")], [LLM_TEXT])
deadline = time.monotonic() + 5
while not linked and time.monotonic() < deadline:
time.sleep(0.01)
self.assertEqual(linked, [prompt_id])
if __name__ == "__main__":
unittest.main()
-285
View File
@@ -1,285 +0,0 @@
"""Tests for tests/run_suite.py, the CI test runner and its hang watchdog (#350).
Each test writes a temporary folder of tiny test modules and runs
``python -m tests.run_suite -s <folder> -v`` in a subprocess from the repo
root. Every subprocess call has a timeout, so a broken watchdog fails this
file instead of hanging the suite. Probes that should hang wait a bounded
time, so a watchdog that never fires shows up as a plain assertion failure.
"""
import os
import subprocess
import sys
import tempfile
import textwrap
import unittest
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parent.parent
SUBPROCESS_TIMEOUT = 60
WATCHDOG_ENV = ("PM_TEST_HANG_LIMIT", "PM_TEST_SHUTDOWN_LIMIT")
def run_suite(folder, extra_env=None, *args, verbosity=("-v",)):
"""Run the runner on ``folder`` and return the completed process."""
env = {k: v for k, v in os.environ.items() if k not in WATCHDOG_ENV}
# Windows writes a piped stderr in the ANSI codepage; ask for the UTF-8
# this helper decodes, so non-ASCII values and paths survive.
env["PYTHONIOENCODING"] = "utf-8"
env.update(extra_env or {})
command = [sys.executable, "-m", "tests.run_suite", "-s", str(folder)]
return subprocess.run(
[*command, *verbosity, *args],
cwd=REPO_ROOT,
env=env,
capture_output=True,
text=True,
encoding="utf-8",
errors="replace",
timeout=SUBPROCESS_TIMEOUT,
)
class RunSuiteTests(unittest.TestCase):
def setUp(self):
self._tmp = tempfile.TemporaryDirectory(prefix="pm-run-suite-")
self.addCleanup(self._tmp.cleanup)
self.folder = Path(self._tmp.name)
def write(self, name, source):
(self.folder / name).write_text(textwrap.dedent(source), encoding="utf-8")
def assertDumped(self, result, test_id):
detail = result.stderr[-4000:]
self.assertNotEqual(result.returncode, 0, detail)
self.assertIn(f"start: {test_id}", result.stderr, detail)
self.assertIn("Timeout (", result.stderr, detail)
self.assertIn("Thread 0x", result.stderr, detail)
def test_passing_test_exits_zero_and_names_each_start(self):
self.write(
"test_probe_pass.py",
"""
import unittest
class Probe(unittest.TestCase):
def test_ok(self):
pass
""",
)
result = run_suite(self.folder)
self.assertEqual(result.returncode, 0, result.stderr)
self.assertIn("start: test_probe_pass.Probe.test_ok", result.stderr)
def test_start_line_is_written_at_quiet_verbosity_too(self):
self.write(
"test_probe_quiet.py",
"""
import unittest
class Probe(unittest.TestCase):
def test_ok(self):
pass
""",
)
result = run_suite(self.folder, None, "-q")
self.assertEqual(result.returncode, 0, result.stderr)
self.assertIn("start: test_probe_quiet.Probe.test_ok", result.stderr)
def test_start_line_begins_its_own_line_at_default_verbosity(self):
self.write(
"test_probe_dots.py",
"""
import unittest
class Probe(unittest.TestCase):
def test_a(self):
pass
def test_b(self):
pass
""",
)
result = run_suite(self.folder, verbosity=())
self.assertEqual(result.returncode, 0, result.stderr)
lines = result.stderr.splitlines()
for name in ("test_a", "test_b"):
self.assertIn(f"start: test_probe_dots.Probe.{name}", lines)
def test_warnings_are_shown_as_with_unittest_discover(self):
self.write(
"test_probe_warn.py",
"""
import unittest
import warnings
class Probe(unittest.TestCase):
def test_warns(self):
warnings.warn("probe-deprecated", DeprecationWarning)
""",
)
env = {"PYTHONWARNINGS": ""}
result = run_suite(self.folder, env)
self.assertEqual(result.returncode, 0, result.stderr)
self.assertIn("probe-deprecated", result.stderr)
def test_failing_test_exits_one(self):
self.write(
"test_probe_fail.py",
"""
import unittest
class Probe(unittest.TestCase):
def test_fails(self):
self.fail("expected")
""",
)
result = run_suite(self.folder)
self.assertEqual(result.returncode, 1, result.stderr)
def test_pattern_option_selects_modules(self):
self.write(
"check_probe.py",
"""
import unittest
class Probe(unittest.TestCase):
def test_fails(self):
self.fail("selected by -p")
""",
)
self.assertEqual(run_suite(self.folder).returncode, 5)
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)
def test_repo_root_modules_import_as_with_unittest_discover(self):
self.write(
"test_probe_import.py",
"""
import unittest
from utils.hashing import generate_prompt_hash
class Probe(unittest.TestCase):
def test_import(self):
same = generate_prompt_hash(" A ") == generate_prompt_hash("a")
self.assertTrue(same)
""",
)
result = run_suite(self.folder)
self.assertEqual(result.returncode, 0, result.stderr)
def test_hung_test_dumps_every_thread_and_exits(self):
self.write(
"test_probe_hang.py",
"""
import time
import unittest
class Probe(unittest.TestCase):
def test_hangs(self):
time.sleep(30)
""",
)
result = run_suite(self.folder, {"PM_TEST_HANG_LIMIT": "2"})
self.assertDumped(result, "test_probe_hang.Probe.test_hangs")
def test_hang_in_async_teardown_dumps_and_exits(self):
self.write(
"test_probe_teardown.py",
"""
import threading
import unittest
class Probe(unittest.IsolatedAsyncioTestCase):
async def asyncTearDown(self):
threading.Event().wait(30)
async def test_body_passes(self):
pass
""",
)
result = run_suite(self.folder, {"PM_TEST_HANG_LIMIT": "2"})
self.assertDumped(result, "test_probe_teardown.Probe.test_body_passes")
def test_hang_in_event_loop_shutdown_after_ok_dumps_and_exits(self):
# The #349 shape: the test passes ("ok" is printed), then the loop's
# shutdown waits for a default-executor thread that never finishes.
self.write(
"test_probe_loop.py",
"""
import asyncio
import threading
import unittest
class Probe(unittest.IsolatedAsyncioTestCase):
async def test_leaves_executor_busy(self):
loop = asyncio.get_running_loop()
loop.run_in_executor(None, threading.Event().wait, 30)
""",
)
result = run_suite(self.folder, {"PM_TEST_HANG_LIMIT": "2"})
self.assertDumped(result, "test_probe_loop.Probe.test_leaves_executor_busy")
start = result.stderr.index("start: test_probe_loop")
self.assertIn("ok", result.stderr[start:], result.stderr[-4000:])
def test_hang_at_interpreter_exit_dumps_and_exits(self):
self.write(
"test_probe_exit.py",
"""
import threading
import unittest
class Probe(unittest.TestCase):
def test_leaves_thread(self):
threading.Thread(target=threading.Event().wait, args=(30,)).start()
""",
)
result = run_suite(self.folder, {"PM_TEST_SHUTDOWN_LIMIT": "2"})
detail = result.stderr[-4000:]
self.assertNotEqual(result.returncode, 0, detail)
self.assertIn("ok", result.stderr, detail)
self.assertIn("Timeout (", result.stderr, detail)
def test_invalid_limit_is_a_usage_error(self):
for name in WATCHDOG_ENV:
for value in ("abc", "0", "-3", "1.5", "\u00b2"):
with self.subTest(name=name, value=value):
result = run_suite(self.folder, {name: value})
self.assertEqual(result.returncode, 2, result.stderr)
lines = result.stderr.strip().splitlines()
self.assertEqual(len(lines), 1, result.stderr)
self.assertIn(name, lines[0])
self.assertIn(repr(value), lines[0])
if __name__ == "__main__":
unittest.main()
-8
View File
@@ -85,14 +85,6 @@ class TestConvertOutdated(ConvertCase):
)
self.assertEqual(self.api._jobs.list(), [])
async def test_conversion_keeps_the_listing(self):
self._legacy("a")
await self.client.get("/prompt_manager/images/output?mode=workspace")
files = self.api._root_listings[str(self.output_dir)].files
_summary, job = await self._convert()
self.assertEqual(job.state, "finished")
self.assertIs(self.api._root_listings[str(self.output_dir)].files, files)
async def test_converts_from_the_source_and_removes_the_old_preview(self):
old = self._legacy("a")
original = (self.output_dir / "a.png").read_bytes()
-14
View File
@@ -345,20 +345,6 @@ class TestThumbnailJob(ThumbnailJobCase):
sizes[quality] = thumb.stat().st_size
self.assertLess(sizes["low"], sizes["high"])
async def test_a_job_keeps_the_listing_and_the_next_page_shows_the_preview(self):
make_png(self.output_dir / "a.png", size=(64, 64))
url = "/prompt_manager/images/output?mode=workspace"
before = (await (await self.client.get(url)).json())["images"][0]
files = self.api._root_listings[str(self.output_dir)].files
status, body = await self._start(mode="missing")
self.assertEqual(status, 202)
await self._finished(body["job"]["id"])
# Previews are not part of the listing (#348): no rewalk is forced.
self.assertIs(self.api._root_listings[str(self.output_dir)].files, files)
after = (await (await self.client.get(url)).json())["images"][0]
self.assertNotIn("v=", before["thumbnail_url"])
self.assertIn("v=", after["thumbnail_url"])
async def test_no_output_folder_fails_the_job(self):
self.api._get_all_output_dirs = lambda: []
_, body = await self._start()
+1 -10
View File
@@ -477,20 +477,11 @@ class TestPendingFirstUse(unittest.TestCase):
class FakeTracker:
def __init__(self):
self.calls = []
self.runs = []
def set_current_prompt(
self, prompt_text, additional_data=None, push_to_queue=True, **run
):
def set_current_prompt(self, prompt_text, additional_data=None, push_to_queue=True):
self.calls.append({"text": prompt_text, "push_to_queue": push_to_queue})
self.runs.append(run)
return "exec-1"
def mark_empty(self, run_key=None, run_id=None, push_to_queue=False):
self.calls.append(
{"empty": True, "run_key": run_key, "push_to_queue": push_to_queue}
)
class TestNodeRoles(DbTestCase):
def _node(self):
+2 -28
View File
@@ -1,6 +1,5 @@
// Structural checks on .github/workflows/viewerjs-check.yml, plus the time limits of every CI workflow.
// The tooling has no YAML parser, so these read the workflow's fixed two-space layout: jobs at two
// spaces, their keys at four, steps at six.
// Structural checks on .github/workflows/viewerjs-check.yml. The tooling has no YAML parser, so these
// read the workflow's fixed two-space layout: jobs at two spaces, their keys at four, steps at six.
import test from 'node:test';
import assert from 'node:assert/strict';
import { execFileSync } from 'node:child_process';
@@ -222,28 +221,3 @@ test('browser job timeouts keep a real margin over measured runs', () => {
assert.ok(minutes >= 20, `${id} allows at least 20 minutes, since runner speed varies`);
}
});
// Every job is bounded (#350): a hang ends at a measured limit instead of GitHub's 6-hour default.
const BOUNDED_WORKFLOWS = ['test.yml', 'code-quality.yml', 'viewerjs-check.yml'];
const readWorkflow = name => readFileSync(new URL(`../../.github/workflows/${name}`, import.meta.url), 'utf8');
test('every job in the CI workflows has a numeric timeout-minutes', () => {
for (const name of BOUNDED_WORKFLOWS) {
const found = jobs(readWorkflow(name));
assert.ok(found.size > 0, `${name} has jobs`);
for (const [id, body] of found) {
const limit = /^ {4}timeout-minutes: (\d+)\s*$/m.exec(body);
assert.ok(limit, `${name} job ${id} sets a numeric timeout-minutes`);
assert.ok(Number(limit[1]) >= 5, `${name} job ${id} allows at least 5 minutes`);
}
}
});
test('the Python matrix and coverage run the suite through the watchdog runner', () => {
const found = jobs(readWorkflow('test.yml'));
assert.match(found.get('test'), /^ {8}run: python -m tests\.run_suite -v$/m);
assert.match(found.get('coverage'), /^ {10}python -m coverage run -m tests\.run_suite$/m);
for (const id of ['test', 'coverage']) {
assert.doesNotMatch(found.get(id), /unittest discover/, `${id} no longer runs the bare unittest discover`);
}
});
+1 -4
View File
@@ -29,18 +29,15 @@ test('execution script discovers only the new matrix and preserves all 78 origin
assert.doesNotMatch(stock, /execution-history/);
assert.doesNotMatch(stock, /execution-outputs/);
assert.doesNotMatch(stock, /library-location/); // it starts its own fresh host (#249 S6b)
assert.doesNotMatch(stock, /image-attribution/); // needs the owned fixture pack (#342)
const executions = list('tools/vendor/playwright.execution.config.mjs');
assert.match(executions, /\[classic-start\]/);
assert.match(executions, /\[nodes2-start\]/);
// + the fresh-host library location case, + #342 image attribution, + #343 linked text
assert.match(executions, /Total: 90 tests in 5 files/);
assert.match(executions, /Total: 86 tests in 4 files/); // + the fresh-host library location case, + #343 linked text
for (const project of ['classic-start', 'nodes2-start']) {
const cases = executions.split('\n').filter(line => line.includes(`[${project}]`));
assert.equal(cases.filter(line => line.includes('execution-history.spec.mjs')).length, 15);
assert.equal(cases.filter(line => line.includes('execution-outputs.spec.mjs')).length, 8);
assert.equal(cases.filter(line => line.includes('library-location.spec.mjs')).length, 1);
assert.equal(cases.filter(line => line.includes('image-attribution.spec.mjs')).length, 2);
assert.equal(cases.filter(line => line.includes('execution-text-link.spec.mjs')).length, 19);
}
assert.doesNotMatch(executions, /prompt-manager.spec/);
-330
View File
@@ -1,330 +0,0 @@
"""Review, and on request repair, image links that nothing certain supports (#342).
Before #342 the image monitor linked an image it could not identify to the
newest prompt in the database or to a stale tracker entry. This tool judges the
selected images again by today's rules, using only what is certain:
- the image's own metadata, as the monitor reads it;
- the output scan's resolver, and its legacy browser-parser floor (#170);
- recorded execution history (``database.run_attribution``).
Each image gets one verdict:
``keep`` a certain source names the prompt it is linked to
``relink`` every certain source names one other stored prompt
``unlink`` nothing certain names a prompt (the link was a guess)
``conflict`` certain sources disagree: decide by hand
``missing`` the file is gone: leave it to Maintenance's missing links
Dry run by default: the database is copied (SQLite backup from a read-only
connection) into a temporary folder and judged there, so the library is never
written. ``--apply`` writes the ``relink`` and ``unlink`` verdicts of the
selected images, after saving a backup beside the database
(``<db>.pre-relink_<stamp>``). Only ids named on the command line are ever
changed, and only while each is still linked to the prompt it was judged on.
What ``--apply`` does (APPLY_EFFECTS below, printed by every run):
- ``relink`` points the image's link record at the other prompt and drops its
pair links (they named the old positive).
- ``unlink`` removes only the image's link record (its ``generated_images`` row
with the metadata snapshot stored in it, and its pair links). That is how the
app represents an unlinked image, exactly like a file the scan could not
attribute: the file stays on disk and in the gallery as unlinked, and its
favorite star and tags (kept by file, not by link) are untouched.
A batch item linked through the monitor's queue (PromptSearchList) can't be
proven after the fact and reads ``unlink``: review before applying.
python3 tools/review_image_links.py <db> --output-dir <ComfyUI output> \\
--ids 58005-58010
python3 tools/review_image_links.py <db> --output-dir <dir> --prompt 5278
python3 tools/review_image_links.py <db> --output-dir <dir> \\
--ids 58005-58010 --apply
Stop ComfyUI before ``--apply``. Run from this repository with its venv; the
package's logger may write to its usual log file.
"""
import argparse
import datetime
import json
import os
import sqlite3
import sys
import tempfile
from contextlib import closing
from pathlib import Path
_ROOT = Path(__file__).resolve().parents[1]
if str(_ROOT) not in sys.path:
sys.path.insert(0, str(_ROOT))
from database.maintenance import _delete_ids # noqa: E402
from database.operations import PromptDatabase # noqa: E402
from database.pair_ingestion import blank_positive # noqa: E402
from database.run_attribution import stored_run_prompt # noqa: E402
from utils.generation_params import ( # noqa: E402
PLACEHOLDERS,
GenerationParamsError,
extract_generation_params,
)
from utils.hashing import generate_prompt_hash # noqa: E402
from utils.image_monitor import ImageGenerationHandler # noqa: E402
from utils.prompt_pair_resolver import resolve_prompt_pair # noqa: E402
VERDICTS = ("keep", "relink", "unlink", "conflict", "missing")
APPLY_EFFECTS = (
"--apply, for ids named with --ids only, while each is still linked to the "
"prompt it was judged on:\n"
" relink: points the image's link record at the other prompt and drops "
"its pair links;\n"
" unlink: removes only the image's link record (its stored metadata "
"snapshot and pair links go with it); the file stays on disk and in the "
"gallery as unlinked, with its favorite and tags.\n"
)
# Ids per query (well under SQLite's bound-parameter limit)
ID_BATCH = 500
def parse_ids(text):
"""``"1,4-6"`` -> ``[1, 4, 5, 6]``; raises ValueError on anything else."""
ids = set()
for part in (text or "").split(","):
part = part.strip()
if not part:
continue
low, _, high = part.partition("-")
first, last = int(low), int(high or low)
if first < 1 or last < first or last - first > 100_000:
raise ValueError(f"Bad id range: {part}")
ids.update(range(first, last + 1))
return sorted(ids)
def _stored_id(db, text):
if not isinstance(text, str) or not text.strip():
return None
row = db.get_prompt_by_hash(generate_prompt_hash(text.strip()))
return row["id"] if row else None
def _parser_text(metadata):
"""The legacy browser parser's positive prompt (the scan's floor), or None."""
records = {key: metadata.get(key) for key in ("prompt", "workflow")}
try:
positive = extract_generation_params(records)["positivePrompt"]
except (GenerationParamsError, ValueError, TypeError, KeyError):
return None
if not isinstance(positive, str) or positive == PLACEHOLDERS["positivePrompt"]:
return None
return positive.strip() or None
def certain_prompts(db, handler, path, metadata, output_dir):
"""``{source: prompt id}`` of every certain source that names a stored prompt."""
found = {}
if metadata:
state, prompt = handler._find_prompt_from_metadata(metadata)
if state == "linked" and prompt:
found["metadata"] = prompt["id"]
resolution = resolve_prompt_pair(metadata)
if resolution["state"] == "resolved" and not blank_positive(resolution):
positive = resolution["positive"]
scanned = _stored_id(db, positive.get("storage_text", positive["text"]))
else:
scanned = _stored_id(db, _parser_text(metadata))
if scanned is not None:
found["scan"] = scanned
recorded = stored_run_prompt(db, path, output_dir)
if recorded is not None:
found["history"] = recorded
return found
def judge(db, handler, row, output_dir):
"""One image's verdict dict."""
path = row["image_path"]
verdict = {
"image_id": row["id"],
"file": path,
"linked_prompt": row["prompt_id"],
"sources": {},
"target": None,
}
if not os.path.isfile(path):
return {**verdict, "verdict": "missing"}
try:
metadata = handler.metadata_extractor.extract_metadata(path)
except Exception:
metadata = None
sources = certain_prompts(db, handler, path, metadata, output_dir)
named = set(sources.values())
verdict["sources"] = sources
if not named:
return {**verdict, "verdict": "unlink"}
if len(named) > 1:
return {**verdict, "verdict": "conflict"}
(target,) = named
if target == row["prompt_id"]:
return {**verdict, "verdict": "keep", "target": target}
return {**verdict, "verdict": "relink", "target": target}
def selected_rows(db, ids, prompt_id):
"""The named images and every image of ``prompt_id``, by id."""
conn = db.model.get_connection()
columns = "SELECT id, prompt_id, image_path FROM generated_images "
rows = {}
for start in range(0, len(ids), ID_BATCH):
batch = ids[start : start + ID_BATCH]
marks = ",".join("?" for _ in batch)
query = columns + f"WHERE id IN ({marks})" # only ? marks
rows.update((row["id"], row) for row in conn.execute(query, batch))
if prompt_id is not None:
found = conn.execute(columns + "WHERE prompt_id = ?", (prompt_id,))
rows.update((row["id"], row) for row in found)
return [rows[key] for key in sorted(rows)]
def review(db, ids, prompt_id, output_dir):
handler = ImageGenerationHandler(db, None)
return [
judge(db, handler, row, output_dir) for row in selected_rows(db, ids, prompt_id)
]
def apply(db, verdicts, ids):
"""Write the relink/unlink verdicts of explicitly named ids; returns counts.
An image whose link changed since it was judged (now another prompt, or
gone) is skipped.
"""
allowed = set(ids)
chosen = [
v
for v in verdicts
if v["verdict"] in ("relink", "unlink") and v["image_id"] in allowed
]
with db._pair_transaction() as conn:
current = [
item
for item in chosen
if (
row := conn.execute(
"SELECT prompt_id FROM generated_images WHERE id=?",
(item["image_id"],),
).fetchone()
)
and row[0] == item["linked_prompt"]
]
relink = [item for item in current if item["verdict"] == "relink"]
for item in relink:
# A pair link records the old positive: it goes with the old link
conn.execute(
"DELETE FROM image_prompt_pairs WHERE image_id=?", (item["image_id"],)
)
conn.execute(
"UPDATE generated_images SET prompt_id=? WHERE id=? AND prompt_id=?",
(item["target"], item["image_id"], item["linked_prompt"]),
)
removed = _delete_ids(
conn,
"generated_images",
[item["image_id"] for item in current if item["verdict"] == "unlink"],
)
return {
"relinked": len(relink),
"unlinked": removed,
"skipped": len(chosen) - len(current),
}
def backup_database(source, target):
"""SQLite online backup of ``source`` (opened read-only) into ``target``."""
uri = Path(os.path.abspath(source)).as_uri() + "?mode=ro"
with (
closing(sqlite3.connect(uri, uri=True)) as src,
closing(sqlite3.connect(target)) as dst,
):
src.backup(dst)
def _print(verdicts, out):
for item in verdicts:
sources = ", ".join(f"{k}={v}" for k, v in sorted(item["sources"].items()))
target = f" -> {item['target']}" if item["verdict"] == "relink" else ""
linked = item["linked_prompt"]
out.write(
f"{item['image_id']:>8} {item['verdict']:<8} prompt {linked}{target}"
f" [{sources or 'nothing certain'}] {item['file']}\n"
)
counts = {name: sum(v["verdict"] == name for v in verdicts) for name in VERDICTS}
out.write(" ".join(f"{name}: {count}" for name, count in counts.items()) + "\n")
def main(argv=None, out=sys.stdout):
parser = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0])
parser.add_argument("database", help="the library database (prompts.db)")
parser.add_argument(
"--output-dir", required=True, help="ComfyUI's output directory (history)"
)
parser.add_argument("--ids", default="", help="image ids, e.g. 58005-58010,58020")
parser.add_argument("--prompt", type=int, help="every image linked to this prompt")
parser.add_argument(
"--apply", action="store_true", help="write relink/unlink for --ids images"
)
parser.add_argument("--json", action="store_true", help="print JSON verdicts")
args = parser.parse_args(argv)
try:
ids = parse_ids(args.ids)
except ValueError as error:
parser.error(str(error))
if not ids and args.prompt is None:
parser.error("select images with --ids and/or --prompt")
if args.apply and not ids:
parser.error("--apply changes only images named with --ids")
if not os.path.isfile(args.database):
parser.error(f"no database at {args.database}")
if not args.apply:
with tempfile.TemporaryDirectory() as scratch:
copy = os.path.join(scratch, "review.db")
backup_database(args.database, copy)
db = PromptDatabase(copy)
try:
verdicts = review(db, ids, args.prompt, args.output_dir)
finally:
db.close_all()
_report(verdicts, args.json, out)
out.write(APPLY_EFFECTS)
out.write("Dry run: nothing was changed.\n")
return 0
stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
backup = f"{args.database}.pre-relink_{stamp}"
backup_database(args.database, backup)
out.write(f"Backup: {backup}\n")
db = PromptDatabase(args.database)
try:
verdicts = review(db, ids, args.prompt, args.output_dir)
_report(verdicts, args.json, out)
result = apply(db, verdicts, ids)
finally:
db.close_all()
out.write(
f"Applied: {result['relinked']} relinked, {result['unlinked']} unlinked, "
f"{result['skipped']} skipped (link changed since review)\n"
)
return 0
def _report(verdicts, as_json, out):
if as_json:
out.write(json.dumps(verdicts, indent=2, sort_keys=True) + "\n")
else:
_print(verdicts, out)
if __name__ == "__main__":
sys.exit(main())
+1 -2
View File
@@ -19,8 +19,7 @@ const servers = [{
export default defineConfig({
testDir: '../../tests/comfyui',
testIgnore: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs', 'execution-text-link.spec.mjs', 'library-location.spec.mjs',
'image-attribution.spec.mjs'],
testIgnore: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs', 'execution-text-link.spec.mjs', 'library-location.spec.mjs'],
outputDir: './test-results/comfyui',
timeout: 60000,
workers: 1,
-81
View File
@@ -1,81 +0,0 @@
// Maintainer tool, run by hand (never in CI): how long the workspace Images
// page takes to show its first thumbnail (#348, TESTING.md "Images listing
// cache (#348)").
//
// Once: npm ci --prefix tools/vendor
// npx --prefix tools/vendor playwright install chromium
// Then: node tools/vendor/perf/images_load_probe.mjs http://127.0.0.1:8188
//
// Prints the time to the first loaded thumbnail, the time to network idle and
// the 10 slowest requests. Each run opens a fresh browser context (no cache).
import { chromium } from '@playwright/test';
const SLOWEST = 10;
const TIMEOUT_MS = 120_000;
function usage() {
console.error('usage: node tools/vendor/perf/images_load_probe.mjs <base-url>');
console.error('example: node tools/vendor/perf/images_load_probe.mjs http://127.0.0.1:8188');
process.exit(2);
}
function parseBase(argument) {
if (!argument) usage();
let url;
try {
url = new URL(argument);
} catch {
usage();
}
if (url.protocol !== 'http:' && url.protocol !== 'https:') usage();
return url.origin + url.pathname.replace(/\/+$/, '');
}
const ms = value => `${Math.round(value)} ms`;
async function probe(base) {
const browser = await chromium.launch();
try {
const context = await browser.newContext({ viewport: { width: 1440, height: 900 } });
const page = await context.newPage();
const requests = [];
page.on('requestfinished', request => {
const timing = request.timing();
requests.push({
url: request.url().replace(base, ''),
status: null,
ms: timing.responseEnd >= 0 ? timing.responseEnd : NaN,
response: request.response(),
});
});
const started = Date.now();
await page.goto(`${base}/prompt_manager/workspace#/images`, { timeout: TIMEOUT_MS });
// The first Images thumbnail the browser has decoded.
await page.waitForFunction(() => [...document.images].some(image =>
image.complete && image.naturalWidth > 0
&& image.src.includes('/prompt_manager/images/serve/')), null,
{ timeout: TIMEOUT_MS, polling: 'raf' });
const firstThumbnail = Date.now() - started;
await page.waitForLoadState('networkidle', { timeout: TIMEOUT_MS });
const networkIdle = Date.now() - started;
for (const entry of requests) {
entry.status = (await entry.response)?.status() ?? 'n/a';
}
const slowest = requests.filter(entry => Number.isFinite(entry.ms))
.sort((a, b) => b.ms - a.ms).slice(0, SLOWEST);
console.log(`first thumbnail: ${ms(firstThumbnail)}`);
console.log(`network idle: ${ms(networkIdle)}`);
console.log(`requests: ${requests.length}`);
console.log(`slowest ${slowest.length}:`);
for (const entry of slowest) {
console.log(` ${ms(entry.ms).padStart(9)} ${String(entry.status).padEnd(3)} ${entry.url}`);
}
} finally {
await browser.close();
}
}
probe(parseBase(process.argv[2])).catch(error => {
console.error(`probe failed: ${error.message}`);
process.exit(1);
});
+1 -1
View File
@@ -6,7 +6,7 @@ if (host.href !== 'http://127.0.0.1:8196/') throw new Error('Execution tests req
export default defineConfig({
testDir: '../../tests/comfyui',
testMatch: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs', 'execution-text-link.spec.mjs', 'library-location.spec.mjs', 'image-attribution.spec.mjs'],
testMatch: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs', 'execution-text-link.spec.mjs', 'library-location.spec.mjs'],
outputDir: './test-results/executions/cases',
timeout: 60000,
workers: 1,
+83 -291
View File
@@ -18,9 +18,8 @@ Typical usage:
The system automatically:
- Detects new image files in monitored directories
- Extracts ComfyUI workflow metadata from PNG chunks
- Links images to tracked prompts only when the image's own graph proves the
tracked entry belongs to its run (``utils.run_keys``)
- Leaves an image unlinked when nothing certain identifies its prompt (#342)
- Links images to active prompts using the prompt tracker
- Handles fallback linking when no active prompt is available
- Provides status information and monitoring control
"""
@@ -28,8 +27,7 @@ import os
import queue
import time
import threading
from collections import OrderedDict
from typing import Callable, Optional, Dict, Any, List, Tuple
from typing import Callable, Optional, Dict, Any, Tuple
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler
@@ -38,38 +36,18 @@ from .metadata_extractor import ComfyUIMetadataExtractor
from .logging_config import get_logger
from .prompt_graph import resolve_text, run_prompt_nodes
from .prompt_pair_resolver import resolve_prompt_pair
from .run_keys import image_run_keys
from .usage_tracking import survive_switch
try:
from ..database.pair_ingestion import ingest_resolved_pair
from ..database.run_attribution import (
host_output_directory,
in_output_folder,
stored_run_prompt,
subscribe_outputs_sealed,
)
except ImportError:
from database.pair_ingestion import ingest_resolved_pair
from database.run_attribution import (
host_output_directory,
in_output_folder,
stored_run_prompt,
subscribe_outputs_sealed,
)
# Upper bound on images waiting to be processed (a burst beyond this is skipped)
MAX_PENDING_IMAGES = 1000
# A file counts as fully written once its size is unchanged across two reads
SETTLE_INTERVAL_SECONDS = 0.1
SETTLE_ATTEMPTS = 20
# Keyless images (no embedded API graph) left unlinked wait for their run's
# recorded outputs for one history-only re-check (#342): at most this many,
# each for at most this long.
MAX_AWAITING_HISTORY = 1000
AWAIT_HISTORY_SECONDS = 3600
# ``schedule`` snapshot meaning "history-only re-check" (not a new image)
HISTORY_ONLY = object()
class ImageGenerationHandler(FileSystemEventHandler):
@@ -106,10 +84,6 @@ class ImageGenerationHandler(FileSystemEventHandler):
# ``queue_capture_thumbnail``, wired at start-up; #249). It only
# queues: the preview is written on its own worker. None: no queue.
self.thumbnail_queue: Optional[Callable[[str], str]] = None
# Told about every settled output, with its resolved path (the API's
# ``notify_output_settled``, wired at start-up; #348): the Images
# listing walks that root again. None: nobody is told.
self.output_listener: Optional[Callable[[str], None]] = None
# Read from GalleryConfig if available, otherwise use defaults
try:
@@ -135,14 +109,6 @@ class ImageGenerationHandler(FileSystemEventHandler):
self._work: "queue.Queue[Optional[str]]" = queue.Queue()
self._stop = threading.Event()
self._worker: Optional[threading.Thread] = None
# Keyless unlinked images awaiting their run's recorded outputs:
# normalized path -> (path, monotonic deadline), oldest first.
self._awaiting: "OrderedDict[str, Tuple[str, float]]" = OrderedDict()
# Sealed re-checks ``schedule`` refused (the path still pending, or the
# pending set full): retried each time a pending image finishes.
self._history_due: "OrderedDict[str, None]" = OrderedDict()
self._awaiting_lock = threading.Lock()
self._sealed_subscribed = False
def on_created(self, event):
"""Handle filesystem creation events.
@@ -160,17 +126,15 @@ class ImageGenerationHandler(FileSystemEventHandler):
if event.is_directory or not self.is_image_file(event.src_path):
return
self.logger.info(f"New image detected: {event.src_path}")
# Snapshot the tracked runs NOW, before the delay: a later run may
# execute the same node again before this image is processed.
self.schedule(event.src_path, self._keyed_snapshot())
def _keyed_snapshot(self):
"""The tracker's latest entry per run key, or None when unavailable."""
try:
return self.prompt_tracker.keyed_snapshot()
except Exception as error:
self.logger.warning(f"Could not snapshot tracked prompts: {error}")
return None
# Snapshot prompt context NOW before the delay — in batch workflows
# the tracker advances to the next prompt before images are processed.
prompt_snapshot = self.prompt_tracker.get_current_prompt()
if prompt_snapshot:
self.logger.debug(
f"Snapshot prompt {prompt_snapshot.get('id', '?')} for "
f"{os.path.basename(event.src_path)}"
)
self.schedule(event.src_path, prompt_snapshot)
def schedule(self, image_path: str, prompt_snapshot=None) -> bool:
"""Queue an image for processing on the worker thread, once per path.
@@ -210,9 +174,6 @@ class ImageGenerationHandler(FileSystemEventHandler):
and exits; the next ``schedule`` starts a fresh worker with its own
queue. Safe to call repeatedly and before anything was scheduled.
"""
with self._awaiting_lock:
self._awaiting.clear()
self._history_due.clear()
with self._pending_lock:
self._pending.clear()
worker, self._worker = self._worker, None
@@ -263,16 +224,11 @@ class ImageGenerationHandler(FileSystemEventHandler):
f"Image never settled or vanished, skipping: {image_path}"
)
return
try:
self.process_new_image(image_path, prompt_snapshot=prompt_snapshot)
self._queue_thumbnail(image_path)
finally:
# Linked or not, even if linking raised: the file exists.
self._report_output(image_path)
self.process_new_image(image_path, prompt_snapshot=prompt_snapshot)
self._queue_thumbnail(image_path)
finally:
with self._pending_lock:
self._pending.pop(image_path, None)
self._drain_history_due()
def _queue_thumbnail(self, image_path: str) -> None:
"""Queue a settled output's preview, best effort (linked or not).
@@ -293,25 +249,6 @@ class ImageGenerationHandler(FileSystemEventHandler):
return
self.logger.debug(f"Preview {outcome}: {os.path.basename(image_path)}")
def _report_output(self, image_path: str) -> None:
"""Tell the output listener a settled output exists, best effort.
The path is resolved here, on the worker thread, so the listener's
event loop does no file I/O. A failure is logged, never raised: the
output is then listed by the listing's 30 s refresh.
"""
listener = getattr(self, "output_listener", None)
if listener is None:
return
try:
listener(os.path.realpath(image_path))
except Exception:
self.logger.warning(
f"{os.path.basename(image_path)} could not be reported to the "
"Images listing",
exc_info=True,
)
def wait_until_settled(self, image_path: str, stop=None) -> bool:
"""True once the file's size is non-zero and unchanged across two reads.
@@ -353,20 +290,15 @@ class ImageGenerationHandler(FileSystemEventHandler):
This method handles the complete processing pipeline for a new image:
1. Verifies the file still exists
2. Extracts ComfyUI metadata from the image
3. Identifies the prompt from the image's metadata, else from a tracked
run the image's graph proves (``_find_prompt_from_run``)
4. Links the image to that prompt, or leaves it unlinked when nothing
certain identifies it (logged at info; #342)
3. Tries to identify the correct prompt via metadata-based lookup
4. Falls back to the prompt snapshot captured at event time
5. Links the image to the appropriate prompt in the database
Args:
image_path: Full path to the newly created image file
prompt_snapshot: The tracker's keyed snapshot taken when the image
appeared (``PromptTracker.keyed_snapshot``); None reads it now
prompt_snapshot: Prompt context captured at file-creation time (optional)
"""
try:
if prompt_snapshot is HISTORY_ONLY:
self._recheck_history(image_path)
return
self.logger.info(f"Processing image: {image_path}")
if not os.path.exists(image_path):
@@ -397,17 +329,43 @@ class ImageGenerationHandler(FileSystemEventHandler):
f"Metadata match: prompt {current_prompt['id']} for "
f"{os.path.basename(image_path)}"
)
else:
current_prompt = self._find_prompt_from_run(
image_path, metadata, prompt_snapshot
)
# Strategy 2: Pop from batch queue. Batch items (text supplied by
# PromptSearchList) share one graph, so metadata can't tell them apart;
# they are queued in encoding order and images save in the same order.
if not current_prompt:
# Nothing certain: never guess from recency or a stale entry (#342)
current_prompt = self.prompt_tracker.pop_next_prompt()
if current_prompt:
self.logger.info(
f"Queue match: prompt {current_prompt['id']} for "
f"{os.path.basename(image_path)}"
)
# Strategy 3: Use snapshot captured at file-creation time
if not current_prompt and prompt_snapshot:
current_prompt = prompt_snapshot
self.logger.info(
"No certain prompt for "
f"{os.path.basename(image_path)}; left unlinked"
f"Snapshot match: prompt {current_prompt.get('id', 'unknown')}"
)
return
# Strategy 4: Check live prompt tracker
if not current_prompt:
current_prompt = self.prompt_tracker.get_current_prompt()
if current_prompt:
self.logger.info(
f"Live tracker match: prompt {current_prompt['id']}"
)
# Strategy 5: Fallback to most recent prompt in DB
if not current_prompt:
current_prompt = self._get_fallback_prompt()
if current_prompt:
self.logger.debug(f"Fallback match: prompt {current_prompt['id']}")
else:
self.logger.warning(
f"No prompt context available, skipping image: {image_path}"
)
return
# Extend timeout if we have an active execution
if current_prompt.get("execution_id"):
@@ -435,193 +393,6 @@ class ImageGenerationHandler(FileSystemEventHandler):
self.logger.error(traceback.format_exc())
def _find_prompt_from_run(
self, image_path: str, metadata, prompt_snapshot=None
) -> Optional[Dict[str, Any]]:
"""The prompt a tracked run proves for an image whose metadata can't say.
Only an entry registered under the run key of the image's own positive
node counts (``utils.run_keys``): never the newest prompt in the
database and never an entry of another graph (#342).
The snapshot taken when the image appeared names the key's latest run.
1. Batch queue: the oldest queued entry of that key from that run.
Batch items (text from PromptSearchList) share one key and are
linked in order; an empty item's placeholder leaves its image
unlinked; entries queued by later runs stay for their own images.
2. The snapshot's entry itself, when its run is known and registered
exactly one entry for the key (a batch's latest item says nothing
about the other images). This covers a cached node that did not
run again.
3. Recorded execution history: the run that wrote this file names one
stored prompt (``database.run_attribution``).
"""
keys = image_run_keys(metadata)
name = os.path.basename(image_path)
if not keys:
# No tracked run can be proven: history now, and once more after
# the run's outputs are recorded (``_on_outputs_sealed``). A file
# outside ComfyUI's output folder can never be attributed, so it
# does not take a waiting slot.
if not in_output_folder(image_path, host_output_directory()):
return None
self._await_history(image_path)
prompt_id = self._history_prompt(image_path)
if prompt_id is None:
return None
self._stop_awaiting(image_path)
self.logger.info(f"History match: prompt {prompt_id} for {name}")
return {"id": prompt_id, "from_history": True}
if len(keys) == 1:
(key,) = keys.values()
if prompt_snapshot is None:
prompt_snapshot = self._keyed_snapshot()
# No entry of the key when the image appeared: no tracked run made it
entry = (prompt_snapshot or {}).get(key)
if entry is not None:
context, run_id, count = entry
queued = self.prompt_tracker.pop_for_key(key, run_id)
if queued is not None:
if queued.get("placeholder"):
self.logger.info(f"Queue entry for {name} has no prompt text")
return None
self.logger.info(f"Queue match: prompt {queued['id']} for {name}")
return queued
if context.get("placeholder"):
# Same rule as a popped placeholder: that run had no prompt
self.logger.info(f"Latest run for {name} had no prompt text")
return None
if run_id and count == 1:
self.logger.info(f"Run match: prompt {context['id']} for {name}")
return context
else:
self.logger.debug(f"{len(keys)} positive prompt nodes for {name}")
prompt_id = self._history_prompt(image_path)
if prompt_id is None:
return None
self.logger.info(f"History match: prompt {prompt_id} for {name}")
return {"id": prompt_id, "from_history": True}
@staticmethod
def _awaiting_key(image_path: str) -> str:
"""One spelling for a waiting file, through symlinks and junctions."""
try:
resolved = os.path.realpath(image_path)
except (OSError, ValueError):
resolved = os.path.abspath(image_path)
return os.path.normcase(resolved)
def _prune_awaiting(self) -> None:
"""Drop expired and surplus entries (``_awaiting_lock`` held)."""
now = time.monotonic()
for key in [k for k, (_p, due) in self._awaiting.items() if due <= now]:
del self._awaiting[key]
while len(self._awaiting) > MAX_AWAITING_HISTORY:
self._awaiting.popitem(last=False)
def _await_history(self, image_path: str) -> None:
"""Re-check ``image_path`` once, history only, when its run is sealed."""
if not self._sealed_subscribed:
subscribe_outputs_sealed(self._on_outputs_sealed)
self._sealed_subscribed = True
key = self._awaiting_key(image_path)
with self._awaiting_lock:
self._awaiting.pop(key, None)
self._awaiting[key] = (
image_path,
time.monotonic() + AWAIT_HISTORY_SECONDS,
)
self._prune_awaiting()
def _stop_awaiting(self, image_path: str) -> None:
with self._awaiting_lock:
self._awaiting.pop(self._awaiting_key(image_path), None)
def awaiting_history(self) -> List[str]:
"""Normalized paths still waiting for their run's recorded outputs."""
with self._awaiting_lock:
self._prune_awaiting()
return list(self._awaiting)
def _on_outputs_sealed(self, run_id: str, references: List[dict]) -> None:
"""A run's outputs were recorded: re-check its waiting images once.
Called on the history writer's thread, so it only schedules the work.
"""
root = host_output_directory()
if not root:
return
due = []
with self._awaiting_lock:
self._prune_awaiting()
for reference in references:
if reference.get("folder_type") != "output":
continue
key = self._awaiting_key(
os.path.join(
root, reference.get("subfolder", ""), reference["filename"]
)
)
waiting = self._awaiting.pop(key, None)
if waiting is not None:
due.append(waiting[0])
for image_path in due:
self._schedule_recheck(image_path)
def _schedule_recheck(self, image_path: str) -> None:
"""Schedule the history-only re-check, or keep it due for later."""
if self.schedule(image_path, HISTORY_ONLY):
return
with self._awaiting_lock:
self._history_due.pop(image_path, None)
self._history_due[image_path] = None
while len(self._history_due) > MAX_AWAITING_HISTORY:
self._history_due.popitem(last=False)
def _drain_history_due(self) -> None:
"""Schedule due re-checks now that a pending slot or path is free."""
with self._awaiting_lock:
if not self._history_due:
return
due = list(self._history_due)
for image_path in due:
if self.schedule(image_path, HISTORY_ONLY):
with self._awaiting_lock:
self._history_due.pop(image_path, None)
def _recheck_history(self, image_path: str) -> None:
"""The one history-only re-check of a keyless image (#342)."""
name = os.path.basename(image_path)
if not os.path.exists(image_path):
return
prompt_id = self._history_prompt(image_path)
if prompt_id is None:
self.logger.info(
f"No certain prompt for {name} after its run ended; left unlinked"
)
return
self.logger.info(f"History match after its run ended: {prompt_id} for {name}")
try:
metadata = self.metadata_extractor.extract_metadata(image_path)
except Exception:
metadata = None
self.link_image_to_prompt(
image_path,
{"id": prompt_id, "from_history": True},
metadata or {"file_info": self.get_basic_file_info(image_path)},
)
def _history_prompt(self, image_path: str) -> Optional[int]:
"""The stored prompt recorded history attributes to this file, if one."""
try:
return stored_run_prompt(
self.db_manager, image_path, host_output_directory()
)
except Exception as error:
self.logger.warning(f"History attribution failed: {error}")
return None
def _find_prompt_from_metadata(self, metadata) -> Tuple[str, Optional[Dict]]:
"""Identify the prompt that produced an image from its embedded metadata.
@@ -771,6 +542,35 @@ class ImageGenerationHandler(FileSystemEventHandler):
self.logger.error(f"Error getting file info: {e}")
return {}
def _get_fallback_prompt(self) -> Optional[Dict[str, Any]]:
"""Get the most recent prompt from database as fallback.
When no active prompt is available from the tracker, this method attempts
to find the most recently created prompt in the database to use as a fallback
for image linking.
Returns:
Dictionary containing prompt information with 'fallback' flag set to True,
or None if no recent prompt is available
"""
try:
recent = self.db_manager.get_recent_prompts(limit=1)
# Paginated dict ({"prompts": [...]}) or a plain list
recent_prompts = (
recent.get("prompts") if isinstance(recent, dict) else recent
)
if recent_prompts:
prompt = recent_prompts[0]
return {
"id": prompt["id"],
"text": prompt["text"],
"timestamp": prompt.get("created_at"),
"fallback": True,
}
except Exception as e:
self.logger.error(f"Error getting fallback prompt: {e}")
return None
def link_image_to_prompt(
self, image_path: str, prompt_context: Dict, metadata: Dict
):
@@ -815,9 +615,10 @@ class ImageGenerationHandler(FileSystemEventHandler):
if image_id is None:
return
if image_id: # 0: nothing new was linked (logged by the database)
fallback_note = " (fallback)" if prompt_context.get("fallback") else ""
self.logger.debug(
f"Successfully linked image {image_id} to prompt "
f"{prompt_context['id']}"
f"{prompt_context['id']}{fallback_note}"
)
except Exception as e:
self.logger.error(f"Failed to link image to prompt: {e}")
@@ -875,14 +676,6 @@ class ImageMonitor:
self.logger = get_logger("prompt_manager.image_monitor")
self.review_recorder: Optional[Callable[[str], str]] = None
self.thumbnail_queue: Optional[Callable[[str], str]] = None
self.output_listener: Optional[Callable[[str], None]] = None
def set_output_listener(self, listener: Optional[Callable[[str], None]]) -> None:
"""Report every settled output's resolved path to ``listener`` (None
stops it), on the running handler and every later one (#348)."""
self.output_listener = listener
if self.handler is not None:
self.handler.output_listener = listener
def set_thumbnail_queue(
self, queue_preview: Optional[Callable[[str], str]]
@@ -960,7 +753,6 @@ class ImageMonitor:
handler = ImageGenerationHandler(self.db_manager, self.prompt_tracker)
handler.review_recorder = self.review_recorder
handler.thumbnail_queue = self.thumbnail_queue
handler.output_listener = self.output_listener
observer = Observer()
scheduled = []
for output_dir in valid_directories:
+9 -151
View File
@@ -28,8 +28,7 @@ Or using the context manager:
import threading
import time
import uuid
from collections import OrderedDict
from typing import Any, Dict, Optional, Tuple
from typing import Optional, Dict, Any
# Import logging system
try:
@@ -104,8 +103,6 @@ class PromptTracker:
prompt_text: str,
additional_data: Optional[Dict[str, Any]] = None,
push_to_queue: bool = True,
run_key: Optional[str] = None,
run_id: Optional[str] = None,
) -> str:
"""
Set the current prompt for this thread for image tracking.
@@ -123,10 +120,6 @@ class PromptTracker:
push_to_queue: Also push onto the FIFO batch queue. Only batch items whose
text comes from another node need it; typed prompts are linked from
the image's own metadata instead.
run_key: The node's run key (``utils.run_keys.prompt_node_key``). The
image monitor links an image to a tracked prompt only when the
image's own graph gives the same key (#342).
run_id: The ComfyUI prompt id executing the node, when the host says.
Returns:
Unique execution ID for this prompt execution
@@ -164,8 +157,6 @@ class PromptTracker:
"timestamp": time.time(),
"thread_id": threading.current_thread().ident,
"additional_data": additional_data or {},
"run_key": run_key,
"run_id": run_id,
}
# Store in thread-local storage
@@ -175,20 +166,20 @@ class PromptTracker:
with self.lock:
self.active_prompts[execution_id] = execution_context
self._remember_keyed(execution_context)
# Push to batch queue for ordered image linking
if push_to_queue:
with self._queue_lock:
self._prompt_queue.append(execution_context)
self.logger.debug(
f"Queue size: {len(self._prompt_queue)} after push "
f"for prompt {prompt_id}"
)
self.logger.debug(
f"Set current prompt: {execution_id} -> {prompt_text[:50]}... "
f"(thread: {threading.current_thread().ident})"
)
self.logger.debug(f"Active prompts count: {len(self.active_prompts)}")
# Push to batch queue for ordered image linking. Last, with nothing
# after it that can raise: a caller that marks a failed execution
# (#342) must never add a second queue slot for the same item.
if push_to_queue:
with self._queue_lock:
self._prompt_queue.append(execution_context)
return execution_id
def get_current_prompt(self) -> Optional[Dict[str, Any]]:
@@ -234,99 +225,6 @@ class PromptTracker:
)
return None
def mark_empty(
self,
run_key: Optional[str],
run_id: Optional[str] = None,
push_to_queue: bool = False,
) -> None:
"""Record that a positive node ran with nothing to save (#342).
The run key's latest entry becomes a placeholder, so no earlier prompt of
that key can be linked to this run's images; with ``push_to_queue`` the
placeholder also holds this item's place in the batch queue, so a later
item's image never takes the wrong queue entry.
"""
if not run_key:
return
placeholder = {
"id": None,
"placeholder": True,
"execution_id": self.generate_execution_id(),
"text": "",
"timestamp": time.time(),
"additional_data": {},
"run_key": run_key,
"run_id": run_id,
}
self._remember_keyed(placeholder)
if push_to_queue:
with self._queue_lock:
self._prompt_queue.append(placeholder)
def _keyed_runs(self) -> "KeyedRuns":
keyed = getattr(self, "_keyed", None)
if keyed is None:
with _KEYED_INIT_LOCK:
keyed = getattr(self, "_keyed", None)
if keyed is None:
keyed = self._keyed = KeyedRuns()
return keyed
def _remember_keyed(self, context: Dict[str, Any]) -> None:
if context.get("run_key"):
self._keyed_runs().record(context)
def keyed_snapshot(self) -> Dict[str, Tuple[Dict[str, Any], Optional[str], int]]:
"""``{run key: (latest entry, its run id, entries of that run)}`` now.
The image monitor captures it when an image appears, before later runs
can move on.
"""
return self._keyed_runs().snapshot(self.prompt_timeout)
def pop_for_key(
self, run_key: Optional[str], run_id: Optional[str] = None
) -> Optional[Dict[str, Any]]:
"""Pop the oldest queued entry of ``run_key`` (FIFO), from ``run_id`` if set.
``run_id`` is the run the key's latest entry came from when the image
appeared (``keyed_snapshot``): entries of that key queued by earlier
runs ahead of it are stale and dropped, while entries of later runs
(queued after the image appeared) stay for their own images. Entries
of other keys stay queued; expired entries are discarded.
"""
if not run_key:
return None
now = time.time()
with self._queue_lock:
live = [
ctx
for ctx in self._prompt_queue
if now - ctx.get("timestamp", 0) < self.prompt_timeout
]
index = next(
(
i
for i, ctx in enumerate(live)
if ctx.get("run_key") == run_key
and (run_id is None or ctx.get("run_id") == run_id)
),
None,
)
if index is None:
self._prompt_queue = live
return None
found = live[index]
# Same-key entries ahead of it came from earlier runs: stale
stale_or_found = {
i for i in range(index + 1) if live[i].get("run_key") == run_key
}
self._prompt_queue = [
ctx for i, ctx in enumerate(live) if i not in stale_or_found
]
return found
def pop_next_prompt(self) -> Optional[Dict[str, Any]]:
"""Pop the next prompt from the batch queue (FIFO).
@@ -565,46 +463,6 @@ class PromptTracker:
}
_KEYED_INIT_LOCK = threading.Lock()
# Run keys remembered for the image monitor (oldest forgotten first)
MAX_KEYED_RUNS = 256
class KeyedRuns:
"""The latest tracked entry per run key, and how many its run produced.
More than one entry for a key in one run (a batch fed by a list) means the
latest one alone cannot say which item an image belongs to.
"""
def __init__(self, max_keys: int = MAX_KEYED_RUNS):
self._max = max_keys
self._latest: "OrderedDict[str, Tuple[Dict[str, Any], Optional[str], int]]"
self._latest = OrderedDict()
self._lock = threading.Lock()
def record(self, context: Dict[str, Any]) -> None:
key, run_id = context["run_key"], context.get("run_id")
with self._lock:
previous = self._latest.pop(key, None)
same_run = (
previous is not None and run_id is not None and previous[1] == run_id
)
count = previous[2] + 1 if same_run else 1
self._latest[key] = (context, run_id, count)
while len(self._latest) > self._max:
self._latest.popitem(last=False)
def snapshot(self, timeout: float) -> Dict[str, Tuple[Dict[str, Any], Any, int]]:
now = time.time()
with self._lock:
return {
key: entry
for key, entry in self._latest.items()
if now - entry[0].get("timestamp", 0) < timeout
}
class PromptExecutionContext:
"""Context manager for prompt executions.
-129
View File
@@ -1,129 +0,0 @@
"""Prove that a tracked prompt belongs to the run that saved an image (#342).
A PromptManager node registers the text it executed under a *run key*: a digest
of the node and everything upstream of it in the queued API graph. The image
embeds the same graph, so the image monitor can compute the same key for the
image's positive prompt node. Equal keys mean the node ran with exactly the
inputs (and upstream nodes) that produced this image.
Equal keys do not always mean equal text. ComfyUI's cache signature also folds
in each node's ``IS_CHANGED``, which the key cannot see: an upstream node that
forces a re-run (a random line, an "always run" LLM) can give two runs one key
and two texts. The key is therefore never used alone: the monitor also requires
the entry to come from the run the snapshot taken when the image appeared names
(``PromptTracker.pop_for_key``, ``ImageGenerationHandler._find_prompt_from_run``).
Pure functions; the only host dependency (the executing run id) is optional.
"""
import hashlib
import json
from typing import Any, Dict, List, Optional
try:
from .prompt_graph import MAX_VISITS, run_prompt_nodes
except ImportError:
from utils.prompt_graph import MAX_VISITS, run_prompt_nodes
RUN_KEY_VERSION = 1
def is_link(value: Any) -> bool:
"""ComfyUI's own test (``comfy_execution.graph_utils.is_link``), exactly.
A two-item list of a ``str`` node id and an ``int`` or ``float`` slot (a
``bool`` slot passes too, as it does in ComfyUI). Tuples are values.
"""
if not isinstance(value, list):
return False
if len(value) != 2:
return False
if not isinstance(value[0], str):
return False
if not isinstance(value[1], int) and not isinstance(value[1], float):
return False
return True
def _links(inputs: Any) -> List[str]:
if not isinstance(inputs, dict):
return []
return [value[0] for value in inputs.values() if is_link(value)]
def prompt_node_key(graph: Any, node_id: Any) -> Optional[str]:
"""Digest of ``node_id`` and its whole upstream closure, or None.
Only ``class_type`` and ``inputs`` count (titles and other UI fields do not
change what runs). None when the graph or node is missing, an upstream link
points at a missing node, the closure exceeds ``MAX_VISITS`` nodes, or a
value cannot be serialized: no key means no proof.
"""
if not isinstance(graph, dict) or node_id is None:
return None
start = str(node_id)
closure: Dict[str, Any] = {}
stack = [start]
while stack:
current = stack.pop()
if current in closure:
continue
node = graph.get(current)
if not isinstance(node, dict):
return None
if len(closure) >= MAX_VISITS:
return None
inputs = node.get("inputs")
closure[current] = [node.get("class_type"), inputs]
stack.extend(_links(inputs))
try:
encoded = json.dumps(
[RUN_KEY_VERSION, start, closure],
sort_keys=True,
separators=(",", ":"),
ensure_ascii=True,
)
except (TypeError, ValueError):
return None
return hashlib.sha256(encoded.encode("ascii")).hexdigest()
def image_run_keys(metadata: Any) -> Dict[str, str]:
"""``{node_id: run key}`` for each positive PromptManager node of an image.
Empty without an embedded API graph, and empty when any positive node has
no key: one unkeyable node must not leave another looking like the only
positive prompt (a two-pass workflow would link the wrong pass).
"""
graph = metadata.get("prompt") if isinstance(metadata, dict) else None
if isinstance(graph, str):
try:
graph = json.loads(graph)
except (TypeError, ValueError, RecursionError):
return {}
if not isinstance(graph, dict):
return {}
keys = {}
for node_id in run_prompt_nodes(graph):
key = prompt_node_key(graph, node_id)
if key is None:
return {}
keys[node_id] = key
return keys
def executing_run_id() -> Optional[str]:
"""The ComfyUI prompt id currently executing on this thread, if the host says.
None outside ComfyUI or on hosts without ``get_executing_context``.
"""
try:
from comfy_execution.utils import get_executing_context
except Exception:
return None
try:
context = get_executing_context()
except Exception:
return None
run_id = getattr(context, "prompt_id", None)
return run_id if isinstance(run_id, str) and run_id else None
+2 -5
View File
@@ -50,8 +50,7 @@ export function cardMedia(image, lazy = false) {
/**
* One card for `image`. `checked` (a boolean) adds the D8 selection checkbox, `favoriteBusy` paints a pending star
* (D7, shown wherever the image reports `favorite`), `media: false` leaves the open button empty for a file that
* cannot be served, and `showInfo: false` hides the title and file line (D97). A `linked_only` image (LoRA media in
* LoraManager's folders, #355) has no Tag action.
* cannot be served, and `showInfo: false` hides the title and file line (D97).
*/
export function createImageCard(image, { checked = null, favoriteBusy = false, media = true, showInfo = true } = {}) {
const card = el('article', undefined, 'media-card');
@@ -88,9 +87,7 @@ export function createImageCard(image, { checked = null, favoriteBusy = false, m
const inspectButton = createButton('', 'icon-btn');
inspectButton.append(createWorkspaceIcon('info'));
inspectButton.setAttribute('aria-label', `Inspect ${image.filename}`);
actions.append(inspectButton);
// #355: LoRA-linked media lives in LoraManager's folders; tags are kept for output files only.
if (image.linked_only !== true) actions.append(createTagButton(image));
actions.append(inspectButton, createTagButton(image));
info.append(text, actions);
card.append(art, info);
return card;
+11 -43
View File
@@ -16,9 +16,6 @@
* - "LoRA reference data" → "Review import" (D107; the concept's button only simulated a job): the server
* reviews the import (`POST /lora/import/preview`, read-only) and the shared review dialog shows what it replaces,
* what saved pairs block, what it downloads and the backup box; Start runs the "LoRA reference import" Activity job.
* #355: its first row, "LoraManager", names the settings source, storage mode and library the server read (the
* settings file is the row's hover title); "Found" counts metadata files, LoRAs and LoRAs with local images. When
* LoraManager's settings cannot be used, the plan has no token and the review says why and offers only Close.
*/
import { createButton } from '../shared/toolkit/core.js';
import { createBridgeClient } from '../shared/toolkit/prompt-bridge.js';
@@ -58,9 +55,6 @@ export const KEY_COPY = Object.freeze({ saved: 'A key is saved.', none: 'No key
export const NOTHING_TO_IMPORT = 'There is nothing to import; nothing was changed.';
export const BLOCKED_NOTE = 'Delete or change the saved pairs that use these prompts, then review the import again. Nothing was changed.';
export const CACHE_NOTE = 'Example images downloaded before a failure stay in the image cache.';
export const LAYOUT_TERM = 'LoraManager';
/** #355: the close-only note for LoraManager settings the server cannot use. */
export const LAYOUT_UNUSABLE = reason => `LoraManager settings: not usable: ${reason}. Nothing was changed.`;
const number = value => value.toLocaleString('en-US');
@@ -181,38 +175,17 @@ export function keyText(values, changes) {
return values.lora_has_key ? KEY_COPY.saved : KEY_COPY.none;
}
/** #355: the "LoraManager" row: which settings, storage mode and library the server read; null without a layout. */
export function layoutText(layout) {
if (!layout || typeof layout !== 'object') return null;
if (layout.error) return `Not usable: ${layout.error}`;
return `${layout.source} settings · ${layout.mode} storage · library ${layout.library}`;
}
/** The note of a review that offers only Close, or null when the plan can start. */
export function closeNote(plan) {
if (!plan.token) {
if (plan.layout?.error) return LAYOUT_UNUSABLE(plan.layout.error);
return plan.blocked_by_pairs ? BLOCKED_NOTE : NOTHING_TO_IMPORT;
}
return !plan.with_metadata && !plan.current_imports ? NOTHING_TO_IMPORT : null;
}
/** D107: the import review's rows from the server's plan. */
export function importGrid(plan) {
const downloads = Number.isSafeInteger(plan.planned_downloads) ? plan.planned_downloads : 0;
const network = downloads
? `Downloads up to ${number(downloads)} example ${noun(downloads, 'image')} from CivitAI${plan.network?.api_key ? ' with your saved key' : ''}`
: 'No downloads';
const count = value => (Number.isSafeInteger(value) ? value : 0);
const found = count(plan.lora_files);
const readable = count(plan.with_metadata);
const local = count(plan.with_local_media);
const layout = layoutText(plan.layout);
const found = Number.isSafeInteger(plan.lora_files) ? plan.lora_files : 0;
const readable = Number.isSafeInteger(plan.with_metadata) ? plan.with_metadata : 0;
return [
...(layout ? [[LAYOUT_TERM, layout]] : []),
['Scope', plan.scope?.label || `${number(readable)} ${noun(readable, 'LoRA')}`],
['Found', `${number(found)} metadata ${noun(found, 'file')} · ${number(readable)} ${noun(readable, 'LoRA')} · `
+ `${number(local)} with local images`],
['Found', `${number(found)} metadata ${noun(found, 'file')} · ${number(readable)} readable`],
['Protected', protectedText(plan.protected)],
['Network', network],
['Backup', backupText(plan.backup)],
@@ -232,18 +205,11 @@ function namesList(plan) {
return [list];
}
/** The settings file (home shown as `~`) as the LoraManager row's hover title. */
function settingsTitle(dd, plan) {
if (typeof plan.layout?.settings_file === 'string') dd.title = plan.layout.settings_file;
}
function closeOnly({ scope, trigger, plan, note }) {
const grid = el('dl', undefined, 'meta-grid');
for (const [term, value] of importGrid(plan)) {
const entry = el('div');
const dd = el('dd', value);
if (term === LAYOUT_TERM) settingsTitle(dd, plan);
entry.append(el('dt', term), dd);
entry.append(el('dt', term), el('dd', value));
grid.append(entry);
}
const close = createButton('Close');
@@ -257,19 +223,21 @@ function closeOnly({ scope, trigger, plan, note }) {
/** Show the reviewed import; Start applies exactly this token (D107). */
export function showImportPlan(plan, { api, store, scope, trigger = null, reviewAgain, notify = showWorkspaceToast,
navigate = navigateTo, backup = () => downloadDatabaseBackup() }) {
const note = closeNote(plan);
if (note) {
closeOnly({ scope, trigger, plan, note });
if (!plan.token) {
closeOnly({ scope, trigger, plan, note: plan.blocked_by_pairs ? BLOCKED_NOTE : NOTHING_TO_IMPORT });
return;
}
if (!plan.with_metadata && !plan.current_imports) {
closeOnly({ scope, trigger, plan, note: NOTHING_TO_IMPORT });
return;
}
const removesAll = !plan.with_metadata;
const sentence = [planNote(plan), removesAll ? 'No LoRA metadata was found: this removes every imported LoRA prompt.' : '', CACHE_NOTE]
.filter(Boolean).join(' ');
const review = openReviewDialog({ scope, trigger, title: IMPORT_TITLE, heading: 'LoRA reference data', grid: importGrid(plan),
openReviewDialog({ scope, trigger, title: IMPORT_TITLE, heading: 'LoRA reference data', grid: importGrid(plan),
options: namesList(plan), sentence, backup: { copy: BACKUP_COPY, onBackup: backup }, cancelLabel: 'Cancel',
primary: { label: IMPORT_START, danger: removesAll },
onStart: () => applyPlan({ api, store, scope, token: plan.token, reviewAgain, notify, navigate, url: IMPORT_APPLY_URL }) });
if (typeof plan.layout?.settings_file === 'string') review.setTitle(LAYOUT_TERM, plan.layout.settings_file);
}
/** Review the import (read-only) and show its plan; the pressed button reads "Reviewing…" until it answered. */
+1 -5
View File
@@ -429,7 +429,7 @@ function requireGrid(grid) {
* primary (`ca-primary`, or the concept's `danger` when `primary.danger`). The primary closes the dialog, then calls
* `onStart()`; `onDismiss()` runs once when the dialog closes any other way (Cancel, ×, Escape, a replacing dialog, the
* page leaving), so a staged review is discarded (S9). Returns `{setValue(term, text)}` for a grid value that follows
* an option, and `{setTitle(term, text)}` for a value's hover title.
* an option.
*/
export function openReviewDialog({ scope, title, heading = null, grid, options = [], sentence = null, backup = null,
cancelLabel, primary, onStart, trigger = null, onDismiss = null }) {
@@ -475,10 +475,6 @@ export function openReviewDialog({ scope, title, heading = null, grid, options =
const dd = values.get(term);
if (dd) dd.textContent = text;
},
setTitle(term, text) {
const dd = values.get(term);
if (dd) dd.title = text;
},
};
}